lance_encoding/encodings/logical/
list.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{ops::Range, sync::Arc};
5
6use arrow_array::{Array, ArrayRef, LargeListArray, ListArray, cast::AsArray, make_array};
7use arrow_schema::DataType;
8use futures::future::BoxFuture;
9use lance_arrow::deepcopy::deep_copy_nulls;
10use lance_arrow::list::ListArrayExt;
11use lance_core::Result;
12
13use crate::{
14    decoder::{
15        DecodedArray, FilterExpression, ScheduledScanLine, SchedulerContext,
16        StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
17        StructuralSchedulingJob,
18    },
19    encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers},
20    repdef::RepDefBuilder,
21};
22
23/// A structural encoder for list fields
24///
25/// The list's offsets are added to the rep/def builder
26/// and the list array's values are passed to the child encoder
27///
28/// The values will have any garbage values removed and will be trimmed
29/// to only include the values that are actually used.
30pub struct ListStructuralEncoder {
31    keep_original_array: bool,
32    child: Box<dyn FieldEncoder>,
33}
34
35impl ListStructuralEncoder {
36    pub fn new(keep_original_array: bool, child: Box<dyn FieldEncoder>) -> Self {
37        Self {
38            keep_original_array,
39            child,
40        }
41    }
42}
43
44impl FieldEncoder for ListStructuralEncoder {
45    fn maybe_encode(
46        &mut self,
47        array: ArrayRef,
48        external_buffers: &mut OutOfLineBuffers,
49        mut repdef: RepDefBuilder,
50        row_number: u64,
51        num_rows: u64,
52    ) -> Result<Vec<EncodeTask>> {
53        let values = if let Some(list_arr) = array.as_list_opt::<i32>() {
54            let has_garbage_values = if self.keep_original_array {
55                repdef.add_offsets(list_arr.offsets().clone(), array.nulls().cloned())
56            } else {
57                // there is no need to deep copy offsets, because offset buffers will be cast to a common type (i64).
58                repdef.add_offsets(list_arr.offsets().clone(), deep_copy_nulls(array.nulls()))
59            };
60            if has_garbage_values {
61                list_arr.filter_garbage_nulls().trimmed_values()
62            } else {
63                list_arr.trimmed_values()
64            }
65        } else if let Some(list_arr) = array.as_list_opt::<i64>() {
66            let has_garbage_values = if self.keep_original_array {
67                repdef.add_offsets(list_arr.offsets().clone(), array.nulls().cloned())
68            } else {
69                repdef.add_offsets(list_arr.offsets().clone(), deep_copy_nulls(array.nulls()))
70            };
71            if has_garbage_values {
72                list_arr.filter_garbage_nulls().trimmed_values()
73            } else {
74                list_arr.trimmed_values()
75            }
76        } else {
77            panic!("List encoder used for non-list data")
78        };
79        self.child
80            .maybe_encode(values, external_buffers, repdef, row_number, num_rows)
81    }
82
83    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
84        self.child.flush(external_buffers)
85    }
86
87    fn num_columns(&self) -> u32 {
88        self.child.num_columns()
89    }
90
91    fn finish(
92        &mut self,
93        external_buffers: &mut OutOfLineBuffers,
94    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
95        self.child.finish(external_buffers)
96    }
97}
98
99#[derive(Debug)]
100pub struct StructuralListScheduler {
101    child: Box<dyn StructuralFieldScheduler>,
102}
103
104impl StructuralListScheduler {
105    pub fn new(child: Box<dyn StructuralFieldScheduler>) -> Self {
106        Self { child }
107    }
108}
109
110impl StructuralFieldScheduler for StructuralListScheduler {
111    fn schedule_ranges<'a>(
112        &'a self,
113        ranges: &[Range<u64>],
114        filter: &FilterExpression,
115    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
116        let child = self.child.schedule_ranges(ranges, filter)?;
117
118        Ok(Box::new(StructuralListSchedulingJob::new(child)))
119    }
120
121    fn initialize<'a>(
122        &'a mut self,
123        filter: &'a FilterExpression,
124        context: &'a SchedulerContext,
125    ) -> BoxFuture<'a, Result<()>> {
126        self.child.initialize(filter, context)
127    }
128}
129
130/// Scheduling job for list data
131///
132/// Scheduling is handled by the primitive encoder and nothing special
133/// happens here.
134#[derive(Debug)]
135struct StructuralListSchedulingJob<'a> {
136    child: Box<dyn StructuralSchedulingJob + 'a>,
137}
138
139impl<'a> StructuralListSchedulingJob<'a> {
140    fn new(child: Box<dyn StructuralSchedulingJob + 'a>) -> Self {
141        Self { child }
142    }
143}
144
145impl StructuralSchedulingJob for StructuralListSchedulingJob<'_> {
146    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
147        self.child.schedule_next(context)
148    }
149}
150
151#[derive(Debug)]
152pub struct StructuralListDecoder {
153    child: Box<dyn StructuralFieldDecoder>,
154    data_type: DataType,
155}
156
157impl StructuralListDecoder {
158    pub fn new(child: Box<dyn StructuralFieldDecoder>, data_type: DataType) -> Self {
159        Self { child, data_type }
160    }
161}
162
163impl StructuralFieldDecoder for StructuralListDecoder {
164    fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> {
165        self.child.accept_page(child)
166    }
167
168    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
169        let child_task = self.child.drain(num_rows)?;
170        Ok(Box::new(StructuralListDecodeTask::new(
171            child_task,
172            self.data_type.clone(),
173        )))
174    }
175
176    fn data_type(&self) -> &DataType {
177        &self.data_type
178    }
179}
180
181#[derive(Debug)]
182struct StructuralListDecodeTask {
183    child_task: Box<dyn StructuralDecodeArrayTask>,
184    data_type: DataType,
185}
186
187impl StructuralListDecodeTask {
188    fn new(child_task: Box<dyn StructuralDecodeArrayTask>, data_type: DataType) -> Self {
189        Self {
190            child_task,
191            data_type,
192        }
193    }
194}
195
196impl StructuralDecodeArrayTask for StructuralListDecodeTask {
197    fn decode(self: Box<Self>) -> Result<DecodedArray> {
198        let DecodedArray {
199            array,
200            mut repdef,
201            data_size,
202        } = self.child_task.decode()?;
203        match &self.data_type {
204            DataType::List(child_field) => {
205                let (offsets, validity) = repdef.unravel_offsets::<i32>()?;
206                let array = if !child_field.is_nullable() && array.null_count() == array.len() {
207                    make_array(array.into_data().into_builder().nulls(None).build()?)
208                } else {
209                    array
210                };
211                let list_array = ListArray::try_new(child_field.clone(), offsets, array, validity)?;
212
213                Ok(DecodedArray {
214                    array: Arc::new(list_array),
215                    repdef,
216                    data_size,
217                })
218            }
219            DataType::LargeList(child_field) => {
220                let (offsets, validity) = repdef.unravel_offsets::<i64>()?;
221                let list_array =
222                    LargeListArray::try_new(child_field.clone(), offsets, array, validity)?;
223                Ok(DecodedArray {
224                    array: Arc::new(list_array),
225                    repdef,
226                    data_size,
227                })
228            }
229            _ => panic!("List decoder did not have a list field"),
230        }
231    }
232}
233
234#[cfg(test)]
235mod tests {
236
237    use std::{collections::HashMap, sync::Arc};
238
239    use crate::constants::{
240        STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK,
241    };
242    use arrow_array::{
243        Array, ArrayRef, BooleanArray, DictionaryArray, LargeStringArray, ListArray, StructArray,
244        UInt8Array, UInt64Array,
245        builder::{
246            Int32Builder, Int64Builder, LargeListBuilder, ListBuilder, StringBuilder, UInt32Builder,
247        },
248    };
249
250    use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
251    use arrow_schema::{DataType, Field, Fields};
252    use rstest::rstest;
253
254    use crate::{
255        testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data},
256        version::LanceFileVersion,
257    };
258
259    fn make_list_type(inner_type: DataType) -> DataType {
260        DataType::List(Arc::new(Field::new("item", inner_type, true)))
261    }
262
263    fn make_large_list_type(inner_type: DataType) -> DataType {
264        DataType::LargeList(Arc::new(Field::new("item", inner_type, true)))
265    }
266
267    #[rstest]
268    #[test_log::test(tokio::test)]
269    async fn test_list(
270        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
271        structural_encoding: &str,
272    ) {
273        let mut field_metadata = HashMap::new();
274        field_metadata.insert(
275            STRUCTURAL_ENCODING_META_KEY.to_string(),
276            structural_encoding.into(),
277        );
278        let field =
279            Field::new("", make_list_type(DataType::Int32), true).with_metadata(field_metadata);
280        check_basic_random(field).await;
281    }
282
283    #[rstest]
284    #[test_log::test(tokio::test)]
285    async fn test_deeply_nested_lists(
286        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
287        structural_encoding: &str,
288    ) {
289        let mut field_metadata = HashMap::new();
290        field_metadata.insert(
291            STRUCTURAL_ENCODING_META_KEY.to_string(),
292            structural_encoding.into(),
293        );
294        let field = Field::new("item", DataType::Int32, true).with_metadata(field_metadata);
295        for _ in 0..5 {
296            let field = Field::new("", make_list_type(field.data_type().clone()), true);
297            check_basic_random(field).await;
298        }
299    }
300
301    #[test_log::test(tokio::test)]
302    async fn test_large_list() {
303        let field = Field::new("", make_large_list_type(DataType::Int32), true);
304        check_basic_random(field).await;
305    }
306
307    #[test_log::test(tokio::test)]
308    async fn test_nested_strings() {
309        let field = Field::new("", make_list_type(DataType::Utf8), true);
310        check_basic_random(field).await;
311    }
312
313    #[test_log::test(tokio::test)]
314    async fn test_nested_list() {
315        let field = Field::new("", make_list_type(make_list_type(DataType::Int32)), true);
316        check_basic_random(field).await;
317    }
318
319    #[test_log::test(tokio::test)]
320    async fn test_list_struct_list() {
321        let struct_type = DataType::Struct(Fields::from(vec![Field::new(
322            "inner_str",
323            DataType::Utf8,
324            false,
325        )]));
326
327        let field = Field::new("", make_list_type(struct_type), true);
328        check_basic_random(field).await;
329    }
330
331    #[test_log::test(tokio::test)]
332    async fn test_list_struct_empty() {
333        let fields = Fields::from(vec![Field::new("inner", DataType::UInt64, true)]);
334        let items = UInt64Array::from(Vec::<u64>::new());
335        let structs = StructArray::new(fields, vec![Arc::new(items)], None);
336        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0; 2 * 1024 * 1024 + 1]));
337        let lists = ListArray::new(
338            Arc::new(Field::new("item", structs.data_type().clone(), true)),
339            offsets,
340            Arc::new(structs),
341            None,
342        );
343
344        check_round_trip_encoding_of_data(
345            vec![Arc::new(lists)],
346            &TestCases::default(),
347            HashMap::new(),
348        )
349        .await;
350    }
351
352    #[rstest]
353    #[test_log::test(tokio::test)]
354    async fn test_simple_list(
355        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
356        structural_encoding: &str,
357    ) {
358        let items_builder = Int32Builder::new();
359        let mut list_builder = ListBuilder::new(items_builder);
360        list_builder.append_value([Some(1), Some(2), Some(3)]);
361        list_builder.append_value([Some(4), Some(5)]);
362        list_builder.append_null();
363        list_builder.append_value([Some(6), Some(7), Some(8)]);
364        let list_array = list_builder.finish();
365
366        let mut field_metadata = HashMap::new();
367        field_metadata.insert(
368            STRUCTURAL_ENCODING_META_KEY.to_string(),
369            structural_encoding.into(),
370        );
371
372        let test_cases = TestCases::default()
373            .with_range(0..2)
374            .with_range(0..3)
375            .with_range(1..3)
376            .with_indices(vec![1, 3])
377            .with_indices(vec![2]);
378        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
379            .await;
380    }
381
382    #[rstest]
383    #[test_log::test(tokio::test)]
384    async fn test_simple_nested_list_ends_with_null(
385        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
386        structural_encoding: &str,
387    ) {
388        use arrow_array::Int32Array;
389
390        let values = Int32Array::from(vec![1, 2, 3, 4, 5]);
391        let inner_offsets = ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 4, 5, 5]);
392        let inner_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
393        let outer_offsets = ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 4, 5, 6, 6]);
394        let outer_validity = BooleanBuffer::from(vec![true, true, true, true, true, true, false]);
395
396        let inner_list = ListArray::new(
397            Arc::new(Field::new("item", DataType::Int32, true)),
398            OffsetBuffer::new(inner_offsets),
399            Arc::new(values),
400            Some(NullBuffer::new(inner_validity)),
401        );
402        let outer_list = ListArray::new(
403            Arc::new(Field::new(
404                "item",
405                DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
406                true,
407            )),
408            OffsetBuffer::new(outer_offsets),
409            Arc::new(inner_list),
410            Some(NullBuffer::new(outer_validity)),
411        );
412
413        let mut field_metadata = HashMap::new();
414        field_metadata.insert(
415            STRUCTURAL_ENCODING_META_KEY.to_string(),
416            structural_encoding.into(),
417        );
418
419        let test_cases = TestCases::default()
420            .with_range(0..2)
421            .with_range(0..3)
422            .with_range(5..7)
423            .with_indices(vec![1, 6])
424            .with_indices(vec![6])
425            .with_min_file_version(LanceFileVersion::V2_1);
426        check_round_trip_encoding_of_data(vec![Arc::new(outer_list)], &test_cases, field_metadata)
427            .await;
428    }
429
430    #[rstest]
431    #[test_log::test(tokio::test)]
432    async fn test_simple_string_list(
433        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
434        structural_encoding: &str,
435    ) {
436        let items_builder = StringBuilder::new();
437        let mut list_builder = ListBuilder::new(items_builder);
438        list_builder.append_value([Some("a"), Some("bc"), Some("def")]);
439        list_builder.append_value([Some("gh"), None]);
440        list_builder.append_null();
441        list_builder.append_value([Some("ijk"), Some("lmnop"), Some("qrs")]);
442        let list_array = list_builder.finish();
443
444        let mut field_metadata = HashMap::new();
445        field_metadata.insert(
446            STRUCTURAL_ENCODING_META_KEY.to_string(),
447            structural_encoding.into(),
448        );
449
450        let test_cases = TestCases::default()
451            .with_range(0..2)
452            .with_range(0..3)
453            .with_range(1..3)
454            .with_indices(vec![1, 3])
455            .with_indices(vec![2])
456            .with_min_file_version(LanceFileVersion::V2_1);
457        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
458            .await;
459    }
460
461    #[rstest]
462    #[test_log::test(tokio::test)]
463    async fn test_simple_string_list_no_null(
464        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
465        structural_encoding: &str,
466    ) {
467        let items_builder = StringBuilder::new();
468        let mut list_builder = ListBuilder::new(items_builder);
469        list_builder.append_value([Some("a"), Some("bc"), Some("def")]);
470        list_builder.append_value([Some("gh"), Some("zxy")]);
471        list_builder.append_value([Some("gh"), Some("z")]);
472        list_builder.append_value([Some("ijk"), Some("lmnop"), Some("qrs")]);
473        let list_array = list_builder.finish();
474
475        let mut field_metadata = HashMap::new();
476        field_metadata.insert(
477            STRUCTURAL_ENCODING_META_KEY.to_string(),
478            structural_encoding.into(),
479        );
480
481        let test_cases = TestCases::default()
482            .with_range(0..2)
483            .with_range(0..3)
484            .with_range(1..3)
485            .with_indices(vec![1, 3])
486            .with_indices(vec![2])
487            .with_min_file_version(LanceFileVersion::V2_1);
488        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
489            .await;
490    }
491
492    #[rstest]
493    #[test_log::test(tokio::test)]
494    async fn test_simple_sliced_list(
495        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
496        structural_encoding: &str,
497    ) {
498        let items_builder = Int32Builder::new();
499        let mut list_builder = ListBuilder::new(items_builder);
500        list_builder.append_value([Some(1), Some(2), Some(3)]);
501        list_builder.append_value([Some(4), Some(5)]);
502        list_builder.append_null();
503        list_builder.append_value([Some(6), Some(7), Some(8)]);
504        let list_array = list_builder.finish();
505
506        let list_array = list_array.slice(1, 2);
507
508        let mut field_metadata = HashMap::new();
509        field_metadata.insert(
510            STRUCTURAL_ENCODING_META_KEY.to_string(),
511            structural_encoding.into(),
512        );
513
514        let test_cases = TestCases::default()
515            .with_range(0..2)
516            .with_range(1..2)
517            .with_indices(vec![0])
518            .with_indices(vec![1])
519            .with_min_file_version(LanceFileVersion::V2_1);
520        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
521            .await;
522    }
523
524    #[test_log::test(tokio::test)]
525    async fn test_simple_list_dict() {
526        let values = LargeStringArray::from_iter_values(["a", "bb", "ccc"]);
527        let indices = UInt8Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
528        let dict_array = DictionaryArray::new(indices, Arc::new(values));
529        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 6, 9]));
530        let list_array = ListArray::new(
531            Arc::new(Field::new("item", dict_array.data_type().clone(), true)),
532            offsets,
533            Arc::new(dict_array),
534            None,
535        );
536
537        let test_cases = TestCases::default()
538            .with_range(0..2)
539            .with_range(1..3)
540            .with_range(2..4)
541            .with_indices(vec![1])
542            .with_indices(vec![2]);
543        check_round_trip_encoding_of_data(
544            vec![Arc::new(list_array)],
545            &test_cases,
546            HashMap::default(),
547        )
548        .await;
549    }
550
551    #[test_log::test(tokio::test)]
552    async fn test_simple_list_all_null() {
553        let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
554        let offsets = ScalarBuffer::<i32>::from(vec![0, 5, 8, 10]);
555        let offsets = OffsetBuffer::new(offsets);
556        let list_validity = NullBuffer::new(BooleanBuffer::from(vec![false, false, false]));
557
558        // The list array is nullable but the items are not.  Then, all lists are null.
559        let list_arr = ListArray::new(
560            Arc::new(Field::new("item", DataType::UInt64, false)),
561            offsets,
562            Arc::new(items),
563            Some(list_validity),
564        );
565
566        let test_cases = TestCases::default()
567            .with_range(0..3)
568            .with_range(1..2)
569            .with_indices(vec![1])
570            .with_indices(vec![2])
571            .with_min_file_version(LanceFileVersion::V2_1);
572        check_round_trip_encoding_of_data(
573            vec![Arc::new(list_arr)],
574            &test_cases,
575            HashMap::default(),
576        )
577        .await;
578    }
579
580    #[rstest]
581    #[test_log::test(tokio::test)]
582    async fn test_list_with_garbage_nulls(
583        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
584        structural_encoding: &str,
585    ) {
586        // In Arrow, list nulls are allowed to be non-empty, with masked garbage values
587        // Here we make a list with a null row in the middle with 3 garbage values
588        let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
589        let offsets = ScalarBuffer::<i32>::from(vec![0, 5, 8, 10]);
590        let offsets = OffsetBuffer::new(offsets);
591        let list_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
592        let list_arr = ListArray::new(
593            Arc::new(Field::new("item", DataType::UInt64, true)),
594            offsets,
595            Arc::new(items),
596            Some(list_validity),
597        );
598
599        let mut field_metadata = HashMap::new();
600        field_metadata.insert(
601            STRUCTURAL_ENCODING_META_KEY.to_string(),
602            structural_encoding.into(),
603        );
604
605        let test_cases = TestCases::default()
606            .with_range(0..3)
607            .with_range(1..2)
608            .with_indices(vec![1])
609            .with_indices(vec![2])
610            .with_min_file_version(LanceFileVersion::V2_1);
611        check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, field_metadata)
612            .await;
613    }
614
615    #[rstest]
616    #[test_log::test(tokio::test)]
617    async fn test_simple_two_page_list(
618        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
619        structural_encoding: &str,
620    ) {
621        // This is a simple pre-defined list that spans two pages.  This test is useful for
622        // debugging the repetition index
623
624        let items_builder = Int64Builder::new();
625        let mut list_builder = ListBuilder::new(items_builder);
626        for i in 0..512 {
627            list_builder.append_value([Some(i), Some(i * 2)]);
628        }
629        let list_array_1 = list_builder.finish();
630
631        let items_builder = Int64Builder::new();
632        let mut list_builder = ListBuilder::new(items_builder);
633        for i in 0..512 {
634            let i = i + 512;
635            list_builder.append_value([Some(i), Some(i * 2)]);
636        }
637        let list_array_2 = list_builder.finish();
638
639        let mut metadata = HashMap::new();
640        metadata.insert(
641            STRUCTURAL_ENCODING_META_KEY.to_string(),
642            structural_encoding.into(),
643        );
644
645        let test_cases = TestCases::default()
646            .with_min_file_version(LanceFileVersion::V2_1)
647            .with_page_sizes(vec![100])
648            .with_range(800..900);
649        check_round_trip_encoding_of_data(
650            vec![Arc::new(list_array_1), Arc::new(list_array_2)],
651            &test_cases,
652            metadata,
653        )
654        .await;
655    }
656
657    #[test_log::test(tokio::test)]
658    async fn test_simple_large_list() {
659        let items_builder = Int32Builder::new();
660        let mut list_builder = LargeListBuilder::new(items_builder);
661        list_builder.append_value([Some(1), Some(2), Some(3)]);
662        list_builder.append_value([Some(4), Some(5)]);
663        list_builder.append_null();
664        list_builder.append_value([Some(6), Some(7), Some(8)]);
665        let list_array = list_builder.finish();
666
667        let test_cases = TestCases::default()
668            .with_range(0..2)
669            .with_range(0..3)
670            .with_range(1..3)
671            .with_indices(vec![1, 3]);
672        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
673            .await;
674    }
675
676    #[rstest]
677    #[test_log::test(tokio::test)]
678    async fn test_empty_lists(
679        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
680        structural_encoding: &str,
681    ) {
682        let mut field_metadata = HashMap::new();
683        field_metadata.insert(
684            STRUCTURAL_ENCODING_META_KEY.to_string(),
685            structural_encoding.into(),
686        );
687
688        // Scenario 1: Some lists are empty
689
690        let values = [vec![Some(1), Some(2), Some(3)], vec![], vec![None]];
691        // Test empty list at beginning, middle, and end
692        for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] {
693            let items_builder = Int32Builder::new();
694            let mut list_builder = ListBuilder::new(items_builder);
695            for idx in order {
696                list_builder.append_value(values[idx].clone());
697            }
698            let list_array = Arc::new(list_builder.finish());
699            let test_cases = TestCases::default()
700                .with_indices(vec![1])
701                .with_indices(vec![0])
702                .with_indices(vec![2])
703                .with_indices(vec![0, 1]);
704            check_round_trip_encoding_of_data(
705                vec![list_array.clone()],
706                &test_cases,
707                field_metadata.clone(),
708            )
709            .await;
710            let test_cases = test_cases.with_batch_size(1);
711            check_round_trip_encoding_of_data(
712                vec![list_array],
713                &test_cases,
714                field_metadata.clone(),
715            )
716            .await;
717        }
718
719        // Scenario 2: All lists are empty
720
721        // When encoding a list of empty lists there are no items to encode
722        // which is strange and we want to ensure we handle it
723        let items_builder = Int32Builder::new();
724        let mut list_builder = ListBuilder::new(items_builder);
725        list_builder.append(true);
726        list_builder.append_null();
727        list_builder.append(true);
728        let list_array = Arc::new(list_builder.finish());
729
730        let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
731        check_round_trip_encoding_of_data(
732            vec![list_array.clone()],
733            &test_cases,
734            field_metadata.clone(),
735        )
736        .await;
737        let test_cases = test_cases.with_batch_size(1);
738        check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone())
739            .await;
740
741        // Scenario 2B: All lists are empty (but now with strings)
742
743        // When encoding a list of empty lists there are no items to encode
744        // which is strange and we want to ensure we handle it
745        let items_builder = StringBuilder::new();
746        let mut list_builder = ListBuilder::new(items_builder);
747        list_builder.append(true);
748        list_builder.append_null();
749        list_builder.append(true);
750        let list_array = Arc::new(list_builder.finish());
751
752        let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
753        check_round_trip_encoding_of_data(
754            vec![list_array.clone()],
755            &test_cases,
756            field_metadata.clone(),
757        )
758        .await;
759        let test_cases = test_cases.with_batch_size(1);
760        check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone())
761            .await;
762
763        // Scenario 3: All lists are null
764
765        let items_builder = Int32Builder::new();
766        let mut list_builder = ListBuilder::new(items_builder);
767        list_builder.append_null();
768        list_builder.append_null();
769        list_builder.append_null();
770        let list_array = Arc::new(list_builder.finish());
771
772        let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
773        check_round_trip_encoding_of_data(
774            vec![list_array.clone()],
775            &test_cases,
776            field_metadata.clone(),
777        )
778        .await;
779        let test_cases = test_cases.with_batch_size(1);
780        check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone())
781            .await;
782
783        // Scenario 4: All lists are null and inside a struct (only valid for 2.1 since 2.0 doesn't
784        // support null structs)
785        let items_builder = Int32Builder::new();
786        let mut list_builder = ListBuilder::new(items_builder);
787        list_builder.append_null();
788        list_builder.append_null();
789        list_builder.append_null();
790        let list_array = Arc::new(list_builder.finish());
791
792        let struct_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
793        let struct_array = Arc::new(StructArray::new(
794            Fields::from(vec![Field::new(
795                "lists",
796                list_array.data_type().clone(),
797                true,
798            )]),
799            vec![list_array],
800            Some(struct_validity),
801        ));
802
803        let test_cases = TestCases::default()
804            .with_range(0..2)
805            .with_indices(vec![1])
806            .with_min_file_version(LanceFileVersion::V2_1);
807        check_round_trip_encoding_of_data(
808            vec![struct_array.clone()],
809            &test_cases,
810            field_metadata.clone(),
811        )
812        .await;
813        let test_cases = test_cases.with_batch_size(1);
814        check_round_trip_encoding_of_data(vec![struct_array], &test_cases, field_metadata.clone())
815            .await;
816    }
817
818    #[test_log::test(tokio::test)]
819    async fn test_empty_list_list() {
820        let items_builder = Int32Builder::new();
821        let list_builder = ListBuilder::new(items_builder);
822        let mut outer_list_builder = ListBuilder::new(list_builder);
823        outer_list_builder.append_null();
824        outer_list_builder.append_null();
825        outer_list_builder.append_null();
826        let list_array = Arc::new(outer_list_builder.finish());
827
828        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
829        check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
830    }
831
832    #[test_log::test(tokio::test)]
833    #[ignore] // This test is quite slow in debug mode
834    async fn test_jumbo_list() {
835        // This is an overflow test.  We have a list of lists where each list
836        // has 1Mi items.  We encode 5000 of these lists and so we have over 4Gi in the
837        // offsets range
838        let items = BooleanArray::new_null(1024 * 1024);
839        let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 1024 * 1024]));
840        let list_arr = Arc::new(ListArray::new(
841            Arc::new(Field::new("item", DataType::Boolean, true)),
842            offsets,
843            Arc::new(items),
844            None,
845        )) as ArrayRef;
846        let arrs = vec![list_arr; 5000];
847
848        // We can't validate because our validation relies on concatenating all input arrays
849        let test_cases = TestCases::default().without_validation();
850        check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await;
851    }
852
853    // Regression test for issue with ListArray encoding when crossing 1024 value boundary
854    // This test reproduces the bug where rows_avail assertion fails in schedule_instructions
855    // when encoding a ListArray with specific size patterns that cross the 1024 value boundary
856    #[tokio::test]
857    async fn test_fuzz_issue_4466() {
858        // This specific pattern of list sizes triggers the bug when total values cross 1024
859        // 94 lists total 1009 values (passes), 95 lists total 1025 values (fails)
860        let list_sizes = vec![
861            13, 18, 12, 7, 14, 12, 6, 13, 18, 8, // 0-9: 119 values
862            6, 11, 17, 12, 8, 19, 5, 6, 10, 13, // 10-19: 107 values
863            8, 6, 10, 4, 8, 16, 14, 12, 18, 9, // 20-29: 105 values
864            17, 8, 14, 18, 15, 3, 2, 4, 5, 1, // 30-39: 82 values
865            3, 13, 1, 2, 10, 4, 10, 18, 7, 14, // 40-49: 75 values
866            18, 13, 9, 17, 3, 13, 10, 14, 8, 19, // 50-59: 125 values
867            17, 10, 5, 11, 6, 15, 10, 18, 18, 20, // 60-69: 130 values
868            16, 11, 12, 15, 7, 9, 3, 10, 20, 5, // 70-79: 102 values
869            2, 3, 17, 4, 8, 12, 15, 6, 3, 20, // 80-89: 90 values
870            15, 20, 1, 19, 16, // 90-94: 71 values
871        ];
872
873        // Build the ListArray
874        let mut list_builder = ListBuilder::new(Int32Builder::new());
875        let mut total_values = 0;
876
877        for size in &list_sizes {
878            for i in 0..*size {
879                list_builder.values().append_value(i);
880            }
881            list_builder.append(true);
882            total_values += size;
883        }
884
885        let list_array = Arc::new(list_builder.finish());
886
887        // Verify we have the expected number of values
888        assert_eq!(list_array.len(), 95);
889        assert_eq!(total_values, 1025);
890
891        // This should trigger the assertion failure at primitive.rs:1362
892        // debug_assert!(rows_avail > 0)
893        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
894
895        // The bug manifests when encoding this specific pattern
896        // Expected: successful round-trip encoding
897        // Actual: panic at primitive.rs:1362 - assertion failed: rows_avail > 0
898        check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
899    }
900
901    #[rstest]
902    #[test_log::test(tokio::test)]
903    async fn test_sparse_large_string_list(
904        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
905        structural_encoding: &str,
906    ) {
907        // 2.5 million rows, mostly empty lists. ~100 lists have 10 short strings each.
908        let num_rows = 2_500_000u32;
909        let num_non_empty = 100u32;
910        let strings_per_list = 10;
911
912        let items_builder = StringBuilder::new();
913        let mut list_builder = ListBuilder::new(items_builder);
914
915        // Spread non-empty lists evenly across the range
916        let step = num_rows / num_non_empty;
917        let mut next_non_empty = step / 2;
918
919        for i in 0..num_rows {
920            if i == next_non_empty {
921                let vals: Vec<Option<&str>> = (0..strings_per_list)
922                    .map(|j| match j % 4 {
923                        0 => Some("a"),
924                        1 => Some("bb"),
925                        2 => Some("ccc"),
926                        _ => Some("d"),
927                    })
928                    .collect();
929                list_builder.append_value(vals);
930                next_non_empty = next_non_empty.saturating_add(step);
931            } else {
932                list_builder.append_value([] as [Option<&str>; 0]);
933            }
934        }
935        let list_array = list_builder.finish();
936
937        let mut field_metadata = HashMap::new();
938        field_metadata.insert(
939            STRUCTURAL_ENCODING_META_KEY.to_string(),
940            structural_encoding.into(),
941        );
942
943        let test_cases = TestCases::default()
944            .with_range(0..1000)
945            .with_range(0..num_rows as u64)
946            .with_indices(vec![0, (step / 2) as u64, num_rows as u64 - 1])
947            .with_max_file_version(LanceFileVersion::V2_2);
948        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
949            .await;
950    }
951
952    /// Builds the HNSW-flush repro shape: a dense prefix where every row has
953    /// `NEIGHBORS_PER_ROW` distinct values, followed by a long tail of empty
954    /// lists. Mirrors `HNSW::schema()` `__neighbors` / `__dists` columns:
955    /// dense level-0 lists, then ~6× as many mostly-empty higher-level rows.
956    fn make_hnsw_shaped_list_u32() -> ListArray {
957        const DENSE_ROWS: u32 = 40_000;
958        const NEIGHBORS_PER_ROW: u32 = 32;
959        const EMPTY_TAIL_ROWS: u32 = 240_000;
960
961        let mut list_builder = ListBuilder::new(UInt32Builder::new());
962        let mut next_val: u32 = 0;
963        for _ in 0..DENSE_ROWS {
964            for _ in 0..NEIGHBORS_PER_ROW {
965                list_builder.values().append_value(next_val);
966                next_val = next_val.wrapping_add(1);
967            }
968            list_builder.append(true);
969        }
970        for _ in 0..EMPTY_TAIL_ROWS {
971            list_builder.append(true);
972        }
973        list_builder.finish()
974    }
975
976    /// Reproduces the HNSW-shaped variable-length `List` miniblock bug at v2.2
977    /// **on the auto-routing path** (no `STRUCTURAL_ENCODING` metadata): a
978    /// dense prefix followed by a long tail of empty lists. Globally the data
979    /// looks dense, so the unfixed `repdef_too_sparse_for_miniblock` heuristic
980    /// picks miniblock; the final chunk's level count then overflows the u16
981    /// stored in the chunk header and the read drops rows. After the heuristic
982    /// fix, this shape correctly routes to fullzip and the round-trip is
983    /// lossless.
984    #[test_log::test(tokio::test)]
985    async fn test_list_hnsw_shape_auto_routes_around_miniblock_overflow_v2_2() {
986        let list_array = make_hnsw_shaped_list_u32();
987        let dense_rows: u64 = 40_000;
988        let total_rows = list_array.len() as u64;
989
990        let field_metadata = HashMap::new();
991
992        let test_cases = TestCases::default()
993            .with_range(0..1000)
994            .with_range(dense_rows.saturating_sub(8)..(dense_rows + 8))
995            .with_range(0..total_rows)
996            .with_indices(vec![0, dense_rows - 1, dense_rows, total_rows - 1])
997            .with_min_file_version(LanceFileVersion::V2_2)
998            .with_max_file_version(LanceFileVersion::V2_2);
999        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
1000            .await;
1001    }
1002
1003    /// Companion to the auto-routing test: even when the user *explicitly*
1004    /// sets `STRUCTURAL_ENCODING_MINIBLOCK` on the HNSW shape, the heuristic
1005    /// must still detect the per-chunk `num_levels: u16` overflow and override
1006    /// the request with fullzip — silently overriding a corrupt-encoding
1007    /// preference is the safe behaviour given v2.2 is a stable on-disk format
1008    /// whose chunk header width cannot widen.
1009    ///
1010    /// (The codec-side `u16::try_from` safety net at the chunk-build site is
1011    /// dormant on this shape precisely because the heuristic intercepts it;
1012    /// the safety net is there in case a future shape sneaks past the
1013    /// heuristic.)
1014    #[test_log::test(tokio::test)]
1015    async fn test_forced_miniblock_hnsw_shape_routed_to_fullzip_v2_2() {
1016        let list_array = make_hnsw_shaped_list_u32();
1017        let total_rows = list_array.len() as u64;
1018
1019        let mut field_metadata = HashMap::new();
1020        field_metadata.insert(
1021            STRUCTURAL_ENCODING_META_KEY.to_string(),
1022            STRUCTURAL_ENCODING_MINIBLOCK.into(),
1023        );
1024
1025        let test_cases = TestCases::default()
1026            .with_range(0..total_rows)
1027            .with_min_file_version(LanceFileVersion::V2_2)
1028            .with_max_file_version(LanceFileVersion::V2_2);
1029        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
1030            .await;
1031    }
1032}
lance_encoding/encodings/logical/list.rs

lance_encoding/encodings/logical/
list.rs
