lance_encoding/encoder/
structural.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Version-free structural field encoder builders.
5
6use std::sync::Arc;
7
8use arrow_schema::DataType;
9use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt};
10
11pub use crate::encodings::logical::primitive::PrimitivePageEncoding;
12
13use crate::encodings::logical::{
14    blob::{BlobStructuralEncoder, BlobV2StructuralEncoder},
15    fixed_size_list::FixedSizeListStructuralEncoder,
16    list::ListStructuralEncoder,
17    map::MapStructuralEncoder,
18    primitive::PrimitiveStructuralEncoder,
19    r#struct::StructStructuralEncoder,
20};
21
22use super::{ColumnIndexSequence, FieldEncoder, FieldEncodingContext};
23
24/// Encode primitive leaves, primitive fixed-size lists, dictionaries, and
25/// packed or empty structs using one concrete primitive page grammar.
26#[derive(Debug, Clone)]
27pub struct PrimitiveFieldEncoding {
28    page_encodings: Arc<[PrimitivePageEncoding]>,
29}
30
31impl PrimitiveFieldEncoding {
32    /// Create a primitive field mechanism from ordered executable page behaviors.
33    pub fn new(page_encodings: impl IntoIterator<Item = PrimitivePageEncoding>) -> Self {
34        Self {
35            page_encodings: page_encodings.into_iter().collect(),
36        }
37    }
38
39    fn is_primitive_type(data_type: &DataType) -> bool {
40        match data_type {
41            DataType::FixedSizeList(inner, _) => Self::is_primitive_type(inner.data_type()),
42            _ => matches!(
43                data_type,
44                DataType::Boolean
45                    | DataType::Date32
46                    | DataType::Date64
47                    | DataType::Decimal128(_, _)
48                    | DataType::Decimal256(_, _)
49                    | DataType::Duration(_)
50                    | DataType::Float16
51                    | DataType::Float32
52                    | DataType::Float64
53                    | DataType::Int16
54                    | DataType::Int32
55                    | DataType::Int64
56                    | DataType::Int8
57                    | DataType::Interval(_)
58                    | DataType::Null
59                    | DataType::Time32(_)
60                    | DataType::Time64(_)
61                    | DataType::Timestamp(_, _)
62                    | DataType::UInt16
63                    | DataType::UInt32
64                    | DataType::UInt64
65                    | DataType::UInt8
66                    | DataType::FixedSizeBinary(_)
67                    | DataType::Binary
68                    | DataType::LargeBinary
69                    | DataType::Utf8
70                    | DataType::LargeUtf8,
71            ),
72        }
73    }
74
75    /// Compose a primitive encoder at an already allocated physical column.
76    /// File grammars use this when another logical type, such as a Blob,
77    /// supplies the physical descriptor field instead of the logical field.
78    fn create_at(
79        &self,
80        field: Field,
81        column_index: u32,
82        context: &FieldEncodingContext<'_>,
83    ) -> Result<Box<dyn FieldEncoder>> {
84        Ok(Box::new(PrimitiveStructuralEncoder::try_new(
85            context.options,
86            self.page_encodings.clone(),
87            column_index,
88            field,
89            Arc::new(context.root_field_metadata.clone()),
90        )?))
91    }
92
93    /// Create a primitive field encoder when this mechanism recognizes `field`.
94    pub fn try_create(
95        &self,
96        field: &Field,
97        column_index: &mut ColumnIndexSequence,
98        context: &FieldEncodingContext<'_>,
99    ) -> Result<Option<Box<dyn FieldEncoder>>> {
100        if field.is_blob() {
101            return Ok(None);
102        }
103
104        let data_type = field.data_type();
105        let is_primitive = Self::is_primitive_type(&data_type);
106        let is_packed_or_empty_struct = matches!(
107            &data_type,
108            DataType::Struct(fields) if field.is_packed_struct() || fields.is_empty()
109        );
110        let is_primitive_dictionary = matches!(
111            &data_type,
112            DataType::Dictionary(_, value_type) if Self::is_primitive_type(value_type)
113        );
114
115        if !is_primitive && !is_packed_or_empty_struct && !is_primitive_dictionary {
116            if let DataType::Dictionary(_, value_type) = data_type {
117                return Err(Error::not_supported_source(
118                    format!(
119                        "cannot encode a dictionary column whose value type is a logical type ({})",
120                        value_type
121                    )
122                    .into(),
123                ));
124            }
125            return Ok(None);
126        }
127
128        Ok(Some(self.create_at(
129            field.clone(),
130            column_index.next_column_index(field.id as u32),
131            context,
132        )?))
133    }
134}
135
136/// Create the original binary blob descriptor when `field` matches.
137pub fn try_create_binary_blob(
138    primitive: &PrimitiveFieldEncoding,
139    field: &Field,
140    column_index: &mut ColumnIndexSequence,
141    context: &FieldEncodingContext<'_>,
142) -> Result<Option<Box<dyn FieldEncoder>>> {
143    if !field.is_blob() || !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) {
144        return Ok(None);
145    }
146    let descriptor_column_index = column_index.next_column_index(field.id as u32);
147    Ok(Some(Box::new(BlobStructuralEncoder::new(
148        field,
149        |descriptor_field| primitive.create_at(descriptor_field, descriptor_column_index, context),
150    )?)))
151}
152
153/// Create the structural blob descriptor when `field` matches.
154pub fn try_create_structural_blob(
155    primitive: &PrimitiveFieldEncoding,
156    field: &Field,
157    column_index: &mut ColumnIndexSequence,
158    context: &FieldEncodingContext<'_>,
159) -> Result<Option<Box<dyn FieldEncoder>>> {
160    if !field.is_blob() || !matches!(field.data_type(), DataType::Struct(_)) {
161        return Ok(None);
162    }
163    let descriptor_column_index = column_index.next_column_index(field.id as u32);
164    Ok(Some(Box::new(BlobV2StructuralEncoder::new(
165        field,
166        |descriptor_field| primitive.create_at(descriptor_field, descriptor_column_index, context),
167    )?)))
168}
169
170/// Create a variable-size list encoder when `field` matches.
171pub fn try_create_list(
172    field: &Field,
173    column_index: &mut ColumnIndexSequence,
174    context: &FieldEncodingContext<'_>,
175) -> Result<Option<Box<dyn FieldEncoder>>> {
176    if !matches!(
177        field.data_type(),
178        DataType::List(_) | DataType::LargeList(_)
179    ) {
180        return Ok(None);
181    }
182    let child = field.children.first().expect_ok()?;
183    let child_encoder = context
184        .strategy
185        .create_field_encoder(child, column_index, context)?;
186    Ok(Some(Box::new(ListStructuralEncoder::new(
187        context.options.keep_original_array,
188        child_encoder,
189    ))))
190}
191
192/// Create a fixed-size-list encoder whose child is a struct when applicable.
193pub fn try_create_structural_fixed_size_list(
194    field: &Field,
195    column_index: &mut ColumnIndexSequence,
196    context: &FieldEncodingContext<'_>,
197) -> Result<Option<Box<dyn FieldEncoder>>> {
198    if !matches!(
199        field.data_type(),
200        DataType::FixedSizeList(inner, _) if matches!(inner.data_type(), DataType::Struct(_))
201    ) {
202        return Ok(None);
203    }
204    let child = field.children.first().expect_ok()?;
205    let child_encoder = context
206        .strategy
207        .create_field_encoder(child, column_index, context)?;
208    Ok(Some(Box::new(FixedSizeListStructuralEncoder::new(
209        context.options.keep_original_array,
210        child_encoder,
211    ))))
212}
213
214/// Create an Arrow map encoder when `field` matches.
215pub fn try_create_map(
216    field: &Field,
217    column_index: &mut ColumnIndexSequence,
218    context: &FieldEncodingContext<'_>,
219) -> Result<Option<Box<dyn FieldEncoder>>> {
220    let DataType::Map(_, keys_sorted) = field.data_type() else {
221        return Ok(None);
222    };
223    if keys_sorted {
224        return Err(Error::not_supported_source(
225            format!(
226                "Map data type is not supported with keys_sorted=true now, current value is {}",
227                keys_sorted
228            )
229            .into(),
230        ));
231    }
232    let entries_child = field
233        .children
234        .first()
235        .ok_or_else(|| Error::schema("Map should have an entries child".to_string()))?;
236    let DataType::Struct(struct_fields) = entries_child.data_type() else {
237        return Err(Error::schema(
238            "Map entries field must be a Struct<key, value>".to_string(),
239        ));
240    };
241    if struct_fields.len() < 2 {
242        return Err(Error::schema(
243            "Map entries struct must contain both key and value fields".to_string(),
244        ));
245    }
246    let key_field = &struct_fields[0];
247    if key_field.is_nullable() {
248        return Err(Error::schema(format!(
249            "Map key field '{}' must be non-nullable according to Arrow Map specification",
250            key_field.name()
251        )));
252    }
253    let child_encoder =
254        context
255            .strategy
256            .create_field_encoder(entries_child, column_index, context)?;
257    Ok(Some(Box::new(MapStructuralEncoder::new(
258        context.options.keep_original_array,
259        child_encoder,
260    ))))
261}
262
263/// Create a non-packed, non-empty struct encoder when `field` matches.
264pub fn try_create_struct(
265    field: &Field,
266    column_index: &mut ColumnIndexSequence,
267    context: &FieldEncodingContext<'_>,
268) -> Result<Option<Box<dyn FieldEncoder>>> {
269    let DataType::Struct(fields) = field.data_type() else {
270        return Ok(None);
271    };
272    if field.is_blob() || field.is_packed_struct() || fields.is_empty() {
273        return Ok(None);
274    }
275    let children_encoders = field
276        .children
277        .iter()
278        .map(|child| {
279            context
280                .strategy
281                .create_field_encoder(child, column_index, context)
282        })
283        .collect::<Result<Vec<_>>>()?;
284    Ok(Some(Box::new(StructStructuralEncoder::new(
285        context.options.keep_original_array,
286        children_encoders,
287    ))))
288}
lance_encoding/encoder/structural.rs

lance_encoding/encoder/
structural.rs
