1use crate::column::chunker::CdcChunk;
44use crate::column::writer::LevelDataRef;
45use crate::errors::{ParquetError, Result};
46use arrow_array::cast::AsArray;
47use arrow_array::types::RunEndIndexType;
48use arrow_array::{Array, ArrayRef, Int32Array, OffsetSizeTrait, RunArray, downcast_run_array};
49use arrow_buffer::bit_iterator::BitIndexIterator;
50use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
51use arrow_schema::{DataType, Field};
52use std::ops::Range;
53use std::sync::Arc;
54
55fn expand_ree_array(array: &ArrayRef) -> Result<ArrayRef> {
60 downcast_run_array!(
61 array => expand_typed_ree(array),
62 _ => unreachable!("expand_ree_array called on non-REE array"),
63 )
64}
65
66fn expand_typed_ree<R: RunEndIndexType>(run_array: &RunArray<R>) -> Result<ArrayRef> {
67 let run_ends = run_array.run_ends();
68 let values = run_array.values();
69 let len = run_array.len();
70 let indices: Int32Array = (0..len)
71 .map(|i| run_ends.get_physical_index(i) as i32)
72 .collect();
73 arrow_select::take::take(values.as_ref(), &indices, None)
74 .map_err(|e| arrow_err!("Failed to expand REE array: {}", e))
75}
76
77pub(crate) fn calculate_array_levels(array: &ArrayRef, field: &Field) -> Result<Vec<ArrayLevels>> {
80 let mut builder = LevelInfoBuilder::try_new(field, Default::default(), array)?;
81 builder.write(0..array.len());
82 Ok(builder.finish())
83}
84
85fn is_leaf(data_type: &DataType) -> bool {
88 matches!(
89 data_type,
90 DataType::Null
91 | DataType::Boolean
92 | DataType::Int8
93 | DataType::Int16
94 | DataType::Int32
95 | DataType::Int64
96 | DataType::UInt8
97 | DataType::UInt16
98 | DataType::UInt32
99 | DataType::UInt64
100 | DataType::Float16
101 | DataType::Float32
102 | DataType::Float64
103 | DataType::Utf8
104 | DataType::Utf8View
105 | DataType::LargeUtf8
106 | DataType::Timestamp(_, _)
107 | DataType::Date32
108 | DataType::Date64
109 | DataType::Time32(_)
110 | DataType::Time64(_)
111 | DataType::Duration(_)
112 | DataType::Interval(_)
113 | DataType::Binary
114 | DataType::LargeBinary
115 | DataType::BinaryView
116 | DataType::Decimal32(_, _)
117 | DataType::Decimal64(_, _)
118 | DataType::Decimal128(_, _)
119 | DataType::Decimal256(_, _)
120 | DataType::FixedSizeBinary(_)
121 )
122}
123
124#[derive(Debug, Default, Clone, Copy)]
126struct LevelContext {
127 rep_level: i16,
129 def_level: i16,
131}
132
133#[derive(Debug)]
135enum LevelInfoBuilder {
136 Primitive(ArrayLevels),
138 List(
140 Box<LevelInfoBuilder>, LevelContext, OffsetBuffer<i32>, Option<NullBuffer>, bool, ),
146 LargeList(
148 Box<LevelInfoBuilder>, LevelContext, OffsetBuffer<i64>, Option<NullBuffer>, bool, ),
154 FixedSizeList(
156 Box<LevelInfoBuilder>, LevelContext, usize, Option<NullBuffer>, ),
161 ListView(
163 Box<LevelInfoBuilder>, LevelContext, ScalarBuffer<i32>, ScalarBuffer<i32>, Option<NullBuffer>, ),
169 LargeListView(
171 Box<LevelInfoBuilder>, LevelContext, ScalarBuffer<i64>, ScalarBuffer<i64>, Option<NullBuffer>, ),
177 Struct(Vec<LevelInfoBuilder>, LevelContext, Option<NullBuffer>),
179}
180
181const BULK_FILL_MIN_LEN: usize = 64;
187
188impl LevelInfoBuilder {
189 fn try_new(field: &Field, parent_ctx: LevelContext, array: &ArrayRef) -> Result<Self> {
191 if !Self::types_compatible(field.data_type(), array.data_type()) {
192 return Err(arrow_err!(format!(
193 "Incompatible type. Field '{}' has type {}, array has type {}",
194 field.name(),
195 field.data_type(),
196 array.data_type(),
197 )));
198 }
199
200 let is_nullable = field.is_nullable();
201
202 match array.data_type() {
203 d if is_leaf(d) => {
204 let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone());
205 Ok(Self::Primitive(levels))
206 }
207 DataType::Dictionary(_, v) if is_leaf(v.as_ref()) => {
208 let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone());
209 Ok(Self::Primitive(levels))
210 }
211 DataType::RunEndEncoded(_, value_field) => {
212 let flat = expand_ree_array(array)?;
213 let flat_field = Field::new(
214 field.name(),
215 value_field.data_type().clone(),
216 field.is_nullable(),
217 );
218 Self::try_new(&flat_field, parent_ctx, &flat)
219 }
220 DataType::Struct(children) => {
221 let array = array.as_struct();
222 let def_level = match is_nullable {
223 true => parent_ctx.def_level + 1,
224 false => parent_ctx.def_level,
225 };
226
227 let ctx = LevelContext {
228 rep_level: parent_ctx.rep_level,
229 def_level,
230 };
231
232 let children = children
233 .iter()
234 .zip(array.columns())
235 .map(|(f, a)| Self::try_new(f, ctx, a))
236 .collect::<Result<_>>()?;
237
238 Ok(Self::Struct(children, ctx, array.nulls().cloned()))
239 }
240 DataType::List(child)
241 | DataType::LargeList(child)
242 | DataType::Map(child, _)
243 | DataType::FixedSizeList(child, _)
244 | DataType::ListView(child)
245 | DataType::LargeListView(child) => {
246 let def_level = match is_nullable {
247 true => parent_ctx.def_level + 2,
248 false => parent_ctx.def_level + 1,
249 };
250
251 let ctx = LevelContext {
252 rep_level: parent_ctx.rep_level + 1,
253 def_level,
254 };
255
256 Ok(match field.data_type() {
257 DataType::List(_) => {
258 let list = array.as_list();
259 let child = Self::try_new(child.as_ref(), ctx, list.values())?;
260 let is_last = child.child_has_no_nested_rep();
261 let offsets = list.offsets().clone();
262 Self::List(
263 Box::new(child),
264 ctx,
265 offsets,
266 list.nulls().cloned(),
267 is_last,
268 )
269 }
270 DataType::LargeList(_) => {
271 let list = array.as_list();
272 let child = Self::try_new(child.as_ref(), ctx, list.values())?;
273 let is_last = child.child_has_no_nested_rep();
274 let offsets = list.offsets().clone();
275 let nulls = list.nulls().cloned();
276 Self::LargeList(Box::new(child), ctx, offsets, nulls, is_last)
277 }
278 DataType::Map(_, _) => {
279 let map = array.as_map();
280 let entries = Arc::new(map.entries().clone()) as ArrayRef;
281 let child = Self::try_new(child.as_ref(), ctx, &entries)?;
282 let is_last = child.child_has_no_nested_rep();
283 let offsets = map.offsets().clone();
284 Self::List(Box::new(child), ctx, offsets, map.nulls().cloned(), is_last)
285 }
286 DataType::FixedSizeList(_, size) => {
287 let list = array.as_fixed_size_list();
288 let child = Self::try_new(child.as_ref(), ctx, list.values())?;
289 let nulls = list.nulls().cloned();
290 Self::FixedSizeList(Box::new(child), ctx, *size as _, nulls)
291 }
292 DataType::ListView(_) => {
293 let list = array.as_list_view();
294 let child = Self::try_new(child.as_ref(), ctx, list.values())?;
295 let offsets = list.offsets().clone();
296 let sizes = list.sizes().clone();
297 let nulls = list.nulls().cloned();
298 Self::ListView(Box::new(child), ctx, offsets, sizes, nulls)
299 }
300 DataType::LargeListView(_) => {
301 let list = array.as_list_view();
302 let child = Self::try_new(child.as_ref(), ctx, list.values())?;
303 let offsets = list.offsets().clone();
304 let sizes = list.sizes().clone();
305 let nulls = list.nulls().cloned();
306 Self::LargeListView(Box::new(child), ctx, offsets, sizes, nulls)
307 }
308 _ => unreachable!(),
309 })
310 }
311 d => Err(nyi_err!("Datatype {} is not yet supported", d)),
312 }
313 }
314
315 fn finish(self) -> Vec<ArrayLevels> {
318 match self {
319 LevelInfoBuilder::Primitive(v) => vec![v],
320 LevelInfoBuilder::List(v, _, _, _, _)
321 | LevelInfoBuilder::LargeList(v, _, _, _, _)
322 | LevelInfoBuilder::FixedSizeList(v, _, _, _)
323 | LevelInfoBuilder::ListView(v, _, _, _, _)
324 | LevelInfoBuilder::LargeListView(v, _, _, _, _) => v.finish(),
325 LevelInfoBuilder::Struct(v, _, _) => v.into_iter().flat_map(|l| l.finish()).collect(),
326 }
327 }
328
329 fn write(&mut self, range: Range<usize>) {
331 match self {
332 LevelInfoBuilder::Primitive(info) => Self::write_leaf(info, range),
333 LevelInfoBuilder::List(child, ctx, offsets, nulls, is_last) => {
334 Self::write_list(child, ctx, offsets, nulls.as_ref(), range, *is_last)
335 }
336 LevelInfoBuilder::LargeList(child, ctx, offsets, nulls, is_last) => {
337 Self::write_list(child, ctx, offsets, nulls.as_ref(), range, *is_last)
338 }
339 LevelInfoBuilder::FixedSizeList(child, ctx, size, nulls) => {
340 Self::write_fixed_size_list(child, ctx, *size, nulls.as_ref(), range)
341 }
342 LevelInfoBuilder::ListView(child, ctx, offsets, sizes, nulls) => {
343 Self::write_list_view(child, ctx, offsets, sizes, nulls.as_ref(), range)
344 }
345 LevelInfoBuilder::LargeListView(child, ctx, offsets, sizes, nulls) => {
346 Self::write_list_view(child, ctx, offsets, sizes, nulls.as_ref(), range)
347 }
348 LevelInfoBuilder::Struct(children, ctx, nulls) => {
349 Self::write_struct(children, ctx, nulls.as_ref(), range)
350 }
351 }
352 }
353
354 fn child_has_no_nested_rep(&self) -> bool {
358 match self {
359 LevelInfoBuilder::Primitive(_) => true,
360 LevelInfoBuilder::Struct(children, _, _) => {
361 children.iter().all(|c| c.child_has_no_nested_rep())
362 }
363 _ => false,
364 }
365 }
366
367 fn write_list<O: OffsetSizeTrait>(
371 child: &mut LevelInfoBuilder,
372 ctx: &LevelContext,
373 offsets: &[O],
374 nulls: Option<&NullBuffer>,
375 range: Range<usize>,
376 is_last_level: bool,
377 ) {
378 if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) {
380 let count = range.end - range.start;
381 child.visit_leaves(|leaf| {
382 leaf.extend_uniform_levels(ctx.def_level - 2, ctx.rep_level - 1, count);
383 });
384 return;
385 }
386
387 if is_last_level {
390 Self::write_list_direct(child, ctx, offsets, nulls, range);
391 } else {
392 Self::write_list_scan(child, ctx, offsets, nulls, range);
393 }
394 }
395
396 fn write_list_direct<O: OffsetSizeTrait>(
400 child: &mut LevelInfoBuilder,
401 ctx: &LevelContext,
402 offsets: &[O],
403 nulls: Option<&NullBuffer>,
404 range: Range<usize>,
405 ) {
406 let list_start_rep = ctx.rep_level - 1;
407
408 let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets: &[O]| {
409 debug_assert!(run_offsets.len() >= 2);
410 let values_start = run_offsets[0].as_usize();
411 let values_end = run_offsets[run_offsets.len() - 1].as_usize();
412 debug_assert!(values_end > values_start);
413
414 child.write(values_start..values_end);
415
416 child.visit_leaves(|leaf| {
421 debug_assert_eq!(leaf.max_rep_level, ctx.rep_level);
422 let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
423 let batch_len = values_end - values_start;
424 let batch_base = rep_levels.len() - batch_len;
425 for slot_offset in run_offsets.iter().take(run_offsets.len() - 1) {
426 let pos = batch_base + (slot_offset.as_usize() - values_start);
427 rep_levels[pos] = list_start_rep;
428 }
429 });
430 };
431
432 Self::write_list_impl(child, ctx, offsets, nulls, range, emit_non_empty_run);
433 }
434
435 fn write_list_scan<O: OffsetSizeTrait>(
442 child: &mut LevelInfoBuilder,
443 ctx: &LevelContext,
444 offsets: &[O],
445 nulls: Option<&NullBuffer>,
446 range: Range<usize>,
447 ) {
448 let list_start_rep = ctx.rep_level - 1;
449
450 let emit_non_empty_run = |child: &mut LevelInfoBuilder, run_offsets: &[O]| {
451 debug_assert!(run_offsets.len() >= 2);
452 let values_start = run_offsets[0].as_usize();
453 let values_end = run_offsets[run_offsets.len() - 1].as_usize();
454 debug_assert!(values_end > values_start);
455
456 child.write(values_start..values_end);
457
458 child.visit_leaves(|leaf| {
459 let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
460
461 if leaf.max_rep_level == ctx.rep_level {
462 let batch_len = values_end - values_start;
466 let batch_base = rep_levels.len() - batch_len;
467 for slot_offset in run_offsets.iter().take(run_offsets.len() - 1) {
468 let pos = batch_base + (slot_offset.as_usize() - values_start);
469 rep_levels[pos] = list_start_rep;
470 }
471 } else {
472 let mut slot_bounds = run_offsets[..run_offsets.len() - 1].iter().rev();
475 let mut next_stamp_at = values_end - slot_bounds.next().unwrap().as_usize();
476 let mut seen = 0usize;
477
478 for rep in rep_levels.iter_mut().rev() {
479 debug_assert!(*rep >= ctx.rep_level);
482 if *rep <= ctx.rep_level {
488 seen += 1;
489 if seen == next_stamp_at {
490 *rep = list_start_rep;
491 match slot_bounds.next() {
492 Some(offset) => next_stamp_at = values_end - offset.as_usize(),
493 None => break,
494 }
495 }
496 }
497 }
498 }
499 });
500 };
501
502 Self::write_list_impl(child, ctx, offsets, nulls, range, emit_non_empty_run);
503 }
504
505 fn write_list_impl<O: OffsetSizeTrait>(
509 child: &mut LevelInfoBuilder,
510 ctx: &LevelContext,
511 offsets: &[O],
512 nulls: Option<&NullBuffer>,
513 range: Range<usize>,
514 mut emit_non_empty_run: impl FnMut(&mut LevelInfoBuilder, &[O]),
515 ) {
516 let null_offset = range.start;
517 let offsets = &offsets[range.start..range.end + 1];
518 let list_start_rep = ctx.rep_level - 1;
519
520 let emit_nulls = |child: &mut LevelInfoBuilder, count: usize| {
521 child.visit_leaves(|leaf| {
522 leaf.append_rep_level_run(list_start_rep, count);
523 leaf.append_def_level_run(ctx.def_level - 2, count);
524 });
525 };
526
527 let emit_empties = |child: &mut LevelInfoBuilder, count: usize| {
528 child.visit_leaves(|leaf| {
529 leaf.append_rep_level_run(list_start_rep, count);
530 leaf.append_def_level_run(ctx.def_level - 1, count);
531 });
532 };
533
534 #[derive(Clone, Copy, PartialEq)]
536 enum SlotKind {
537 Null,
538 Empty,
539 NonEmpty,
540 }
541
542 let num_slots = offsets.len() - 1;
543 if num_slots == 0 {
544 return;
545 }
546
547 macro_rules! classify {
548 ($i:expr, $nulls:expr) => {
549 if !$nulls.is_valid($i + null_offset) {
550 SlotKind::Null
551 } else if offsets[$i] == offsets[$i + 1] {
552 SlotKind::Empty
553 } else {
554 SlotKind::NonEmpty
555 }
556 };
557 }
558
559 macro_rules! flush_run {
560 ($kind:expr, $start:expr, $end:expr) => {
561 match $kind {
562 SlotKind::Null => emit_nulls(child, $end - $start),
563 SlotKind::Empty => emit_empties(child, $end - $start),
564 SlotKind::NonEmpty => emit_non_empty_run(child, &offsets[$start..$end + 1]),
565 }
566 };
567 }
568
569 match nulls {
570 Some(nulls) if nulls.null_count() > 0 => {
573 let mut run_kind = classify!(0, nulls);
574 let mut run_start: usize = 0;
575 for i in 1..num_slots {
576 let kind = classify!(i, nulls);
577 if kind != run_kind {
578 flush_run!(run_kind, run_start, i);
579 run_kind = kind;
580 run_start = i;
581 }
582 }
583 flush_run!(run_kind, run_start, num_slots);
584 }
585 _ => {
586 let mut run_kind = if offsets[0] == offsets[1] {
587 SlotKind::Empty
588 } else {
589 SlotKind::NonEmpty
590 };
591 let mut run_start: usize = 0;
592 for i in 1..num_slots {
593 let kind = if offsets[i] == offsets[i + 1] {
594 SlotKind::Empty
595 } else {
596 SlotKind::NonEmpty
597 };
598 if kind != run_kind {
599 flush_run!(run_kind, run_start, i);
600 run_kind = kind;
601 run_start = i;
602 }
603 }
604 flush_run!(run_kind, run_start, num_slots);
605 }
606 }
607 }
608
609 fn write_list_view<O: OffsetSizeTrait>(
611 child: &mut LevelInfoBuilder,
612 ctx: &LevelContext,
613 offsets: &[O],
614 sizes: &[O],
615 nulls: Option<&NullBuffer>,
616 range: Range<usize>,
617 ) {
618 let offsets = &offsets[range.clone()];
619 let sizes = &sizes[range.clone()];
620
621 let write_non_null_slice =
622 |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| {
623 child.write(start_idx..end_idx);
624 child.visit_leaves(|leaf| {
625 let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
626 let mut rev = rep_levels.iter_mut().rev();
627 let mut remaining = end_idx - start_idx;
628
629 loop {
630 let next = rev.next().unwrap();
631 if *next > ctx.rep_level {
632 continue;
634 }
635
636 remaining -= 1;
637 if remaining == 0 {
638 *next = ctx.rep_level - 1;
639 break;
640 }
641 }
642 })
643 };
644
645 let write_empty_slice = |child: &mut LevelInfoBuilder| {
646 child.visit_leaves(|leaf| {
647 leaf.append_rep_level_run(ctx.rep_level - 1, 1);
648 leaf.append_def_level_run(ctx.def_level - 1, 1);
649 })
650 };
651
652 let write_null_slice = |child: &mut LevelInfoBuilder| {
653 child.visit_leaves(|leaf| {
654 leaf.append_rep_level_run(ctx.rep_level - 1, 1);
655 leaf.append_def_level_run(ctx.def_level - 2, 1);
656 })
657 };
658
659 match nulls {
660 Some(nulls) => {
661 let null_offset = range.start;
662 for (idx, (offset, size)) in offsets.iter().zip(sizes.iter()).enumerate() {
664 let is_valid = nulls.is_valid(idx + null_offset);
665 let start_idx = offset.as_usize();
666 let size = size.as_usize();
667 let end_idx = start_idx + size;
668 if !is_valid {
669 write_null_slice(child)
670 } else if size == 0 {
671 write_empty_slice(child)
672 } else {
673 write_non_null_slice(child, start_idx, end_idx)
674 }
675 }
676 }
677 None => {
678 for (offset, size) in offsets.iter().zip(sizes.iter()) {
679 let start_idx = offset.as_usize();
680 let size = size.as_usize();
681 let end_idx = start_idx + size;
682 if size == 0 {
683 write_empty_slice(child)
684 } else {
685 write_non_null_slice(child, start_idx, end_idx)
686 }
687 }
688 }
689 }
690 }
691
692 fn write_struct(
694 children: &mut [LevelInfoBuilder],
695 ctx: &LevelContext,
696 nulls: Option<&NullBuffer>,
697 range: Range<usize>,
698 ) {
699 let write_null = |children: &mut [LevelInfoBuilder], range: Range<usize>| {
700 let len = range.end - range.start;
701 for child in children {
702 child.visit_leaves(|info| {
703 info.extend_uniform_levels(ctx.def_level - 1, ctx.rep_level, len);
704 })
705 }
706 };
707
708 if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) {
710 write_null(children, range);
711 return;
712 }
713
714 let write_non_null = |children: &mut [LevelInfoBuilder], range: Range<usize>| {
715 for child in children {
716 child.write(range.clone())
717 }
718 };
719
720 match nulls {
721 Some(validity) => {
722 let mut last_non_null_idx = None;
723 let mut last_null_idx = None;
724
725 for i in range.clone() {
727 match validity.is_valid(i) {
728 true => {
729 if let Some(last_idx) = last_null_idx.take() {
730 write_null(children, last_idx..i)
731 }
732 last_non_null_idx.get_or_insert(i);
733 }
734 false => {
735 if let Some(last_idx) = last_non_null_idx.take() {
736 write_non_null(children, last_idx..i)
737 }
738 last_null_idx.get_or_insert(i);
739 }
740 }
741 }
742
743 if let Some(last_idx) = last_null_idx.take() {
744 write_null(children, last_idx..range.end)
745 }
746
747 if let Some(last_idx) = last_non_null_idx.take() {
748 write_non_null(children, last_idx..range.end)
749 }
750 }
751 None => write_non_null(children, range),
752 }
753 }
754
755 fn write_fixed_size_list(
757 child: &mut LevelInfoBuilder,
758 ctx: &LevelContext,
759 fixed_size: usize,
760 nulls: Option<&NullBuffer>,
761 range: Range<usize>,
762 ) {
763 if nulls.is_some_and(|nulls| nulls.null_count() == nulls.len()) {
765 let count = range.end - range.start;
766 child.visit_leaves(|leaf| {
767 leaf.extend_uniform_levels(ctx.def_level - 2, ctx.rep_level - 1, count);
768 });
769 return;
770 }
771
772 let write_non_null = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| {
773 let values_start = start_idx * fixed_size;
774 let values_end = end_idx * fixed_size;
775 child.write(values_start..values_end);
776
777 child.visit_leaves(|leaf| {
778 let rep_levels = leaf.rep_levels.materialize_mut().unwrap();
779
780 let row_indices = (0..fixed_size)
781 .rev()
782 .cycle()
783 .take(values_end - values_start);
784
785 rep_levels
787 .iter_mut()
788 .rev()
789 .filter(|&&mut r| r == ctx.rep_level)
791 .zip(row_indices)
792 .for_each(|(r, idx)| {
793 if idx == 0 {
794 *r = ctx.rep_level - 1;
795 }
796 });
797 })
798 };
799
800 let write_empty = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| {
802 let len = end_idx - start_idx;
803 child.visit_leaves(|leaf| {
804 leaf.append_rep_level_run(ctx.rep_level - 1, len);
805 leaf.append_def_level_run(ctx.def_level - 1, len);
806 })
807 };
808
809 let write_rows = |child: &mut LevelInfoBuilder, start_idx: usize, end_idx: usize| {
810 if fixed_size > 0 {
811 write_non_null(child, start_idx, end_idx)
812 } else {
813 write_empty(child, start_idx, end_idx)
814 }
815 };
816
817 match nulls {
818 Some(nulls) => {
819 let mut start_idx = None;
820 for idx in range.clone() {
821 if nulls.is_valid(idx) {
822 start_idx.get_or_insert(idx);
824 } else {
825 if let Some(start) = start_idx.take() {
827 write_rows(child, start, idx);
828 }
829 child.visit_leaves(|leaf| {
831 leaf.append_rep_level_run(ctx.rep_level - 1, 1);
832 leaf.append_def_level_run(ctx.def_level - 2, 1);
833 })
834 }
835 }
836 if let Some(start) = start_idx.take() {
838 write_rows(child, start, range.end);
839 }
840 }
841 None => write_rows(child, range.start, range.end),
843 }
844 }
845
846 fn write_leaf(info: &mut ArrayLevels, range: Range<usize>) {
848 let len = range.end - range.start;
849
850 if let Some(nulls) = &info.logical_nulls
852 && !matches!(info.def_levels, LevelData::Absent)
853 && nulls.null_count() == nulls.len()
854 {
855 info.extend_uniform_levels(info.max_def_level - 1, info.max_rep_level, len);
856 return;
857 }
858
859 if matches!(info.def_levels, LevelData::Absent) {
860 info.non_null_indices.extend(range.clone());
861 } else {
862 let max_def_level = info.max_def_level;
863 match &info.logical_nulls {
864 Some(nulls) => {
865 assert!(range.end <= nulls.len());
866 if len >= BULK_FILL_MIN_LEN && nulls.null_count() * 2 >= nulls.len() {
871 let range_nulls = nulls.slice(range.start, len);
872 let valid_in_range = len - range_nulls.null_count();
873 let null_def_level = max_def_level - 1;
874 let buf = info
875 .def_levels
876 .materialize_mut()
877 .expect("definition levels present");
878 let base = buf.len();
879 buf.resize(base + len, null_def_level);
880 for i in range_nulls.valid_indices() {
881 buf[base + i] = max_def_level;
882 }
883 info.non_null_indices.reserve(valid_in_range);
884 info.non_null_indices
885 .extend(range_nulls.valid_indices().map(|i| i + range.start));
886 } else {
887 let bits = nulls.inner();
888 info.def_levels.extend_from_iter(range.clone().map(|i| {
889 let valid = unsafe { bits.value_unchecked(i) };
891 max_def_level - (!valid as i16)
892 }));
893 info.non_null_indices.reserve(len);
894 info.non_null_indices.extend(
895 BitIndexIterator::new(bits.inner(), bits.offset() + range.start, len)
896 .map(|i| i + range.start),
897 );
898 }
899 }
900 None => {
901 info.append_def_level_run(max_def_level, len);
902 info.non_null_indices.reserve(len);
903 info.non_null_indices.extend(range.clone());
904 }
905 }
906 }
907
908 if !matches!(info.rep_levels, LevelData::Absent) {
909 info.append_rep_level_run(info.max_rep_level, len);
910 }
911 }
912
913 fn visit_leaves(&mut self, visit: impl Fn(&mut ArrayLevels) + Copy) {
915 match self {
916 LevelInfoBuilder::Primitive(info) => visit(info),
917 LevelInfoBuilder::List(c, _, _, _, _)
918 | LevelInfoBuilder::LargeList(c, _, _, _, _)
919 | LevelInfoBuilder::FixedSizeList(c, _, _, _)
920 | LevelInfoBuilder::ListView(c, _, _, _, _)
921 | LevelInfoBuilder::LargeListView(c, _, _, _, _) => c.visit_leaves(visit),
922 LevelInfoBuilder::Struct(children, _, _) => {
923 for c in children {
924 c.visit_leaves(visit)
925 }
926 }
927 }
928 }
929
930 fn types_compatible(a: &DataType, b: &DataType) -> bool {
936 if a.equals_datatype(b) {
938 return true;
939 }
940
941 let (a, b) = match (a, b) {
943 (DataType::Dictionary(_, va), DataType::Dictionary(_, vb)) => {
944 (va.as_ref(), vb.as_ref())
945 }
946 (DataType::Dictionary(_, v), b) => (v.as_ref(), b),
947 (a, DataType::Dictionary(_, v)) => (a, v.as_ref()),
948 _ => (a, b),
949 };
950
951 if a == b {
954 return true;
955 }
956
957 match a {
960 DataType::Utf8 => matches!(b, DataType::LargeUtf8 | DataType::Utf8View),
962 DataType::Utf8View => matches!(b, DataType::LargeUtf8 | DataType::Utf8),
963 DataType::LargeUtf8 => matches!(b, DataType::Utf8 | DataType::Utf8View),
964
965 DataType::Binary => matches!(b, DataType::LargeBinary | DataType::BinaryView),
967 DataType::BinaryView => matches!(b, DataType::LargeBinary | DataType::Binary),
968 DataType::LargeBinary => matches!(b, DataType::Binary | DataType::BinaryView),
969
970 _ => false,
972 }
973 }
974}
975
976#[derive(Debug, Clone)]
979pub(crate) enum LevelData {
980 Absent,
981 Materialized(Vec<i16>),
982 Uniform { value: i16, count: usize },
983}
984
985impl PartialEq for LevelData {
988 fn eq(&self, other: &Self) -> bool {
989 match (self, other) {
990 (Self::Absent, Self::Absent) => true,
991 (Self::Materialized(a), Self::Materialized(b)) => a == b,
992 (Self::Uniform { value: v, count: n }, Self::Materialized(b))
993 | (Self::Materialized(b), Self::Uniform { value: v, count: n }) => {
994 b.len() == *n && b.iter().all(|x| x == v)
995 }
996 (
997 Self::Uniform {
998 value: v1,
999 count: n1,
1000 },
1001 Self::Uniform {
1002 value: v2,
1003 count: n2,
1004 },
1005 ) => v1 == v2 && n1 == n2,
1006 _ => false,
1007 }
1008 }
1009}
1010
1011impl Eq for LevelData {}
1012
1013impl LevelData {
1014 fn new(present: bool) -> Self {
1015 match present {
1016 true => Self::Materialized(Vec::new()),
1017 false => Self::Absent,
1018 }
1019 }
1020
1021 pub(crate) fn as_ref(&self) -> LevelDataRef<'_> {
1022 match self {
1023 Self::Absent => LevelDataRef::Absent,
1024 Self::Materialized(values) => LevelDataRef::Materialized(values),
1025 Self::Uniform { value, count } => LevelDataRef::Uniform {
1026 value: *value,
1027 count: *count,
1028 },
1029 }
1030 }
1031
1032 pub(crate) fn slice(&self, offset: usize, len: usize) -> Self {
1033 match self {
1034 Self::Absent => Self::Absent,
1035 Self::Materialized(values) => Self::Materialized(values[offset..offset + len].to_vec()),
1036 Self::Uniform { value, .. } => Self::Uniform {
1037 value: *value,
1038 count: len,
1039 },
1040 }
1041 }
1042
1043 fn append_run(&mut self, value: i16, count: usize) {
1044 if count == 0 {
1045 return;
1046 }
1047
1048 match self {
1049 Self::Absent => {}
1052 Self::Materialized(values) if values.is_empty() => {
1055 *self = Self::Uniform { value, count };
1056 }
1057 Self::Materialized(values) => values.extend(std::iter::repeat_n(value, count)),
1059 Self::Uniform {
1062 value: uniform_value,
1063 count: uniform_count,
1064 } if *uniform_value == value => {
1065 *uniform_count += count;
1066 }
1067 Self::Uniform { .. } => {
1070 let values = self.materialize_mut().unwrap();
1071 values.extend(std::iter::repeat_n(value, count));
1072 }
1073 }
1074 }
1075
1076 fn extend_from_iter<I>(&mut self, iter: I)
1077 where
1078 I: IntoIterator<Item = i16>,
1079 {
1080 if let Some(values) = self.materialize_mut() {
1081 values.extend(iter);
1082 }
1083 }
1084
1085 fn materialize_mut(&mut self) -> Option<&mut Vec<i16>> {
1088 match self {
1089 Self::Absent => None,
1090 Self::Materialized(values) => Some(values),
1091 Self::Uniform { value, count } => {
1092 let values = vec![*value; *count];
1093 *self = Self::Materialized(values);
1094 match self {
1095 Self::Materialized(values) => Some(values),
1096 _ => unreachable!(),
1097 }
1098 }
1099 }
1100 }
1101}
1102
1103#[derive(Debug, Clone)]
1104pub(crate) struct ArrayLevels {
1105 def_levels: LevelData,
1109
1110 rep_levels: LevelData,
1114
1115 non_null_indices: Vec<usize>,
1118
1119 max_def_level: i16,
1121
1122 max_rep_level: i16,
1124
1125 array: ArrayRef,
1127
1128 logical_nulls: Option<NullBuffer>,
1130}
1131
1132impl PartialEq for ArrayLevels {
1133 fn eq(&self, other: &Self) -> bool {
1134 self.def_levels == other.def_levels
1135 && self.rep_levels == other.rep_levels
1136 && self.non_null_indices == other.non_null_indices
1137 && self.max_def_level == other.max_def_level
1138 && self.max_rep_level == other.max_rep_level
1139 && self.array.as_ref() == other.array.as_ref()
1140 && self.logical_nulls.as_ref() == other.logical_nulls.as_ref()
1141 }
1142}
1143impl Eq for ArrayLevels {}
1144
1145impl ArrayLevels {
1146 fn new(ctx: LevelContext, is_nullable: bool, array: ArrayRef) -> Self {
1147 let max_rep_level = ctx.rep_level;
1148 let max_def_level = match is_nullable {
1149 true => ctx.def_level + 1,
1150 false => ctx.def_level,
1151 };
1152
1153 let logical_nulls = array.logical_nulls();
1154
1155 Self {
1156 def_levels: LevelData::new(max_def_level != 0),
1157 rep_levels: LevelData::new(max_rep_level != 0),
1158 non_null_indices: vec![],
1159 max_def_level,
1160 max_rep_level,
1161 array,
1162 logical_nulls,
1163 }
1164 }
1165
1166 pub fn array(&self) -> &ArrayRef {
1167 &self.array
1168 }
1169
1170 pub(crate) fn def_level_data(&self) -> &LevelData {
1171 &self.def_levels
1172 }
1173
1174 pub(crate) fn rep_level_data(&self) -> &LevelData {
1175 &self.rep_levels
1176 }
1177
1178 pub fn non_null_indices(&self) -> &[usize] {
1179 &self.non_null_indices
1180 }
1181
1182 pub(crate) fn slice_for_chunk(&self, chunk: &CdcChunk) -> Self {
1188 let def_levels = self.def_levels.slice(chunk.level_offset, chunk.num_levels);
1189 let rep_levels = self.rep_levels.slice(chunk.level_offset, chunk.num_levels);
1190
1191 let nni = &self.non_null_indices[chunk.value_offset..chunk.value_offset + chunk.num_values];
1193 let start = nni.first().copied().unwrap_or(0);
1198 let end = nni.last().map_or(0, |&i| i + 1);
1199 let non_null_indices = nni.iter().map(|&idx| idx - start).collect();
1201 let array = self.array.slice(start, end - start);
1203 let logical_nulls = array.logical_nulls();
1204
1205 Self {
1206 def_levels,
1207 rep_levels,
1208 non_null_indices,
1209 max_def_level: self.max_def_level,
1210 max_rep_level: self.max_rep_level,
1211 array,
1212 logical_nulls,
1213 }
1214 }
1215
1216 fn extend_uniform_levels(&mut self, def_val: i16, rep_val: i16, count: usize) {
1218 self.def_levels.append_run(def_val, count);
1219 self.rep_levels.append_run(rep_val, count);
1220 }
1221
1222 fn append_def_level_run(&mut self, value: i16, count: usize) {
1223 self.def_levels.append_run(value, count);
1224 }
1225
1226 fn append_rep_level_run(&mut self, value: i16, count: usize) {
1227 self.rep_levels.append_run(value, count);
1228 }
1229}
1230
1231#[cfg(test)]
1232mod tests {
1233 use super::*;
1234 use crate::column::chunker::CdcChunk;
1235
1236 use arrow_array::builder::*;
1237 use arrow_array::types::Int32Type;
1238 use arrow_array::*;
1239 use arrow_buffer::{Buffer, ToByteSlice};
1240 use arrow_cast::display::array_value_to_string;
1241 use arrow_data::{ArrayData, ArrayDataBuilder};
1242 use arrow_schema::{Fields, Schema};
1243
1244 #[test]
1245 fn test_calculate_array_levels_twitter_example() {
1246 let leaf_type = Field::new_list_field(DataType::Int32, false);
1250 let inner_type = DataType::List(Arc::new(leaf_type));
1251 let inner_field = Field::new("l2", inner_type.clone(), false);
1252 let outer_type = DataType::List(Arc::new(inner_field));
1253 let outer_field = Field::new("l1", outer_type.clone(), false);
1254
1255 let primitives = Int32Array::from_iter(0..10);
1256
1257 let offsets = Buffer::from_iter([0_i32, 3, 7, 8, 10]);
1259 let inner_list = ArrayDataBuilder::new(inner_type)
1260 .len(4)
1261 .add_buffer(offsets)
1262 .add_child_data(primitives.to_data())
1263 .build()
1264 .unwrap();
1265
1266 let offsets = Buffer::from_iter([0_i32, 2, 4]);
1267 let outer_list = ArrayDataBuilder::new(outer_type)
1268 .len(2)
1269 .add_buffer(offsets)
1270 .add_child_data(inner_list)
1271 .build()
1272 .unwrap();
1273 let outer_list = make_array(outer_list);
1274
1275 let levels = calculate_array_levels(&outer_list, &outer_field).unwrap();
1276 assert_eq!(levels.len(), 1);
1277
1278 let expected = ArrayLevels {
1279 def_levels: LevelData::Materialized(vec![2; 10]),
1280 rep_levels: LevelData::Materialized(vec![0, 2, 2, 1, 2, 2, 2, 0, 1, 2]),
1281 non_null_indices: vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
1282 max_def_level: 2,
1283 max_rep_level: 2,
1284 array: Arc::new(primitives),
1285 logical_nulls: None,
1286 };
1287 assert_eq!(&levels[0], &expected);
1288 }
1289
1290 #[test]
1291 fn test_calculate_one_level_1() {
1292 let array = Arc::new(Int32Array::from_iter(0..10)) as ArrayRef;
1294 let field = Field::new_list_field(DataType::Int32, false);
1295
1296 let levels = calculate_array_levels(&array, &field).unwrap();
1297 assert_eq!(levels.len(), 1);
1298
1299 let expected_levels = ArrayLevels {
1300 def_levels: LevelData::Absent,
1301 rep_levels: LevelData::Absent,
1302 non_null_indices: (0..10).collect(),
1303 max_def_level: 0,
1304 max_rep_level: 0,
1305 array,
1306 logical_nulls: None,
1307 };
1308 assert_eq!(&levels[0], &expected_levels);
1309 }
1310
1311 #[test]
1312 fn test_calculate_one_level_2() {
1313 let array = Arc::new(Int32Array::from_iter([
1315 Some(0),
1316 None,
1317 Some(0),
1318 Some(0),
1319 None,
1320 ])) as ArrayRef;
1321 let field = Field::new_list_field(DataType::Int32, true);
1322
1323 let levels = calculate_array_levels(&array, &field).unwrap();
1324 assert_eq!(levels.len(), 1);
1325
1326 let logical_nulls = array.logical_nulls();
1327 let expected_levels = ArrayLevels {
1328 def_levels: LevelData::Materialized(vec![1, 0, 1, 1, 0]),
1329 rep_levels: LevelData::Absent,
1330 non_null_indices: vec![0, 2, 3],
1331 max_def_level: 1,
1332 max_rep_level: 0,
1333 array,
1334 logical_nulls,
1335 };
1336 assert_eq!(&levels[0], &expected_levels);
1337 }
1338
1339 #[test]
1340 fn test_calculate_array_levels_1() {
1341 let leaf_field = Field::new_list_field(DataType::Int32, false);
1342 let list_type = DataType::List(Arc::new(leaf_field));
1343
1344 let leaf_array = Int32Array::from_iter(0..5);
1348 let offsets = Buffer::from_iter(0_i32..6);
1350 let list = ArrayDataBuilder::new(list_type.clone())
1351 .len(5)
1352 .add_buffer(offsets)
1353 .add_child_data(leaf_array.to_data())
1354 .build()
1355 .unwrap();
1356 let list = make_array(list);
1357
1358 let list_field = Field::new("list", list_type.clone(), false);
1359 let levels = calculate_array_levels(&list, &list_field).unwrap();
1360 assert_eq!(levels.len(), 1);
1361
1362 let expected_levels = ArrayLevels {
1363 def_levels: LevelData::Materialized(vec![1; 5]),
1364 rep_levels: LevelData::Materialized(vec![0; 5]),
1365 non_null_indices: (0..5).collect(),
1366 max_def_level: 1,
1367 max_rep_level: 1,
1368 array: Arc::new(leaf_array),
1369 logical_nulls: None,
1370 };
1371 assert_eq!(&levels[0], &expected_levels);
1372
1373 let leaf_array = Int32Array::from_iter([0, 0, 2, 2, 3, 3, 3, 3, 4, 4, 4]);
1382 let offsets = Buffer::from_iter([0_i32, 2, 2, 4, 8, 11]);
1383 let list = ArrayDataBuilder::new(list_type.clone())
1384 .len(5)
1385 .add_buffer(offsets)
1386 .add_child_data(leaf_array.to_data())
1387 .null_bit_buffer(Some(Buffer::from([0b00011101])))
1388 .build()
1389 .unwrap();
1390 let list = make_array(list);
1391
1392 let list_field = Field::new("list", list_type, true);
1393 let levels = calculate_array_levels(&list, &list_field).unwrap();
1394 assert_eq!(levels.len(), 1);
1395
1396 let expected_levels = ArrayLevels {
1397 def_levels: LevelData::Materialized(vec![2, 2, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2]),
1398 rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1]),
1399 non_null_indices: (0..11).collect(),
1400 max_def_level: 2,
1401 max_rep_level: 1,
1402 array: Arc::new(leaf_array),
1403 logical_nulls: None,
1404 };
1405 assert_eq!(&levels[0], &expected_levels);
1406 }
1407
1408 #[test]
1409 fn test_calculate_array_levels_2() {
1410 let leaf = Int32Array::from_iter(0..11);
1424 let leaf_field = Field::new("leaf", DataType::Int32, false);
1425
1426 let list_type = DataType::List(Arc::new(leaf_field));
1427 let list = ArrayData::builder(list_type.clone())
1428 .len(5)
1429 .add_child_data(leaf.to_data())
1430 .add_buffer(Buffer::from_iter([0_i32, 2, 2, 4, 8, 11]))
1431 .build()
1432 .unwrap();
1433
1434 let list = make_array(list);
1435 let list_field = Arc::new(Field::new("list", list_type, true));
1436
1437 let struct_array =
1438 StructArray::from((vec![(list_field, list)], Buffer::from([0b00011010])));
1439 let array = Arc::new(struct_array) as ArrayRef;
1440
1441 let struct_field = Field::new("struct", array.data_type().clone(), true);
1442
1443 let levels = calculate_array_levels(&array, &struct_field).unwrap();
1444 assert_eq!(levels.len(), 1);
1445
1446 let expected_levels = ArrayLevels {
1447 def_levels: LevelData::Materialized(vec![0, 2, 0, 3, 3, 3, 3, 3, 3, 3]),
1448 rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 1, 1, 1, 0, 1, 1]),
1449 non_null_indices: (4..11).collect(),
1450 max_def_level: 3,
1451 max_rep_level: 1,
1452 array: Arc::new(leaf),
1453 logical_nulls: None,
1454 };
1455
1456 assert_eq!(&levels[0], &expected_levels);
1457
1458 let leaf = Int32Array::from_iter(100..122);
1467 let leaf_field = Field::new("leaf", DataType::Int32, true);
1468
1469 let l1_type = DataType::List(Arc::new(leaf_field));
1470 let offsets = Buffer::from_iter([0_i32, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]);
1471 let l1 = ArrayData::builder(l1_type.clone())
1472 .len(11)
1473 .add_child_data(leaf.to_data())
1474 .add_buffer(offsets)
1475 .build()
1476 .unwrap();
1477
1478 let l1_field = Field::new("l1", l1_type, true);
1479 let l2_type = DataType::List(Arc::new(l1_field));
1480 let l2 = ArrayData::builder(l2_type)
1481 .len(5)
1482 .add_child_data(l1)
1483 .add_buffer(Buffer::from_iter([0, 2, 2, 4, 8, 11]))
1484 .build()
1485 .unwrap();
1486
1487 let l2 = make_array(l2);
1488 let l2_field = Field::new("l2", l2.data_type().clone(), true);
1489
1490 let levels = calculate_array_levels(&l2, &l2_field).unwrap();
1491 assert_eq!(levels.len(), 1);
1492
1493 let expected_levels = ArrayLevels {
1494 def_levels: LevelData::Materialized(vec![
1495 5, 5, 5, 5, 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
1496 ]),
1497 rep_levels: LevelData::Materialized(vec![
1498 0, 2, 1, 2, 0, 0, 2, 1, 2, 0, 2, 1, 2, 1, 2, 1, 2, 0, 2, 1, 2, 1, 2,
1499 ]),
1500 non_null_indices: (0..22).collect(),
1501 max_def_level: 5,
1502 max_rep_level: 2,
1503 array: Arc::new(leaf),
1504 logical_nulls: None,
1505 };
1506
1507 assert_eq!(&levels[0], &expected_levels);
1508 }
1509
1510 #[test]
1511 fn test_calculate_array_levels_nested_list() {
1512 let leaf_field = Field::new("leaf", DataType::Int32, false);
1513 let list_type = DataType::List(Arc::new(leaf_field));
1514
1515 let leaf = Int32Array::from_iter([0; 4]);
1523 let list = ArrayData::builder(list_type.clone())
1524 .len(4)
1525 .add_buffer(Buffer::from_iter(0_i32..5))
1526 .add_child_data(leaf.to_data())
1527 .build()
1528 .unwrap();
1529 let list = make_array(list);
1530
1531 let list_field = Field::new("list", list_type.clone(), false);
1532 let levels = calculate_array_levels(&list, &list_field).unwrap();
1533 assert_eq!(levels.len(), 1);
1534
1535 let expected_levels = ArrayLevels {
1536 def_levels: LevelData::Materialized(vec![1; 4]),
1537 rep_levels: LevelData::Materialized(vec![0; 4]),
1538 non_null_indices: (0..4).collect(),
1539 max_def_level: 1,
1540 max_rep_level: 1,
1541 array: Arc::new(leaf),
1542 logical_nulls: None,
1543 };
1544 assert_eq!(&levels[0], &expected_levels);
1545
1546 let leaf = Int32Array::from_iter(0..8);
1551 let list = ArrayData::builder(list_type.clone())
1552 .len(4)
1553 .add_buffer(Buffer::from_iter([0_i32, 0, 3, 5, 7]))
1554 .null_bit_buffer(Some(Buffer::from([0b00001110])))
1555 .add_child_data(leaf.to_data())
1556 .build()
1557 .unwrap();
1558 let list = make_array(list);
1559 let list_field = Arc::new(Field::new("list", list_type, true));
1560
1561 let struct_array = StructArray::from(vec![(list_field, list)]);
1562 let array = Arc::new(struct_array) as ArrayRef;
1563
1564 let struct_field = Field::new("struct", array.data_type().clone(), true);
1565 let levels = calculate_array_levels(&array, &struct_field).unwrap();
1566 assert_eq!(levels.len(), 1);
1567
1568 let expected_levels = ArrayLevels {
1569 def_levels: LevelData::Materialized(vec![1, 3, 3, 3, 3, 3, 3, 3]),
1570 rep_levels: LevelData::Materialized(vec![0, 0, 1, 1, 0, 1, 0, 1]),
1571 non_null_indices: (0..7).collect(),
1572 max_def_level: 3,
1573 max_rep_level: 1,
1574 array: Arc::new(leaf),
1575 logical_nulls: None,
1576 };
1577 assert_eq!(&levels[0], &expected_levels);
1578
1579 let leaf = Int32Array::from_iter(201..216);
1587 let leaf_field = Field::new("leaf", DataType::Int32, false);
1588 let list_1_type = DataType::List(Arc::new(leaf_field));
1589 let list_1 = ArrayData::builder(list_1_type.clone())
1590 .len(7)
1591 .add_buffer(Buffer::from_iter([0_i32, 1, 3, 3, 6, 10, 10, 15]))
1592 .add_child_data(leaf.to_data())
1593 .build()
1594 .unwrap();
1595
1596 let list_1_field = Field::new("l1", list_1_type, true);
1597 let list_2_type = DataType::List(Arc::new(list_1_field));
1598 let list_2 = ArrayData::builder(list_2_type.clone())
1599 .len(4)
1600 .add_buffer(Buffer::from_iter([0_i32, 0, 3, 5, 7]))
1601 .null_bit_buffer(Some(Buffer::from([0b00001110])))
1602 .add_child_data(list_1)
1603 .build()
1604 .unwrap();
1605
1606 let list_2 = make_array(list_2);
1607 let list_2_field = Arc::new(Field::new("list_2", list_2_type, true));
1608
1609 let struct_array =
1610 StructArray::from((vec![(list_2_field, list_2)], Buffer::from([0b00001111])));
1611 let struct_field = Field::new("struct", struct_array.data_type().clone(), true);
1612
1613 let array = Arc::new(struct_array) as ArrayRef;
1614 let levels = calculate_array_levels(&array, &struct_field).unwrap();
1615 assert_eq!(levels.len(), 1);
1616
1617 let expected_levels = ArrayLevels {
1618 def_levels: LevelData::Materialized(vec![
1619 1, 5, 5, 5, 4, 5, 5, 5, 5, 5, 5, 5, 4, 5, 5, 5, 5, 5,
1620 ]),
1621 rep_levels: LevelData::Materialized(vec![
1622 0, 0, 1, 2, 1, 0, 2, 2, 1, 2, 2, 2, 0, 1, 2, 2, 2, 2,
1623 ]),
1624 non_null_indices: (0..15).collect(),
1625 max_def_level: 5,
1626 max_rep_level: 2,
1627 array: Arc::new(leaf),
1628 logical_nulls: None,
1629 };
1630 assert_eq!(&levels[0], &expected_levels);
1631 }
1632
1633 #[test]
1634 fn test_calculate_nested_struct_levels() {
1635 let c = Int32Array::from_iter([Some(1), None, Some(3), None, Some(5), Some(6)]);
1645 let leaf = Arc::new(c) as ArrayRef;
1646 let c_field = Arc::new(Field::new("c", DataType::Int32, true));
1647 let b = StructArray::from(((vec![(c_field, leaf.clone())]), Buffer::from([0b00110111])));
1648
1649 let b_field = Arc::new(Field::new("b", b.data_type().clone(), true));
1650 let a = StructArray::from((
1651 (vec![(b_field, Arc::new(b) as ArrayRef)]),
1652 Buffer::from([0b00101111]),
1653 ));
1654
1655 let a_field = Field::new("a", a.data_type().clone(), true);
1656 let a_array = Arc::new(a) as ArrayRef;
1657
1658 let levels = calculate_array_levels(&a_array, &a_field).unwrap();
1659 assert_eq!(levels.len(), 1);
1660
1661 let logical_nulls = leaf.logical_nulls();
1662 let expected_levels = ArrayLevels {
1663 def_levels: LevelData::Materialized(vec![3, 2, 3, 1, 0, 3]),
1664 rep_levels: LevelData::Absent,
1665 non_null_indices: vec![0, 2, 5],
1666 max_def_level: 3,
1667 max_rep_level: 0,
1668 array: leaf,
1669 logical_nulls,
1670 };
1671 assert_eq!(&levels[0], &expected_levels);
1672 }
1673
1674 #[test]
1675 fn list_single_column() {
1676 let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
1679 let a_value_offsets = arrow::buffer::Buffer::from_iter([0_i32, 1, 3, 3, 6, 10]);
1680 let a_list_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1681 let a_list_data = ArrayData::builder(a_list_type.clone())
1682 .len(5)
1683 .add_buffer(a_value_offsets)
1684 .null_bit_buffer(Some(Buffer::from([0b00011011])))
1685 .add_child_data(a_values.to_data())
1686 .build()
1687 .unwrap();
1688
1689 assert_eq!(a_list_data.null_count(), 1);
1690
1691 let a = ListArray::from(a_list_data);
1692
1693 let item_field = Field::new_list_field(a_list_type, true);
1694 let mut builder = levels(&item_field, a);
1695 builder.write(2..4);
1696 let levels = builder.finish();
1697
1698 assert_eq!(levels.len(), 1);
1699
1700 let list_level = &levels[0];
1701
1702 let expected_level = ArrayLevels {
1703 def_levels: LevelData::Materialized(vec![0, 3, 3, 3]),
1704 rep_levels: LevelData::Materialized(vec![0, 0, 1, 1]),
1705 non_null_indices: vec![3, 4, 5],
1706 max_def_level: 3,
1707 max_rep_level: 1,
1708 array: Arc::new(a_values),
1709 logical_nulls: None,
1710 };
1711 assert_eq!(list_level, &expected_level);
1712 }
1713
1714 #[test]
1715 fn mixed_struct_list() {
1716 let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
1720 let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
1721 let struct_field_g = Arc::new(Field::new(
1722 "g",
1723 DataType::List(Arc::new(Field::new("items", DataType::Int16, false))),
1724 false,
1725 ));
1726 let struct_field_e = Arc::new(Field::new(
1727 "e",
1728 DataType::Struct(vec![struct_field_f.clone(), struct_field_g.clone()].into()),
1729 true,
1730 ));
1731 let schema = Schema::new(vec![
1732 Field::new("a", DataType::Int32, false),
1733 Field::new("b", DataType::Int32, true),
1734 Field::new(
1735 "c",
1736 DataType::Struct(vec![struct_field_d.clone(), struct_field_e.clone()].into()),
1737 true, ),
1739 ]);
1740
1741 let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
1743 let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
1744 let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
1745 let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
1746
1747 let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
1748
1749 let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
1752
1753 let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
1755 .len(5)
1756 .add_buffer(g_value_offsets)
1757 .add_child_data(g_value.into_data())
1758 .build()
1759 .unwrap();
1760 let g = ListArray::from(g_list_data);
1761
1762 let e = StructArray::from(vec![
1763 (struct_field_f, Arc::new(f.clone()) as ArrayRef),
1764 (struct_field_g, Arc::new(g) as ArrayRef),
1765 ]);
1766
1767 let c = StructArray::from(vec![
1768 (struct_field_d, Arc::new(d.clone()) as ArrayRef),
1769 (struct_field_e, Arc::new(e) as ArrayRef),
1770 ]);
1771
1772 let batch = RecordBatch::try_new(
1774 Arc::new(schema),
1775 vec![Arc::new(a.clone()), Arc::new(b.clone()), Arc::new(c)],
1776 )
1777 .unwrap();
1778
1779 let mut levels = vec![];
1782 batch
1783 .columns()
1784 .iter()
1785 .zip(batch.schema().fields())
1786 .for_each(|(array, field)| {
1787 let mut array_levels = calculate_array_levels(array, field).unwrap();
1788 levels.append(&mut array_levels);
1789 });
1790 assert_eq!(levels.len(), 5);
1791
1792 let list_level = &levels[0];
1794
1795 let expected_level = ArrayLevels {
1796 def_levels: LevelData::Absent,
1797 rep_levels: LevelData::Absent,
1798 non_null_indices: vec![0, 1, 2, 3, 4],
1799 max_def_level: 0,
1800 max_rep_level: 0,
1801 array: Arc::new(a),
1802 logical_nulls: None,
1803 };
1804 assert_eq!(list_level, &expected_level);
1805
1806 let list_level = levels.get(1).unwrap();
1808
1809 let b_logical_nulls = b.logical_nulls();
1810 let expected_level = ArrayLevels {
1811 def_levels: LevelData::Materialized(vec![1, 0, 0, 1, 1]),
1812 rep_levels: LevelData::Absent,
1813 non_null_indices: vec![0, 3, 4],
1814 max_def_level: 1,
1815 max_rep_level: 0,
1816 array: Arc::new(b),
1817 logical_nulls: b_logical_nulls,
1818 };
1819 assert_eq!(list_level, &expected_level);
1820
1821 let list_level = levels.get(2).unwrap();
1823
1824 let d_logical_nulls = d.logical_nulls();
1825 let expected_level = ArrayLevels {
1826 def_levels: LevelData::Materialized(vec![1, 1, 1, 2, 1]),
1827 rep_levels: LevelData::Absent,
1828 non_null_indices: vec![3],
1829 max_def_level: 2,
1830 max_rep_level: 0,
1831 array: Arc::new(d),
1832 logical_nulls: d_logical_nulls,
1833 };
1834 assert_eq!(list_level, &expected_level);
1835
1836 let list_level = levels.get(3).unwrap();
1838
1839 let f_logical_nulls = f.logical_nulls();
1840 let expected_level = ArrayLevels {
1841 def_levels: LevelData::Materialized(vec![3, 2, 3, 2, 3]),
1842 rep_levels: LevelData::Absent,
1843 non_null_indices: vec![0, 2, 4],
1844 max_def_level: 3,
1845 max_rep_level: 0,
1846 array: Arc::new(f),
1847 logical_nulls: f_logical_nulls,
1848 };
1849 assert_eq!(list_level, &expected_level);
1850 }
1851
1852 #[test]
1853 fn test_null_vs_nonnull_struct() {
1854 let offset_field = Arc::new(Field::new("offset", DataType::Int32, true));
1856 let schema = Schema::new(vec![Field::new(
1857 "some_nested_object",
1858 DataType::Struct(vec![offset_field.clone()].into()),
1859 false,
1860 )]);
1861
1862 let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
1864
1865 let some_nested_object =
1866 StructArray::from(vec![(offset_field, Arc::new(offset) as ArrayRef)]);
1867
1868 let batch =
1870 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
1871
1872 let struct_null_level =
1873 calculate_array_levels(batch.column(0), batch.schema().field(0)).unwrap();
1874
1875 let offset_field = Arc::new(Field::new("offset", DataType::Int32, true));
1878 let schema = Schema::new(vec![Field::new(
1879 "some_nested_object",
1880 DataType::Struct(vec![offset_field.clone()].into()),
1881 true,
1882 )]);
1883
1884 let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
1886
1887 let some_nested_object =
1888 StructArray::from(vec![(offset_field, Arc::new(offset) as ArrayRef)]);
1889
1890 let batch =
1892 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
1893
1894 let struct_non_null_level =
1895 calculate_array_levels(batch.column(0), batch.schema().field(0)).unwrap();
1896
1897 if struct_non_null_level == struct_null_level {
1899 panic!("Levels should not be equal, to reflect the difference in struct nullness");
1900 }
1901 }
1902
1903 #[test]
1904 fn test_map_array() {
1905 let json_content = r#"
1907 {"stocks":{"long": "$AAA", "short": "$BBB"}}
1908 {"stocks":{"long": "$CCC", "short": null}}
1909 {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
1910 "#;
1911 let entries_struct_type = DataType::Struct(Fields::from(vec![
1912 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1913 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
1914 ]));
1915 let stocks_field = Field::new(
1916 "stocks",
1917 DataType::Map(
1918 Arc::new(Field::new(
1919 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1920 entries_struct_type,
1921 false,
1922 )),
1923 false,
1924 ),
1925 false,
1927 );
1928 let schema = Arc::new(Schema::new(vec![stocks_field]));
1929 let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
1930 let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
1931
1932 let batch = reader.next().unwrap().unwrap();
1933
1934 let mut levels = vec![];
1936 batch
1937 .columns()
1938 .iter()
1939 .zip(batch.schema().fields())
1940 .for_each(|(array, field)| {
1941 let mut array_levels = calculate_array_levels(array, field).unwrap();
1942 levels.append(&mut array_levels);
1943 });
1944 assert_eq!(levels.len(), 2);
1945
1946 let map = batch.column(0).as_map();
1947 let map_keys_logical_nulls = map.keys().logical_nulls();
1948
1949 let list_level = &levels[0];
1951
1952 let expected_level = ArrayLevels {
1953 def_levels: LevelData::Materialized(vec![1; 7]),
1954 rep_levels: LevelData::Materialized(vec![0, 1, 0, 1, 0, 1, 1]),
1955 non_null_indices: vec![0, 1, 2, 3, 4, 5, 6],
1956 max_def_level: 1,
1957 max_rep_level: 1,
1958 array: map.keys().clone(),
1959 logical_nulls: map_keys_logical_nulls,
1960 };
1961 assert_eq!(list_level, &expected_level);
1962
1963 let list_level = levels.get(1).unwrap();
1965 let map_values_logical_nulls = map.values().logical_nulls();
1966
1967 let expected_level = ArrayLevels {
1968 def_levels: LevelData::Materialized(vec![2, 2, 2, 1, 2, 1, 2]),
1969 rep_levels: LevelData::Materialized(vec![0, 1, 0, 1, 0, 1, 1]),
1970 non_null_indices: vec![0, 1, 2, 4, 6],
1971 max_def_level: 2,
1972 max_rep_level: 1,
1973 array: map.values().clone(),
1974 logical_nulls: map_values_logical_nulls,
1975 };
1976 assert_eq!(list_level, &expected_level);
1977 }
1978
1979 #[test]
1980 fn test_list_of_struct() {
1981 let int_field = Field::new("a", DataType::Int32, true);
1983 let fields = Fields::from([Arc::new(int_field)]);
1984 let item_field = Field::new_list_field(DataType::Struct(fields.clone()), true);
1985 let list_field = Field::new("list", DataType::List(Arc::new(item_field)), true);
1986
1987 let int_builder = Int32Builder::with_capacity(10);
1988 let struct_builder = StructBuilder::new(fields, vec![Box::new(int_builder)]);
1989 let mut list_builder = ListBuilder::new(struct_builder);
1990
1991 let values = list_builder.values();
1995 values
1996 .field_builder::<Int32Builder>(0)
1997 .unwrap()
1998 .append_value(1);
1999 values.append(true);
2000 list_builder.append(true);
2001
2002 list_builder.append(true);
2004
2005 list_builder.append(false);
2007
2008 let values = list_builder.values();
2010 values
2011 .field_builder::<Int32Builder>(0)
2012 .unwrap()
2013 .append_null();
2014 values.append(false);
2015 values
2016 .field_builder::<Int32Builder>(0)
2017 .unwrap()
2018 .append_null();
2019 values.append(false);
2020 list_builder.append(true);
2021
2022 let values = list_builder.values();
2024 values
2025 .field_builder::<Int32Builder>(0)
2026 .unwrap()
2027 .append_null();
2028 values.append(true);
2029 list_builder.append(true);
2030
2031 let values = list_builder.values();
2033 values
2034 .field_builder::<Int32Builder>(0)
2035 .unwrap()
2036 .append_value(2);
2037 values.append(true);
2038 list_builder.append(true);
2039
2040 let array = Arc::new(list_builder.finish());
2041
2042 let values = array.values().as_struct().column(0).clone();
2043 let values_len = values.len();
2044 assert_eq!(values_len, 5);
2045
2046 let schema = Arc::new(Schema::new(vec![list_field]));
2047
2048 let rb = RecordBatch::try_new(schema, vec![array]).unwrap();
2049
2050 let levels = calculate_array_levels(rb.column(0), rb.schema().field(0)).unwrap();
2051 let list_level = &levels[0];
2052
2053 let logical_nulls = values.logical_nulls();
2054 let expected_level = ArrayLevels {
2055 def_levels: LevelData::Materialized(vec![4, 1, 0, 2, 2, 3, 4]),
2056 rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 1, 0, 0]),
2057 non_null_indices: vec![0, 4],
2058 max_def_level: 4,
2059 max_rep_level: 1,
2060 array: values,
2061 logical_nulls,
2062 };
2063
2064 assert_eq!(list_level, &expected_level);
2065 }
2066
2067 #[test]
2068 fn test_struct_mask_list() {
2069 let inner = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2071 Some(vec![Some(1), Some(2)]),
2072 Some(vec![None]),
2073 Some(vec![]),
2074 Some(vec![Some(3), None]), Some(vec![Some(4), Some(5)]),
2076 None, None,
2078 ]);
2079 let values = inner.values().clone();
2080
2081 assert_eq!(inner.values().len(), 7);
2083
2084 let field = Arc::new(Field::new("list", inner.data_type().clone(), true));
2085 let array = Arc::new(inner) as ArrayRef;
2086 let nulls = Buffer::from([0b01010111]);
2087 let struct_a = StructArray::from((vec![(field, array)], nulls));
2088
2089 let field = Field::new("struct", struct_a.data_type().clone(), true);
2090 let array = Arc::new(struct_a) as ArrayRef;
2091 let levels = calculate_array_levels(&array, &field).unwrap();
2092
2093 assert_eq!(levels.len(), 1);
2094
2095 let logical_nulls = values.logical_nulls();
2096 let expected_level = ArrayLevels {
2097 def_levels: LevelData::Materialized(vec![4, 4, 3, 2, 0, 4, 4, 0, 1]),
2098 rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 0, 0, 1, 0, 0]),
2099 non_null_indices: vec![0, 1, 5, 6],
2100 max_def_level: 4,
2101 max_rep_level: 1,
2102 array: values,
2103 logical_nulls,
2104 };
2105
2106 assert_eq!(&levels[0], &expected_level);
2107 }
2108
2109 #[test]
2110 fn test_list_mask_struct() {
2111 let a1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2115 Some(vec![None]), Some(vec![]), Some(vec![Some(3), None]),
2118 Some(vec![Some(4), Some(5), None, Some(6)]), None,
2120 None,
2121 ]);
2122 let a1_values = a1.values().clone();
2123 let a1 = Arc::new(a1) as ArrayRef;
2124
2125 let a2 = Arc::new(Int32Array::from_iter(vec![
2126 Some(1), Some(2), None,
2129 Some(4), Some(5),
2131 None,
2132 ])) as ArrayRef;
2133 let a2_values = a2.clone();
2134
2135 let field_a1 = Arc::new(Field::new("list", a1.data_type().clone(), true));
2136 let field_a2 = Arc::new(Field::new("integers", a2.data_type().clone(), true));
2137
2138 let nulls = Buffer::from([0b00110111]);
2139 let struct_a = Arc::new(StructArray::from((
2140 vec![(field_a1, a1), (field_a2, a2)],
2141 nulls,
2142 ))) as ArrayRef;
2143
2144 let offsets = Buffer::from_iter([0_i32, 0, 2, 2, 3, 5, 5]);
2145 let nulls = Buffer::from([0b00111100]);
2146
2147 let list_type = DataType::List(Arc::new(Field::new(
2148 "struct",
2149 struct_a.data_type().clone(),
2150 true,
2151 )));
2152
2153 let data = ArrayDataBuilder::new(list_type.clone())
2154 .len(6)
2155 .null_bit_buffer(Some(nulls))
2156 .add_buffer(offsets)
2157 .add_child_data(struct_a.into_data())
2158 .build()
2159 .unwrap();
2160
2161 let list = make_array(data);
2162 let list_field = Field::new("col", list_type, true);
2163
2164 let expected = vec![
2165 String::new(),
2166 String::new(),
2167 "[]".to_string(),
2168 "[{list: [3, ], integers: }]".to_string(),
2169 "[, {list: , integers: 5}]".to_string(),
2170 "[]".to_string(),
2171 ];
2172
2173 let actual: Vec<_> = (0..6)
2174 .map(|x| array_value_to_string(&list, x).unwrap())
2175 .collect();
2176 assert_eq!(actual, expected);
2177
2178 let levels = calculate_array_levels(&list, &list_field).unwrap();
2179
2180 assert_eq!(levels.len(), 2);
2181
2182 let a1_logical_nulls = a1_values.logical_nulls();
2183 let expected_level = ArrayLevels {
2184 def_levels: LevelData::Materialized(vec![0, 0, 1, 6, 5, 2, 3, 1]),
2185 rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 2, 0, 1, 0]),
2186 non_null_indices: vec![1],
2187 max_def_level: 6,
2188 max_rep_level: 2,
2189 array: a1_values,
2190 logical_nulls: a1_logical_nulls,
2191 };
2192
2193 assert_eq!(&levels[0], &expected_level);
2194
2195 let a2_logical_nulls = a2_values.logical_nulls();
2196 let expected_level = ArrayLevels {
2197 def_levels: LevelData::Materialized(vec![0, 0, 1, 3, 2, 4, 1]),
2198 rep_levels: LevelData::Materialized(vec![0, 0, 0, 0, 0, 1, 0]),
2199 non_null_indices: vec![4],
2200 max_def_level: 4,
2201 max_rep_level: 1,
2202 array: a2_values,
2203 logical_nulls: a2_logical_nulls,
2204 };
2205
2206 assert_eq!(&levels[1], &expected_level);
2207 }
2208
2209 #[test]
2210 fn test_fixed_size_list() {
2211 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 2);
2213 builder.values().append_slice(&[1, 2]);
2214 builder.append(true);
2215 builder.values().append_slice(&[3, 4]);
2216 builder.append(false);
2217 builder.values().append_slice(&[5, 6]);
2218 builder.append(false);
2219 builder.values().append_slice(&[7, 8]);
2220 builder.append(true);
2221 builder.values().append_slice(&[9, 10]);
2222 builder.append(false);
2223 let a = builder.finish();
2224 let values = a.values().clone();
2225
2226 let item_field = Field::new_list_field(a.data_type().clone(), true);
2227 let mut builder = levels(&item_field, a);
2228 builder.write(1..4);
2229 let levels = builder.finish();
2230
2231 assert_eq!(levels.len(), 1);
2232
2233 let list_level = &levels[0];
2234
2235 let logical_nulls = values.logical_nulls();
2236 let expected_level = ArrayLevels {
2237 def_levels: LevelData::Materialized(vec![0, 0, 3, 3]),
2238 rep_levels: LevelData::Materialized(vec![0, 0, 0, 1]),
2239 non_null_indices: vec![6, 7],
2240 max_def_level: 3,
2241 max_rep_level: 1,
2242 array: values,
2243 logical_nulls,
2244 };
2245 assert_eq!(list_level, &expected_level);
2246 }
2247
2248 #[test]
2249 fn test_fixed_size_list_of_struct() {
2250 let field_a = Field::new("a", DataType::Int32, true);
2252 let field_b = Field::new("b", DataType::Int64, false);
2253 let fields = Fields::from([Arc::new(field_a), Arc::new(field_b)]);
2254 let item_field = Field::new_list_field(DataType::Struct(fields.clone()), true);
2255 let list_field = Field::new(
2256 "list",
2257 DataType::FixedSizeList(Arc::new(item_field), 2),
2258 true,
2259 );
2260
2261 let builder_a = Int32Builder::with_capacity(10);
2262 let builder_b = Int64Builder::with_capacity(10);
2263 let struct_builder =
2264 StructBuilder::new(fields, vec![Box::new(builder_a), Box::new(builder_b)]);
2265 let mut list_builder = FixedSizeListBuilder::new(struct_builder, 2);
2266
2267 let values = list_builder.values();
2276 values
2278 .field_builder::<Int32Builder>(0)
2279 .unwrap()
2280 .append_value(1);
2281 values
2282 .field_builder::<Int64Builder>(1)
2283 .unwrap()
2284 .append_value(2);
2285 values.append(true);
2286 values
2288 .field_builder::<Int32Builder>(0)
2289 .unwrap()
2290 .append_null();
2291 values
2292 .field_builder::<Int64Builder>(1)
2293 .unwrap()
2294 .append_value(0);
2295 values.append(false);
2296 list_builder.append(true);
2297
2298 let values = list_builder.values();
2300 values
2302 .field_builder::<Int32Builder>(0)
2303 .unwrap()
2304 .append_null();
2305 values
2306 .field_builder::<Int64Builder>(1)
2307 .unwrap()
2308 .append_value(0);
2309 values.append(false);
2310 values
2312 .field_builder::<Int32Builder>(0)
2313 .unwrap()
2314 .append_null();
2315 values
2316 .field_builder::<Int64Builder>(1)
2317 .unwrap()
2318 .append_value(0);
2319 values.append(false);
2320 list_builder.append(false);
2321
2322 let values = list_builder.values();
2324 values
2326 .field_builder::<Int32Builder>(0)
2327 .unwrap()
2328 .append_null();
2329 values
2330 .field_builder::<Int64Builder>(1)
2331 .unwrap()
2332 .append_value(0);
2333 values.append(false);
2334 values
2336 .field_builder::<Int32Builder>(0)
2337 .unwrap()
2338 .append_null();
2339 values
2340 .field_builder::<Int64Builder>(1)
2341 .unwrap()
2342 .append_value(0);
2343 values.append(false);
2344 list_builder.append(true);
2345
2346 let values = list_builder.values();
2348 values
2350 .field_builder::<Int32Builder>(0)
2351 .unwrap()
2352 .append_null();
2353 values
2354 .field_builder::<Int64Builder>(1)
2355 .unwrap()
2356 .append_value(3);
2357 values.append(true);
2358 values
2360 .field_builder::<Int32Builder>(0)
2361 .unwrap()
2362 .append_value(2);
2363 values
2364 .field_builder::<Int64Builder>(1)
2365 .unwrap()
2366 .append_value(4);
2367 values.append(true);
2368 list_builder.append(true);
2369
2370 let array = Arc::new(list_builder.finish());
2371
2372 assert_eq!(array.values().len(), 8);
2373 assert_eq!(array.len(), 4);
2374
2375 let struct_values = array.values().as_struct();
2376 let values_a = struct_values.column(0).clone();
2377 let values_b = struct_values.column(1).clone();
2378
2379 let schema = Arc::new(Schema::new(vec![list_field]));
2380 let rb = RecordBatch::try_new(schema, vec![array]).unwrap();
2381
2382 let levels = calculate_array_levels(rb.column(0), rb.schema().field(0)).unwrap();
2383 let a_levels = &levels[0];
2384 let b_levels = &levels[1];
2385
2386 let values_a_logical_nulls = values_a.logical_nulls();
2388 let expected_a = ArrayLevels {
2389 def_levels: LevelData::Materialized(vec![4, 2, 0, 2, 2, 3, 4]),
2390 rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1]),
2391 non_null_indices: vec![0, 7],
2392 max_def_level: 4,
2393 max_rep_level: 1,
2394 array: values_a,
2395 logical_nulls: values_a_logical_nulls,
2396 };
2397 let values_b_logical_nulls = values_b.logical_nulls();
2399 let expected_b = ArrayLevels {
2400 def_levels: LevelData::Materialized(vec![3, 2, 0, 2, 2, 3, 3]),
2401 rep_levels: LevelData::Materialized(vec![0, 1, 0, 0, 1, 0, 1]),
2402 non_null_indices: vec![0, 6, 7],
2403 max_def_level: 3,
2404 max_rep_level: 1,
2405 array: values_b,
2406 logical_nulls: values_b_logical_nulls,
2407 };
2408
2409 assert_eq!(a_levels, &expected_a);
2410 assert_eq!(b_levels, &expected_b);
2411 }
2412
2413 #[test]
2414 fn test_fixed_size_list_empty() {
2415 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 0);
2416 builder.append(true);
2417 builder.append(false);
2418 builder.append(true);
2419 let array = builder.finish();
2420 let values = array.values().clone();
2421
2422 let item_field = Field::new_list_field(array.data_type().clone(), true);
2423 let mut builder = levels(&item_field, array);
2424 builder.write(0..3);
2425 let levels = builder.finish();
2426
2427 assert_eq!(levels.len(), 1);
2428
2429 let list_level = &levels[0];
2430
2431 let logical_nulls = values.logical_nulls();
2432 let expected_level = ArrayLevels {
2433 def_levels: LevelData::Materialized(vec![1, 0, 1]),
2434 rep_levels: LevelData::Materialized(vec![0, 0, 0]),
2435 non_null_indices: vec![],
2436 max_def_level: 3,
2437 max_rep_level: 1,
2438 array: values,
2439 logical_nulls,
2440 };
2441 assert_eq!(list_level, &expected_level);
2442 }
2443
2444 #[test]
2445 fn test_fixed_size_list_of_var_lists() {
2446 let mut builder = FixedSizeListBuilder::new(ListBuilder::new(Int32Builder::new()), 2);
2448 builder.values().append_value([Some(1), None, Some(3)]);
2449 builder.values().append_null();
2450 builder.append(true);
2451 builder.values().append_value([Some(4)]);
2452 builder.values().append_value([]);
2453 builder.append(true);
2454 builder.values().append_value([Some(5), Some(6)]);
2455 builder.values().append_value([None, None]);
2456 builder.append(true);
2457 builder.values().append_null();
2458 builder.values().append_null();
2459 builder.append(false);
2460 let a = builder.finish();
2461 let values = a.values().as_list::<i32>().values().clone();
2462
2463 let item_field = Field::new_list_field(a.data_type().clone(), true);
2464 let mut builder = levels(&item_field, a);
2465 builder.write(0..4);
2466 let levels = builder.finish();
2467
2468 let logical_nulls = values.logical_nulls();
2469 let expected_level = ArrayLevels {
2470 def_levels: LevelData::Materialized(vec![5, 4, 5, 2, 5, 3, 5, 5, 4, 4, 0]),
2471 rep_levels: LevelData::Materialized(vec![0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0]),
2472 non_null_indices: vec![0, 2, 3, 4, 5],
2473 max_def_level: 5,
2474 max_rep_level: 2,
2475 array: values,
2476 logical_nulls,
2477 };
2478
2479 assert_eq!(levels[0], expected_level);
2480 }
2481
2482 #[test]
2483 fn test_null_dictionary_values() {
2484 let values = Int32Array::new(
2485 vec![1, 2, 3, 4].into(),
2486 Some(NullBuffer::from(vec![true, false, true, true])),
2487 );
2488 let keys = Int32Array::new(
2489 vec![1, 54, 2, 0].into(),
2490 Some(NullBuffer::from(vec![true, false, true, true])),
2491 );
2492 let dict = DictionaryArray::new(keys, Arc::new(values));
2494
2495 let item_field = Field::new_list_field(dict.data_type().clone(), true);
2496
2497 let mut builder = levels(&item_field, dict.clone());
2498 builder.write(0..4);
2499 let levels = builder.finish();
2500
2501 let logical_nulls = dict.logical_nulls();
2502 let expected_level = ArrayLevels {
2503 def_levels: LevelData::Materialized(vec![0, 0, 1, 1]),
2504 rep_levels: LevelData::Absent,
2505 non_null_indices: vec![2, 3],
2506 max_def_level: 1,
2507 max_rep_level: 0,
2508 array: Arc::new(dict),
2509 logical_nulls,
2510 };
2511 assert_eq!(levels[0], expected_level);
2512 }
2513
2514 #[test]
2515 fn mismatched_types() {
2516 let array = Arc::new(Int32Array::from_iter(0..10)) as ArrayRef;
2517 let field = Field::new_list_field(DataType::Float64, false);
2518
2519 let err = LevelInfoBuilder::try_new(&field, Default::default(), &array)
2520 .unwrap_err()
2521 .to_string();
2522
2523 assert_eq!(
2524 err,
2525 "Arrow: Incompatible type. Field 'item' has type Float64, array has type Int32",
2526 );
2527 }
2528
2529 fn levels<T: Array + 'static>(field: &Field, array: T) -> LevelInfoBuilder {
2530 let v = Arc::new(array) as ArrayRef;
2531 LevelInfoBuilder::try_new(field, Default::default(), &v).unwrap()
2532 }
2533
2534 #[test]
2535 fn test_slice_for_chunk_flat() {
2536 let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6]));
2541 let logical_nulls = array.logical_nulls();
2542 let levels = ArrayLevels {
2543 def_levels: LevelData::Absent,
2544 rep_levels: LevelData::Absent,
2545 non_null_indices: vec![0, 1, 2, 3, 4, 5],
2546 max_def_level: 0,
2547 max_rep_level: 0,
2548 array,
2549 logical_nulls,
2550 };
2551 let sliced = levels.slice_for_chunk(&CdcChunk {
2552 level_offset: 0,
2553 num_levels: 0,
2554 value_offset: 2,
2555 num_values: 3,
2556 });
2557 assert!(matches!(sliced.def_levels, LevelData::Absent));
2558 assert!(matches!(sliced.rep_levels, LevelData::Absent));
2559 assert_eq!(sliced.non_null_indices, vec![0, 1, 2]);
2560 assert_eq!(sliced.array.len(), 3);
2561
2562 let array: ArrayRef = Arc::new(Int32Array::from(vec![
2568 Some(1),
2569 None,
2570 Some(3),
2571 None,
2572 Some(5),
2573 Some(6),
2574 ]));
2575 let logical_nulls = array.logical_nulls();
2576 let levels = ArrayLevels {
2577 def_levels: LevelData::Materialized(vec![1, 0, 1, 0, 1, 1]),
2578 rep_levels: LevelData::Absent,
2579 non_null_indices: vec![0, 2, 4, 5],
2580 max_def_level: 1,
2581 max_rep_level: 0,
2582 array,
2583 logical_nulls,
2584 };
2585 let sliced = levels.slice_for_chunk(&CdcChunk {
2586 level_offset: 1,
2587 num_levels: 3,
2588 value_offset: 1,
2589 num_values: 1,
2590 });
2591 assert_eq!(sliced.def_levels, LevelData::Materialized(vec![0, 1, 0]));
2592 assert!(matches!(sliced.rep_levels, LevelData::Absent));
2593 assert_eq!(sliced.non_null_indices, vec![0]); assert_eq!(sliced.array.len(), 1);
2595 }
2596
2597 #[test]
2598 fn test_slice_for_chunk_nested_with_nulls() {
2599 let array: ArrayRef = Arc::new(Int32Array::from(vec![
2618 Some(1), None, None, Some(2), None, None, None, None, Some(4), Some(5), ]));
2629 let logical_nulls = array.logical_nulls();
2630 let levels = ArrayLevels {
2631 def_levels: LevelData::Materialized(vec![3, 0, 3, 2, 0, 3, 3]),
2632 rep_levels: LevelData::Materialized(vec![0, 0, 0, 1, 0, 0, 1]),
2633 non_null_indices: vec![0, 3, 8, 9],
2634 max_def_level: 3,
2635 max_rep_level: 1,
2636 array,
2637 logical_nulls,
2638 };
2639
2640 let chunk0 = levels.slice_for_chunk(&CdcChunk {
2642 level_offset: 0,
2643 num_levels: 2,
2644 value_offset: 0,
2645 num_values: 1,
2646 });
2647 assert_eq!(chunk0.non_null_indices, vec![0]);
2648 assert_eq!(chunk0.array.len(), 1);
2649
2650 let chunk1 = levels.slice_for_chunk(&CdcChunk {
2652 level_offset: 2,
2653 num_levels: 3,
2654 value_offset: 1,
2655 num_values: 1,
2656 });
2657 assert_eq!(chunk1.non_null_indices, vec![0]);
2658 assert_eq!(chunk1.array.len(), 1);
2659
2660 let chunk2 = levels.slice_for_chunk(&CdcChunk {
2662 level_offset: 5,
2663 num_levels: 2,
2664 value_offset: 2,
2665 num_values: 2,
2666 });
2667 assert_eq!(chunk2.non_null_indices, vec![0, 1]);
2668 assert_eq!(chunk2.array.len(), 2);
2669 }
2670
2671 #[test]
2672 fn test_slice_for_chunk_all_null() {
2673 let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, None, Some(4)]));
2675 let logical_nulls = array.logical_nulls();
2676 let levels = ArrayLevels {
2677 def_levels: LevelData::Materialized(vec![1, 0, 0, 1]),
2678 rep_levels: LevelData::Absent,
2679 non_null_indices: vec![0, 3],
2680 max_def_level: 1,
2681 max_rep_level: 0,
2682 array,
2683 logical_nulls,
2684 };
2685 let sliced = levels.slice_for_chunk(&CdcChunk {
2687 level_offset: 1,
2688 num_levels: 2,
2689 value_offset: 1,
2690 num_values: 0,
2691 });
2692 assert_eq!(sliced.def_levels, LevelData::Materialized(vec![0, 0]));
2693 assert_eq!(sliced.non_null_indices, Vec::<usize>::new());
2694 assert_eq!(sliced.array.len(), 0);
2695 }
2696
2697 #[test]
2698 fn test_all_null_list() {
2699 let item_field = Arc::new(Field::new_list_field(DataType::Int32, true));
2705 let list = ListArray::new_null(item_field, 4);
2706 let values = list.values().clone();
2707 let field = Field::new("list", list.data_type().clone(), true);
2708 let array = Arc::new(list) as ArrayRef;
2709
2710 let levels = calculate_array_levels(&array, &field).unwrap();
2711 assert_eq!(levels.len(), 1);
2712
2713 let logical_nulls = values.logical_nulls();
2714 let expected = ArrayLevels {
2715 def_levels: LevelData::Uniform { value: 0, count: 4 },
2716 rep_levels: LevelData::Uniform { value: 0, count: 4 },
2717 non_null_indices: vec![],
2718 max_def_level: 3,
2719 max_rep_level: 1,
2720 array: values,
2721 logical_nulls,
2722 };
2723 assert_eq!(&levels[0], &expected);
2724 }
2725
2726 #[test]
2727 fn test_all_null_fixed_size_list() {
2728 let item_field = Arc::new(Field::new_list_field(DataType::Int32, true));
2734 let list = FixedSizeListArray::new_null(item_field, 2, 3);
2735 let values = list.values().clone();
2736 let field = Field::new("list", list.data_type().clone(), true);
2737 let array = Arc::new(list) as ArrayRef;
2738
2739 let levels = calculate_array_levels(&array, &field).unwrap();
2740 assert_eq!(levels.len(), 1);
2741
2742 let logical_nulls = values.logical_nulls();
2743 let expected = ArrayLevels {
2744 def_levels: LevelData::Uniform { value: 0, count: 3 },
2745 rep_levels: LevelData::Uniform { value: 0, count: 3 },
2746 non_null_indices: vec![],
2747 max_def_level: 3,
2748 max_rep_level: 1,
2749 array: values,
2750 logical_nulls,
2751 };
2752 assert_eq!(&levels[0], &expected);
2753 }
2754
2755 #[test]
2756 fn test_all_null_struct() {
2757 let c = Int32Array::from(vec![None::<i32>; 4]);
2764 let leaf = Arc::new(c) as ArrayRef;
2765 let c_field = Arc::new(Field::new("c", DataType::Int32, true));
2766 let a = StructArray::from((vec![(c_field, leaf.clone())], Buffer::from([0b00000000])));
2767 let a_field = Field::new("a", a.data_type().clone(), true);
2768 let a_array = Arc::new(a) as ArrayRef;
2769
2770 let levels = calculate_array_levels(&a_array, &a_field).unwrap();
2771 assert_eq!(levels.len(), 1);
2772
2773 let expected = ArrayLevels {
2774 def_levels: LevelData::Uniform { value: 0, count: 4 },
2775 rep_levels: LevelData::Absent,
2776 non_null_indices: vec![],
2777 max_def_level: 2,
2778 max_rep_level: 0,
2779 array: leaf,
2780 logical_nulls: Some(NullBuffer::new_null(4)),
2781 };
2782 assert_eq!(&levels[0], &expected);
2783 }
2784
2785 #[test]
2786 fn test_all_null_nested_struct() {
2787 let c = Int32Array::from(vec![None::<i32>; 3]);
2793 let leaf = Arc::new(c) as ArrayRef;
2794 let c_field = Arc::new(Field::new("c", DataType::Int32, true));
2795 let b = StructArray::from((vec![(c_field, leaf.clone())], Buffer::from([0b00000000])));
2796 let b_field = Arc::new(Field::new("b", b.data_type().clone(), true));
2797 let a = StructArray::from((
2798 vec![(b_field, Arc::new(b) as ArrayRef)],
2799 Buffer::from([0b00000000]),
2800 ));
2801 let a_field = Field::new("a", a.data_type().clone(), true);
2802 let a_array = Arc::new(a) as ArrayRef;
2803
2804 let levels = calculate_array_levels(&a_array, &a_field).unwrap();
2805 assert_eq!(levels.len(), 1);
2806
2807 let expected = ArrayLevels {
2808 def_levels: LevelData::Uniform { value: 0, count: 3 },
2809 rep_levels: LevelData::Absent,
2810 non_null_indices: vec![],
2811 max_def_level: 3,
2812 max_rep_level: 0,
2813 array: leaf,
2814 logical_nulls: Some(NullBuffer::new_null(3)),
2815 };
2816 assert_eq!(&levels[0], &expected);
2817 }
2818
2819 #[test]
2820 fn test_all_null_struct_multiple_children() {
2821 let c1 = Arc::new(Int32Array::from(vec![None::<i32>; 2])) as ArrayRef;
2827 let c2 = Arc::new(Int32Array::from(vec![None::<i32>; 2])) as ArrayRef;
2828 let c1_field = Arc::new(Field::new("c1", DataType::Int32, true));
2829 let c2_field = Arc::new(Field::new("c2", DataType::Int32, true));
2830 let a = StructArray::from((
2831 vec![(c1_field, c1.clone()), (c2_field, c2.clone())],
2832 Buffer::from([0b00000000]),
2833 ));
2834 let a_field = Field::new("a", a.data_type().clone(), true);
2835 let a_array = Arc::new(a) as ArrayRef;
2836
2837 let levels = calculate_array_levels(&a_array, &a_field).unwrap();
2838 assert_eq!(levels.len(), 2);
2839
2840 for (i, leaf) in [c1, c2].into_iter().enumerate() {
2841 let expected = ArrayLevels {
2842 def_levels: LevelData::Uniform { value: 0, count: 2 },
2843 rep_levels: LevelData::Absent,
2844 non_null_indices: vec![],
2845 max_def_level: 2,
2846 max_rep_level: 0,
2847 array: leaf,
2848 logical_nulls: Some(NullBuffer::new_null(2)),
2849 };
2850 assert_eq!(&levels[i], &expected, "leaf {i} mismatch");
2851 }
2852 }
2853}