1use std::{
5 collections::{BinaryHeap, VecDeque},
6 ops::Range,
7 sync::Arc,
8};
9
10use super::{
11 fixed_size_list::StructuralFixedSizeListDecoder, list::StructuralListDecoder,
12 map::StructuralMapDecoder, primitive::StructuralPrimitiveFieldDecoder,
13};
14use crate::{
15 decoder::{
16 DecodedArray, FilterExpression, LoadedPageShard, NextDecodeTask, PageEncoding,
17 ScheduledScanLine, SchedulerContext, StructuralDecodeArrayTask, StructuralFieldDecoder,
18 StructuralFieldScheduler, StructuralSchedulingJob,
19 },
20 encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers},
21 format::pb,
22 repdef::{CompositeRepDefUnraveler, RepDefBuilder},
23};
24use arrow_array::{Array, ArrayRef, StructArray, cast::AsArray};
25use arrow_schema::{DataType, Fields};
26use futures::{
27 FutureExt, StreamExt, TryStreamExt,
28 future::BoxFuture,
29 stream::{FuturesOrdered, FuturesUnordered},
30};
31use itertools::Itertools;
32use lance_arrow::FieldExt;
33use lance_arrow::{deepcopy::deep_copy_nulls, r#struct::StructArrayExt};
34use lance_core::{Error, Result, datatypes::validate_fixed_size_list_dimensions};
35use log::trace;
36
37#[derive(Debug)]
38struct StructuralSchedulingJobWithStatus<'a> {
39 col_idx: u32,
40 col_name: &'a str,
41 job: Box<dyn StructuralSchedulingJob + 'a>,
42 rows_scheduled: u64,
43 rows_remaining: u64,
44 ready_scan_lines: VecDeque<ScheduledScanLine>,
45}
46
47impl PartialEq for StructuralSchedulingJobWithStatus<'_> {
48 fn eq(&self, other: &Self) -> bool {
49 self.col_idx == other.col_idx
50 }
51}
52
53impl Eq for StructuralSchedulingJobWithStatus<'_> {}
54
55impl PartialOrd for StructuralSchedulingJobWithStatus<'_> {
56 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
57 Some(self.cmp(other))
58 }
59}
60
61impl Ord for StructuralSchedulingJobWithStatus<'_> {
62 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
63 other.rows_scheduled.cmp(&self.rows_scheduled)
65 }
66}
67
68#[derive(Debug)]
75struct RepDefStructSchedulingJob<'a> {
76 children: BinaryHeap<StructuralSchedulingJobWithStatus<'a>>,
78 rows_scheduled: u64,
79 num_rows: u64,
80}
81
82impl<'a> RepDefStructSchedulingJob<'a> {
83 fn new(
84 scheduler: &'a StructuralStructScheduler,
85 children: Vec<Box<dyn StructuralSchedulingJob + 'a>>,
86 num_rows: u64,
87 ) -> Self {
88 let children = children
89 .into_iter()
90 .enumerate()
91 .map(|(idx, job)| StructuralSchedulingJobWithStatus {
92 col_idx: idx as u32,
93 col_name: scheduler.child_fields[idx].name(),
94 job,
95 rows_scheduled: 0,
96 rows_remaining: num_rows,
97 ready_scan_lines: VecDeque::new(),
98 })
99 .collect::<BinaryHeap<_>>();
100 Self {
101 children,
102 rows_scheduled: 0,
103 num_rows,
104 }
105 }
106}
107
108impl StructuralSchedulingJob for RepDefStructSchedulingJob<'_> {
109 fn schedule_next(
110 &mut self,
111 mut context: &mut SchedulerContext,
112 ) -> Result<Vec<ScheduledScanLine>> {
113 if self.children.is_empty() {
114 if self.rows_scheduled == self.num_rows {
116 return Ok(Vec::new());
117 }
118 self.rows_scheduled = self.num_rows;
119 return Ok(vec![ScheduledScanLine {
120 decoders: Vec::new(),
121 rows_scheduled: self.num_rows,
122 }]);
123 }
124
125 let mut decoders = Vec::new();
126 let old_rows_scheduled = self.rows_scheduled;
127 while old_rows_scheduled == self.rows_scheduled {
130 if self.children.is_empty() {
131 return Ok(Vec::new());
133 }
134 let mut next_child = self.children.pop().unwrap();
135 if next_child.ready_scan_lines.is_empty() {
136 let scoped = context.push(next_child.col_name, next_child.col_idx);
137 let child_scans = next_child.job.schedule_next(scoped.context)?;
138 context = scoped.pop();
139 if child_scans.is_empty() {
140 continue;
142 }
143 next_child.ready_scan_lines.extend(child_scans);
144 }
145 let child_scan = next_child.ready_scan_lines.pop_front().unwrap();
146 trace!(
147 "Scheduled {} rows for child {}",
148 child_scan.rows_scheduled, next_child.col_idx
149 );
150 next_child.rows_scheduled += child_scan.rows_scheduled;
151 next_child.rows_remaining -= child_scan.rows_scheduled;
152 decoders.extend(child_scan.decoders);
153 self.children.push(next_child);
154 self.rows_scheduled = self.children.peek().unwrap().rows_scheduled;
155 }
156 let struct_rows_scheduled = self.rows_scheduled - old_rows_scheduled;
157 Ok(vec![ScheduledScanLine {
158 decoders,
159 rows_scheduled: struct_rows_scheduled,
160 }])
161 }
162}
163
164#[derive(Debug)]
175pub struct StructuralStructScheduler {
176 children: Vec<Box<dyn StructuralFieldScheduler>>,
177 child_fields: Fields,
178}
179
180impl StructuralStructScheduler {
181 pub fn new(children: Vec<Box<dyn StructuralFieldScheduler>>, child_fields: Fields) -> Self {
182 Self {
183 children,
184 child_fields,
185 }
186 }
187}
188
189impl StructuralFieldScheduler for StructuralStructScheduler {
190 fn schedule_ranges<'a>(
191 &'a self,
192 ranges: &[Range<u64>],
193 filter: &FilterExpression,
194 ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
195 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
196
197 let child_schedulers = self
198 .children
199 .iter()
200 .map(|child| child.schedule_ranges(ranges, filter))
201 .collect::<Result<Vec<_>>>()?;
202
203 Ok(Box::new(RepDefStructSchedulingJob::new(
204 self,
205 child_schedulers,
206 num_rows,
207 )))
208 }
209
210 fn initialize<'a>(
211 &'a mut self,
212 filter: &'a FilterExpression,
213 context: &'a SchedulerContext,
214 ) -> BoxFuture<'a, Result<()>> {
215 let children_initialization = self
216 .children
217 .iter_mut()
218 .map(|child| child.initialize(filter, context))
219 .collect::<FuturesUnordered<_>>();
220 async move {
221 children_initialization
222 .map(|res| res.map(|_| ()))
223 .try_collect::<Vec<_>>()
224 .await?;
225 Ok(())
226 }
227 .boxed()
228 }
229}
230
231#[derive(Debug)]
232pub struct StructuralStructDecoder {
233 children: Vec<Box<dyn StructuralFieldDecoder>>,
234 data_type: DataType,
235 child_fields: Fields,
236 is_root: bool,
238}
239
240impl StructuralStructDecoder {
241 pub fn new(fields: Fields, should_validate: bool, is_root: bool) -> Result<Self> {
242 let children = fields
243 .iter()
244 .map(|field| Self::field_to_decoder(field, should_validate))
245 .collect::<Result<Vec<_>>>()?;
246 let data_type = DataType::Struct(fields.clone());
247 Ok(Self {
248 data_type,
249 children,
250 child_fields: fields,
251 is_root,
252 })
253 }
254
255 fn field_to_decoder(
256 field: &Arc<arrow_schema::Field>,
257 should_validate: bool,
258 ) -> Result<Box<dyn StructuralFieldDecoder>> {
259 match field.data_type() {
260 DataType::Struct(fields) => {
261 if field.is_packed_struct() || field.is_blob() {
262 let decoder =
263 StructuralPrimitiveFieldDecoder::new(&field.clone(), should_validate);
264 Ok(Box::new(decoder))
265 } else {
266 Ok(Box::new(Self::new(fields.clone(), should_validate, false)?))
267 }
268 }
269 DataType::List(child_field) | DataType::LargeList(child_field) => {
270 let child_decoder = Self::field_to_decoder(child_field, should_validate)?;
271 Ok(Box::new(StructuralListDecoder::new(
272 child_decoder,
273 field.data_type().clone(),
274 )))
275 }
276 DataType::FixedSizeList(child_field, _)
277 if matches!(child_field.data_type(), DataType::Struct(_)) =>
278 {
279 validate_fixed_size_list_dimensions(field.name(), field.data_type())?;
285 let child_decoder = Self::field_to_decoder(child_field, should_validate)?;
287 Ok(Box::new(StructuralFixedSizeListDecoder::new(
288 child_decoder,
289 field.data_type().clone(),
290 )))
291 }
292 DataType::Map(entries_field, keys_sorted) => {
293 if *keys_sorted {
294 return Err(Error::not_supported_source(
295 "Map data type with keys_sorted=true is not supported yet"
296 .to_string()
297 .into(),
298 ));
299 }
300 let child_decoder = Self::field_to_decoder(entries_field, should_validate)?;
301 Ok(Box::new(StructuralMapDecoder::new(
302 child_decoder,
303 field.data_type().clone(),
304 )))
305 }
306 DataType::RunEndEncoded(_, _) => todo!(),
307 DataType::ListView(_) | DataType::LargeListView(_) => todo!(),
308 DataType::Union(_, _) => todo!(),
309 _ => Ok(Box::new(StructuralPrimitiveFieldDecoder::new(
310 field,
311 should_validate,
312 ))),
313 }
314 }
315
316 pub fn drain_batch_task(&mut self, num_rows: u64) -> Result<NextDecodeTask> {
317 let array_drain = self.drain(num_rows)?;
318 Ok(NextDecodeTask {
319 num_rows,
320 task: Box::new(array_drain),
321 })
322 }
323}
324
325impl StructuralFieldDecoder for StructuralStructDecoder {
326 fn accept_page(&mut self, mut child: LoadedPageShard) -> Result<()> {
327 let child_idx = child.path.pop_front().unwrap();
329 self.children[child_idx as usize].accept_page(child)?;
331 Ok(())
332 }
333
334 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
335 let child_tasks = self
336 .children
337 .iter_mut()
338 .map(|child| child.drain(num_rows))
339 .collect::<Result<Vec<_>>>()?;
340 Ok(Box::new(RepDefStructDecodeTask {
341 children: child_tasks,
342 child_fields: self.child_fields.clone(),
343 is_root: self.is_root,
344 num_rows,
345 }))
346 }
347
348 fn data_type(&self) -> &DataType {
349 &self.data_type
350 }
351}
352
353#[derive(Debug)]
354struct RepDefStructDecodeTask {
355 children: Vec<Box<dyn StructuralDecodeArrayTask>>,
356 child_fields: Fields,
357 is_root: bool,
358 num_rows: u64,
359}
360
361impl StructuralDecodeArrayTask for RepDefStructDecodeTask {
362 fn decode(self: Box<Self>) -> Result<DecodedArray> {
363 if self.children.is_empty() {
364 return Ok(DecodedArray {
365 array: Arc::new(StructArray::new_empty_fields(self.num_rows as usize, None)),
366 repdef: CompositeRepDefUnraveler::new(vec![]),
367 data_size: 0,
368 });
369 }
370
371 let arrays = self
372 .children
373 .into_iter()
374 .map(|task| task.decode())
375 .collect::<Result<Vec<_>>>()?;
376 let mut children = Vec::with_capacity(arrays.len());
377 let mut repdefs = Vec::with_capacity(arrays.len());
378 let mut data_size = 0u64;
379 let mut arrays_iter = arrays.into_iter();
380 let first_array = arrays_iter.next().ok_or_else(|| {
381 Error::internal("Struct decoder unexpectedly has no child arrays".to_string())
382 })?;
383 let length = first_array.array.len();
384
385 repdefs.push(first_array.repdef);
387 data_size += first_array.data_size;
388 children.push(first_array.array);
389
390 for array in arrays_iter {
391 if length != array.array.len() {
392 return Err(Error::invalid_input_source(
393 format!(
394 "Struct child array length {} does not match sibling length {}",
395 array.array.len(),
396 length
397 )
398 .into(),
399 ));
400 }
401 data_size += array.data_size;
402 children.push(array.array);
403 repdefs.push(array.repdef);
404 }
405
406 let primary_repdef = repdefs
411 .iter()
412 .position(CompositeRepDefUnraveler::has_sparse)
413 .unwrap_or(0);
414 let mut repdef = repdefs.swap_remove(primary_repdef);
415 if repdef.has_sparse() {
416 for sibling in repdefs {
417 if sibling.has_sparse() {
418 repdef.add_compatibility_check(sibling);
419 }
420 }
421 }
422
423 let validity = if self.is_root {
424 repdef.ensure_exhausted()?;
425 None
426 } else {
427 repdef.unravel_validity(length)?
428 };
429
430 let array = StructArray::try_new(self.child_fields, children, validity)
431 .map_err(|e| Error::invalid_input_source(e.to_string().into()))?;
432 Ok(DecodedArray {
433 array: Arc::new(array),
434 repdef,
435 data_size,
436 })
437 }
438}
439
440pub struct StructStructuralEncoder {
445 keep_original_array: bool,
446 children: Vec<Box<dyn FieldEncoder>>,
447}
448
449impl StructStructuralEncoder {
450 pub fn new(keep_original_array: bool, children: Vec<Box<dyn FieldEncoder>>) -> Self {
451 Self {
452 keep_original_array,
453 children,
454 }
455 }
456}
457
458impl FieldEncoder for StructStructuralEncoder {
459 fn maybe_encode(
460 &mut self,
461 array: ArrayRef,
462 external_buffers: &mut OutOfLineBuffers,
463 mut repdef: RepDefBuilder,
464 row_number: u64,
465 num_rows: u64,
466 ) -> Result<Vec<EncodeTask>> {
467 let struct_array = array.as_struct();
468 let mut struct_array = struct_array.normalize_slicing()?;
469 if let Some(validity) = struct_array.nulls() {
470 if self.keep_original_array {
471 repdef.add_validity_bitmap(validity.clone())
472 } else {
473 repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap())
474 }
475 struct_array = struct_array.pushdown_nulls()?;
476 } else {
477 repdef.add_no_null(struct_array.len());
478 }
479 let child_tasks = self
480 .children
481 .iter_mut()
482 .zip(struct_array.columns().iter())
483 .map(|(encoder, arr)| {
484 encoder.maybe_encode(
485 arr.clone(),
486 external_buffers,
487 repdef.clone(),
488 row_number,
489 num_rows,
490 )
491 })
492 .collect::<Result<Vec<_>>>()?;
493 Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
494 }
495
496 fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
497 self.children
498 .iter_mut()
499 .map(|encoder| encoder.flush(external_buffers))
500 .flatten_ok()
501 .collect::<Result<Vec<_>>>()
502 }
503
504 fn num_columns(&self) -> u32 {
505 self.children
506 .iter()
507 .map(|child| child.num_columns())
508 .sum::<u32>()
509 }
510
511 fn finish(
512 &mut self,
513 external_buffers: &mut OutOfLineBuffers,
514 ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
515 let mut child_columns = self
516 .children
517 .iter_mut()
518 .map(|child| child.finish(external_buffers))
519 .collect::<FuturesOrdered<_>>();
520 async move {
521 let mut encoded_columns = Vec::with_capacity(child_columns.len());
522 while let Some(child_cols) = child_columns.next().await {
523 encoded_columns.extend(child_cols?);
524 }
525 Ok(encoded_columns)
526 }
527 .boxed()
528 }
529}
530
531pub struct StructFieldEncoder {
532 children: Vec<Box<dyn FieldEncoder>>,
533 column_index: u32,
534 num_rows_seen: u64,
535}
536
537impl StructFieldEncoder {
538 pub fn new(children: Vec<Box<dyn FieldEncoder>>, column_index: u32) -> Self {
539 Self {
540 children,
541 column_index,
542 num_rows_seen: 0,
543 }
544 }
545}
546
547impl FieldEncoder for StructFieldEncoder {
548 fn maybe_encode(
549 &mut self,
550 array: ArrayRef,
551 external_buffers: &mut OutOfLineBuffers,
552 repdef: RepDefBuilder,
553 row_number: u64,
554 num_rows: u64,
555 ) -> Result<Vec<EncodeTask>> {
556 self.num_rows_seen += array.len() as u64;
557 let struct_array = array.as_struct();
558 let child_tasks = self
559 .children
560 .iter_mut()
561 .zip(struct_array.columns().iter())
562 .map(|(encoder, arr)| {
563 encoder.maybe_encode(
564 arr.clone(),
565 external_buffers,
566 repdef.clone(),
567 row_number,
568 num_rows,
569 )
570 })
571 .collect::<Result<Vec<_>>>()?;
572 Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
573 }
574
575 fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
576 let child_tasks = self
577 .children
578 .iter_mut()
579 .map(|encoder| encoder.flush(external_buffers))
580 .collect::<Result<Vec<_>>>()?;
581 Ok(child_tasks.into_iter().flatten().collect::<Vec<_>>())
582 }
583
584 fn num_columns(&self) -> u32 {
585 self.children
586 .iter()
587 .map(|child| child.num_columns())
588 .sum::<u32>()
589 + 1
590 }
591
592 fn finish(
593 &mut self,
594 external_buffers: &mut OutOfLineBuffers,
595 ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
596 let mut child_columns = self
597 .children
598 .iter_mut()
599 .map(|child| child.finish(external_buffers))
600 .collect::<FuturesOrdered<_>>();
601 let num_rows_seen = self.num_rows_seen;
602 let column_index = self.column_index;
603 async move {
604 let mut columns = Vec::new();
605 let mut header = EncodedColumn::default();
607 header.final_pages.push(EncodedPage {
608 data: Vec::new(),
609 description: PageEncoding::Legacy(pb::ArrayEncoding {
610 array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct(
611 pb::SimpleStruct {},
612 )),
613 }),
614 num_rows: num_rows_seen,
615 column_idx: column_index,
616 row_number: 0, });
618 columns.push(header);
619 while let Some(child_cols) = child_columns.next().await {
621 columns.extend(child_cols?);
622 }
623 Ok(columns)
624 }
625 .boxed()
626 }
627}
628
629#[cfg(test)]
630mod tests {
631
632 use std::{collections::HashMap, sync::Arc};
633
634 use arrow_array::{
635 Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, StructArray,
636 builder::{Int32Builder, ListBuilder},
637 };
638 use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
639 use arrow_schema::{DataType, Field, Fields};
640
641 use super::StructuralStructDecoder;
642 use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data};
643
644 #[test]
645 fn test_zero_dimension_fsl_decoder_errors() {
646 let item_fields = Fields::from(vec![Field::new("x", DataType::Int32, true)]);
650 let fields = Fields::from(vec![Field::new(
651 "vecs",
652 DataType::FixedSizeList(
653 Arc::new(Field::new("item", DataType::Struct(item_fields), true)),
654 0,
655 ),
656 true,
657 )]);
658
659 let err = StructuralStructDecoder::new(fields, false, true).unwrap_err();
660 assert!(matches!(err, lance_core::Error::Schema { .. }));
661 assert!(
662 err.to_string()
663 .contains("dimension must be a positive integer"),
664 "unexpected error: {}",
665 err
666 );
667 }
668
669 #[test_log::test(tokio::test)]
670 async fn test_simple_struct() {
671 let data_type = DataType::Struct(Fields::from(vec![
672 Field::new("a", DataType::Int32, false),
673 Field::new("b", DataType::Int32, false),
674 ]));
675 let field = Field::new("", data_type, false);
676 check_basic_random(field).await;
677 }
678
679 #[test_log::test(tokio::test)]
680 async fn test_nullable_struct() {
681 let inner_fields = Fields::from(vec![
699 Field::new("x", DataType::Int32, false),
700 Field::new("y", DataType::Int32, true),
701 ]);
702 let inner_struct = DataType::Struct(inner_fields.clone());
703 let outer_fields = Fields::from(vec![
704 Field::new("score", DataType::Int32, true),
705 Field::new("location", inner_struct, true),
706 ]);
707
708 let x_vals = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
709 let y_vals = Int32Array::from(vec![Some(6), None, Some(8), Some(9), Some(10)]);
710 let scores = Int32Array::from(vec![None, Some(12), Some(13), Some(14), Some(15)]);
711
712 let location_validity = NullBuffer::from(vec![true, true, true, false, true]);
713 let locations = StructArray::new(
714 inner_fields,
715 vec![Arc::new(x_vals), Arc::new(y_vals)],
716 Some(location_validity),
717 );
718
719 let rows_validity = NullBuffer::from(vec![true, true, true, true, false]);
720 let rows = StructArray::new(
721 outer_fields,
722 vec![Arc::new(scores), Arc::new(locations)],
723 Some(rows_validity),
724 );
725
726 let test_cases = TestCases::default().with_structural_encodings();
727
728 check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await;
729 }
730
731 #[test_log::test(tokio::test)]
732 async fn test_simple_masked_nonempty_list() {
733 let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4), Some(5), Some(6)]);
736 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 5]));
737 let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
738 let list_array = ListArray::new(
739 Arc::new(Field::new("item", DataType::Int32, true)),
740 offsets,
741 Arc::new(items),
742 Some(NullBuffer::new(list_validity)),
743 );
744 let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
745 let struct_array = StructArray::new(
746 Fields::from(vec![Field::new(
747 "inner_list",
748 list_array.data_type().clone(),
749 true,
750 )]),
751 vec![Arc::new(list_array)],
752 Some(NullBuffer::new(struct_validity)),
753 );
754 check_round_trip_encoding_of_data(
755 vec![Arc::new(struct_array)],
756 &TestCases::default().with_structural_encodings(),
757 HashMap::new(),
758 )
759 .await;
760 }
761
762 #[test_log::test(tokio::test)]
763 async fn test_simple_struct_list() {
764 let items = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
767 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 4, 4, 4]));
768 let list_validity = BooleanBuffer::from(vec![true, true, true, true, false, true]);
769 let list_array = ListArray::new(
770 Arc::new(Field::new("item", DataType::Int32, true)),
771 offsets,
772 Arc::new(items),
773 Some(NullBuffer::new(list_validity)),
774 );
775 let struct_validity = BooleanBuffer::from(vec![true, true, true, true, true, false]);
776 let struct_array = StructArray::new(
777 Fields::from(vec![Field::new(
778 "inner_list",
779 list_array.data_type().clone(),
780 true,
781 )]),
782 vec![Arc::new(list_array)],
783 Some(NullBuffer::new(struct_validity)),
784 );
785 check_round_trip_encoding_of_data(
786 vec![Arc::new(struct_array)],
787 &TestCases::default().with_structural_encodings(),
788 HashMap::new(),
789 )
790 .await;
791 }
792
793 #[test_log::test(tokio::test)]
794 async fn test_struct_list() {
795 let data_type = DataType::Struct(Fields::from(vec![
796 Field::new(
797 "inner_list",
798 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
799 true,
800 ),
801 Field::new("outer_int", DataType::Int32, true),
802 ]));
803 let field = Field::new("row", data_type, false);
804 check_basic_random(field).await;
805 }
806
807 #[test_log::test(tokio::test)]
808 async fn test_empty_struct() {
809 let data_type = DataType::Struct(Fields::from(Vec::<Field>::default()));
812 let field = Field::new("row", data_type, false);
813 check_basic_random(field).await;
814 }
815
816 #[test_log::test(tokio::test)]
817 async fn test_complicated_struct() {
818 let data_type = DataType::Struct(Fields::from(vec![
819 Field::new("int", DataType::Int32, true),
820 Field::new(
821 "inner",
822 DataType::Struct(Fields::from(vec![
823 Field::new("inner_int", DataType::Int32, true),
824 Field::new(
825 "inner_list",
826 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
827 true,
828 ),
829 ])),
830 true,
831 ),
832 Field::new("outer_binary", DataType::Binary, true),
833 ]));
834 let field = Field::new("row", data_type, false);
835 check_basic_random(field).await;
836 }
837
838 #[test_log::test(tokio::test)]
839 async fn test_list_of_struct_with_null_struct_element() {
840 use arrow_array::StringArray;
843
844 let tag_array = StringArray::from(vec![
845 Some("valid"),
846 Some("null_struct"),
847 Some("valid"),
848 Some("valid"),
849 ]);
850 let struct_fields = Fields::from(vec![Field::new("tag", DataType::Utf8, true)]);
851 let struct_validity = NullBuffer::from(vec![false, true, false, false]);
853 let struct_array = StructArray::new(
854 struct_fields.clone(),
855 vec![Arc::new(tag_array)],
856 Some(struct_validity),
857 );
858
859 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4]));
860 let list_field = Field::new("item", DataType::Struct(struct_fields), true);
861 let list_array =
862 ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
863
864 check_round_trip_encoding_of_data(
865 vec![Arc::new(list_array)],
866 &TestCases::default().with_u32_structural_encodings(),
867 HashMap::new(),
868 )
869 .await;
870 }
871
872 #[test_log::test(tokio::test)]
873 async fn test_list_of_struct_with_all_null_child() {
874 use arrow_array::StringArray;
875
876 let a_array = StringArray::from(vec![Some("w"), Some("x"), Some("y"), Some("z")]);
877 let b_array = StringArray::from(vec![None::<&str>, None, None, None]);
878 let struct_fields = Fields::from(vec![
879 Field::new("a", DataType::Utf8, true),
880 Field::new("b", DataType::Utf8, true),
881 ]);
882 let struct_array = StructArray::new(
883 struct_fields.clone(),
884 vec![Arc::new(a_array), Arc::new(b_array)],
885 None,
886 );
887
888 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 4]));
889 let list_field = Field::new("item", DataType::Struct(struct_fields), true);
890 let list_array =
891 ListArray::new(Arc::new(list_field), offsets, Arc::new(struct_array), None);
892
893 check_round_trip_encoding_of_data(
894 vec![Arc::new(list_array)],
895 &TestCases::default()
896 .with_range(0..1)
897 .with_range(1..2)
898 .with_indices(vec![0])
899 .with_indices(vec![1])
900 .with_structural_encodings(),
901 HashMap::new(),
902 )
903 .await;
904 }
905
906 #[test_log::test(tokio::test)]
907 async fn test_list_of_struct_with_constant_child_and_empty_lists() {
908 let item_fields = Fields::from(vec![
909 Field::new("maneuver", DataType::Int64, true),
910 Field::new("remaining_dist", DataType::Float64, true),
911 ]);
912 let item_array = StructArray::new(
913 item_fields.clone(),
914 vec![
915 Arc::new(Int64Array::from(vec![Some(3), Some(3), Some(3)])),
916 Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])),
917 ],
918 None,
919 );
920
921 let list_field = Field::new("item", DataType::Struct(item_fields), true);
922 let list_array = ListArray::new(
923 Arc::new(list_field),
924 OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 2, 3, 3, 3, 3])),
925 Arc::new(item_array),
926 None,
927 );
928
929 let data_fields = Fields::from(vec![Field::new(
930 "maneuvers",
931 list_array.data_type().clone(),
932 true,
933 )]);
934 let data_array = StructArray::new(data_fields.clone(), vec![Arc::new(list_array)], None);
935 let row_validity = NullBuffer::from(vec![true, true, true, true, true, false]);
936 let row_array = StructArray::new(
937 Fields::from(vec![Field::new(
938 "data",
939 DataType::Struct(data_fields),
940 true,
941 )]),
942 vec![Arc::new(data_array)],
943 Some(row_validity),
944 );
945
946 check_round_trip_encoding_of_data(
947 vec![Arc::new(row_array)],
948 &TestCases::default().with_structural_encodings(),
949 HashMap::new(),
950 )
951 .await;
952 }
953
954 #[test_log::test(tokio::test)]
955 async fn test_ragged_scheduling() {
956 let items_builder = Int32Builder::new();
960 let mut list_builder = ListBuilder::new(items_builder);
961 for _ in 0..10000 {
962 list_builder.append_null();
963 }
964 let list_array = Arc::new(list_builder.finish());
965 let int_array = Arc::new(Int32Array::from_iter_values(0..10000));
966 let fields = vec![
967 Field::new("", list_array.data_type().clone(), true),
968 Field::new("", int_array.data_type().clone(), true),
969 ];
970 let struct_array = Arc::new(StructArray::new(
971 Fields::from(fields),
972 vec![list_array, int_array],
973 None,
974 )) as ArrayRef;
975 let struct_arrays = (0..10000)
976 .step_by(437)
978 .map(|offset| struct_array.slice(offset, 437.min(10000 - offset)))
979 .collect::<Vec<_>>();
980 check_round_trip_encoding_of_data(struct_arrays, &TestCases::default(), HashMap::new())
981 .await;
982 }
983}