1use 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
23pub 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 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 requested_ranges: Option<&'a [Range<u64>]>,
124 filter: &'a FilterExpression,
125 context: &'a SchedulerContext,
126 ) -> BoxFuture<'a, Result<()>> {
127 self.child.initialize(requested_ranges, filter, context)
131 }
132}
133
134#[derive(Debug)]
139struct StructuralListSchedulingJob<'a> {
140 child: Box<dyn StructuralSchedulingJob + 'a>,
141}
142
143impl<'a> StructuralListSchedulingJob<'a> {
144 fn new(child: Box<dyn StructuralSchedulingJob + 'a>) -> Self {
145 Self { child }
146 }
147}
148
149impl StructuralSchedulingJob for StructuralListSchedulingJob<'_> {
150 fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
151 self.child.schedule_next(context)
152 }
153}
154
155#[derive(Debug)]
156pub struct StructuralListDecoder {
157 child: Box<dyn StructuralFieldDecoder>,
158 data_type: DataType,
159}
160
161impl StructuralListDecoder {
162 pub fn new(child: Box<dyn StructuralFieldDecoder>, data_type: DataType) -> Self {
163 Self { child, data_type }
164 }
165}
166
167impl StructuralFieldDecoder for StructuralListDecoder {
168 fn accept_page(&mut self, child: crate::decoder::LoadedPageShard) -> Result<()> {
169 self.child.accept_page(child)
170 }
171
172 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
173 let child_task = self.child.drain(num_rows)?;
174 Ok(Box::new(StructuralListDecodeTask::new(
175 child_task,
176 self.data_type.clone(),
177 )))
178 }
179
180 fn data_type(&self) -> &DataType {
181 &self.data_type
182 }
183
184 fn plan_decoded_bytes(&self, _rows_remaining: u64) -> lance_core::Result<[u64; 8]> {
185 Err(lance_core::Error::not_supported(
186 "plan_decoded_bytes is not yet supported for fields with repetition (list/repeated types)"
187 .to_string(),
188 ))
189 }
190}
191
192#[derive(Debug)]
193struct StructuralListDecodeTask {
194 child_task: Box<dyn StructuralDecodeArrayTask>,
195 data_type: DataType,
196}
197
198impl StructuralListDecodeTask {
199 fn new(child_task: Box<dyn StructuralDecodeArrayTask>, data_type: DataType) -> Self {
200 Self {
201 child_task,
202 data_type,
203 }
204 }
205}
206
207impl StructuralDecodeArrayTask for StructuralListDecodeTask {
208 fn decode(self: Box<Self>) -> Result<DecodedArray> {
209 let DecodedArray {
210 array,
211 mut repdef,
212 data_size,
213 } = self.child_task.decode()?;
214 match &self.data_type {
215 DataType::List(child_field) => {
216 let (offsets, validity) = repdef.unravel_offsets::<i32>()?;
217 let array = if !child_field.is_nullable() && array.null_count() == array.len() {
218 make_array(array.into_data().into_builder().nulls(None).build()?)
219 } else {
220 array
221 };
222 let list_array = ListArray::try_new(child_field.clone(), offsets, array, validity)?;
223
224 Ok(DecodedArray {
225 array: Arc::new(list_array),
226 repdef,
227 data_size,
228 })
229 }
230 DataType::LargeList(child_field) => {
231 let (offsets, validity) = repdef.unravel_offsets::<i64>()?;
232 let list_array =
233 LargeListArray::try_new(child_field.clone(), offsets, array, validity)?;
234 Ok(DecodedArray {
235 array: Arc::new(list_array),
236 repdef,
237 data_size,
238 })
239 }
240 _ => panic!("List decoder did not have a list field"),
241 }
242 }
243}
244
245#[cfg(test)]
246mod tests {
247
248 use std::{collections::HashMap, sync::Arc};
249
250 use crate::constants::{
251 STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK,
252 };
253 use arrow_array::{
254 Array, ArrayRef, BooleanArray, DictionaryArray, LargeStringArray, ListArray, StringArray,
255 StructArray, UInt8Array, UInt64Array,
256 builder::{
257 Int32Builder, Int64Builder, LargeListBuilder, ListBuilder, StringBuilder, UInt32Builder,
258 },
259 };
260
261 use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
262 use arrow_schema::{DataType, Field, Fields};
263 use lance_datagen::{RowCount, Seed, array, gen_batch};
264 use rstest::rstest;
265
266 use crate::testing::{
267 TestCases, TestEncoding, check_round_trip_encoding_of_data, create_test_field_encoder,
268 test_encoding_strategy,
269 };
270
271 fn make_list_type(inner_type: DataType) -> DataType {
272 DataType::List(Arc::new(Field::new("item", inner_type, true)))
273 }
274
275 fn make_large_list_type(inner_type: DataType) -> DataType {
276 DataType::LargeList(Arc::new(Field::new("item", inner_type, true)))
277 }
278
279 #[derive(Clone, Copy)]
280 enum NullPattern {
281 None,
282 Mixed,
283 All,
284 }
285
286 async fn check_nested_type(
287 data_type: DataType,
288 null_pattern: NullPattern,
289 encoding: TestEncoding,
290 ) {
291 check_nested_type_with_metadata(data_type, null_pattern, encoding, HashMap::new()).await;
292 }
293
294 async fn check_nested_type_with_metadata(
295 data_type: DataType,
296 null_pattern: NullPattern,
297 encoding: TestEncoding,
298 field_metadata: HashMap<String, String>,
299 ) {
300 let null_rate = match null_pattern {
301 NullPattern::None => None,
302 NullPattern::Mixed => Some(0.5),
303 NullPattern::All => Some(1.0),
304 };
305 let make_batch = |seed, rows| {
306 let mut generator = gen_batch()
307 .with_seed(Seed::from(seed))
308 .anon_col(array::rand_type(&data_type));
309 if let Some(null_rate) = null_rate {
310 generator.with_random_nulls(null_rate);
311 }
312 generator
313 .into_batch_rows(RowCount::from(rows))
314 .unwrap()
315 .column(0)
316 .clone()
317 };
318
319 let first = make_batch(0, 513).slice(1, 512);
323 let second = make_batch(1, 513);
324 let test_cases = TestCases::default()
325 .with_page_sizes(vec![4096])
326 .with_encoding(encoding)
327 .with_batch_size(257)
328 .with_range(510..515)
329 .with_indices(vec![0, 511, 512, 1024]);
330
331 check_round_trip_encoding_of_data(vec![first, second], &test_cases, field_metadata).await;
332 }
333
334 async fn try_encode_v22_pages(
335 array: ArrayRef,
336 ) -> lance_core::Result<Vec<crate::encoder::EncodedPage>> {
337 try_encode_v22_pages_with_metadata(array, HashMap::new()).await
338 }
339
340 async fn try_encode_v22_pages_with_metadata(
341 array: ArrayRef,
342 field_metadata: HashMap<String, String>,
343 ) -> lance_core::Result<Vec<crate::encoder::EncodedPage>> {
344 let arrow_field =
345 Field::new("", array.data_type().clone(), true).with_metadata(field_metadata);
346 let lance_field = lance_core::datatypes::Field::try_from(&arrow_field).unwrap();
347 let encoding_strategy = test_encoding_strategy(TestEncoding::StructuralU32);
348 let mut column_index_seq = crate::encoder::ColumnIndexSequence::default();
349 let encoding_options = crate::encoder::EncodingOptions::default();
350 let mut encoder = create_test_field_encoder(
351 encoding_strategy.as_ref(),
352 &lance_field,
353 &mut column_index_seq,
354 &encoding_options,
355 )
356 .unwrap();
357 let mut external_buffers =
358 crate::encoder::OutOfLineBuffers::new(0, crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT);
359 let num_rows = array.len() as u64;
360 let mut pages = Vec::new();
361 for task in encoder
362 .maybe_encode(
363 array,
364 &mut external_buffers,
365 crate::repdef::RepDefBuilder::default(),
366 0,
367 num_rows,
368 )
369 .unwrap()
370 {
371 pages.push(task.await?);
372 }
373 for task in encoder.flush(&mut external_buffers).unwrap() {
374 pages.push(task.await?);
375 }
376 Ok(pages)
377 }
378
379 async fn encode_v22_pages(array: ArrayRef) -> Vec<crate::encoder::EncodedPage> {
380 try_encode_v22_pages(array).await.unwrap()
381 }
382
383 fn assert_split_miniblock_layout(
384 pages: &[crate::encoder::EncodedPage],
385 min_miniblock_pages: usize,
386 expect_structural_only_page: bool,
387 ) {
388 let mut miniblock_pages = 0;
389 let mut fullzip_pages = 0;
390 let mut structural_only_pages = 0;
391
392 for page in pages {
393 let crate::decoder::PageEncoding::Structural(layout) = &page.description else {
394 continue;
395 };
396 match layout.layout.as_ref().unwrap() {
397 crate::format::pb21::page_layout::Layout::MiniBlockLayout(_) => {
398 miniblock_pages += 1;
399 }
400 crate::format::pb21::page_layout::Layout::FullZipLayout(_) => {
401 fullzip_pages += 1;
402 }
403 crate::format::pb21::page_layout::Layout::ConstantLayout(layout) => {
404 if layout.inline_value.is_none()
405 && (layout.num_rep_values > 0 || layout.num_def_values > 0)
406 {
407 structural_only_pages += 1;
408 }
409 }
410 crate::format::pb21::page_layout::Layout::BlobLayout(_) => {}
411 crate::format::pb21::page_layout::Layout::SparseLayout(_) => {}
412 }
413 }
414
415 assert!(
416 miniblock_pages >= min_miniblock_pages,
417 "expected at least {min_miniblock_pages} mini-block pages, got {miniblock_pages}"
418 );
419 assert_eq!(
420 fullzip_pages, 0,
421 "split list pages should not fall back to full-zip"
422 );
423 assert_eq!(
424 structural_only_pages > 0,
425 expect_structural_only_page,
426 "structural-only page presence did not match expectation; got {structural_only_pages}"
427 );
428 }
429
430 fn assert_has_fullzip_layout(pages: &[crate::encoder::EncodedPage]) {
431 let has_fullzip = pages.iter().any(|page| {
432 let crate::decoder::PageEncoding::Structural(layout) = &page.description else {
433 return false;
434 };
435 matches!(
436 layout.layout.as_ref().unwrap(),
437 crate::format::pb21::page_layout::Layout::FullZipLayout(_)
438 )
439 });
440 assert!(has_fullzip, "expected at least one full-zip page");
441 }
442
443 #[rstest]
444 #[test_log::test(tokio::test)]
445 async fn test_list(
446 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
447 structural_encoding: &str,
448 #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)]
449 null_pattern: NullPattern,
450 #[values(
451 TestEncoding::Array,
452 TestEncoding::StructuralU16,
453 TestEncoding::StructuralU32,
454 TestEncoding::StructuralSparse
455 )]
456 encoding: TestEncoding,
457 ) {
458 let mut field_metadata = HashMap::new();
459 field_metadata.insert(
460 STRUCTURAL_ENCODING_META_KEY.to_string(),
461 structural_encoding.into(),
462 );
463 check_nested_type_with_metadata(
464 make_list_type(DataType::Int32),
465 null_pattern,
466 encoding,
467 field_metadata,
468 )
469 .await;
470 }
471
472 #[rstest]
473 #[test_log::test(tokio::test)]
474 async fn test_deeply_nested_lists(
475 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
476 structural_encoding: &str,
477 #[values(1, 2, 3, 4, 5)] depth: usize,
478 #[values(
479 TestEncoding::Array,
480 TestEncoding::StructuralU16,
481 TestEncoding::StructuralU32,
482 TestEncoding::StructuralSparse
483 )]
484 test_encoding: TestEncoding,
485 ) {
486 let mut data_type = DataType::Int32;
487 for _ in 0..depth {
488 data_type = make_list_type(data_type);
489 }
490
491 let mut generator = gen_batch()
492 .with_seed(Seed::from(depth as u64))
493 .anon_col(array::rand_type(&data_type));
494 generator.with_random_nulls(0.2);
495 let source = generator
496 .into_batch_rows(RowCount::from(1026))
497 .unwrap()
498 .column(0)
499 .clone();
500
501 let data = vec![source.slice(1, 512), source.slice(513, 513)];
504 let test_cases = TestCases::default()
505 .with_page_sizes(vec![4096])
506 .with_encoding(test_encoding)
507 .with_batch_size(257)
508 .with_range(510..515)
509 .with_indices(vec![0, 511, 512, 1024]);
510 let field_metadata = HashMap::from([(
511 STRUCTURAL_ENCODING_META_KEY.to_string(),
512 structural_encoding.into(),
513 )]);
514
515 check_round_trip_encoding_of_data(data, &test_cases, field_metadata).await;
516 }
517
518 #[rstest]
519 #[test_log::test(tokio::test)]
520 async fn test_large_list(
521 #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)]
522 null_pattern: NullPattern,
523 #[values(
524 TestEncoding::Array,
525 TestEncoding::StructuralU16,
526 TestEncoding::StructuralU32,
527 TestEncoding::StructuralSparse
528 )]
529 encoding: TestEncoding,
530 ) {
531 check_nested_type(
532 make_large_list_type(DataType::Int32),
533 null_pattern,
534 encoding,
535 )
536 .await;
537 }
538
539 #[rstest]
540 #[test_log::test(tokio::test)]
541 async fn test_nested_strings(
542 #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)]
543 null_pattern: NullPattern,
544 #[values(
545 TestEncoding::Array,
546 TestEncoding::StructuralU16,
547 TestEncoding::StructuralU32,
548 TestEncoding::StructuralSparse
549 )]
550 encoding: TestEncoding,
551 ) {
552 check_nested_type(make_list_type(DataType::Utf8), null_pattern, encoding).await;
553 }
554
555 #[rstest]
556 #[test_log::test(tokio::test)]
557 async fn test_nested_list(
558 #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)]
559 null_pattern: NullPattern,
560 #[values(
561 TestEncoding::Array,
562 TestEncoding::StructuralU16,
563 TestEncoding::StructuralU32,
564 TestEncoding::StructuralSparse
565 )]
566 encoding: TestEncoding,
567 ) {
568 check_nested_type(
569 make_list_type(make_list_type(DataType::Int32)),
570 null_pattern,
571 encoding,
572 )
573 .await;
574 }
575
576 #[rstest]
595 #[test_log::test(tokio::test)]
596 async fn test_multipage_nested_float_list(
597 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
598 structural_encoding: &str,
599 #[values(
600 TestEncoding::StructuralU16,
601 TestEncoding::StructuralU32,
602 TestEncoding::StructuralSparse
603 )]
604 test_encoding: TestEncoding,
605 ) {
606 use arrow_array::Float32Array;
607
608 let inner_per_row: usize = 3;
610 let inner_len: usize = 768;
611 let chunk_rows: &[usize] = &[150, 149];
615
616 let make_chunk = |start_row: usize, num_rows: usize| -> Arc<dyn Array> {
617 let total_inner = num_rows * inner_per_row;
618 let total_values = total_inner * inner_len;
619 let values = Float32Array::from(
620 (0..total_values)
621 .map(|i| (start_row + i) as f32)
622 .collect::<Vec<_>>(),
623 );
624 let inner_offsets = ScalarBuffer::<i32>::from(
625 (0..=total_inner)
626 .map(|i| (i * inner_len) as i32)
627 .collect::<Vec<_>>(),
628 );
629 let inner_list = ListArray::new(
630 Arc::new(Field::new("item", DataType::Float32, true)),
631 OffsetBuffer::new(inner_offsets),
632 Arc::new(values),
633 None,
634 );
635 let outer_offsets = ScalarBuffer::<i32>::from(
636 (0..=num_rows)
637 .map(|i| (i * inner_per_row) as i32)
638 .collect::<Vec<_>>(),
639 );
640 Arc::new(ListArray::new(
641 Arc::new(Field::new(
642 "item",
643 DataType::List(Arc::new(Field::new("item", DataType::Float32, true))),
644 true,
645 )),
646 OffsetBuffer::new(outer_offsets),
647 Arc::new(inner_list),
648 None,
649 ))
650 };
651
652 let mut start = 0;
653 let chunks: Vec<Arc<dyn Array>> = chunk_rows
654 .iter()
655 .map(|&n| {
656 let c = make_chunk(start, n);
657 start += n;
658 c
659 })
660 .collect();
661
662 let mut field_metadata = HashMap::new();
663 field_metadata.insert(
664 STRUCTURAL_ENCODING_META_KEY.to_string(),
665 structural_encoding.into(),
666 );
667
668 let test_cases = TestCases::default()
669 .with_page_sizes(vec![1024 * 1024])
670 .with_encoding(test_encoding)
671 .with_batch_size(151)
672 .with_range(148..152)
673 .with_indices(vec![0, 149, 150, 298]);
674 check_round_trip_encoding_of_data(chunks, &test_cases, field_metadata).await;
675 }
676
677 #[rstest]
678 #[test_log::test(tokio::test)]
679 async fn test_list_struct_list(
680 #[values(NullPattern::None, NullPattern::Mixed, NullPattern::All)]
681 null_pattern: NullPattern,
682 #[values(
683 TestEncoding::Array,
684 TestEncoding::StructuralU16,
685 TestEncoding::StructuralU32,
686 TestEncoding::StructuralSparse
687 )]
688 encoding: TestEncoding,
689 ) {
690 let struct_type = DataType::Struct(Fields::from(vec![Field::new(
691 "inner_str",
692 DataType::Utf8,
693 false,
694 )]));
695
696 check_nested_type(make_list_type(struct_type), null_pattern, encoding).await;
697 }
698
699 #[rstest]
700 #[test_log::test(tokio::test)]
701 async fn test_list_struct_empty(
702 #[values(
703 TestEncoding::Array,
704 TestEncoding::StructuralU16,
705 TestEncoding::StructuralU32,
706 TestEncoding::StructuralSparse
707 )]
708 encoding: TestEncoding,
709 ) {
710 let fields = Fields::from(vec![Field::new("inner", DataType::UInt64, true)]);
711 let items = UInt64Array::from(Vec::<u64>::new());
712 let structs = StructArray::new(fields, vec![Arc::new(items)], None);
713 let num_rows = 2 * 1024 * 1024 / size_of::<i32>() + 1;
716 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0; num_rows + 1]));
717 let lists = ListArray::new(
718 Arc::new(Field::new("item", structs.data_type().clone(), true)),
719 offsets,
720 Arc::new(structs),
721 None,
722 );
723
724 check_round_trip_encoding_of_data(
725 vec![Arc::new(lists)],
726 &TestCases::default()
727 .with_page_sizes(vec![1024 * 1024])
728 .with_encoding(encoding),
729 HashMap::new(),
730 )
731 .await;
732 }
733
734 #[rstest]
735 #[test_log::test(tokio::test)]
736 async fn test_simple_list(
737 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
738 structural_encoding: &str,
739 ) {
740 let items_builder = Int32Builder::new();
741 let mut list_builder = ListBuilder::new(items_builder);
742 list_builder.append_value([Some(1), Some(2), Some(3)]);
743 list_builder.append_value([Some(4), Some(5)]);
744 list_builder.append_null();
745 list_builder.append_value([Some(6), Some(7), Some(8)]);
746 let list_array = list_builder.finish();
747
748 let mut field_metadata = HashMap::new();
749 field_metadata.insert(
750 STRUCTURAL_ENCODING_META_KEY.to_string(),
751 structural_encoding.into(),
752 );
753
754 let test_cases = TestCases::default()
755 .with_range(0..2)
756 .with_range(0..3)
757 .with_range(1..3)
758 .with_indices(vec![1, 3])
759 .with_indices(vec![2]);
760 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
761 .await;
762 }
763
764 #[rstest]
765 #[test_log::test(tokio::test)]
766 async fn test_simple_nested_list_ends_with_null(
767 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
768 structural_encoding: &str,
769 ) {
770 use arrow_array::Int32Array;
771
772 let values = Int32Array::from(vec![1, 2, 3, 4, 5]);
773 let inner_offsets = ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 4, 5, 5]);
774 let inner_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
775 let outer_offsets = ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 4, 5, 6, 6]);
776 let outer_validity = BooleanBuffer::from(vec![true, true, true, true, true, true, false]);
777
778 let inner_list = ListArray::new(
779 Arc::new(Field::new("item", DataType::Int32, true)),
780 OffsetBuffer::new(inner_offsets),
781 Arc::new(values),
782 Some(NullBuffer::new(inner_validity)),
783 );
784 let outer_list = ListArray::new(
785 Arc::new(Field::new(
786 "item",
787 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
788 true,
789 )),
790 OffsetBuffer::new(outer_offsets),
791 Arc::new(inner_list),
792 Some(NullBuffer::new(outer_validity)),
793 );
794
795 let mut field_metadata = HashMap::new();
796 field_metadata.insert(
797 STRUCTURAL_ENCODING_META_KEY.to_string(),
798 structural_encoding.into(),
799 );
800
801 let test_cases = TestCases::default()
802 .with_range(0..2)
803 .with_range(0..3)
804 .with_range(5..7)
805 .with_indices(vec![1, 6])
806 .with_indices(vec![6])
807 .with_structural_encodings();
808 check_round_trip_encoding_of_data(vec![Arc::new(outer_list)], &test_cases, field_metadata)
809 .await;
810 }
811
812 #[rstest]
813 #[test_log::test(tokio::test)]
814 async fn test_simple_string_list(
815 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
816 structural_encoding: &str,
817 ) {
818 let items_builder = StringBuilder::new();
819 let mut list_builder = ListBuilder::new(items_builder);
820 list_builder.append_value([Some("a"), Some("bc"), Some("def")]);
821 list_builder.append_value([Some("gh"), None]);
822 list_builder.append_null();
823 list_builder.append_value([Some("ijk"), Some("lmnop"), Some("qrs")]);
824 let list_array = list_builder.finish();
825
826 let mut field_metadata = HashMap::new();
827 field_metadata.insert(
828 STRUCTURAL_ENCODING_META_KEY.to_string(),
829 structural_encoding.into(),
830 );
831
832 let test_cases = TestCases::default()
833 .with_range(0..2)
834 .with_range(0..3)
835 .with_range(1..3)
836 .with_indices(vec![1, 3])
837 .with_indices(vec![2])
838 .with_structural_encodings();
839 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
840 .await;
841 }
842
843 #[rstest]
844 #[test_log::test(tokio::test)]
845 async fn test_simple_string_list_no_null(
846 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
847 structural_encoding: &str,
848 ) {
849 let items_builder = StringBuilder::new();
850 let mut list_builder = ListBuilder::new(items_builder);
851 list_builder.append_value([Some("a"), Some("bc"), Some("def")]);
852 list_builder.append_value([Some("gh"), Some("zxy")]);
853 list_builder.append_value([Some("gh"), Some("z")]);
854 list_builder.append_value([Some("ijk"), Some("lmnop"), Some("qrs")]);
855 let list_array = list_builder.finish();
856
857 let mut field_metadata = HashMap::new();
858 field_metadata.insert(
859 STRUCTURAL_ENCODING_META_KEY.to_string(),
860 structural_encoding.into(),
861 );
862
863 let test_cases = TestCases::default()
864 .with_range(0..2)
865 .with_range(0..3)
866 .with_range(1..3)
867 .with_indices(vec![1, 3])
868 .with_indices(vec![2])
869 .with_structural_encodings();
870 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
871 .await;
872 }
873
874 #[rstest]
875 #[test_log::test(tokio::test)]
876 async fn test_simple_sliced_list(
877 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
878 structural_encoding: &str,
879 ) {
880 let items_builder = Int32Builder::new();
881 let mut list_builder = ListBuilder::new(items_builder);
882 list_builder.append_value([Some(1), Some(2), Some(3)]);
883 list_builder.append_value([Some(4), Some(5)]);
884 list_builder.append_null();
885 list_builder.append_value([Some(6), Some(7), Some(8)]);
886 let list_array = list_builder.finish();
887
888 let list_array = list_array.slice(1, 2);
889
890 let mut field_metadata = HashMap::new();
891 field_metadata.insert(
892 STRUCTURAL_ENCODING_META_KEY.to_string(),
893 structural_encoding.into(),
894 );
895
896 let test_cases = TestCases::default()
897 .with_range(0..2)
898 .with_range(1..2)
899 .with_indices(vec![0])
900 .with_indices(vec![1])
901 .with_structural_encodings();
902 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
903 .await;
904 }
905
906 #[test_log::test(tokio::test)]
907 async fn test_simple_list_dict() {
908 let values = LargeStringArray::from_iter_values(["a", "bb", "ccc"]);
909 let indices = UInt8Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
910 let dict_array = DictionaryArray::new(indices, Arc::new(values));
911 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 6, 9]));
912 let list_array = ListArray::new(
913 Arc::new(Field::new("item", dict_array.data_type().clone(), true)),
914 offsets,
915 Arc::new(dict_array),
916 None,
917 );
918
919 let test_cases = TestCases::default()
920 .with_range(0..2)
921 .with_range(1..3)
922 .with_range(2..4)
923 .with_indices(vec![1])
924 .with_indices(vec![2]);
925 check_round_trip_encoding_of_data(
926 vec![Arc::new(list_array)],
927 &test_cases,
928 HashMap::default(),
929 )
930 .await;
931 }
932
933 #[test_log::test(tokio::test)]
934 async fn test_simple_list_all_null() {
935 let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
936 let offsets = ScalarBuffer::<i32>::from(vec![0, 5, 8, 10]);
937 let offsets = OffsetBuffer::new(offsets);
938 let list_validity = NullBuffer::new(BooleanBuffer::from(vec![false, false, false]));
939
940 let list_arr = ListArray::new(
942 Arc::new(Field::new("item", DataType::UInt64, false)),
943 offsets,
944 Arc::new(items),
945 Some(list_validity),
946 );
947
948 let test_cases = TestCases::default()
949 .with_range(0..3)
950 .with_range(1..2)
951 .with_indices(vec![1])
952 .with_indices(vec![2])
953 .with_structural_encodings();
954 check_round_trip_encoding_of_data(
955 vec![Arc::new(list_arr)],
956 &test_cases,
957 HashMap::default(),
958 )
959 .await;
960 }
961
962 #[rstest]
963 #[test_log::test(tokio::test)]
964 async fn test_list_with_garbage_nulls(
965 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
966 structural_encoding: &str,
967 ) {
968 let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
971 let offsets = ScalarBuffer::<i32>::from(vec![0, 5, 8, 10]);
972 let offsets = OffsetBuffer::new(offsets);
973 let list_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
974 let list_arr = ListArray::new(
975 Arc::new(Field::new("item", DataType::UInt64, true)),
976 offsets,
977 Arc::new(items),
978 Some(list_validity),
979 );
980
981 let mut field_metadata = HashMap::new();
982 field_metadata.insert(
983 STRUCTURAL_ENCODING_META_KEY.to_string(),
984 structural_encoding.into(),
985 );
986
987 let test_cases = TestCases::default()
988 .with_range(0..3)
989 .with_range(1..2)
990 .with_indices(vec![1])
991 .with_indices(vec![2])
992 .with_structural_encodings();
993 check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, field_metadata)
994 .await;
995 }
996
997 #[rstest]
998 #[test_log::test(tokio::test)]
999 async fn test_simple_two_page_list(
1000 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
1001 structural_encoding: &str,
1002 ) {
1003 let items_builder = Int64Builder::new();
1007 let mut list_builder = ListBuilder::new(items_builder);
1008 for i in 0..512 {
1009 list_builder.append_value([Some(i), Some(i * 2)]);
1010 }
1011 let list_array_1 = list_builder.finish();
1012
1013 let items_builder = Int64Builder::new();
1014 let mut list_builder = ListBuilder::new(items_builder);
1015 for i in 0..512 {
1016 let i = i + 512;
1017 list_builder.append_value([Some(i), Some(i * 2)]);
1018 }
1019 let list_array_2 = list_builder.finish();
1020
1021 let mut metadata = HashMap::new();
1022 metadata.insert(
1023 STRUCTURAL_ENCODING_META_KEY.to_string(),
1024 structural_encoding.into(),
1025 );
1026
1027 let test_cases = TestCases::default()
1028 .with_structural_encodings()
1029 .with_page_sizes(vec![100])
1030 .with_range(800..900);
1031 check_round_trip_encoding_of_data(
1032 vec![Arc::new(list_array_1), Arc::new(list_array_2)],
1033 &test_cases,
1034 metadata,
1035 )
1036 .await;
1037 }
1038
1039 #[test_log::test(tokio::test)]
1040 async fn test_simple_large_list() {
1041 let items_builder = Int32Builder::new();
1042 let mut list_builder = LargeListBuilder::new(items_builder);
1043 list_builder.append_value([Some(1), Some(2), Some(3)]);
1044 list_builder.append_value([Some(4), Some(5)]);
1045 list_builder.append_null();
1046 list_builder.append_value([Some(6), Some(7), Some(8)]);
1047 let list_array = list_builder.finish();
1048
1049 let test_cases = TestCases::default()
1050 .with_range(0..2)
1051 .with_range(0..3)
1052 .with_range(1..3)
1053 .with_indices(vec![1, 3]);
1054 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
1055 .await;
1056 }
1057
1058 #[rstest]
1059 #[test_log::test(tokio::test)]
1060 async fn test_empty_lists(
1061 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
1062 structural_encoding: &str,
1063 ) {
1064 let mut field_metadata = HashMap::new();
1065 field_metadata.insert(
1066 STRUCTURAL_ENCODING_META_KEY.to_string(),
1067 structural_encoding.into(),
1068 );
1069
1070 let values = [vec![Some(1), Some(2), Some(3)], vec![], vec![None]];
1073 for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] {
1075 let items_builder = Int32Builder::new();
1076 let mut list_builder = ListBuilder::new(items_builder);
1077 for idx in order {
1078 list_builder.append_value(values[idx].clone());
1079 }
1080 let list_array = Arc::new(list_builder.finish());
1081 let test_cases = TestCases::default()
1082 .with_indices(vec![1])
1083 .with_indices(vec![0])
1084 .with_indices(vec![2])
1085 .with_indices(vec![0, 1]);
1086 check_round_trip_encoding_of_data(
1087 vec![list_array.clone()],
1088 &test_cases,
1089 field_metadata.clone(),
1090 )
1091 .await;
1092 let test_cases = test_cases.with_batch_size(1);
1093 check_round_trip_encoding_of_data(
1094 vec![list_array],
1095 &test_cases,
1096 field_metadata.clone(),
1097 )
1098 .await;
1099 }
1100
1101 let items_builder = Int32Builder::new();
1106 let mut list_builder = ListBuilder::new(items_builder);
1107 list_builder.append(true);
1108 list_builder.append_null();
1109 list_builder.append(true);
1110 let list_array = Arc::new(list_builder.finish());
1111
1112 let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
1113 check_round_trip_encoding_of_data(
1114 vec![list_array.clone()],
1115 &test_cases,
1116 field_metadata.clone(),
1117 )
1118 .await;
1119 let test_cases = test_cases.with_batch_size(1);
1120 check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone())
1121 .await;
1122
1123 let items_builder = StringBuilder::new();
1128 let mut list_builder = ListBuilder::new(items_builder);
1129 list_builder.append(true);
1130 list_builder.append_null();
1131 list_builder.append(true);
1132 let list_array = Arc::new(list_builder.finish());
1133
1134 let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
1135 check_round_trip_encoding_of_data(
1136 vec![list_array.clone()],
1137 &test_cases,
1138 field_metadata.clone(),
1139 )
1140 .await;
1141 let test_cases = test_cases.with_batch_size(1);
1142 check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone())
1143 .await;
1144
1145 let items_builder = Int32Builder::new();
1148 let mut list_builder = ListBuilder::new(items_builder);
1149 list_builder.append_null();
1150 list_builder.append_null();
1151 list_builder.append_null();
1152 let list_array = Arc::new(list_builder.finish());
1153
1154 let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
1155 check_round_trip_encoding_of_data(
1156 vec![list_array.clone()],
1157 &test_cases,
1158 field_metadata.clone(),
1159 )
1160 .await;
1161 let test_cases = test_cases.with_batch_size(1);
1162 check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata.clone())
1163 .await;
1164
1165 let items_builder = Int32Builder::new();
1168 let mut list_builder = ListBuilder::new(items_builder);
1169 list_builder.append_null();
1170 list_builder.append_null();
1171 list_builder.append_null();
1172 let list_array = Arc::new(list_builder.finish());
1173
1174 let struct_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
1175 let struct_array = Arc::new(StructArray::new(
1176 Fields::from(vec![Field::new(
1177 "lists",
1178 list_array.data_type().clone(),
1179 true,
1180 )]),
1181 vec![list_array],
1182 Some(struct_validity),
1183 ));
1184
1185 let test_cases = TestCases::default()
1186 .with_range(0..2)
1187 .with_indices(vec![1])
1188 .with_structural_encodings();
1189 check_round_trip_encoding_of_data(
1190 vec![struct_array.clone()],
1191 &test_cases,
1192 field_metadata.clone(),
1193 )
1194 .await;
1195 let test_cases = test_cases.with_batch_size(1);
1196 check_round_trip_encoding_of_data(vec![struct_array], &test_cases, field_metadata.clone())
1197 .await;
1198 }
1199
1200 #[test_log::test(tokio::test)]
1201 async fn test_empty_list_list() {
1202 let items_builder = Int32Builder::new();
1203 let list_builder = ListBuilder::new(items_builder);
1204 let mut outer_list_builder = ListBuilder::new(list_builder);
1205 outer_list_builder.append_null();
1206 outer_list_builder.append_null();
1207 outer_list_builder.append_null();
1208 let list_array = Arc::new(outer_list_builder.finish());
1209
1210 let test_cases = TestCases::default().with_structural_encodings();
1211 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1212 }
1213
1214 #[test_log::test(tokio::test)]
1215 #[ignore] async fn test_jumbo_list() {
1217 let items = BooleanArray::new_null(1024 * 1024);
1221 let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 1024 * 1024]));
1222 let list_arr = Arc::new(ListArray::new(
1223 Arc::new(Field::new("item", DataType::Boolean, true)),
1224 offsets,
1225 Arc::new(items),
1226 None,
1227 )) as ArrayRef;
1228 let arrs = vec![list_arr; 5000];
1229
1230 let test_cases = TestCases::default().without_validation();
1232 check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await;
1233 }
1234
1235 #[tokio::test]
1239 async fn test_fuzz_issue_4466() {
1240 let list_sizes = vec![
1243 13, 18, 12, 7, 14, 12, 6, 13, 18, 8, 6, 11, 17, 12, 8, 19, 5, 6, 10, 13, 8, 6, 10, 4, 8, 16, 14, 12, 18, 9, 17, 8, 14, 18, 15, 3, 2, 4, 5, 1, 3, 13, 1, 2, 10, 4, 10, 18, 7, 14, 18, 13, 9, 17, 3, 13, 10, 14, 8, 19, 17, 10, 5, 11, 6, 15, 10, 18, 18, 20, 16, 11, 12, 15, 7, 9, 3, 10, 20, 5, 2, 3, 17, 4, 8, 12, 15, 6, 3, 20, 15, 20, 1, 19, 16, ];
1254
1255 let mut list_builder = ListBuilder::new(Int32Builder::new());
1257 let mut total_values = 0;
1258
1259 for size in &list_sizes {
1260 for i in 0..*size {
1261 list_builder.values().append_value(i);
1262 }
1263 list_builder.append(true);
1264 total_values += size;
1265 }
1266
1267 let list_array = Arc::new(list_builder.finish());
1268
1269 assert_eq!(list_array.len(), 95);
1271 assert_eq!(total_values, 1025);
1272
1273 let test_cases = TestCases::default().with_structural_encodings();
1276
1277 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1281 }
1282
1283 #[rstest]
1284 #[test_log::test(tokio::test)]
1285 async fn test_sparse_large_string_list(
1286 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
1287 structural_encoding: &str,
1288 ) {
1289 let levels_per_chunk =
1292 crate::encodings::logical::primitive::miniblock::max_repdef_levels_per_chunk(2);
1293 let num_rows = (levels_per_chunk * 3) as u32;
1294 let num_non_empty = 100u32;
1295 let strings_per_list = 10;
1296
1297 let items_builder = StringBuilder::new();
1298 let mut list_builder = ListBuilder::new(items_builder);
1299
1300 let step = num_rows / num_non_empty;
1302 let mut next_non_empty = step / 2;
1303
1304 for i in 0..num_rows {
1305 if i == next_non_empty {
1306 let vals: Vec<Option<&str>> = (0..strings_per_list)
1307 .map(|j| match j % 4 {
1308 0 => Some("a"),
1309 1 => Some("bb"),
1310 2 => Some("ccc"),
1311 _ => Some("d"),
1312 })
1313 .collect();
1314 list_builder.append_value(vals);
1315 next_non_empty = next_non_empty.saturating_add(step);
1316 } else {
1317 list_builder.append_value([] as [Option<&str>; 0]);
1318 }
1319 }
1320 let list_array = list_builder.finish();
1321
1322 let mut field_metadata = HashMap::new();
1323 field_metadata.insert(
1324 STRUCTURAL_ENCODING_META_KEY.to_string(),
1325 structural_encoding.into(),
1326 );
1327
1328 let list_array = Arc::new(list_array) as ArrayRef;
1329 let pages = try_encode_v22_pages_with_metadata(list_array.clone(), field_metadata.clone())
1330 .await
1331 .unwrap();
1332 if structural_encoding == STRUCTURAL_ENCODING_MINIBLOCK {
1333 assert_split_miniblock_layout(&pages, 2, true);
1334 }
1335
1336 let chunk_boundary = levels_per_chunk;
1337 let test_cases = TestCases::default()
1338 .with_range(chunk_boundary - 2..chunk_boundary + 2)
1339 .with_indices(vec![
1340 0,
1341 (step / 2) as u64,
1342 chunk_boundary - 1,
1343 chunk_boundary,
1344 num_rows as u64 - 1,
1345 ])
1346 .with_batch_size(64 * 1024)
1347 .with_page_sizes(vec![1024 * 1024])
1348 .with_encoding(TestEncoding::StructuralU32);
1349 check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata).await;
1350 }
1351
1352 #[test_log::test(tokio::test)]
1353 async fn test_sparse_boolean_list_uses_miniblock() {
1354 let levels_per_chunk =
1358 crate::encodings::logical::primitive::miniblock::max_repdef_levels_per_chunk(2);
1359 let num_rows = (levels_per_chunk + 1) as usize;
1362 let booleans_per_list = 8usize;
1363
1364 let mut offsets = Vec::with_capacity(num_rows + 1);
1365 let mut values = Vec::with_capacity(2 * booleans_per_list);
1366 offsets.push(0i32);
1367
1368 for row in 0..num_rows {
1369 if row == 0 || row == num_rows - 1 {
1370 values.extend((0..booleans_per_list).map(|idx| idx % 2 == 0));
1371 }
1372 offsets.push(values.len() as i32);
1373 }
1374
1375 let items = BooleanArray::from(values);
1376 let list_array = ListArray::new(
1377 Arc::new(Field::new("item", DataType::Boolean, true)),
1378 OffsetBuffer::new(ScalarBuffer::from(offsets)),
1379 Arc::new(items),
1380 None,
1381 );
1382
1383 let test_cases = TestCases::default()
1384 .with_range(0..1000)
1385 .with_indices(vec![0, levels_per_chunk / 2, num_rows as u64 - 1])
1386 .with_batch_size(64 * 1024)
1387 .with_page_sizes(vec![1024 * 1024])
1388 .with_encoding(TestEncoding::StructuralU32);
1389 let list_array = Arc::new(list_array) as ArrayRef;
1390 let pages = encode_v22_pages(list_array.clone()).await;
1391 assert_split_miniblock_layout(&pages, 2, false);
1392 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1393 }
1394
1395 #[test_log::test(tokio::test)]
1396 async fn test_sparse_boolean_list_with_long_empty_prefix() {
1397 let empty_prefix_rows = 70_000usize;
1398 let trailing_empty_rows = 9usize;
1399 let booleans_per_list = 8usize;
1400 let num_rows = empty_prefix_rows + 1 + trailing_empty_rows;
1401
1402 let mut offsets = Vec::with_capacity(num_rows + 1);
1403 offsets.extend(std::iter::repeat_n(0i32, empty_prefix_rows + 1));
1404 let values = (0..booleans_per_list)
1405 .map(|idx| idx % 2 == 0)
1406 .collect::<Vec<_>>();
1407 offsets.push(values.len() as i32);
1408 offsets.extend(std::iter::repeat_n(
1409 values.len() as i32,
1410 trailing_empty_rows,
1411 ));
1412
1413 let items = BooleanArray::from(values);
1414 let list_array = ListArray::new(
1415 Arc::new(Field::new("item", DataType::Boolean, true)),
1416 OffsetBuffer::new(ScalarBuffer::from(offsets)),
1417 Arc::new(items),
1418 None,
1419 );
1420
1421 let test_cases = TestCases::default()
1422 .with_range(0..num_rows as u64)
1423 .with_indices(vec![0, empty_prefix_rows as u64, num_rows as u64 - 1])
1424 .with_dense_encodings();
1425 let list_array = Arc::new(list_array) as ArrayRef;
1426 let pages = encode_v22_pages(list_array.clone()).await;
1427 assert_split_miniblock_layout(&pages, 1, true);
1428 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1429 }
1430
1431 #[rstest]
1432 #[test_log::test(tokio::test)]
1433 async fn test_sparse_boolean_list_with_long_null_prefix(
1434 #[values(
1435 TestEncoding::Array,
1436 TestEncoding::StructuralU16,
1437 TestEncoding::StructuralU32
1438 )]
1439 encoding: TestEncoding,
1440 ) {
1441 let null_prefix_rows = 70_000usize;
1442 let trailing_empty_rows = 9usize;
1443 let booleans_per_list = 8usize;
1444 let num_rows = null_prefix_rows + 1 + trailing_empty_rows;
1445
1446 let mut offsets = Vec::with_capacity(num_rows + 1);
1447 offsets.extend(std::iter::repeat_n(0i32, null_prefix_rows + 1));
1448 let values = (0..booleans_per_list)
1449 .map(|idx| idx % 2 == 0)
1450 .collect::<Vec<_>>();
1451 offsets.push(values.len() as i32);
1452 offsets.extend(std::iter::repeat_n(
1453 values.len() as i32,
1454 trailing_empty_rows,
1455 ));
1456 let validity = BooleanBuffer::from_iter((0..num_rows).map(|row| row >= null_prefix_rows));
1457
1458 let items = BooleanArray::from(values);
1459 let list_array = ListArray::new(
1460 Arc::new(Field::new("item", DataType::Boolean, true)),
1461 OffsetBuffer::new(ScalarBuffer::from(offsets)),
1462 Arc::new(items),
1463 Some(NullBuffer::new(validity)),
1464 );
1465
1466 let test_cases = TestCases::default()
1467 .with_range(0..num_rows as u64)
1468 .with_indices(vec![0, null_prefix_rows as u64, num_rows as u64 - 1])
1469 .with_encoding(encoding);
1470 let list_array = Arc::new(list_array) as ArrayRef;
1471 let pages = encode_v22_pages(list_array.clone()).await;
1472 assert_split_miniblock_layout(&pages, 1, true);
1473 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1474 }
1475
1476 #[test_log::test(tokio::test)]
1477 async fn test_sparse_boolean_list_with_amortized_long_empty_prefix() {
1478 let empty_prefix_rows = 62_000usize;
1479 let booleans_per_list = 8_192usize;
1480 let num_rows = empty_prefix_rows + 1;
1481
1482 let mut offsets = Vec::with_capacity(num_rows + 1);
1483 offsets.extend(std::iter::repeat_n(0i32, empty_prefix_rows + 1));
1484 let values = (0..booleans_per_list)
1485 .map(|idx| idx % 2 == 0)
1486 .collect::<Vec<_>>();
1487 offsets.push(values.len() as i32);
1488
1489 let items = BooleanArray::from(values);
1490 let list_array = ListArray::new(
1491 Arc::new(Field::new("item", DataType::Boolean, true)),
1492 OffsetBuffer::new(ScalarBuffer::from(offsets)),
1493 Arc::new(items),
1494 None,
1495 );
1496
1497 let test_cases = TestCases::default()
1498 .with_range(0..num_rows as u64)
1499 .with_indices(vec![0, empty_prefix_rows as u64])
1500 .with_dense_encodings();
1501 let list_array = Arc::new(list_array) as ArrayRef;
1502 let pages = encode_v22_pages(list_array.clone()).await;
1503 assert_split_miniblock_layout(&pages, 1, true);
1504 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1505 }
1506
1507 fn unsplittable_nested_list(items: ArrayRef, empty_inner_lists: usize) -> ArrayRef {
1508 let mut inner_offsets = vec![0i32; empty_inner_lists + 1];
1509 inner_offsets.push(items.len() as i32);
1510 let inner_list = ListArray::new(
1511 Arc::new(Field::new("item", items.data_type().clone(), true)),
1512 OffsetBuffer::new(ScalarBuffer::from(inner_offsets)),
1513 items,
1514 None,
1515 );
1516 Arc::new(ListArray::new(
1517 Arc::new(Field::new("item", inner_list.data_type().clone(), true)),
1518 OffsetBuffer::new(ScalarBuffer::from(vec![0i32, empty_inner_lists as i32 + 1])),
1519 Arc::new(inner_list),
1520 None,
1521 ))
1522 }
1523
1524 #[rstest]
1525 #[case::boolean(Arc::new(BooleanArray::from(vec![true, false, true, false, true, false, true, false])))]
1526 #[case::string(Arc::new(StringArray::from(vec!["value", "other"])))]
1527 #[test_log::test(tokio::test)]
1528 async fn test_nested_sparse_single_row_falls_back_to_fullzip(#[case] items: ArrayRef) {
1529 let list = unsplittable_nested_list(items, 70_000);
1530 let pages = encode_v22_pages(list.clone()).await;
1531 assert_has_fullzip_layout(&pages);
1532
1533 let test_cases = TestCases::default()
1534 .with_range(0..1)
1535 .with_indices(vec![0])
1536 .with_dense_encodings();
1537 check_round_trip_encoding_of_data(vec![list], &test_cases, HashMap::new()).await;
1538 }
1539
1540 fn make_hnsw_shaped_list_u32() -> ListArray {
1545 const DENSE_ROWS: u32 = 5_000;
1546 const NEIGHBORS_PER_ROW: u32 = 32;
1547 const EMPTY_TAIL_ROWS: u32 = 70_000;
1548
1549 let mut list_builder = ListBuilder::new(UInt32Builder::new());
1550 let mut next_val: u32 = 0;
1551 for _ in 0..DENSE_ROWS {
1552 for _ in 0..NEIGHBORS_PER_ROW {
1553 list_builder.values().append_value(next_val);
1554 next_val = next_val.wrapping_add(1);
1555 }
1556 list_builder.append(true);
1557 }
1558 for _ in 0..EMPTY_TAIL_ROWS {
1559 list_builder.append(true);
1560 }
1561 list_builder.finish()
1562 }
1563
1564 #[test_log::test(tokio::test)]
1574 async fn test_list_hnsw_shape_splits_to_miniblock_v2_2() {
1575 let list_array = make_hnsw_shaped_list_u32();
1576 let dense_rows: u64 = 5_000;
1577 let total_rows = list_array.len() as u64;
1578
1579 let test_cases = TestCases::default()
1580 .with_range(0..1000)
1581 .with_range(dense_rows.saturating_sub(8)..(dense_rows + 8))
1582 .with_range(0..total_rows)
1583 .with_indices(vec![0, dense_rows - 1, dense_rows, total_rows - 1])
1584 .with_encoding(TestEncoding::StructuralU32);
1585 let list_array = Arc::new(list_array) as ArrayRef;
1586 let pages = encode_v22_pages(list_array.clone()).await;
1587 assert_split_miniblock_layout(&pages, 1, true);
1588 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
1589 }
1590
1591 #[test_log::test(tokio::test)]
1597 async fn test_forced_miniblock_hnsw_shape_splits_to_miniblock_v2_2() {
1598 let list_array = make_hnsw_shaped_list_u32();
1599 let total_rows = list_array.len() as u64;
1600
1601 let mut field_metadata = HashMap::new();
1602 field_metadata.insert(
1603 STRUCTURAL_ENCODING_META_KEY.to_string(),
1604 STRUCTURAL_ENCODING_MINIBLOCK.into(),
1605 );
1606
1607 let test_cases = TestCases::default()
1608 .with_range(0..total_rows)
1609 .with_encoding(TestEncoding::StructuralU32);
1610 let list_array = Arc::new(list_array) as ArrayRef;
1611 let pages = try_encode_v22_pages_with_metadata(list_array.clone(), field_metadata.clone())
1612 .await
1613 .unwrap();
1614 assert_split_miniblock_layout(&pages, 1, true);
1615 check_round_trip_encoding_of_data(vec![list_array], &test_cases, field_metadata).await;
1616 }
1617
1618 #[test]
1619 fn test_plan_decoded_bytes_list_i32_not_supported() {
1620 use crate::decoder::StructuralFieldDecoder;
1625 use crate::encodings::logical::primitive::StructuralPrimitiveFieldDecoder;
1626 use arrow_array::builder::{Int32Builder, ListBuilder};
1627 use arrow_schema::Field;
1628
1629 let mut builder = ListBuilder::new(Int32Builder::new());
1632 let lengths = [0usize, 3, 1, 7, 2, 0, 5];
1633 let mut value = 0i32;
1634 for &len in &lengths {
1635 for _ in 0..len {
1636 builder.values().append_value(value);
1637 value += 1;
1638 }
1639 builder.append(true);
1640 }
1641 let _list_array = builder.finish(); let item_field = Arc::new(Field::new("item", DataType::Int32, true));
1645 let list_type = DataType::List(item_field.clone());
1646
1647 let child = Box::new(StructuralPrimitiveFieldDecoder::new(&item_field, false));
1648 let decoder = super::StructuralListDecoder::new(child, list_type);
1649
1650 let err = decoder
1651 .plan_decoded_bytes(lengths.len() as u64)
1652 .unwrap_err();
1653 assert!(
1654 matches!(err, lance_core::Error::NotSupported { .. }),
1655 "expected NotSupported, got: {err:?}"
1656 );
1657 assert!(
1658 err.to_string().contains("repetition"),
1659 "error should mention repetition, got: {err}"
1660 );
1661 }
1662}