lance_encoding/encodings/logical/
fixed_size_list.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Encoding support for complex FixedSizeList types (FSL with non-primitive children).
5//!
6//! Primitive FSL (e.g., `FixedSizeList<Int32>`) is handled in the physical encoding layer.
7//! This module handles FSL with complex children (Struct, Map, List) which require
8//! structural encoding.
9
10use std::{ops::Range, sync::Arc};
11
12use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait, StructArray, cast::AsArray};
13use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer, ScalarBuffer};
14use arrow_schema::DataType;
15use futures::{FutureExt, future::BoxFuture};
16use lance_arrow::deepcopy::deep_copy_nulls;
17use lance_core::{Error, Result};
18
19use crate::{
20    decoder::{
21        DecodedArray, FilterExpression, ScheduledScanLine, SchedulerContext,
22        StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
23        StructuralSchedulingJob,
24    },
25    encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers},
26    repdef::RepDefBuilder,
27};
28
29/// A structural encoder for complex fixed-size list fields
30///
31/// The FSL's validity is added to the rep/def builder along with the dimension
32/// and the FSL array's values are passed to the child encoder.
33pub struct FixedSizeListStructuralEncoder {
34    keep_original_array: bool,
35    child: Box<dyn FieldEncoder>,
36}
37
38impl FixedSizeListStructuralEncoder {
39    pub fn new(keep_original_array: bool, child: Box<dyn FieldEncoder>) -> Self {
40        Self {
41            keep_original_array,
42            child,
43        }
44    }
45}
46
47impl FieldEncoder for FixedSizeListStructuralEncoder {
48    fn maybe_encode(
49        &mut self,
50        array: ArrayRef,
51        external_buffers: &mut OutOfLineBuffers,
52        mut repdef: RepDefBuilder,
53        row_number: u64,
54        num_rows: u64,
55    ) -> Result<Vec<EncodeTask>> {
56        let fsl_arr = array.as_fixed_size_list_opt().ok_or_else(|| {
57            Error::internal("FixedSizeList encoder used for non-fixed-size-list data".to_string())
58        })?;
59
60        let dimension = fsl_arr.value_length() as usize;
61        let values = fsl_arr.values().clone();
62
63        let validity = if self.keep_original_array {
64            array.nulls().cloned()
65        } else {
66            deep_copy_nulls(array.nulls())
67        };
68        repdef.add_fsl(validity.clone(), dimension, fsl_arr.len());
69
70        // FSL forces child elements to exist even under null rows. Normalize any
71        // nested lists under null FSL rows to null empty lists.
72        let values = if let Some(ref fsl_validity) = validity {
73            if needs_garbage_filtering(values.data_type()) {
74                let is_garbage =
75                    expand_garbage_mask(&fsl_validity_to_garbage_mask(fsl_validity), dimension);
76                filter_fsl_child_garbage(values, &is_garbage)
77            } else {
78                values
79            }
80        } else {
81            values
82        };
83
84        self.child.maybe_encode(
85            values,
86            external_buffers,
87            repdef,
88            row_number,
89            num_rows * dimension as u64,
90        )
91    }
92
93    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
94        self.child.flush(external_buffers)
95    }
96
97    fn num_columns(&self) -> u32 {
98        self.child.num_columns()
99    }
100
101    fn finish(
102        &mut self,
103        external_buffers: &mut OutOfLineBuffers,
104    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
105        self.child.finish(external_buffers)
106    }
107}
108
109/// A scheduler for complex fixed-size list fields
110///
111/// Scales row ranges by the FSL dimension when scheduling child rows,
112/// and scales scheduled rows back when reporting to the parent.
113#[derive(Debug)]
114pub struct StructuralFixedSizeListScheduler {
115    child: Box<dyn StructuralFieldScheduler>,
116    dimension: u64,
117}
118
119impl StructuralFixedSizeListScheduler {
120    pub fn new(child: Box<dyn StructuralFieldScheduler>, dimension: i32) -> Self {
121        Self {
122            child,
123            dimension: dimension as u64,
124        }
125    }
126
127    fn child_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Range<u64>>> {
128        ranges
129            .iter()
130            .map(|range| {
131                let start = range.start.checked_mul(self.dimension).ok_or_else(|| {
132                    Error::invalid_input(format!(
133                        "fixed-size-list child range start overflowed: {} * {}",
134                        range.start, self.dimension
135                    ))
136                })?;
137                let end = range.end.checked_mul(self.dimension).ok_or_else(|| {
138                    Error::invalid_input(format!(
139                        "fixed-size-list child range end overflowed: {} * {}",
140                        range.end, self.dimension
141                    ))
142                })?;
143                Ok(start..end)
144            })
145            .collect()
146    }
147}
148
149impl StructuralFieldScheduler for StructuralFixedSizeListScheduler {
150    fn schedule_ranges<'a>(
151        &'a self,
152        ranges: &[Range<u64>],
153        filter: &FilterExpression,
154    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
155        // Scale ranges by dimension for the child - each FSL row becomes `dimension` child rows
156        let child_ranges = self.child_ranges(ranges)?;
157        let child = self.child.schedule_ranges(&child_ranges, filter)?;
158        Ok(Box::new(StructuralFixedSizeListSchedulingJob::new(
159            child,
160            self.dimension,
161        )))
162    }
163
164    fn initialize<'a>(
165        &'a mut self,
166        requested_ranges: Option<&'a [Range<u64>]>,
167        filter: &'a FilterExpression,
168        context: &'a SchedulerContext,
169    ) -> BoxFuture<'a, Result<()>> {
170        let child_ranges = match requested_ranges.map(|ranges| self.child_ranges(ranges)) {
171            Some(Ok(child_ranges)) => Some(child_ranges),
172            Some(Err(error)) => return std::future::ready(Err(error)).boxed(),
173            None => None,
174        };
175        async move {
176            self.child
177                .initialize(child_ranges.as_deref(), filter, context)
178                .await
179        }
180        .boxed()
181    }
182}
183
184#[derive(Debug)]
185struct StructuralFixedSizeListSchedulingJob<'a> {
186    child: Box<dyn StructuralSchedulingJob + 'a>,
187    dimension: u64,
188}
189
190impl<'a> StructuralFixedSizeListSchedulingJob<'a> {
191    fn new(child: Box<dyn StructuralSchedulingJob + 'a>, dimension: u64) -> Self {
192        Self { child, dimension }
193    }
194}
195
196impl StructuralSchedulingJob for StructuralFixedSizeListSchedulingJob<'_> {
197    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
198        // Get the child's scan lines (scheduled in terms of child struct rows)
199        let child_scan_lines = self.child.schedule_next(context)?;
200
201        // Scale down rows_scheduled by dimension to convert from child rows to FSL rows
202        Ok(child_scan_lines
203            .into_iter()
204            .map(|scan_line| ScheduledScanLine {
205                decoders: scan_line.decoders,
206                rows_scheduled: scan_line.rows_scheduled / self.dimension,
207            })
208            .collect())
209    }
210}
211
212/// A decoder for complex fixed-size list fields
213///
214/// Drains `num_rows * dimension` from the child decoder and reconstructs
215/// the FSL array with validity from the rep/def information.
216#[derive(Debug)]
217pub struct StructuralFixedSizeListDecoder {
218    child: Box<dyn StructuralFieldDecoder>,
219    data_type: DataType,
220}
221
222impl StructuralFixedSizeListDecoder {
223    pub fn new(child: Box<dyn StructuralFieldDecoder>, data_type: DataType) -> Self {
224        Self { child, data_type }
225    }
226}
227
228impl StructuralFieldDecoder for StructuralFixedSizeListDecoder {
229    fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> {
230        self.child.accept_page(child)
231    }
232
233    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
234        // For FixedSizeList, we need to drain num_rows * dimension from the child
235        let dimension = match &self.data_type {
236            DataType::FixedSizeList(_, d) => *d as u64,
237            _ => {
238                return Err(Error::internal(
239                    "FixedSizeListDecoder has non-FSL data type".to_string(),
240                ));
241            }
242        };
243        let child_task = self.child.drain(num_rows * dimension)?;
244        Ok(Box::new(StructuralFixedSizeListDecodeTask::new(
245            child_task,
246            self.data_type.clone(),
247            num_rows,
248        )))
249    }
250
251    fn data_type(&self) -> &DataType {
252        &self.data_type
253    }
254}
255
256#[derive(Debug)]
257struct StructuralFixedSizeListDecodeTask {
258    child_task: Box<dyn StructuralDecodeArrayTask>,
259    data_type: DataType,
260    num_rows: u64,
261}
262
263impl StructuralFixedSizeListDecodeTask {
264    fn new(
265        child_task: Box<dyn StructuralDecodeArrayTask>,
266        data_type: DataType,
267        num_rows: u64,
268    ) -> Self {
269        Self {
270            child_task,
271            data_type,
272            num_rows,
273        }
274    }
275}
276
277impl StructuralDecodeArrayTask for StructuralFixedSizeListDecodeTask {
278    fn decode(self: Box<Self>) -> Result<DecodedArray> {
279        let DecodedArray {
280            array,
281            mut repdef,
282            data_size,
283        } = self.child_task.decode()?;
284        match &self.data_type {
285            DataType::FixedSizeList(child_field, dimension) => {
286                let num_rows = self.num_rows as usize;
287                let validity = repdef.unravel_fsl_validity(num_rows, *dimension as usize)?;
288                let fsl_array = arrow_array::FixedSizeListArray::try_new(
289                    child_field.clone(),
290                    *dimension,
291                    array,
292                    validity,
293                )?;
294                Ok(DecodedArray {
295                    array: Arc::new(fsl_array),
296                    repdef,
297                    data_size,
298                })
299            }
300            _ => Err(Error::internal(
301                "FixedSizeList decoder did not have a fixed-size list field".to_string(),
302            )),
303        }
304    }
305}
306
307// =======================
308// Garbage filtering
309// =======================
310
311/// Returns true if the data type contains any variable-length list-like types
312/// (List, LargeList, ListView, LargeListView, Map) that need garbage filtering.
313fn needs_garbage_filtering(data_type: &DataType) -> bool {
314    match data_type {
315        DataType::List(_)
316        | DataType::LargeList(_)
317        | DataType::ListView(_)
318        | DataType::LargeListView(_)
319        | DataType::Map(_, _) => true,
320        DataType::Struct(fields) => fields
321            .iter()
322            .any(|f| needs_garbage_filtering(f.data_type())),
323        DataType::FixedSizeList(field, _) => needs_garbage_filtering(field.data_type()),
324        _ => false,
325    }
326}
327
328/// Filters garbage (undefined data under null FSL rows) from nested list-like types.
329/// Unlike variable-length lists which can remove null children entirely, FSL children
330/// always exist, so we must clean any nested lists before encoding.
331///
332/// NB: Nested FSL is currently precluded at a higher level in our system. However, this code
333/// supports and tests it.
334fn filter_fsl_child_garbage(array: ArrayRef, is_garbage: &[bool]) -> ArrayRef {
335    debug_assert_eq!(array.len(), is_garbage.len());
336
337    match array.data_type() {
338        DataType::List(_) => filter_list_garbage(array.as_list::<i32>(), is_garbage),
339        DataType::LargeList(_) => filter_list_garbage(array.as_list::<i64>(), is_garbage),
340        DataType::ListView(_) | DataType::LargeListView(_) => {
341            unimplemented!("ListView inside complex FSL is not yet supported")
342        }
343        DataType::Map(_, _) => filter_map_garbage(array.as_map(), is_garbage),
344        DataType::FixedSizeList(_, dim) => {
345            filter_nested_fsl_garbage(array.as_fixed_size_list(), is_garbage, *dim as usize)
346        }
347        DataType::Struct(_) => filter_struct_garbage(array.as_struct(), is_garbage),
348        _ => array,
349    }
350}
351
352fn filter_struct_garbage(struct_arr: &StructArray, is_garbage: &[bool]) -> ArrayRef {
353    let needs_filtering = struct_arr
354        .fields()
355        .iter()
356        .any(|f| needs_garbage_filtering(f.data_type()));
357
358    if !needs_filtering {
359        return Arc::new(struct_arr.clone());
360    }
361
362    let new_columns: Vec<ArrayRef> = struct_arr
363        .columns()
364        .iter()
365        .zip(struct_arr.fields().iter())
366        .map(|(col, field)| {
367            if needs_garbage_filtering(field.data_type()) {
368                filter_fsl_child_garbage(col.clone(), is_garbage)
369            } else {
370                col.clone()
371            }
372        })
373        .collect();
374
375    Arc::new(StructArray::new(
376        struct_arr.fields().clone(),
377        new_columns,
378        struct_arr.nulls().cloned(),
379    ))
380}
381
382fn expand_garbage_mask(is_garbage: &[bool], dimension: usize) -> Vec<bool> {
383    let mut expanded = Vec::with_capacity(is_garbage.len() * dimension);
384    for &garbage in is_garbage {
385        for _ in 0..dimension {
386            expanded.push(garbage);
387        }
388    }
389    expanded
390}
391
392fn fsl_validity_to_garbage_mask(fsl_validity: &NullBuffer) -> Vec<bool> {
393    fsl_validity.iter().map(|valid| !valid).collect()
394}
395
396fn filter_list_garbage<O: OffsetSizeTrait>(
397    list_arr: &GenericListArray<O>,
398    is_garbage: &[bool],
399) -> ArrayRef {
400    debug_assert_eq!(
401        list_arr.len(),
402        is_garbage.len(),
403        "list length must match garbage mask length"
404    );
405
406    let old_offsets = list_arr.offsets();
407    let value_field = match list_arr.data_type() {
408        DataType::List(f) | DataType::LargeList(f) => f.clone(),
409        _ => unreachable!(),
410    };
411
412    let mut new_offsets: Vec<O> = Vec::with_capacity(list_arr.len() + 1);
413    let mut values_to_keep: Vec<usize> = Vec::new();
414    let mut validity_builder = BooleanBufferBuilder::new(list_arr.len());
415    let mut current_offset = O::usize_as(0);
416    new_offsets.push(current_offset);
417    let old_validity = list_arr.nulls();
418
419    for (i, &garbage) in is_garbage.iter().enumerate() {
420        if garbage {
421            new_offsets.push(current_offset);
422            validity_builder.append(false);
423        } else {
424            let start = old_offsets[i].as_usize();
425            let end = old_offsets[i + 1].as_usize();
426            values_to_keep.extend(start..end);
427            current_offset += O::usize_as(end - start);
428            new_offsets.push(current_offset);
429            validity_builder.append(old_validity.map(|v| v.is_valid(i)).unwrap_or(true));
430        }
431    }
432
433    let new_values = if values_to_keep.is_empty() {
434        list_arr.values().slice(0, 0)
435    } else {
436        let indices =
437            arrow_array::UInt64Array::from_iter_values(values_to_keep.iter().map(|&i| i as u64));
438        arrow_select::take::take(list_arr.values().as_ref(), &indices, None)
439            .expect("take should succeed")
440    };
441
442    let new_values = if needs_garbage_filtering(value_field.data_type()) && !new_values.is_empty() {
443        let len = new_values.len();
444        filter_fsl_child_garbage(new_values, &vec![false; len])
445    } else {
446        new_values
447    };
448
449    let new_validity = NullBuffer::new(validity_builder.finish());
450    Arc::new(GenericListArray::new(
451        value_field,
452        OffsetBuffer::new(ScalarBuffer::from(new_offsets)),
453        new_values,
454        Some(new_validity),
455    ))
456}
457
458fn filter_map_garbage(map_arr: &arrow_array::MapArray, is_garbage: &[bool]) -> ArrayRef {
459    debug_assert_eq!(map_arr.len(), is_garbage.len());
460
461    let old_offsets = map_arr.offsets();
462    let entries_field = match map_arr.data_type() {
463        DataType::Map(field, _) => field.clone(),
464        _ => unreachable!(),
465    };
466
467    let mut new_offsets: Vec<i32> = Vec::with_capacity(map_arr.len() + 1);
468    let mut values_to_keep: Vec<usize> = Vec::new();
469    let mut validity_builder = BooleanBufferBuilder::new(map_arr.len());
470    let mut current_offset: i32 = 0;
471    new_offsets.push(current_offset);
472    let old_validity = map_arr.nulls();
473
474    for (i, &garbage) in is_garbage.iter().enumerate() {
475        if garbage {
476            new_offsets.push(current_offset);
477            validity_builder.append(false);
478        } else {
479            let start = old_offsets[i] as usize;
480            let end = old_offsets[i + 1] as usize;
481            values_to_keep.extend(start..end);
482            current_offset += (end - start) as i32;
483            new_offsets.push(current_offset);
484            validity_builder.append(old_validity.map(|v| v.is_valid(i)).unwrap_or(true));
485        }
486    }
487
488    let new_entries: ArrayRef = if values_to_keep.is_empty() {
489        Arc::new(map_arr.entries().slice(0, 0))
490    } else {
491        let indices =
492            arrow_array::UInt64Array::from_iter_values(values_to_keep.iter().map(|&i| i as u64));
493        arrow_select::take::take(map_arr.entries(), &indices, None).expect("take should succeed")
494    };
495
496    let new_entries =
497        if needs_garbage_filtering(entries_field.data_type()) && !new_entries.is_empty() {
498            let len = new_entries.len();
499            filter_fsl_child_garbage(new_entries, &vec![false; len])
500        } else {
501            new_entries
502        };
503
504    let new_validity = NullBuffer::new(validity_builder.finish());
505    let keys_sorted = matches!(map_arr.data_type(), DataType::Map(_, true));
506
507    Arc::new(
508        arrow_array::MapArray::try_new(
509            entries_field,
510            OffsetBuffer::new(ScalarBuffer::from(new_offsets)),
511            new_entries.as_struct().clone(),
512            Some(new_validity),
513            keys_sorted,
514        )
515        .expect("MapArray construction should succeed"),
516    )
517}
518
519/// Filters garbage from nested FSL arrays that contain list-like children.
520fn filter_nested_fsl_garbage(
521    fsl_arr: &arrow_array::FixedSizeListArray,
522    is_garbage: &[bool],
523    dimension: usize,
524) -> ArrayRef {
525    debug_assert_eq!(fsl_arr.len(), is_garbage.len());
526
527    let child_field = match fsl_arr.data_type() {
528        DataType::FixedSizeList(field, _) => field.clone(),
529        _ => unreachable!(),
530    };
531
532    if !needs_garbage_filtering(child_field.data_type()) {
533        return Arc::new(fsl_arr.clone());
534    }
535
536    let child_garbage = expand_garbage_mask(is_garbage, dimension);
537    let new_values = filter_fsl_child_garbage(fsl_arr.values().clone(), &child_garbage);
538
539    Arc::new(arrow_array::FixedSizeListArray::new(
540        child_field,
541        dimension as i32,
542        new_values,
543        fsl_arr.nulls().cloned(),
544    ))
545}
546
547#[cfg(test)]
548mod tests {
549    use std::{collections::HashMap, sync::Arc};
550
551    use arrow_array::{
552        Array, FixedSizeListArray,
553        builder::{Int32Builder, ListBuilder},
554        cast::AsArray,
555    };
556    use arrow_schema::{DataType, Field, Fields};
557    use rstest::rstest;
558
559    use super::filter_nested_fsl_garbage;
560    use crate::{
561        constants::{
562            STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY,
563            STRUCTURAL_ENCODING_MINIBLOCK,
564        },
565        testing::{
566            TestCases, TestEncoding, check_round_trip_encoding_of_data, check_specific_random,
567        },
568    };
569
570    fn make_fsl_struct_type(struct_fields: Fields, dimension: i32) -> DataType {
571        DataType::FixedSizeList(
572            Arc::new(Field::new("item", DataType::Struct(struct_fields), true)),
573            dimension,
574        )
575    }
576
577    fn simple_struct_fields() -> Fields {
578        Fields::from(vec![
579            Field::new("x", DataType::Float64, false),
580            Field::new("y", DataType::Float64, false),
581        ])
582    }
583
584    fn nested_struct_fields() -> Fields {
585        let inner = Fields::from(vec![
586            Field::new("a", DataType::Int32, false),
587            Field::new("b", DataType::Int32, false),
588        ]);
589        Fields::from(vec![
590            Field::new("outer_val", DataType::Float64, false),
591            Field::new("inner", DataType::Struct(inner), true),
592        ])
593    }
594
595    fn nested_struct_with_list_fields() -> Fields {
596        let inner = Fields::from(vec![Field::new(
597            "values",
598            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
599            true,
600        )]);
601        Fields::from(vec![
602            Field::new("id", DataType::Int32, false),
603            Field::new("inner", DataType::Struct(inner), true),
604        ])
605    }
606
607    fn struct_with_list_fields() -> Fields {
608        Fields::from(vec![
609            Field::new("id", DataType::Int32, false),
610            Field::new(
611                "values",
612                DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
613                true,
614            ),
615        ])
616    }
617
618    fn struct_with_large_list_fields() -> Fields {
619        Fields::from(vec![
620            Field::new("id", DataType::Int32, false),
621            Field::new(
622                "values",
623                DataType::LargeList(Arc::new(Field::new("item", DataType::Int64, true))),
624                true,
625            ),
626        ])
627    }
628
629    fn struct_with_nested_fsl_fields() -> Fields {
630        Fields::from(vec![
631            Field::new("id", DataType::Int32, false),
632            Field::new(
633                "vectors",
634                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4),
635                true,
636            ),
637        ])
638    }
639
640    fn struct_with_map_fields() -> Fields {
641        let entries_field = Arc::new(Field::new(
642            "entries",
643            DataType::Struct(Fields::from(vec![
644                Field::new("keys", DataType::Utf8, false),
645                Field::new("values", DataType::Int32, true),
646            ])),
647            false,
648        ));
649        Fields::from(vec![
650            Field::new("id", DataType::Int32, false),
651            Field::new("props", DataType::Map(entries_field, false), true),
652        ])
653    }
654
655    fn make_fsl_of_list() -> DataType {
656        DataType::FixedSizeList(
657            Arc::new(Field::new(
658                "item",
659                DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
660                true,
661            )),
662            2,
663        )
664    }
665
666    fn make_fsl_of_large_list() -> DataType {
667        DataType::FixedSizeList(
668            Arc::new(Field::new(
669                "item",
670                DataType::LargeList(Arc::new(Field::new("item", DataType::Int32, true))),
671                true,
672            )),
673            2,
674        )
675    }
676
677    fn make_fsl_of_map() -> DataType {
678        DataType::FixedSizeList(
679            Arc::new(Field::new(
680                "item",
681                DataType::Map(
682                    Arc::new(Field::new(
683                        "entries",
684                        DataType::Struct(Fields::from(vec![
685                            Field::new("key", DataType::Utf8, false),
686                            Field::new("value", DataType::Int32, true),
687                        ])),
688                        false,
689                    )),
690                    false,
691                ),
692                true,
693            )),
694            2,
695        )
696    }
697
698    fn make_fsl_of_nested_fsl_struct() -> DataType {
699        DataType::FixedSizeList(
700            Arc::new(Field::new(
701                "item",
702                DataType::FixedSizeList(
703                    Arc::new(Field::new(
704                        "item",
705                        DataType::Struct(Fields::from(vec![Field::new(
706                            "x",
707                            DataType::Int32,
708                            true,
709                        )])),
710                        true,
711                    )),
712                    4,
713                ),
714                true,
715            )),
716            2,
717        )
718    }
719
720    #[rstest]
721    #[case::simple(simple_struct_fields(), 2)]
722    #[case::nested_struct(nested_struct_fields(), 2)]
723    #[case::struct_with_list(struct_with_list_fields(), 2)]
724    #[case::struct_with_large_list(struct_with_large_list_fields(), 2)]
725    #[case::nested_struct_with_list(nested_struct_with_list_fields(), 2)]
726    #[case::struct_with_nested_fsl(struct_with_nested_fsl_fields(), 2)]
727    #[case::struct_with_map(struct_with_map_fields(), 2)]
728    #[test_log::test(tokio::test)]
729    async fn test_fsl_struct_random(
730        #[case] struct_fields: Fields,
731        #[case] dimension: i32,
732        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
733        structural_encoding: &str,
734        #[values(TestEncoding::StructuralU32, TestEncoding::StructuralSparse)]
735        encoding: TestEncoding,
736        #[values(4096, 1024 * 1024)] page_size: u64,
737        #[values(false, true)] use_slicing: bool,
738        #[values(1, 5, 10)] ingest_batch_count: u32,
739    ) {
740        let data_type = make_fsl_struct_type(struct_fields, dimension);
741        let mut field_metadata = HashMap::new();
742        field_metadata.insert(
743            STRUCTURAL_ENCODING_META_KEY.to_string(),
744            structural_encoding.into(),
745        );
746        let field = Field::new("", data_type, true).with_metadata(field_metadata);
747        let test_cases = TestCases::basic()
748            .with_encoding(encoding)
749            .with_page_sizes(vec![page_size])
750            .with_slicing_modes([use_slicing])
751            .with_ingest_batch_counts([ingest_batch_count]);
752        check_specific_random(field, test_cases).await;
753    }
754
755    #[rstest]
756    #[case::list(make_fsl_of_list())]
757    #[case::large_list(make_fsl_of_large_list())]
758    #[case::map(make_fsl_of_map())]
759    #[case::nested_fsl_struct(make_fsl_of_nested_fsl_struct())]
760    fn test_unsupported_fsl_child_types_return_error(#[case] data_type: DataType) {
761        let arrow_field = Field::new("test", data_type, true);
762        let err = lance_core::datatypes::Field::try_from(&arrow_field).unwrap_err();
763        assert!(err.to_string().contains("Unsupported data type"));
764    }
765
766    #[test]
767    fn test_filter_nested_fsl_garbage() {
768        // Create FSL<List<Int32>> with dimension 2: [[[1], [2]], [[3], [4]], [[5], [6]]]
769        let mut list_builder = ListBuilder::new(Int32Builder::new());
770        for i in 1..=6 {
771            list_builder.values().append_value(i);
772            list_builder.append(true);
773        }
774        let list_arr = list_builder.finish();
775
776        let fsl_field = Arc::new(Field::new(
777            "item",
778            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
779            true,
780        ));
781        let fsl = FixedSizeListArray::new(fsl_field, 2, Arc::new(list_arr), None);
782
783        // Mark second FSL row as garbage
784        let result = filter_nested_fsl_garbage(&fsl, &[false, true, false], 2);
785        let result = result.as_fixed_size_list();
786
787        // Child lists at positions 2,3 (garbage row 1) should be filtered to null
788        let child_list = result.values().as_list::<i32>();
789        assert_eq!(
790            (0..6).map(|i| child_list.is_valid(i)).collect::<Vec<_>>(),
791            vec![true, true, false, false, true, true]
792        );
793    }
794
795    #[test]
796    fn test_filter_nested_fsl_no_list_child() {
797        // FSL<Int32> - no list child, should return unchanged
798        let fsl_field = Arc::new(Field::new("item", DataType::Int32, true));
799        let values = arrow_array::Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
800        let fsl = FixedSizeListArray::new(fsl_field, 2, Arc::new(values), None);
801
802        let result = filter_nested_fsl_garbage(&fsl, &[false, true, false], 2);
803        // Should return the same array unchanged
804        assert_eq!(result.len(), 3);
805    }
806
807    /// A null-free nested FSL written with a forced full-zip layout goes through the
808    /// per-value path that flattens it. Both layers are nullable, as Arrow and pyarrow
809    /// build them by default, and the data carries no nulls.
810    #[rstest]
811    #[case::outer_wider(22, 3)]
812    #[case::inner_wider(3, 22)]
813    #[tokio::test]
814    async fn test_nested_fsl_full_zip_round_trip(#[case] outer_dim: i32, #[case] inner_dim: i32) {
815        let num_rows = 3;
816        let inner_field = Arc::new(Field::new("item", DataType::Int32, true));
817        let leaf = Arc::new(arrow_array::Int32Array::from_iter_values(
818            0..num_rows * outer_dim * inner_dim,
819        ));
820        let inner = FixedSizeListArray::new(inner_field.clone(), inner_dim, leaf, None);
821        let outer = FixedSizeListArray::new(
822            Arc::new(Field::new(
823                "item",
824                DataType::FixedSizeList(inner_field, inner_dim),
825                true,
826            )),
827            outer_dim,
828            Arc::new(inner),
829            None,
830        );
831
832        check_round_trip_encoding_of_data(
833            vec![Arc::new(outer)],
834            &TestCases::default()
835                .with_structural_encodings()
836                .with_range(0..1)
837                .with_range(1..3)
838                .with_indices(vec![0, 2]),
839            HashMap::from([(
840                STRUCTURAL_ENCODING_META_KEY.to_string(),
841                STRUCTURAL_ENCODING_FULLZIP.to_string(),
842            )]),
843        )
844        .await;
845    }
846}
lance_encoding/encodings/logical/fixed_size_list.rs

lance_encoding/encodings/logical/
fixed_size_list.rs
