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, StringArray, StructArray,
669        builder::{Int32Builder, ListBuilder},
670    };
671    use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
672    use arrow_schema::{DataType, Field, Fields};
673
674    use super::StructuralStructDecoder;
675    use crate::constants::{STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY};
676    use crate::decoder::StructuralFieldDecoder;
677    use crate::testing::{
678        TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data,
679    };
680
681    #[test]
682    fn test_zero_dimension_fsl_decoder_errors() {
683        // Simulates a stored schema declaring a zero-dimension FixedSizeList (writers reject
684        // it but old files may contain one).  Building the decoder must fail cleanly instead
685        // of letting the zero dimension reach the rep/def decimation.
686        let item_fields = Fields::from(vec![Field::new("x", DataType::Int32, true)]);
687        let fields = Fields::from(vec![Field::new(
688            "vecs",
689            DataType::FixedSizeList(
690                Arc::new(Field::new("item", DataType::Struct(item_fields), true)),
691                0,
692            ),
693            true,
694        )]);
695
696        let err = StructuralStructDecoder::new(
697            fields, false, /*is_root=*/ true, /*nullable=*/ false,
698        )
699        .unwrap_err();
700        assert!(matches!(err, lance_core::Error::Schema { .. }));
701        assert!(
702            err.to_string()
703                .contains("dimension must be a positive integer"),
704            "unexpected error: {}",
705            err
706        );
707    }
708
709    #[rstest::rstest]
710    #[test_log::test(tokio::test)]
711    async fn test_simple_struct(
712        #[values(
713            TestEncoding::Array,
714            TestEncoding::StructuralU16,
715            TestEncoding::StructuralU32,
716            TestEncoding::StructuralSparse
717        )]
718        encoding: TestEncoding,
719        #[values(4096, 1024 * 1024)] page_size: u64,
720        #[values(false, true)] use_slicing: bool,
721    ) {
722        let data_type = DataType::Struct(Fields::from(vec![
723            Field::new("a", DataType::Int32, false),
724            Field::new("b", DataType::Int32, false),
725        ]));
726        let field = Field::new("", data_type, false);
727        check_basic_random_case(field, encoding, page_size, use_slicing).await;
728    }
729
730    #[test_log::test(tokio::test)]
731    async fn test_nullable_struct() {
732        // Test data struct<score: int32, location: struct<x: int32, y: int32>>
733        // - score: null
734        //   location:
735        //     x: 1
736        //     y: 6
737        // - score: 12
738        //   location:
739        //     x: 2
740        //     y: null
741        // - score: 13
742        //   location:
743        //     x: 3
744        //     y: 8
745        // - score: 14
746        //   location: null
747        // - null
748        //
749        let inner_fields = Fields::from(vec![
750            Field::new("x", DataType::Int32, false),
751            Field::new("y", DataType::Int32, true),
752        ]);
753        let inner_struct = DataType::Struct(inner_fields.clone());
754        let outer_fields = Fields::from(vec![
755            Field::new("score", DataType::Int32, true),
756            Field::new("location", inner_struct, true),
757        ]);
758
759        let x_vals = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
760        let y_vals = Int32Array::from(vec![Some(6), None, Some(8), Some(9), Some(10)]);
761        let scores = Int32Array::from(vec![None, Some(12), Some(13), Some(14), Some(15)]);
762
763        let location_validity = NullBuffer::from(vec![true, true, true, false, true]);
764        let locations = StructArray::new(
765            inner_fields,
766            vec![Arc::new(x_vals), Arc::new(y_vals)],
767            Some(location_validity),
768        );
769
770        let rows_validity = NullBuffer::from(vec![true, true, true, true, false]);
771        let rows = StructArray::new(
772            outer_fields,
773            vec![Arc::new(scores), Arc::new(locations)],
774            Some(rows_validity),
775        );
776
777        let test_cases = TestCases::default().with_structural_encodings();
778
779        check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await;
780    }
781
782    #[test_log::test(tokio::test)]
783    async fn test_simple_masked_nonempty_list() {
784        // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT]
785        //
786        let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4), Some(5), Some(6)]);
787        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 5]));
788        let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
789        let list_array = ListArray::new(
790            Arc::new(Field::new("item", DataType::Int32, true)),
791            offsets,
792            Arc::new(items),
793            Some(NullBuffer::new(list_validity)),
794        );
795        let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
796        let struct_array = StructArray::new(
797            Fields::from(vec![Field::new(
798                "inner_list",
799                list_array.data_type().clone(),
800                true,
801            )]),
802            vec![Arc::new(list_array)],
803            Some(NullBuffer::new(struct_validity)),
804        );
805        check_round_trip_encoding_of_data(
806            vec![Arc::new(struct_array)],
807            &TestCases::default().with_structural_encodings(),
808            HashMap::new(),
809        )
810        .await;
811    }
812
813    #[test_log::test(tokio::test)]
814    async fn test_simple_struct_list() {
815        // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT]
816        //
817        let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
818        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 4]));
819        let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
820        let list_array = ListArray::new(
821            Arc::new(Field::new("item", DataType::Int32, true)),
822            offsets,
823            Arc::new(items),
824            Some(NullBuffer::new(list_validity)),
825        );
826        let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
827        let struct_array = StructArray::new(
828            Fields::from(vec![Field::new(
829                "inner_list",
830                list_array.data_type().clone(),
831                true,
832            )]),
833            vec![Arc::new(list_array)],
834            Some(NullBuffer::new(struct_validity)),
835        );
836        check_round_trip_encoding_of_data(
837            vec![Arc::new(struct_array)],
838            &TestCases::default().with_structural_encodings(),
839            HashMap::new(),
840        )
841        .await;
842    }
843
844    #[rstest::rstest]
845    #[test_log::test(tokio::test)]
846    async fn test_struct_list(
847        #[values(
848            TestEncoding::Array,
849            TestEncoding::StructuralU16,
850            TestEncoding::StructuralU32,
851            TestEncoding::StructuralSparse
852        )]
853        encoding: TestEncoding,
854        #[values(4096, 1024 * 1024)] page_size: u64,
855        #[values(false, true)] use_slicing: bool,
856    ) {
857        let data_type = DataType::Struct(Fields::from(vec![
858            Field::new(
859                "inner_list",
860                DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
861                true,
862            ),
863            Field::new("outer_int", DataType::Int32, true),
864        ]));
865        let field = Field::new("row", data_type, false);
866        check_basic_random_case(field, encoding, page_size, use_slicing).await;
867    }
868
869    #[rstest::rstest]
870    #[test_log::test(tokio::test)]
871    async fn test_empty_struct(
872        #[values(
873            TestEncoding::Array,
874            TestEncoding::StructuralU16,
875            TestEncoding::StructuralU32,
876            TestEncoding::StructuralSparse
877        )]
878        encoding: TestEncoding,
879        #[values(4096, 1024 * 1024)] page_size: u64,
880        #[values(false, true)] use_slicing: bool,
881    ) {
882        // It's technically legal for a struct to have 0 children, need to
883        // make sure we support that
884        let data_type = DataType::Struct(Fields::from(Vec::<Field>::default()));
885        let field = Field::new("row", data_type, false);
886        check_basic_random_case(field, encoding, page_size, use_slicing).await;
887    }
888
889    #[rstest::rstest]
890    #[test_log::test(tokio::test)]
891    async fn test_complicated_struct(
892        #[values(
893            TestEncoding::Array,
894            TestEncoding::StructuralU16,
895            TestEncoding::StructuralU32,
896            TestEncoding::StructuralSparse
897        )]
898        encoding: TestEncoding,
899        #[values(4096, 1024 * 1024)] page_size: u64,
900        #[values(false, true)] use_slicing: bool,
901    ) {
902        let data_type = DataType::Struct(Fields::from(vec![
903            Field::new("int", DataType::Int32, true),
904            Field::new(
905                "inner",
906                DataType::Struct(Fields::from(vec![
907                    Field::new("inner_int", DataType::Int32, true),
908                    Field::new(
909                        "inner_list",
910                        DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
911                        true,
912                    ),
913                ])),
914                true,
915            ),
916            Field::new("outer_binary", DataType::Binary, true),
917        ]));
918        let field = Field::new("row", data_type, false);
919        check_basic_random_case(field, encoding, page_size, use_slicing).await;
920    }
921
922    #[test_log::test(tokio::test)]
923    async fn test_list_of_struct_with_null_struct_element() {
924        // Regression: a list containing structs where most struct elements are null
925        // causes a length mismatch during decoding with V2_2 encoding.
926        use arrow_array::StringArray;
927
928        let tag_array = StringArray::from(vec![
929            Some("valid"),
930            Some("null_struct"),
931            Some("valid"),
932            Some("valid"),
933        ]);
934        let struct_fields = Fields::from(vec![Field::new("tag", DataType::Utf8, true)]);
935        // 3 out of 4 struct elements are null
936        let struct_validity = NullBuffer::from(vec![false, true, false, false]);
937        let struct_array = StructArray::new(
938            struct_fields.clone(),
939            vec![Arc::new(tag_array)],
940            Some(struct_validity),
941        );
942
943        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4]));
944        let list_field = Field::new("item", DataType::Struct(struct_fields), true);
945        let list_array =
946            ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
947
948        check_round_trip_encoding_of_data(
949            vec![Arc::new(list_array)],
950            &TestCases::default().with_u32_structural_encodings(),
951            HashMap::new(),
952        )
953        .await;
954    }
955
956    #[test_log::test(tokio::test)]
957    async fn test_list_of_struct_with_all_null_child() {
958        use arrow_array::StringArray;
959
960        let a_array = StringArray::from(vec![Some("w"), Some("x"), Some("y"), Some("z")]);
961        let b_array = StringArray::from(vec![None::<&str>, None, None, None]);
962        let struct_fields = Fields::from(vec![
963            Field::new("a", DataType::Utf8, true),
964            Field::new("b", DataType::Utf8, true),
965        ]);
966        let struct_array = StructArray::new(
967            struct_fields.clone(),
968            vec![Arc::new(a_array), Arc::new(b_array)],
969            None,
970        );
971
972        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 4]));
973        let list_field = Field::new("item", DataType::Struct(struct_fields), true);
974        let list_array =
975            ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
976
977        check_round_trip_encoding_of_data(
978            vec![Arc::new(list_array)],
979            &TestCases::default()
980                .with_range(0..1)
981                .with_range(1..2)
982                .with_indices(vec![0])
983                .with_indices(vec![1])
984                .with_structural_encodings(),
985            HashMap::new(),
986        )
987        .await;
988    }
989
990    #[test_log::test(tokio::test)]
991    async fn test_list_of_struct_with_constant_child_and_empty_lists() {
992        let item_fields = Fields::from(vec![
993            Field::new("maneuver", DataType::Int64, true),
994            Field::new("remaining_dist", DataType::Float64, true),
995        ]);
996        let item_array = StructArray::new(
997            item_fields.clone(),
998            vec![
999                Arc::new(Int64Array::from(vec![Some(3), Some(3), Some(3)])),
1000                Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])),
1001            ],
1002            None,
1003        );
1004
1005        let list_field = Field::new("item", DataType::Struct(item_fields), true);
1006        let list_array = ListArray::new(
1007            Arc::new(list_field),
1008            OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 3, 3, 3])),
1009            Arc::new(item_array),
1010            None,
1011        );
1012
1013        let data_fields = Fields::from(vec![Field::new(
1014            "maneuvers",
1015            list_array.data_type().clone(),
1016            true,
1017        )]);
1018        let data_array = StructArray::new(data_fields.clone(), vec![Arc::new(list_array)], None);
1019        let row_validity = NullBuffer::from(vec![true, true, true, true, true, false]);
1020        let row_array = StructArray::new(
1021            Fields::from(vec![Field::new(
1022                "data",
1023                DataType::Struct(data_fields),
1024                true,
1025            )]),
1026            vec![Arc::new(data_array)],
1027            Some(row_validity),
1028        );
1029
1030        check_round_trip_encoding_of_data(
1031            vec![Arc::new(row_array)],
1032            &TestCases::default().with_structural_encodings(),
1033            HashMap::new(),
1034        )
1035        .await;
1036    }
1037
1038    #[test]
1039    fn test_plan_decoded_bytes_struct_non_nullable_primitive() {
1040        // Nullable struct containing a non-nullable i32 child.
1041        // Only the struct-level bitmap contributes; the child has no bitmap.
1042        //
1043        // N == CANDIDATE_BATCH_SIZES[5] == 1024 so that all sizes are
1044        // naturally aligned to Arrow's 64-byte allocation boundary.
1045        use arrow_buffer::NullBuffer;
1046
1047        const N: usize = 1024;
1048        let x_vals = Int32Array::from_iter_values(0..N as i32);
1049        let struct_nulls = NullBuffer::from((0..N).map(|i| i % 2 == 0).collect::<Vec<_>>());
1050        let struct_field = Field::new("x", DataType::Int32, false);
1051        let struct_array = StructArray::new(
1052            Fields::from(vec![struct_field]),
1053            vec![Arc::new(x_vals) as ArrayRef],
1054            Some(struct_nulls),
1055        );
1056        let expected = struct_array.get_buffer_memory_size() as u64;
1057
1058        let child_field = Arc::new(Field::new("x", DataType::Int32, false));
1059        let fields = Fields::from(vec![child_field]);
1060        let decoder = StructuralStructDecoder::new(
1061            fields, false, /*is_root=*/ false, /*nullable=*/ true,
1062        )
1063        .unwrap();
1064
1065        let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
1066        // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N
1067        assert_eq!(
1068            bytes[5], expected,
1069            "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
1070            bytes[5],
1071        );
1072    }
1073
1074    #[test]
1075    fn test_plan_decoded_bytes_struct_nullable_primitive() {
1076        // Nullable struct containing a nullable i32 child.
1077        // Both the struct and the child contribute a validity bitmap.
1078        use arrow_buffer::NullBuffer;
1079
1080        const N: usize = 1024;
1081        let y_vals =
1082            Int32Array::from_iter((0..N as i32).map(|i| if i % 2 == 0 { Some(i) } else { None }));
1083        let struct_nulls = NullBuffer::from((0..N).map(|i| i % 3 != 0).collect::<Vec<_>>());
1084        let struct_field = Field::new("y", DataType::Int32, true);
1085        let struct_array = StructArray::new(
1086            Fields::from(vec![struct_field]),
1087            vec![Arc::new(y_vals) as ArrayRef],
1088            Some(struct_nulls),
1089        );
1090        let expected = struct_array.get_buffer_memory_size() as u64;
1091
1092        let child_field = Arc::new(Field::new("y", DataType::Int32, true));
1093        let fields = Fields::from(vec![child_field]);
1094        let decoder = StructuralStructDecoder::new(
1095            fields, false, /*is_root=*/ false, /*nullable=*/ true,
1096        )
1097        .unwrap();
1098
1099        let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
1100        assert_eq!(
1101            bytes[5], expected,
1102            "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
1103            bytes[5],
1104        );
1105    }
1106
1107    #[test_log::test(tokio::test)]
1108    async fn test_ragged_scheduling() {
1109        // This test covers scheduling when batches straddle page boundaries
1110
1111        // Create a list with 10k nulls
1112        let items_builder = Int32Builder::new();
1113        let mut list_builder = ListBuilder::new(items_builder);
1114        for _ in 0..10000 {
1115            list_builder.append_null();
1116        }
1117        let list_array = Arc::new(list_builder.finish());
1118        let int_array = Arc::new(Int32Array::from_iter_values(0..10000));
1119        let fields = vec![
1120            Field::new("", list_array.data_type().clone(), true),
1121            Field::new("", int_array.data_type().clone(), true),
1122        ];
1123        let struct_array = Arc::new(StructArray::new(
1124            Fields::from(fields),
1125            vec![list_array, int_array],
1126            None,
1127        )) as ArrayRef;
1128        let struct_arrays = (0..10000)
1129            // Intentionally skip in some randomish amount to create more ragged scheduling
1130            .step_by(437)
1131            .map(|offset| struct_array.slice(offset, 437.min(10000 - offset)))
1132            .collect::<Vec<_>>();
1133        check_round_trip_encoding_of_data(struct_arrays, &TestCases::default(), HashMap::new())
1134            .await;
1135    }
1136
1137    /// A null list produces an "invisible" item: a rep/def entry that owns no value in the
1138    /// leaf page.  The full-zip decoder works out which items own a value from
1139    /// `max_visible_def`, and adding a second nullable layer *above* the list (the struct
1140    /// itself) must not move that threshold.  When it does, the decoder hands a value to
1141    /// every null list and each later value shifts by one.
1142    #[rstest::rstest]
1143    #[test_log::test(tokio::test)]
1144    async fn test_full_zip_invisible_items_under_nullable_struct(
1145        #[values(false, true)] variable_width_items: bool,
1146    ) {
1147        // struct<x: list<item>> with both layers nullable:
1148        //   {x: [_, _]} / {x: null} / {x: [_]} / null / {x: [_, _]}
1149        let (item_type, items): (DataType, ArrayRef) = if variable_width_items {
1150            (
1151                DataType::Utf8,
1152                Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])),
1153            )
1154        } else {
1155            (
1156                DataType::Int32,
1157                Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
1158            )
1159        };
1160        let lists = ListArray::new(
1161            Arc::new(Field::new("item", item_type, true)),
1162            OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 2, 3, 3, 5])),
1163            items,
1164            Some(NullBuffer::from(vec![true, false, true, true, true])),
1165        );
1166        let fields = Fields::from(vec![Field::new("x", lists.data_type().clone(), true)]);
1167        let rows = StructArray::new(
1168            fields,
1169            vec![Arc::new(lists)],
1170            Some(NullBuffer::from(vec![true, true, true, false, true])),
1171        );
1172
1173        let metadata = HashMap::from([(
1174            STRUCTURAL_ENCODING_META_KEY.to_string(),
1175            STRUCTURAL_ENCODING_FULLZIP.to_string(),
1176        )]);
1177        let test_cases = TestCases::default().with_structural_encodings();
1178        check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, metadata).await;
1179    }
1180}
lance_encoding/encodings/logical/struct.rs

lance_encoding/encodings/logical/
struct.rs
