Skip to main content

lance_encoding/format/
planar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Translate planar wire metadata to and from shared engine descriptions.
5//! These conversions do not use a stable-format protobuf.
6
7use lance_core::{Error, Result};
8
9use super::{metadata, pb23};
10
11fn missing(label: &str) -> Error {
12    Error::invalid_input(format!("Planar page is missing {label}"))
13}
14
15impl pb23::PlanarLayout {
16    /// Return the sparse structure, if this page uses slot-domain mappings.
17    pub fn sparse(&self) -> Option<&pb23::SparseStructure> {
18        match self.structure.as_ref()?.codec.as_ref()? {
19            pb23::structural_plane::Codec::Sparse(sparse) => Some(sparse),
20            _ => None,
21        }
22    }
23
24    /// Build the execution description for dense mini-block, zipped, or constant values.
25    /// Sparse structure is scheduled independently of its value engine.
26    pub fn dense_execution_layout(&self) -> Result<metadata::PageLayout> {
27        use metadata::page_layout::Layout;
28        use pb23::{structural_plane::Codec as Structure, value_plane::Codec as Values};
29        let structure = self
30            .structure
31            .as_ref()
32            .and_then(|s| s.codec.as_ref())
33            .ok_or_else(|| missing("structure"))?;
34        let Structure::Dense(dense) = structure else {
35            return Err(Error::invalid_input("Expected dense planar structure"));
36        };
37        if dense.layers.is_empty()
38            || dense
39                .layers
40                .iter()
41                .any(|&layer| metadata::RepDefLayer::try_from(layer).is_err() || layer == 0)
42        {
43            return Err(Error::invalid_input(
44                "Planar dense layers are empty or invalid",
45            ));
46        }
47        let values = self
48            .values
49            .as_ref()
50            .and_then(|v| v.codec.as_ref())
51            .ok_or_else(|| missing("values"))?;
52        let layout = match values {
53            Values::Chunked(values) => {
54                let levels = values
55                    .levels
56                    .as_ref()
57                    .ok_or_else(|| missing("dense chunk levels"))?;
58                Layout::MiniBlockLayout(metadata::MiniBlockLayout {
59                    rep_compression: levels.repetition.clone(),
60                    def_compression: levels.definition.clone(),
61                    repetition_index_depth: levels.repetition_index_depth,
62                    value_compression: values.compression.clone(),
63                    dictionary: values.dictionary.clone(),
64                    num_dictionary_items: values.num_dictionary_items,
65                    num_buffers: values.num_buffers,
66                    num_items: values.num_items,
67                    layers: dense.layers.clone(),
68                    has_large_chunk: true,
69                })
70            }
71            Values::Zipped(values) => {
72                let levels = values
73                    .levels
74                    .as_ref()
75                    .ok_or_else(|| missing("dense zip levels"))?;
76                Layout::FullZipLayout(values.full_zip(dense.layers.clone(), Some(levels))?)
77            }
78            Values::Constant(value) => {
79                if value.scalar.is_empty() || value.scalar.len() > 32 {
80                    return Err(Error::invalid_input(
81                        "Planar constant scalar must contain 1..=32 bytes",
82                    ));
83                }
84                if dense
85                    .layers
86                    .iter()
87                    .any(|&layer| layer != metadata::RepDefLayer::RepdefAllValidItem as i32)
88                {
89                    return Err(Error::invalid_input(
90                        "Dense constant requires trivial all-valid structure",
91                    ));
92                }
93                Layout::ConstantLayout(metadata::ConstantLayout {
94                    inline_value: Some(value.scalar.clone()),
95                    layers: dense.layers.clone(),
96                    ..Default::default()
97                })
98            }
99            Values::Indexed(_) => {
100                return Err(Error::invalid_input(
101                    "Indexed values require sparse structure",
102                ));
103            }
104            Values::None(_) => {
105                if dense.layers.as_slice() != [metadata::RepDefLayer::RepdefNullableItem as i32] {
106                    return Err(Error::invalid_input(
107                        "Dense no-values requires a single null leaf layer",
108                    ));
109                }
110                Layout::ConstantLayout(metadata::ConstantLayout {
111                    layers: dense.layers.clone(),
112                    ..Default::default()
113                })
114            }
115        };
116        Ok(metadata::PageLayout {
117            layout: Some(layout),
118        })
119    }
120}
121
122impl pb23::ZippedValues {
123    /// Build a zipped execution description, with optional dense controls.
124    pub(crate) fn full_zip(
125        &self,
126        layers: Vec<i32>,
127        levels: Option<&pb23::DenseZipLevels>,
128    ) -> Result<metadata::FullZipLayout> {
129        let details = match self
130            .addressing
131            .as_ref()
132            .ok_or_else(|| missing("zipped addressing"))?
133        {
134            pb23::zipped_values::Addressing::BitsPerValue(bits) => {
135                metadata::full_zip_layout::Details::BitsPerValue(*bits)
136            }
137            pb23::zipped_values::Addressing::BitsPerOffset(bits) => {
138                metadata::full_zip_layout::Details::BitsPerOffset(*bits)
139            }
140        };
141        Ok(metadata::FullZipLayout {
142            bits_rep: levels.map_or(0, |l| l.bits_repetition),
143            bits_def: levels.map_or(0, |l| l.bits_definition),
144            details: Some(details),
145            num_items: self
146                .num_items
147                .try_into()
148                .map_err(|_| Error::invalid_input("Zipped item count exceeds u32"))?,
149            num_visible_items: self
150                .num_visible_items
151                .try_into()
152                .map_err(|_| Error::invalid_input("Zipped leaf count exceeds u32"))?,
153            value_compression: self.compression.clone(),
154            layers,
155        })
156    }
157}
158
159/// Translate a newly encoded dense page without changing its physical buffers.
160pub(crate) fn from_dense_execution(layout: metadata::PageLayout) -> Result<pb23::PlanarLayout> {
161    use metadata::page_layout::Layout;
162    use pb23::value_plane::Codec;
163    let (layers, values) = match layout.layout.ok_or_else(|| missing("dense layout"))? {
164        Layout::MiniBlockLayout(m) => (
165            m.layers,
166            Codec::Chunked(pb23::ChunkedValues {
167                compression: m.value_compression,
168                levels: Some(pb23::DenseChunkLevels {
169                    repetition: m.rep_compression,
170                    definition: m.def_compression,
171                    repetition_index_depth: m.repetition_index_depth,
172                }),
173                dictionary: m.dictionary,
174                num_dictionary_items: m.num_dictionary_items,
175                num_buffers: m.num_buffers,
176                num_items: m.num_items,
177            }),
178        ),
179        Layout::FullZipLayout(m) => (
180            m.layers,
181            Codec::Zipped(pb23::ZippedValues {
182                compression: m.value_compression,
183                levels: Some(pb23::DenseZipLevels {
184                    bits_repetition: m.bits_rep,
185                    bits_definition: m.bits_def,
186                }),
187                addressing: m.details.map(|d| match d {
188                    metadata::full_zip_layout::Details::BitsPerValue(v) => {
189                        pb23::zipped_values::Addressing::BitsPerValue(v)
190                    }
191                    metadata::full_zip_layout::Details::BitsPerOffset(v) => {
192                        pb23::zipped_values::Addressing::BitsPerOffset(v)
193                    }
194                }),
195                num_items: u64::from(m.num_items),
196                num_visible_items: u64::from(m.num_visible_items),
197            }),
198        ),
199        Layout::ConstantLayout(m) => {
200            let value = match m.inline_value {
201                Some(scalar) => Codec::Constant(pb23::ConstantValue { scalar }),
202                None => Codec::None(pb23::NoValues {}),
203            };
204            (m.layers, value)
205        }
206        _ => {
207            return Err(Error::invalid_input(
208                "Only dense value engines can produce dense planar metadata",
209            ));
210        }
211    };
212    Ok(pb23::PlanarLayout {
213        structure: Some(pb23::StructuralPlane {
214            codec: Some(pb23::structural_plane::Codec::Dense(pb23::DenseStructure {
215                layers,
216            })),
217        }),
218        values: Some(pb23::ValuePlane {
219            codec: Some(values),
220        }),
221    })
222}