Skip to main content

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        filter: &'a FilterExpression,
213        context: &'a SchedulerContext,
214    ) -> BoxFuture<'a, Result<()>> {
215        let children_initialization = self
216            .children
217            .iter_mut()
218            .map(|child| child.initialize(filter, context))
219            .collect::<FuturesUnordered<_>>();
220        async move {
221            children_initialization
222                .map(|res| res.map(|_| ()))
223                .try_collect::<Vec<_>>()
224                .await?;
225            Ok(())
226        }
227        .boxed()
228    }
229}
230
231#[derive(Debug)]
232pub struct StructuralStructDecoder {
233    children: Vec<Box<dyn StructuralFieldDecoder>>,
234    data_type: DataType,
235    child_fields: Fields,
236    // The root decoder is slightly different because it cannot have nulls
237    is_root: bool,
238}
239
240impl StructuralStructDecoder {
241    pub fn new(fields: Fields, should_validate: bool, is_root: bool) -> Result<Self> {
242        let children = fields
243            .iter()
244            .map(|field| Self::field_to_decoder(field, should_validate))
245            .collect::<Result<Vec<_>>>()?;
246        let data_type = DataType::Struct(fields.clone());
247        Ok(Self {
248            data_type,
249            children,
250            child_fields: fields,
251            is_root,
252        })
253    }
254
255    fn field_to_decoder(
256        field: &Arc<arrow_schema::Field>,
257        should_validate: bool,
258    ) -> Result<Box<dyn StructuralFieldDecoder>> {
259        match field.data_type() {
260            DataType::Struct(fields) => {
261                if field.is_packed_struct() || field.is_blob() {
262                    let decoder =
263                        StructuralPrimitiveFieldDecoder::new(&field.clone(), should_validate);
264                    Ok(Box::new(decoder))
265                } else {
266                    Ok(Box::new(Self::new(fields.clone(), should_validate, false)?))
267                }
268            }
269            DataType::List(child_field) | DataType::LargeList(child_field) => {
270                let child_decoder = Self::field_to_decoder(child_field, should_validate)?;
271                Ok(Box::new(StructuralListDecoder::new(
272                    child_decoder,
273                    field.data_type().clone(),
274                )))
275            }
276            DataType::FixedSizeList(child_field, _)
277                if matches!(child_field.data_type(), DataType::Struct(_)) =>
278            {
279                // The scheduler factories run the same guard, but the decoder tree can be
280                // built independently (e.g. `create_decode_stream`) so a zero dimension from
281                // a malformed schema must be rejected here as well.  Draining and unraveling
282                // validity both scale by the dimension and a zero would make that math
283                // degenerate.
284                validate_fixed_size_list_dimensions(field.name(), field.data_type())?;
285                // FixedSizeList containing Struct needs structural decoding
286                let child_decoder = Self::field_to_decoder(child_field, should_validate)?;
287                Ok(Box::new(StructuralFixedSizeListDecoder::new(
288                    child_decoder,
289                    field.data_type().clone(),
290                )))
291            }
292            DataType::Map(entries_field, keys_sorted) => {
293                if *keys_sorted {
294                    return Err(Error::not_supported_source(
295                        "Map data type with keys_sorted=true is not supported yet"
296                            .to_string()
297                            .into(),
298                    ));
299                }
300                let child_decoder = Self::field_to_decoder(entries_field, should_validate)?;
301                Ok(Box::new(StructuralMapDecoder::new(
302                    child_decoder,
303                    field.data_type().clone(),
304                )))
305            }
306            DataType::RunEndEncoded(_, _) => todo!(),
307            DataType::ListView(_) | DataType::LargeListView(_) => todo!(),
308            DataType::Union(_, _) => todo!(),
309            _ => Ok(Box::new(StructuralPrimitiveFieldDecoder::new(
310                field,
311                should_validate,
312            ))),
313        }
314    }
315
316    pub fn drain_batch_task(&mut self, num_rows: u64) -> Result<NextDecodeTask> {
317        let array_drain = self.drain(num_rows)?;
318        Ok(NextDecodeTask {
319            num_rows,
320            task: Box::new(array_drain),
321        })
322    }
323}
324
325impl StructuralFieldDecoder for StructuralStructDecoder {
326    fn accept_page(&mut self, mut child: LoadedPageShard) -> Result<()> {
327        // children with empty path should not be delivered to this method
328        let child_idx = child.path.pop_front().unwrap();
329        // This decoder is intended for one of our children
330        self.children[child_idx as usize].accept_page(child)?;
331        Ok(())
332    }
333
334    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
335        let child_tasks = self
336            .children
337            .iter_mut()
338            .map(|child| child.drain(num_rows))
339            .collect::<Result<Vec<_>>>()?;
340        Ok(Box::new(RepDefStructDecodeTask {
341            children: child_tasks,
342            child_fields: self.child_fields.clone(),
343            is_root: self.is_root,
344            num_rows,
345        }))
346    }
347
348    fn data_type(&self) -> &DataType {
349        &self.data_type
350    }
351}
352
353#[derive(Debug)]
354struct RepDefStructDecodeTask {
355    children: Vec<Box<dyn StructuralDecodeArrayTask>>,
356    child_fields: Fields,
357    is_root: bool,
358    num_rows: u64,
359}
360
361impl StructuralDecodeArrayTask for RepDefStructDecodeTask {
362    fn decode(self: Box<Self>) -> Result<DecodedArray> {
363        if self.children.is_empty() {
364            return Ok(DecodedArray {
365                array: Arc::new(StructArray::new_empty_fields(self.num_rows as usize, None)),
366                repdef: CompositeRepDefUnraveler::new(vec![]),
367                data_size: 0,
368            });
369        }
370
371        let arrays = self
372            .children
373            .into_iter()
374            .map(|task| task.decode())
375            .collect::<Result<Vec<_>>>()?;
376        let mut children = Vec::with_capacity(arrays.len());
377        let mut repdefs = Vec::with_capacity(arrays.len());
378        let mut data_size = 0u64;
379        let mut arrays_iter = arrays.into_iter();
380        let first_array = arrays_iter.next().ok_or_else(|| {
381            Error::internal("Struct decoder unexpectedly has no child arrays".to_string())
382        })?;
383        let length = first_array.array.len();
384
385        // The repdef should be identical across all children at this point
386        repdefs.push(first_array.repdef);
387        data_size += first_array.data_size;
388        children.push(first_array.array);
389
390        for array in arrays_iter {
391            if length != array.array.len() {
392                return Err(Error::invalid_input_source(
393                    format!(
394                        "Struct child array length {} does not match sibling length {}",
395                        array.array.len(),
396                        length
397                    )
398                    .into(),
399                ));
400            }
401            data_size += array.data_size;
402            children.push(array.array);
403            repdefs.push(array.repdef);
404        }
405
406        // Dense rep/def state can retain child-specific repetition information after a child
407        // decoder finishes, so comparing dense siblings is not meaningful. If any child carries
408        // sparse state, keep a sparse child as the canonical structural plan and compare it with
409        // every other sparse sibling so sparse metadata is never silently discarded.
410        let primary_repdef = repdefs
411            .iter()
412            .position(CompositeRepDefUnraveler::has_sparse)
413            .unwrap_or(0);
414        let mut repdef = repdefs.swap_remove(primary_repdef);
415        if repdef.has_sparse() {
416            for sibling in repdefs {
417                if sibling.has_sparse() {
418                    repdef.add_compatibility_check(sibling);
419                }
420            }
421        }
422
423        let validity = if self.is_root {
424            repdef.ensure_exhausted()?;
425            None
426        } else {
427            repdef.unravel_validity(length)?
428        };
429
430        let array = StructArray::try_new(self.child_fields, children, validity)
431            .map_err(|e| Error::invalid_input_source(e.to_string().into()))?;
432        Ok(DecodedArray {
433            array: Arc::new(array),
434            repdef,
435            data_size,
436        })
437    }
438}
439
440/// A structural encoder for struct fields
441///
442/// The struct's validity is added to the rep/def builder
443/// and the builder is cloned to all children.
444pub struct StructStructuralEncoder {
445    keep_original_array: bool,
446    children: Vec<Box<dyn FieldEncoder>>,
447}
448
449impl StructStructuralEncoder {
450    pub fn new(keep_original_array: bool, children: Vec<Box<dyn FieldEncoder>>) -> Self {
451        Self {
452            keep_original_array,
453            children,
454        }
455    }
456}
457
458impl FieldEncoder for StructStructuralEncoder {
459    fn maybe_encode(
460        &mut self,
461        array: ArrayRef,
462        external_buffers: &mut OutOfLineBuffers,
463        mut repdef: RepDefBuilder,
464        row_number: u64,
465        num_rows: u64,
466    ) -> Result<Vec<EncodeTask>> {
467        let struct_array = array.as_struct();
468        let mut struct_array = struct_array.normalize_slicing()?;
469        if let Some(validity) = struct_array.nulls() {
470            if self.keep_original_array {
471                repdef.add_validity_bitmap(validity.clone())
472            } else {
473                repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap())
474            }
475            struct_array = struct_array.pushdown_nulls()?;
476        } else {
477            repdef.add_no_null(struct_array.len());
478        }
479        let child_tasks = self
480            .children
481            .iter_mut()
482            .zip(struct_array.columns().iter())
483            .map(|(encoder, arr)| {
484                encoder.maybe_encode(
485                    arr.clone(),
486                    external_buffers,
487                    repdef.clone(),
488                    row_number,
489                    num_rows,
490                )
491            })
492            .collect::<Result<Vec<_>>>()?;
493        Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
494    }
495
496    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
497        self.children
498            .iter_mut()
499            .map(|encoder| encoder.flush(external_buffers))
500            .flatten_ok()
501            .collect::<Result<Vec<_>>>()
502    }
503
504    fn num_columns(&self) -> u32 {
505        self.children
506            .iter()
507            .map(|child| child.num_columns())
508            .sum::<u32>()
509    }
510
511    fn finish(
512        &mut self,
513        external_buffers: &mut OutOfLineBuffers,
514    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
515        let mut child_columns = self
516            .children
517            .iter_mut()
518            .map(|child| child.finish(external_buffers))
519            .collect::<FuturesOrdered<_>>();
520        async move {
521            let mut encoded_columns = Vec::with_capacity(child_columns.len());
522            while let Some(child_cols) = child_columns.next().await {
523                encoded_columns.extend(child_cols?);
524            }
525            Ok(encoded_columns)
526        }
527        .boxed()
528    }
529}
530
531pub struct StructFieldEncoder {
532    children: Vec<Box<dyn FieldEncoder>>,
533    column_index: u32,
534    num_rows_seen: u64,
535}
536
537impl StructFieldEncoder {
538    pub fn new(children: Vec<Box<dyn FieldEncoder>>, column_index: u32) -> Self {
539        Self {
540            children,
541            column_index,
542            num_rows_seen: 0,
543        }
544    }
545}
546
547impl FieldEncoder for StructFieldEncoder {
548    fn maybe_encode(
549        &mut self,
550        array: ArrayRef,
551        external_buffers: &mut OutOfLineBuffers,
552        repdef: RepDefBuilder,
553        row_number: u64,
554        num_rows: u64,
555    ) -> Result<Vec<EncodeTask>> {
556        self.num_rows_seen += array.len() as u64;
557        let struct_array = array.as_struct();
558        let child_tasks = self
559            .children
560            .iter_mut()
561            .zip(struct_array.columns().iter())
562            .map(|(encoder, arr)| {
563                encoder.maybe_encode(
564                    arr.clone(),
565                    external_buffers,
566                    repdef.clone(),
567                    row_number,
568                    num_rows,
569                )
570            })
571            .collect::<Result<Vec<_>>>()?;
572        Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
573    }
574
575    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
576        let child_tasks = self
577            .children
578            .iter_mut()
579            .map(|encoder| encoder.flush(external_buffers))
580            .collect::<Result<Vec<_>>>()?;
581        Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
582    }
583
584    fn num_columns(&self) -> u32 {
585        self.children
586            .iter()
587            .map(|child| child.num_columns())
588            .sum::<u32>()
589            + 1
590    }
591
592    fn finish(
593        &mut self,
594        external_buffers: &mut OutOfLineBuffers,
595    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
596        let mut child_columns = self
597            .children
598            .iter_mut()
599            .map(|child| child.finish(external_buffers))
600            .collect::<FuturesOrdered<_>>();
601        let num_rows_seen = self.num_rows_seen;
602        let column_index = self.column_index;
603        async move {
604            let mut columns = Vec::new();
605            // Add a column for the struct header
606            let mut header = EncodedColumn::default();
607            header.final_pages.push(EncodedPage {
608                data: Vec::new(),
609                description: PageEncoding::Legacy(pb::ArrayEncoding {
610                    array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct(
611                        pb::SimpleStruct {},
612                    )),
613                }),
614                num_rows: num_rows_seen,
615                column_idx: column_index,
616                row_number: 0, // Not used by legacy encoding
617            });
618            columns.push(header);
619            // Now run finish on the children
620            while let Some(child_cols) = child_columns.next().await {
621                columns.extend(child_cols?);
622            }
623            Ok(columns)
624        }
625        .boxed()
626    }
627}
628
629#[cfg(test)]
630mod tests {
631
632    use std::{collections::HashMap, sync::Arc};
633
634    use arrow_array::{
635        Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, StructArray,
636        builder::{Int32Builder, ListBuilder},
637    };
638    use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
639    use arrow_schema::{DataType, Field, Fields};
640
641    use super::StructuralStructDecoder;
642    use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data};
643
644    #[test]
645    fn test_zero_dimension_fsl_decoder_errors() {
646        // Simulates a stored schema declaring a zero-dimension FixedSizeList (writers reject
647        // it but old files may contain one).  Building the decoder must fail cleanly instead
648        // of letting the zero dimension reach the rep/def decimation.
649        let item_fields = Fields::from(vec![Field::new("x", DataType::Int32, true)]);
650        let fields = Fields::from(vec![Field::new(
651            "vecs",
652            DataType::FixedSizeList(
653                Arc::new(Field::new("item", DataType::Struct(item_fields), true)),
654                0,
655            ),
656            true,
657        )]);
658
659        let err = StructuralStructDecoder::new(fields, false, /*is_root=*/ true).unwrap_err();
660        assert!(matches!(err, lance_core::Error::Schema { .. }));
661        assert!(
662            err.to_string()
663                .contains("dimension must be a positive integer"),
664            "unexpected error: {}",
665            err
666        );
667    }
668
669    #[test_log::test(tokio::test)]
670    async fn test_simple_struct() {
671        let data_type = DataType::Struct(Fields::from(vec![
672            Field::new("a", DataType::Int32, false),
673            Field::new("b", DataType::Int32, false),
674        ]));
675        let field = Field::new("", data_type, false);
676        check_basic_random(field).await;
677    }
678
679    #[test_log::test(tokio::test)]
680    async fn test_nullable_struct() {
681        // Test data struct<score: int32, location: struct<x: int32, y: int32>>
682        // - score: null
683        //   location:
684        //     x: 1
685        //     y: 6
686        // - score: 12
687        //   location:
688        //     x: 2
689        //     y: null
690        // - score: 13
691        //   location:
692        //     x: 3
693        //     y: 8
694        // - score: 14
695        //   location: null
696        // - null
697        //
698        let inner_fields = Fields::from(vec![
699            Field::new("x", DataType::Int32, false),
700            Field::new("y", DataType::Int32, true),
701        ]);
702        let inner_struct = DataType::Struct(inner_fields.clone());
703        let outer_fields = Fields::from(vec![
704            Field::new("score", DataType::Int32, true),
705            Field::new("location", inner_struct, true),
706        ]);
707
708        let x_vals = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
709        let y_vals = Int32Array::from(vec![Some(6), None, Some(8), Some(9), Some(10)]);
710        let scores = Int32Array::from(vec![None, Some(12), Some(13), Some(14), Some(15)]);
711
712        let location_validity = NullBuffer::from(vec![true, true, true, false, true]);
713        let locations = StructArray::new(
714            inner_fields,
715            vec![Arc::new(x_vals), Arc::new(y_vals)],
716            Some(location_validity),
717        );
718
719        let rows_validity = NullBuffer::from(vec![true, true, true, true, false]);
720        let rows = StructArray::new(
721            outer_fields,
722            vec![Arc::new(scores), Arc::new(locations)],
723            Some(rows_validity),
724        );
725
726        let test_cases = TestCases::default().with_structural_encodings();
727
728        check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await;
729    }
730
731    #[test_log::test(tokio::test)]
732    async fn test_simple_masked_nonempty_list() {
733        // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT]
734        //
735        let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4), Some(5), Some(6)]);
736        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 5]));
737        let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
738        let list_array = ListArray::new(
739            Arc::new(Field::new("item", DataType::Int32, true)),
740            offsets,
741            Arc::new(items),
742            Some(NullBuffer::new(list_validity)),
743        );
744        let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
745        let struct_array = StructArray::new(
746            Fields::from(vec![Field::new(
747                "inner_list",
748                list_array.data_type().clone(),
749                true,
750            )]),
751            vec![Arc::new(list_array)],
752            Some(NullBuffer::new(struct_validity)),
753        );
754        check_round_trip_encoding_of_data(
755            vec![Arc::new(struct_array)],
756            &TestCases::default().with_structural_encodings(),
757            HashMap::new(),
758        )
759        .await;
760    }
761
762    #[test_log::test(tokio::test)]
763    async fn test_simple_struct_list() {
764        // [[1, 2], [NULL], [4], [], NULL, NULL-STRUCT]
765        //
766        let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
767        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 4]));
768        let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
769        let list_array = ListArray::new(
770            Arc::new(Field::new("item", DataType::Int32, true)),
771            offsets,
772            Arc::new(items),
773            Some(NullBuffer::new(list_validity)),
774        );
775        let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
776        let struct_array = StructArray::new(
777            Fields::from(vec![Field::new(
778                "inner_list",
779                list_array.data_type().clone(),
780                true,
781            )]),
782            vec![Arc::new(list_array)],
783            Some(NullBuffer::new(struct_validity)),
784        );
785        check_round_trip_encoding_of_data(
786            vec![Arc::new(struct_array)],
787            &TestCases::default().with_structural_encodings(),
788            HashMap::new(),
789        )
790        .await;
791    }
792
793    #[test_log::test(tokio::test)]
794    async fn test_struct_list() {
795        let data_type = DataType::Struct(Fields::from(vec![
796            Field::new(
797                "inner_list",
798                DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
799                true,
800            ),
801            Field::new("outer_int", DataType::Int32, true),
802        ]));
803        let field = Field::new("row", data_type, false);
804        check_basic_random(field).await;
805    }
806
807    #[test_log::test(tokio::test)]
808    async fn test_empty_struct() {
809        // It's technically legal for a struct to have 0 children, need to
810        // make sure we support that
811        let data_type = DataType::Struct(Fields::from(Vec::<Field>::default()));
812        let field = Field::new("row", data_type, false);
813        check_basic_random(field).await;
814    }
815
816    #[test_log::test(tokio::test)]
817    async fn test_complicated_struct() {
818        let data_type = DataType::Struct(Fields::from(vec![
819            Field::new("int", DataType::Int32, true),
820            Field::new(
821                "inner",
822                DataType::Struct(Fields::from(vec![
823                    Field::new("inner_int", DataType::Int32, true),
824                    Field::new(
825                        "inner_list",
826                        DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
827                        true,
828                    ),
829                ])),
830                true,
831            ),
832            Field::new("outer_binary", DataType::Binary, true),
833        ]));
834        let field = Field::new("row", data_type, false);
835        check_basic_random(field).await;
836    }
837
838    #[test_log::test(tokio::test)]
839    async fn test_list_of_struct_with_null_struct_element() {
840        // Regression: a list containing structs where most struct elements are null
841        // causes a length mismatch during decoding with V2_2 encoding.
842        use arrow_array::StringArray;
843
844        let tag_array = StringArray::from(vec![
845            Some("valid"),
846            Some("null_struct"),
847            Some("valid"),
848            Some("valid"),
849        ]);
850        let struct_fields = Fields::from(vec![Field::new("tag", DataType::Utf8, true)]);
851        // 3 out of 4 struct elements are null
852        let struct_validity = NullBuffer::from(vec![false, true, false, false]);
853        let struct_array = StructArray::new(
854            struct_fields.clone(),
855            vec![Arc::new(tag_array)],
856            Some(struct_validity),
857        );
858
859        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4]));
860        let list_field = Field::new("item", DataType::Struct(struct_fields), true);
861        let list_array =
862            ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
863
864        check_round_trip_encoding_of_data(
865            vec![Arc::new(list_array)],
866            &TestCases::default().with_u32_structural_encodings(),
867            HashMap::new(),
868        )
869        .await;
870    }
871
872    #[test_log::test(tokio::test)]
873    async fn test_list_of_struct_with_all_null_child() {
874        use arrow_array::StringArray;
875
876        let a_array = StringArray::from(vec![Some("w"), Some("x"), Some("y"), Some("z")]);
877        let b_array = StringArray::from(vec![None::<&str>, None, None, None]);
878        let struct_fields = Fields::from(vec![
879            Field::new("a", DataType::Utf8, true),
880            Field::new("b", DataType::Utf8, true),
881        ]);
882        let struct_array = StructArray::new(
883            struct_fields.clone(),
884            vec![Arc::new(a_array), Arc::new(b_array)],
885            None,
886        );
887
888        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 4]));
889        let list_field = Field::new("item", DataType::Struct(struct_fields), true);
890        let list_array =
891            ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
892
893        check_round_trip_encoding_of_data(
894            vec![Arc::new(list_array)],
895            &TestCases::default()
896                .with_range(0..1)
897                .with_range(1..2)
898                .with_indices(vec![0])
899                .with_indices(vec![1])
900                .with_structural_encodings(),
901            HashMap::new(),
902        )
903        .await;
904    }
905
906    #[test_log::test(tokio::test)]
907    async fn test_list_of_struct_with_constant_child_and_empty_lists() {
908        let item_fields = Fields::from(vec![
909            Field::new("maneuver", DataType::Int64, true),
910            Field::new("remaining_dist", DataType::Float64, true),
911        ]);
912        let item_array = StructArray::new(
913            item_fields.clone(),
914            vec![
915                Arc::new(Int64Array::from(vec![Some(3), Some(3), Some(3)])),
916                Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])),
917            ],
918            None,
919        );
920
921        let list_field = Field::new("item", DataType::Struct(item_fields), true);
922        let list_array = ListArray::new(
923            Arc::new(list_field),
924            OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 3, 3, 3])),
925            Arc::new(item_array),
926            None,
927        );
928
929        let data_fields = Fields::from(vec![Field::new(
930            "maneuvers",
931            list_array.data_type().clone(),
932            true,
933        )]);
934        let data_array = StructArray::new(data_fields.clone(), vec![Arc::new(list_array)], None);
935        let row_validity = NullBuffer::from(vec![true, true, true, true, true, false]);
936        let row_array = StructArray::new(
937            Fields::from(vec![Field::new(
938                "data",
939                DataType::Struct(data_fields),
940                true,
941            )]),
942            vec![Arc::new(data_array)],
943            Some(row_validity),
944        );
945
946        check_round_trip_encoding_of_data(
947            vec![Arc::new(row_array)],
948            &TestCases::default().with_structural_encodings(),
949            HashMap::new(),
950        )
951        .await;
952    }
953
954    #[test_log::test(tokio::test)]
955    async fn test_ragged_scheduling() {
956        // This test covers scheduling when batches straddle page boundaries
957
958        // Create a list with 10k nulls
959        let items_builder = Int32Builder::new();
960        let mut list_builder = ListBuilder::new(items_builder);
961        for _ in 0..10000 {
962            list_builder.append_null();
963        }
964        let list_array = Arc::new(list_builder.finish());
965        let int_array = Arc::new(Int32Array::from_iter_values(0..10000));
966        let fields = vec![
967            Field::new("", list_array.data_type().clone(), true),
968            Field::new("", int_array.data_type().clone(), true),
969        ];
970        let struct_array = Arc::new(StructArray::new(
971            Fields::from(fields),
972            vec![list_array, int_array],
973            None,
974        )) as ArrayRef;
975        let struct_arrays = (0..10000)
976            // Intentionally skip in some randomish amount to create more ragged scheduling
977            .step_by(437)
978            .map(|offset| struct_array.slice(offset, 437.min(10000 - offset)))
979            .collect::<Vec<_>>();
980        check_round_trip_encoding_of_data(struct_arrays, &TestCases::default(), HashMap::new())
981            .await;
982    }
983}