lance_encoding/encodings/logical/
struct.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{
5    collections::{BinaryHeap, VecDeque},
6    ops::Range,
7    sync::Arc,
8};
9
10use super::{
11    fixed_size_list::StructuralFixedSizeListDecoder, list::StructuralListDecoder,
12    map::StructuralMapDecoder, primitive::StructuralPrimitiveFieldDecoder,
13};
14use crate::{
15    decoder::{
16        DecodedArray, FilterExpression, LoadedPageShard, NextDecodeTask, PageEncoding,
17        ScheduledScanLine, SchedulerContext, StructuralDecodeArrayTask, StructuralFieldDecoder,
18        StructuralFieldScheduler, StructuralSchedulingJob,
19    },
20    encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers},
21    format::pb,
22    repdef::{CompositeRepDefUnraveler, RepDefBuilder},
23};
24use arrow_array::{Array, ArrayRef, StructArray, cast::AsArray};
25use arrow_schema::{DataType, Fields};
26use futures::{
27    FutureExt, StreamExt, TryStreamExt,
28    future::BoxFuture,
29    stream::{FuturesOrdered, FuturesUnordered},
30};
31use itertools::Itertools;
32use lance_arrow::FieldExt;
33use lance_arrow::{deepcopy::deep_copy_nulls, r#struct::StructArrayExt};
34use lance_core::{Error, Result, datatypes::validate_fixed_size_list_dimensions};
35use log::trace;
36
37#[derive(Debug)]
38struct StructuralSchedulingJobWithStatus<'a> {
39    col_idx: u32,
40    col_name: &'a str,
41    job: Box<dyn StructuralSchedulingJob + 'a>,
42    rows_scheduled: u64,
43    rows_remaining: u64,
44    ready_scan_lines: VecDeque<ScheduledScanLine>,
45}
46
47impl PartialEq for StructuralSchedulingJobWithStatus<'_> {
48    fn eq(&self, other: &Self) -> bool {
49        self.col_idx == other.col_idx
50    }
51}
52
53impl Eq for StructuralSchedulingJobWithStatus<'_> {}
54
55impl PartialOrd for StructuralSchedulingJobWithStatus<'_> {
56    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
57        Some(self.cmp(other))
58    }
59}
60
61impl Ord for StructuralSchedulingJobWithStatus<'_> {
62    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
63        // Note this is reversed to make it min-heap
64        other.rows_scheduled.cmp(&self.rows_scheduled)
65    }
66}
67
68/// Scheduling job for struct data
69///
70/// The order in which we schedule the children is important.  We want to schedule the child
71/// with the least amount of data first.
72///
73/// This allows us to decode entire rows as quickly as possible
74#[derive(Debug)]
75struct RepDefStructSchedulingJob<'a> {
76    /// A min-heap whose key is the # of rows currently scheduled
77    children: BinaryHeap<StructuralSchedulingJobWithStatus<'a>>,
78    rows_scheduled: u64,
79    num_rows: u64,
80}
81
82impl<'a> RepDefStructSchedulingJob<'a> {
83    fn new(
84        scheduler: &'a StructuralStructScheduler,
85        children: Vec<Box<dyn StructuralSchedulingJob + 'a>>,
86        num_rows: u64,
87    ) -> Self {
88        let children = children
89            .into_iter()
90            .enumerate()
91            .map(|(idx, job)| StructuralSchedulingJobWithStatus {
92                col_idx: idx as u32,
93                col_name: scheduler.child_fields[idx].name(),
94                job,
95                rows_scheduled: 0,
96                rows_remaining: num_rows,
97                ready_scan_lines: VecDeque::new(),
98            })
99            .collect::<BinaryHeap<_>>();
100        Self {
101            children,
102            rows_scheduled: 0,
103            num_rows,
104        }
105    }
106}
107
108impl StructuralSchedulingJob for RepDefStructSchedulingJob<'_> {
109    fn schedule_next(
110        &mut self,
111        mut context: &mut SchedulerContext,
112    ) -> Result<Vec<ScheduledScanLine>> {
113        if self.children.is_empty() {
114            // Special path for empty structs
115            if self.rows_scheduled == self.num_rows {
116                return Ok(Vec::new());
117            }
118            self.rows_scheduled = self.num_rows;
119            return Ok(vec![ScheduledScanLine {
120                decoders: Vec::new(),
121                rows_scheduled: self.num_rows,
122            }]);
123        }
124
125        let mut decoders = Vec::new();
126        let old_rows_scheduled = self.rows_scheduled;
127        // Schedule as many children as we need to until we have scheduled at least one
128        // complete row
129        while old_rows_scheduled == self.rows_scheduled {
130            if self.children.is_empty() {
131                // Early exit when schedulers are exhausted prematurely (TODO: does this still happen?)
132                return Ok(Vec::new());
133            }
134            let mut next_child = self.children.pop().unwrap();
135            if next_child.ready_scan_lines.is_empty() {
136                let scoped = context.push(next_child.col_name, next_child.col_idx);
137                let child_scans = next_child.job.schedule_next(scoped.context)?;
138                context = scoped.pop();
139                if child_scans.is_empty() {
140                    // Continue without pushing next_child back onto the heap (it is done)
141                    continue;
142                }
143                next_child.ready_scan_lines.extend(child_scans);
144            }
145            let child_scan = next_child.ready_scan_lines.pop_front().unwrap();
146            trace!(
147                "Scheduled {} rows for child {}",
148                child_scan.rows_scheduled, next_child.col_idx
149            );
150            next_child.rows_scheduled += child_scan.rows_scheduled;
151            next_child.rows_remaining -= child_scan.rows_scheduled;
152            decoders.extend(child_scan.decoders);
153            self.children.push(next_child);
154            self.rows_scheduled = self.children.peek().unwrap().rows_scheduled;
155        }
156        let struct_rows_scheduled = self.rows_scheduled - old_rows_scheduled;
157        Ok(vec![ScheduledScanLine {
158            decoders,
159            rows_scheduled: struct_rows_scheduled,
160        }])
161    }
162}
163
164/// A scheduler for structs
165///
166/// The implementation is actually a bit more tricky than one might initially think.  We can't just
167/// go through and schedule each column one after the other.  This would mean our decode can't start
168/// until nearly all the data has arrived (since we need data from each column to yield a batch)
169///
170/// Instead, we schedule in row-major fashion
171///
172/// Note: this scheduler is the starting point for all decoding.  This is because we treat the top-level
173/// record batch as a non-nullable struct.
174#[derive(Debug)]
175pub struct StructuralStructScheduler {
176    children: Vec<Box<dyn StructuralFieldScheduler>>,
177    child_fields: Fields,
178}
179
180impl StructuralStructScheduler {
181    pub fn new(children: Vec<Box<dyn StructuralFieldScheduler>>, child_fields: Fields) -> Self {
182        Self {
183            children,
184            child_fields,
185        }
186    }
187}
188
189impl StructuralFieldScheduler for StructuralStructScheduler {
190    fn schedule_ranges<'a>(
191        &'a self,
192        ranges: &[Range<u64>],
193        filter: &FilterExpression,
194    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
195        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
196
197        let child_schedulers = self
198            .children
199            .iter()
200            .map(|child| child.schedule_ranges(ranges, filter))
201            .collect::<Result<Vec<_>>>()?;
202
203        Ok(Box::new(RepDefStructSchedulingJob::new(
204            self,
205            child_schedulers,
206            num_rows,
207        )))
208    }
209
210    fn initialize<'a>(
211        &'a mut self,
212        requested_ranges: Option<&'a [Range<u64>]>,
213        filter: &'a FilterExpression,
214        context: &'a SchedulerContext,
215    ) -> BoxFuture<'a, Result<()>> {
216        // Struct children share the parent's row coordinates 1:1 (see
217        // `schedule_ranges`), so every child gets the same ranges.
218        let children_initialization = self
219            .children
220            .iter_mut()
221            .map(|child| child.initialize(requested_ranges, filter, context))
222            .collect::<FuturesUnordered<_>>();
223        async move {
224            children_initialization
225                .map(|res| res.map(|_| ()))
226                .try_collect::<Vec<_>>()
227                .await?;
228            Ok(())
229        }
230        .boxed()
231    }
232}
233
234#[derive(Debug)]
235pub struct StructuralStructDecoder {
236    children: Vec<Box<dyn StructuralFieldDecoder>>,
237    data_type: DataType,
238    child_fields: Fields,
239    // The root decoder is slightly different because it cannot have nulls
240    is_root: bool,
241    nullable: bool,
242}
243
244impl StructuralStructDecoder {
245    pub fn new(
246        fields: Fields,
247        should_validate: bool,
248        is_root: bool,
249        nullable: bool,
250    ) -> Result<Self> {
251        let children = fields
252            .iter()
253            .map(|field| Self::field_to_decoder(field, should_validate))
254            .collect::<Result<Vec<_>>>()?;
255        let data_type = DataType::Struct(fields.clone());
256        Ok(Self {
257            data_type,
258            children,
259            child_fields: fields,
260            is_root,
261            nullable: !is_root && nullable,
262        })
263    }
264
265    fn field_to_decoder(
266        field: &Arc<arrow_schema::Field>,
267        should_validate: bool,
268    ) -> Result<Box<dyn StructuralFieldDecoder>> {
269        match field.data_type() {
270            DataType::Struct(fields) => {
271                if field.is_packed_struct() || field.is_blob() {
272                    let decoder =
273                        StructuralPrimitiveFieldDecoder::new(&field.clone(), should_validate);
274                    Ok(Box::new(decoder))
275                } else {
276                    Ok(Box::new(Self::new(
277                        fields.clone(),
278                        should_validate,
279                        false,
280                        field.is_nullable(),
281                    )?))
282                }
283            }
284            DataType::List(child_field) | DataType::LargeList(child_field) => {
285                let child_decoder = Self::field_to_decoder(child_field, should_validate)?;
286                Ok(Box::new(StructuralListDecoder::new(
287                    child_decoder,
288                    field.data_type().clone(),
289                )))
290            }
291            DataType::FixedSizeList(child_field, _)
292                if matches!(child_field.data_type(), DataType::Struct(_)) =>
293            {
294                // The scheduler factories run the same guard, but the decoder tree can be
295                // built independently (e.g. `create_decode_stream`) so a zero dimension from
296                // a malformed schema must be rejected here as well.  Draining and unraveling
297                // validity both scale by the dimension and a zero would make that math
298                // degenerate.
299                validate_fixed_size_list_dimensions(field.name(), field.data_type())?;
300                // FixedSizeList containing Struct needs structural decoding
301                let child_decoder = Self::field_to_decoder(child_field, should_validate)?;
302                Ok(Box::new(StructuralFixedSizeListDecoder::new(
303                    child_decoder,
304                    field.data_type().clone(),
305                )))
306            }
307            DataType::Map(entries_field, keys_sorted) => {
308                if *keys_sorted {
309                    return Err(Error::not_supported_source(
310                        "Map data type with keys_sorted=true is not supported yet"
311                            .to_string()
312                            .into(),
313                    ));
314                }
315                let child_decoder = Self::field_to_decoder(entries_field, should_validate)?;
316                Ok(Box::new(StructuralMapDecoder::new(
317                    child_decoder,
318                    field.data_type().clone(),
319                )))
320            }
321            DataType::RunEndEncoded(_, _) => todo!(),
322            DataType::ListView(_) | DataType::LargeListView(_) => todo!(),
323            DataType::Union(_, _) => todo!(),
324            _ => Ok(Box::new(StructuralPrimitiveFieldDecoder::new(
325                field,
326                should_validate,
327            ))),
328        }
329    }
330
331    pub fn drain_batch_task(&mut self, num_rows: u64) -> Result<NextDecodeTask> {
332        let array_drain = self.drain(num_rows)?;
333        Ok(NextDecodeTask {
334            num_rows,
335            task: Box::new(array_drain),
336        })
337    }
338}
339
340impl StructuralFieldDecoder for StructuralStructDecoder {
341    fn accept_page(&mut self, mut child: LoadedPageShard) -> Result<()> {
342        // children with empty path should not be delivered to this method
343        let child_idx = child.path.pop_front().unwrap();
344        // This decoder is intended for one of our children
345        self.children[child_idx as usize].accept_page(child)?;
346        Ok(())
347    }
348
349    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
350        let child_tasks = self
351            .children
352            .iter_mut()
353            .map(|child| child.drain(num_rows))
354            .collect::<Result<Vec<_>>>()?;
355        Ok(Box::new(RepDefStructDecodeTask {
356            children: child_tasks,
357            child_fields: self.child_fields.clone(),
358            is_root: self.is_root,
359            num_rows,
360        }))
361    }
362
363    fn data_type(&self) -> &DataType {
364        &self.data_type
365    }
366
367    fn plan_decoded_bytes(&self, rows_remaining: u64) -> lance_core::Result<[u64; 8]> {
368        use crate::decoder::CANDIDATE_BATCH_SIZES;
369        let mut out = [0u64; 8];
370        for child in &self.children {
371            let child_bytes = child.plan_decoded_bytes(rows_remaining)?;
372            for (o, c) in out.iter_mut().zip(child_bytes) {
373                *o += c;
374            }
375        }
376        if self.nullable {
377            for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
378                let rows = (candidate as u64).min(rows_remaining);
379                out[i] += rows.div_ceil(8);
380            }
381        }
382        Ok(out)
383    }
384}
385
386#[derive(Debug)]
387struct RepDefStructDecodeTask {
388    children: Vec<Box<dyn StructuralDecodeArrayTask>>,
389    child_fields: Fields,
390    is_root: bool,
391    num_rows: u64,
392}
393
394impl StructuralDecodeArrayTask for RepDefStructDecodeTask {
395    fn decode(self: Box<Self>) -> Result<DecodedArray> {
396        if self.children.is_empty() {
397            return Ok(DecodedArray {
398                array: Arc::new(StructArray::new_empty_fields(self.num_rows as usize, None)),
399                repdef: CompositeRepDefUnraveler::new(vec![]),
400                data_size: 0,
401            });
402        }
403
404        let arrays = self
405            .children
406            .into_iter()
407            .map(|task| task.decode())
408            .collect::<Result<Vec<_>>>()?;
409        let mut children = Vec::with_capacity(arrays.len());
410        let mut repdefs = Vec::with_capacity(arrays.len());
411        let mut data_size = 0u64;
412        let mut arrays_iter = arrays.into_iter();
413        let first_array = arrays_iter.next().ok_or_else(|| {
414            Error::internal("Struct decoder unexpectedly has no child arrays".to_string())
415        })?;
416        let length = first_array.array.len();
417
418        // The repdef should be identical across all children at this point
419        repdefs.push(first_array.repdef);
420        data_size += first_array.data_size;
421        children.push(first_array.array);
422
423        for array in arrays_iter {
424            if length != array.array.len() {
425                return Err(Error::invalid_input_source(
426                    format!(
427                        "Struct child array length {} does not match sibling length {}",
428                        array.array.len(),
429                        length
430                    )
431                    .into(),
432                ));
433            }
434            data_size += array.data_size;
435            children.push(array.array);
436            repdefs.push(array.repdef);
437        }
438
439        // Dense rep/def state can retain child-specific repetition information after a child
440        // decoder finishes, so comparing dense siblings is not meaningful. If any child carries
441        // sparse state, keep a sparse child as the canonical structural plan and compare it with
442        // every other sparse sibling so sparse metadata is never silently discarded.
443        let primary_repdef = repdefs
444            .iter()
445            .position(CompositeRepDefUnraveler::has_sparse)
446            .unwrap_or(0);
447        let mut repdef = repdefs.swap_remove(primary_repdef);
448        if repdef.has_sparse() {
449            for sibling in repdefs {
450                if sibling.has_sparse() {
451                    repdef.add_compatibility_check(sibling);
452                }
453            }
454        }
455
456        let validity = if self.is_root {
457            repdef.ensure_exhausted()?;
458            None
459        } else {
460            repdef.unravel_validity(length)?
461        };
462
463        let array = StructArray::try_new(self.child_fields, children, validity)
464            .map_err(|e| Error::invalid_input_source(e.to_string().into()))?;
465        Ok(DecodedArray {
466            array: Arc::new(array),
467            repdef,
468            data_size,
469        })
470    }
471}
472
473/// A structural encoder for struct fields
474///
475/// The struct's validity is added to the rep/def builder
476/// and the builder is cloned to all children.
477pub struct StructStructuralEncoder {
478    keep_original_array: bool,
479    children: Vec<Box<dyn FieldEncoder>>,
480}
481
482impl StructStructuralEncoder {
483    pub fn new(keep_original_array: bool, children: Vec<Box<dyn FieldEncoder>>) -> Self {
484        Self {
485            keep_original_array,
486            children,
487        }
488    }
489}
490
491impl FieldEncoder for StructStructuralEncoder {
492    fn maybe_encode(
493        &mut self,
494        array: ArrayRef,
495        external_buffers: &mut OutOfLineBuffers,
496        mut repdef: RepDefBuilder,
497        row_number: u64,
498        num_rows: u64,
499    ) -> Result<Vec<EncodeTask>> {
500        let struct_array = array.as_struct();
501        let mut struct_array = struct_array.normalize_slicing()?;
502        if let Some(validity) = struct_array.nulls() {
503            if self.keep_original_array {
504                repdef.add_validity_bitmap(validity.clone())
505            } else {
506                repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap())
507            }
508            struct_array = struct_array.pushdown_nulls()?;
509        } else {
510            repdef.add_no_null(struct_array.len());
511        }
512        let child_tasks = self
513            .children
514            .iter_mut()
515            .zip(struct_array.columns().iter())
516            .map(|(encoder, arr)| {
517                encoder.maybe_encode(
518                    arr.clone(),
519                    external_buffers,
520                    repdef.clone(),
521                    row_number,
522                    num_rows,
523                )
524            })
525            .collect::<Result<Vec<_>>>()?;
526        Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
527    }
528
529    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
530        self.children
531            .iter_mut()
532            .map(|encoder| encoder.flush(external_buffers))
533            .flatten_ok()
534            .collect::<Result<Vec<_>>>()
535    }
536
537    fn num_columns(&self) -> u32 {
538        self.children
539            .iter()
540            .map(|child| child.num_columns())
541            .sum::<u32>()
542    }
543
544    fn finish(
545        &mut self,
546        external_buffers: &mut OutOfLineBuffers,
547    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
548        let mut child_columns = self
549            .children
550            .iter_mut()
551            .map(|child| child.finish(external_buffers))
552            .collect::<FuturesOrdered<_>>();
553        async move {
554            let mut encoded_columns = Vec::with_capacity(child_columns.len());
555            while let Some(child_cols) = child_columns.next().await {
556                encoded_columns.extend(child_cols?);
557            }
558            Ok(encoded_columns)
559        }
560        .boxed()
561    }
562}
563
564pub struct StructFieldEncoder {
565    children: Vec<Box<dyn FieldEncoder>>,
566    column_index: u32,
567    num_rows_seen: u64,
568}
569
570impl StructFieldEncoder {
571    pub fn new(children: Vec<Box<dyn FieldEncoder>>, column_index: u32) -> Self {
572        Self {
573            children,
574            column_index,
575            num_rows_seen: 0,
576        }
577    }
578}
579
580impl FieldEncoder for StructFieldEncoder {
581    fn maybe_encode(
582        &mut self,
583        array: ArrayRef,
584        external_buffers: &mut OutOfLineBuffers,
585        repdef: RepDefBuilder,
586        row_number: u64,
587        num_rows: u64,
588    ) -> Result<Vec<EncodeTask>> {
589        let struct_array = array.as_struct();
590        self.num_rows_seen += array.len() as u64;
591        let child_tasks = self
592            .children
593            .iter_mut()
594            .zip(struct_array.columns().iter())
595            .map(|(encoder, arr)| {
596                encoder.maybe_encode(
597                    arr.clone(),
598                    external_buffers,
599                    repdef.clone(),
600                    row_number,
601                    num_rows,
602                )
603            })
604            .collect::<Result<Vec<_>>>()?;
605        Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
606    }
607
608    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
609        let child_tasks = self
610            .children
611            .iter_mut()
612            .map(|encoder| encoder.flush(external_buffers))
613            .collect::<Result<Vec<_>>>()?;
614        Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
615    }
616
617    fn num_columns(&self) -> u32 {
618        self.children
619            .iter()
620            .map(|child| child.num_columns())
621            .sum::<u32>()
622            + 1
623    }
624
625    fn finish(
626        &mut self,
627        external_buffers: &mut OutOfLineBuffers,
628    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
629        let mut child_columns = self
630            .children
631            .iter_mut()
632            .map(|child| child.finish(external_buffers))
633            .collect::<FuturesOrdered<_>>();
634        let num_rows_seen = self.num_rows_seen;
635        let column_index = self.column_index;
636        async move {
637            let mut columns = Vec::new();
638            // Add a column for the struct header
639            let mut header = EncodedColumn::default();
640            header.final_pages.push(EncodedPage {
641                data: Vec::new(),
642                description: PageEncoding::Legacy(pb::ArrayEncoding {
643                    array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct(
644                        pb::SimpleStruct {},
645                    )),
646                }),
647                num_rows: num_rows_seen,
648                column_idx: column_index,
649                row_number: 0, // Not used by legacy encoding
650            });
651            columns.push(header);
652            // Now run finish on the children
653            while let Some(child_cols) = child_columns.next().await {
654                columns.extend(child_cols?);
655            }
656            Ok(columns)
657        }
658        .boxed()
659    }
660}
661
662#[cfg(test)]
663mod tests {
664
665    use std::{collections::HashMap, sync::Arc};
666
667    use arrow_array::{
668        Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, NullArray, StringArray,
669        StructArray,
670        builder::{Int32Builder, ListBuilder},
671    };
672    use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
673    use arrow_schema::{DataType, Field, Fields};
674
675    use super::StructuralStructDecoder;
676    use crate::constants::{STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY};
677    use crate::decoder::StructuralFieldDecoder;
678    use crate::testing::{
679        TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data,
680    };
681
682    #[test]
683    fn test_zero_dimension_fsl_decoder_errors() {
684        // Simulates a stored schema declaring a zero-dimension FixedSizeList (writers reject
685        // it but old files may contain one).  Building the decoder must fail cleanly instead
686        // of letting the zero dimension reach the rep/def decimation.
687        let item_fields = Fields::from(vec![Field::new("x", DataType::Int32, true)]);
688        let fields = Fields::from(vec![Field::new(
689            "vecs",
690            DataType::FixedSizeList(
691                Arc::new(Field::new("item", DataType::Struct(item_fields), true)),
692                0,
693            ),
694            true,
695        )]);
696
697        let err = StructuralStructDecoder::new(
698            fields, false, /*is_root=*/ true, /*nullable=*/ false,
699        )
700        .unwrap_err();
701        assert!(matches!(err, lance_core::Error::Schema { .. }));
702        assert!(
703            err.to_string()
704                .contains("dimension must be a positive integer"),
705            "unexpected error: {}",
706            err
707        );
708    }
709
710    #[rstest::rstest]
711    #[test_log::test(tokio::test)]
712    async fn test_simple_struct(
713        #[values(
714            TestEncoding::Array,
715            TestEncoding::StructuralU16,
716            TestEncoding::StructuralU32,
717            TestEncoding::StructuralSparse
718        )]
719        encoding: TestEncoding,
720        #[values(4096, 1024 * 1024)] page_size: u64,
721        #[values(false, true)] use_slicing: bool,
722    ) {
723        let data_type = DataType::Struct(Fields::from(vec![
724            Field::new("a", DataType::Int32, false),
725            Field::new("b", DataType::Int32, false),
726        ]));
727        let field = Field::new("", data_type, false);
728        check_basic_random_case(field, encoding, page_size, use_slicing).await;
729    }
730
731    #[test_log::test(tokio::test)]
732    async fn test_nullable_struct() {
733        // Test data struct<score: int32, location: struct<x: int32, y: int32>>
734        // - score: null
735        //   location:
736        //     x: 1
737        //     y: 6
738        // - score: 12
739        //   location:
740        //     x: 2
741        //     y: null
742        // - score: 13
743        //   location:
744        //     x: 3
745        //     y: 8
746        // - score: 14
747        //   location: null
748        // - null
749        //
750        let inner_fields = Fields::from(vec![
751            Field::new("x", DataType::Int32, false),
752            Field::new("y", DataType::Int32, true),
753        ]);
754        let inner_struct = DataType::Struct(inner_fields.clone());
755        let outer_fields = Fields::from(vec![
756            Field::new("score", DataType::Int32, true),
757            Field::new("location", inner_struct, true),
758        ]);
759
760        let x_vals = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
761        let y_vals = Int32Array::from(vec![Some(6), None, Some(8), Some(9), Some(10)]);
762        let scores = Int32Array::from(vec![None, Some(12), Some(13), Some(14), Some(15)]);
763
764        let location_validity = NullBuffer::from(vec![true, true, true, false, true]);
765        let locations = StructArray::new(
766            inner_fields,
767            vec![Arc::new(x_vals), Arc::new(y_vals)],
768            Some(location_validity),
769        );
770
771        let rows_validity = NullBuffer::from(vec![true, true, true, true, false]);
772        let rows = StructArray::new(
773            outer_fields,
774            vec![Arc::new(scores), Arc::new(locations)],
775            Some(rows_validity),
776        );
777
778        let test_cases = TestCases::default().with_structural_encodings();
779
780        check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await;
781    }
782
783    #[test_log::test(tokio::test)]
784    async fn test_nullable_struct_with_null_typed_child() {
785        // A field that is null in every input row is inferred as the Null type, e.g.
786        // `pa.Table.from_pylist([{"s": {"a": 1, "b": None}}, {"s": None}])`.
787        let fields = Fields::from(vec![
788            Field::new("a", DataType::Int32, true),
789            Field::new("b", DataType::Null, true),
790        ]);
791        let rows = StructArray::new(
792            fields,
793            vec![
794                Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4)])),
795                Arc::new(NullArray::new(4)),
796            ],
797            Some(NullBuffer::from(vec![true, false, true, true])),
798        );
799
800        let test_cases = TestCases::default().with_structural_encodings();
801
802        check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await;
803    }
804
805    #[test_log::test(tokio::test)]
806    async fn test_simple_masked_nonempty_list() {
807        // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT]
808        //
809        let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4), Some(5), Some(6)]);
810        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 5]));
811        let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
812        let list_array = ListArray::new(
813            Arc::new(Field::new("item", DataType::Int32, true)),
814            offsets,
815            Arc::new(items),
816            Some(NullBuffer::new(list_validity)),
817        );
818        let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
819        let struct_array = StructArray::new(
820            Fields::from(vec![Field::new(
821                "inner_list",
822                list_array.data_type().clone(),
823                true,
824            )]),
825            vec![Arc::new(list_array)],
826            Some(NullBuffer::new(struct_validity)),
827        );
828        check_round_trip_encoding_of_data(
829            vec![Arc::new(struct_array)],
830            &TestCases::default().with_structural_encodings(),
831            HashMap::new(),
832        )
833        .await;
834    }
835
836    #[test_log::test(tokio::test)]
837    async fn test_simple_struct_list() {
838        // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT]
839        //
840        let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
841        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 4]));
842        let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
843        let list_array = ListArray::new(
844            Arc::new(Field::new("item", DataType::Int32, true)),
845            offsets,
846            Arc::new(items),
847            Some(NullBuffer::new(list_validity)),
848        );
849        let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
850        let struct_array = StructArray::new(
851            Fields::from(vec![Field::new(
852                "inner_list",
853                list_array.data_type().clone(),
854                true,
855            )]),
856            vec![Arc::new(list_array)],
857            Some(NullBuffer::new(struct_validity)),
858        );
859        check_round_trip_encoding_of_data(
860            vec![Arc::new(struct_array)],
861            &TestCases::default().with_structural_encodings(),
862            HashMap::new(),
863        )
864        .await;
865    }
866
867    #[rstest::rstest]
868    #[test_log::test(tokio::test)]
869    async fn test_struct_list(
870        #[values(
871            TestEncoding::Array,
872            TestEncoding::StructuralU16,
873            TestEncoding::StructuralU32,
874            TestEncoding::StructuralSparse
875        )]
876        encoding: TestEncoding,
877        #[values(4096, 1024 * 1024)] page_size: u64,
878        #[values(false, true)] use_slicing: bool,
879    ) {
880        let data_type = DataType::Struct(Fields::from(vec![
881            Field::new(
882                "inner_list",
883                DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
884                true,
885            ),
886            Field::new("outer_int", DataType::Int32, true),
887        ]));
888        let field = Field::new("row", data_type, false);
889        check_basic_random_case(field, encoding, page_size, use_slicing).await;
890    }
891
892    #[rstest::rstest]
893    #[test_log::test(tokio::test)]
894    async fn test_empty_struct(
895        #[values(
896            TestEncoding::Array,
897            TestEncoding::StructuralU16,
898            TestEncoding::StructuralU32,
899            TestEncoding::StructuralSparse
900        )]
901        encoding: TestEncoding,
902        #[values(4096, 1024 * 1024)] page_size: u64,
903        #[values(false, true)] use_slicing: bool,
904    ) {
905        // It's technically legal for a struct to have 0 children, need to
906        // make sure we support that
907        let data_type = DataType::Struct(Fields::from(Vec::<Field>::default()));
908        let field = Field::new("row", data_type, false);
909        check_basic_random_case(field, encoding, page_size, use_slicing).await;
910    }
911
912    #[rstest::rstest]
913    #[test_log::test(tokio::test)]
914    async fn test_complicated_struct(
915        #[values(
916            TestEncoding::Array,
917            TestEncoding::StructuralU16,
918            TestEncoding::StructuralU32,
919            TestEncoding::StructuralSparse
920        )]
921        encoding: TestEncoding,
922        #[values(4096, 1024 * 1024)] page_size: u64,
923        #[values(false, true)] use_slicing: bool,
924    ) {
925        let data_type = DataType::Struct(Fields::from(vec![
926            Field::new("int", DataType::Int32, true),
927            Field::new(
928                "inner",
929                DataType::Struct(Fields::from(vec![
930                    Field::new("inner_int", DataType::Int32, true),
931                    Field::new(
932                        "inner_list",
933                        DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
934                        true,
935                    ),
936                ])),
937                true,
938            ),
939            Field::new("outer_binary", DataType::Binary, true),
940        ]));
941        let field = Field::new("row", data_type, false);
942        check_basic_random_case(field, encoding, page_size, use_slicing).await;
943    }
944
945    #[test_log::test(tokio::test)]
946    async fn test_list_of_struct_with_null_struct_element() {
947        // Regression: a list containing structs where most struct elements are null
948        // causes a length mismatch during decoding with V2_2 encoding.
949        use arrow_array::StringArray;
950
951        let tag_array = StringArray::from(vec![
952            Some("valid"),
953            Some("null_struct"),
954            Some("valid"),
955            Some("valid"),
956        ]);
957        let struct_fields = Fields::from(vec![Field::new("tag", DataType::Utf8, true)]);
958        // 3 out of 4 struct elements are null
959        let struct_validity = NullBuffer::from(vec![false, true, false, false]);
960        let struct_array = StructArray::new(
961            struct_fields.clone(),
962            vec![Arc::new(tag_array)],
963            Some(struct_validity),
964        );
965
966        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4]));
967        let list_field = Field::new("item", DataType::Struct(struct_fields), true);
968        let list_array =
969            ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
970
971        check_round_trip_encoding_of_data(
972            vec![Arc::new(list_array)],
973            &TestCases::default().with_u32_structural_encodings(),
974            HashMap::new(),
975        )
976        .await;
977    }
978
979    #[test_log::test(tokio::test)]
980    async fn test_list_of_struct_with_all_null_child() {
981        use arrow_array::StringArray;
982
983        let a_array = StringArray::from(vec![Some("w"), Some("x"), Some("y"), Some("z")]);
984        let b_array = StringArray::from(vec![None::<&str>, None, None, None]);
985        let struct_fields = Fields::from(vec![
986            Field::new("a", DataType::Utf8, true),
987            Field::new("b", DataType::Utf8, true),
988        ]);
989        let struct_array = StructArray::new(
990            struct_fields.clone(),
991            vec![Arc::new(a_array), Arc::new(b_array)],
992            None,
993        );
994
995        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 4]));
996        let list_field = Field::new("item", DataType::Struct(struct_fields), true);
997        let list_array =
998            ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
999
1000        check_round_trip_encoding_of_data(
1001            vec![Arc::new(list_array)],
1002            &TestCases::default()
1003                .with_range(0..1)
1004                .with_range(1..2)
1005                .with_indices(vec![0])
1006                .with_indices(vec![1])
1007                .with_structural_encodings(),
1008            HashMap::new(),
1009        )
1010        .await;
1011    }
1012
1013    #[test_log::test(tokio::test)]
1014    async fn test_list_of_struct_with_constant_child_and_empty_lists() {
1015        let item_fields = Fields::from(vec![
1016            Field::new("maneuver", DataType::Int64, true),
1017            Field::new("remaining_dist", DataType::Float64, true),
1018        ]);
1019        let item_array = StructArray::new(
1020            item_fields.clone(),
1021            vec![
1022                Arc::new(Int64Array::from(vec![Some(3), Some(3), Some(3)])),
1023                Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])),
1024            ],
1025            None,
1026        );
1027
1028        let list_field = Field::new("item", DataType::Struct(item_fields), true);
1029        let list_array = ListArray::new(
1030            Arc::new(list_field),
1031            OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 3, 3, 3])),
1032            Arc::new(item_array),
1033            None,
1034        );
1035
1036        let data_fields = Fields::from(vec![Field::new(
1037            "maneuvers",
1038            list_array.data_type().clone(),
1039            true,
1040        )]);
1041        let data_array = StructArray::new(data_fields.clone(), vec![Arc::new(list_array)], None);
1042        let row_validity = NullBuffer::from(vec![true, true, true, true, true, false]);
1043        let row_array = StructArray::new(
1044            Fields::from(vec![Field::new(
1045                "data",
1046                DataType::Struct(data_fields),
1047                true,
1048            )]),
1049            vec![Arc::new(data_array)],
1050            Some(row_validity),
1051        );
1052
1053        check_round_trip_encoding_of_data(
1054            vec![Arc::new(row_array)],
1055            &TestCases::default().with_structural_encodings(),
1056            HashMap::new(),
1057        )
1058        .await;
1059    }
1060
1061    #[test]
1062    fn test_plan_decoded_bytes_struct_non_nullable_primitive() {
1063        // Nullable struct containing a non-nullable i32 child.
1064        // Only the struct-level bitmap contributes; the child has no bitmap.
1065        //
1066        // N == CANDIDATE_BATCH_SIZES[5] == 1024 so that all sizes are
1067        // naturally aligned to Arrow's 64-byte allocation boundary.
1068        use arrow_buffer::NullBuffer;
1069
1070        const N: usize = 1024;
1071        let x_vals = Int32Array::from_iter_values(0..N as i32);
1072        let struct_nulls = NullBuffer::from((0..N).map(|i| i % 2 == 0).collect::<Vec<_>>());
1073        let struct_field = Field::new("x", DataType::Int32, false);
1074        let struct_array = StructArray::new(
1075            Fields::from(vec![struct_field]),
1076            vec![Arc::new(x_vals) as ArrayRef],
1077            Some(struct_nulls),
1078        );
1079        let expected = struct_array.get_buffer_memory_size() as u64;
1080
1081        let child_field = Arc::new(Field::new("x", DataType::Int32, false));
1082        let fields = Fields::from(vec![child_field]);
1083        let decoder = StructuralStructDecoder::new(
1084            fields, false, /*is_root=*/ false, /*nullable=*/ true,
1085        )
1086        .unwrap();
1087
1088        let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
1089        // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N
1090        assert_eq!(
1091            bytes[5], expected,
1092            "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
1093            bytes[5],
1094        );
1095    }
1096
1097    #[test]
1098    fn test_plan_decoded_bytes_struct_nullable_primitive() {
1099        // Nullable struct containing a nullable i32 child.
1100        // Both the struct and the child contribute a validity bitmap.
1101        use arrow_buffer::NullBuffer;
1102
1103        const N: usize = 1024;
1104        let y_vals =
1105            Int32Array::from_iter((0..N as i32).map(|i| if i % 2 == 0 { Some(i) } else { None }));
1106        let struct_nulls = NullBuffer::from((0..N).map(|i| i % 3 != 0).collect::<Vec<_>>());
1107        let struct_field = Field::new("y", DataType::Int32, true);
1108        let struct_array = StructArray::new(
1109            Fields::from(vec![struct_field]),
1110            vec![Arc::new(y_vals) as ArrayRef],
1111            Some(struct_nulls),
1112        );
1113        let expected = struct_array.get_buffer_memory_size() as u64;
1114
1115        let child_field = Arc::new(Field::new("y", DataType::Int32, true));
1116        let fields = Fields::from(vec![child_field]);
1117        let decoder = StructuralStructDecoder::new(
1118            fields, false, /*is_root=*/ false, /*nullable=*/ true,
1119        )
1120        .unwrap();
1121
1122        let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
1123        assert_eq!(
1124            bytes[5], expected,
1125            "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
1126            bytes[5],
1127        );
1128    }
1129
1130    #[test_log::test(tokio::test)]
1131    async fn test_ragged_scheduling() {
1132        // This test covers scheduling when batches straddle page boundaries
1133
1134        // Create a list with 10k nulls
1135        let items_builder = Int32Builder::new();
1136        let mut list_builder = ListBuilder::new(items_builder);
1137        for _ in 0..10000 {
1138            list_builder.append_null();
1139        }
1140        let list_array = Arc::new(list_builder.finish());
1141        let int_array = Arc::new(Int32Array::from_iter_values(0..10000));
1142        let fields = vec![
1143            Field::new("", list_array.data_type().clone(), true),
1144            Field::new("", int_array.data_type().clone(), true),
1145        ];
1146        let struct_array = Arc::new(StructArray::new(
1147            Fields::from(fields),
1148            vec![list_array, int_array],
1149            None,
1150        )) as ArrayRef;
1151        let struct_arrays = (0..10000)
1152            // Intentionally skip in some randomish amount to create more ragged scheduling
1153            .step_by(437)
1154            .map(|offset| struct_array.slice(offset, 437.min(10000 - offset)))
1155            .collect::<Vec<_>>();
1156        check_round_trip_encoding_of_data(struct_arrays, &TestCases::default(), HashMap::new())
1157            .await;
1158    }
1159
1160    /// A null list produces an "invisible" item: a rep/def entry that owns no value in the
1161    /// leaf page.  The full-zip decoder works out which items own a value from
1162    /// `max_visible_def`, and adding a second nullable layer *above* the list (the struct
1163    /// itself) must not move that threshold.  When it does, the decoder hands a value to
1164    /// every null list and each later value shifts by one.
1165    #[rstest::rstest]
1166    #[test_log::test(tokio::test)]
1167    async fn test_full_zip_invisible_items_under_nullable_struct(
1168        #[values(false, true)] variable_width_items: bool,
1169    ) {
1170        // struct<x: list<item>> with both layers nullable:
1171        //   {x: [_, _]} / {x: null} / {x: [_]} / null / {x: [_, _]}
1172        let (item_type, items): (DataType, ArrayRef) = if variable_width_items {
1173            (
1174                DataType::Utf8,
1175                Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])),
1176            )
1177        } else {
1178            (
1179                DataType::Int32,
1180                Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
1181            )
1182        };
1183        let lists = ListArray::new(
1184            Arc::new(Field::new("item", item_type, true)),
1185            OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 2, 3, 3, 5])),
1186            items,
1187            Some(NullBuffer::from(vec![true, false, true, true, true])),
1188        );
1189        let fields = Fields::from(vec![Field::new("x", lists.data_type().clone(), true)]);
1190        let rows = StructArray::new(
1191            fields,
1192            vec![Arc::new(lists)],
1193            Some(NullBuffer::from(vec![true, true, true, false, true])),
1194        );
1195
1196        let metadata = HashMap::from([(
1197            STRUCTURAL_ENCODING_META_KEY.to_string(),
1198            STRUCTURAL_ENCODING_FULLZIP.to_string(),
1199        )]);
1200        let test_cases = TestCases::default().with_structural_encodings();
1201        check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, metadata).await;
1202    }
1203}
lance_encoding/encodings/logical/struct.rs

lance_encoding/encodings/logical/
struct.rs
