Skip to main content

parquet/arrow/arrow_writer/
levels.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Parquet definition and repetition levels
19//!
20//! Contains the algorithm for computing definition and repetition levels.
21//! The algorithm works by tracking the slots of an array that should
22//! ultimately be populated when writing to Parquet.
23//! Parquet achieves nesting through definition levels and repetition levels \[1\].
24//! Definition levels specify how many optional fields in the part for the column
25//! are defined.
26//! Repetition levels specify at what repeated field (list) in the path a column
27//! is defined.
28//!
29//! In a nested data structure such as `a.b.c`, one can see levels as defining
30//! whether a record is defined at `a`, `a.b`, or `a.b.c`.
31//! Optional fields are nullable fields, thus if all 3 fields
32//! are nullable, the maximum definition could be = 3 if there are no lists.
33//!
34//! The algorithm in this module computes the necessary information to enable
35//! the writer to keep track of which columns are at which levels, and to extract
36//! the correct values at the correct slots from Arrow arrays.
37//!
38//! It works by walking a record batch's arrays, keeping track of what values
39//! are non-null, their positions and computing what their levels are.
40//!
41//! \[1\] [parquet-format#nested-encoding](https://github.com/apache/parquet-format#nested-encoding)
42
43use 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
55/// Expands a [`DataType::RunEndEncoded`] array into a flat (logical) array of its values type.
56///
57/// use `arrow_select::take` to materialize the  full-length flat array.
58/// This is intentionally simple (O(n)); efficiency can/should be improved
59fn 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
77/// Performs a depth-first scan of the children of `array`, constructing [`ArrayLevels`]
78/// for each leaf column encountered
79pub(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
85/// Returns true if the DataType can be represented as a primitive parquet column,
86/// i.e. a leaf array with no children
87fn 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/// The definition and repetition level of an array within a potentially nested hierarchy
125#[derive(Debug, Default, Clone, Copy)]
126struct LevelContext {
127    /// The current repetition level
128    rep_level: i16,
129    /// The current definition level
130    def_level: i16,
131}
132
133/// A helper to construct [`ArrayLevels`] from a potentially nested [`Field`]
134#[derive(Debug)]
135enum LevelInfoBuilder {
136    /// A primitive, leaf array
137    Primitive(ArrayLevels),
138    /// A list array
139    List(
140        Box<LevelInfoBuilder>, // Child Values
141        LevelContext,          // Context
142        OffsetBuffer<i32>,     // Offsets
143        Option<NullBuffer>,    // Nulls
144        bool,                  // is_last_level (child has no nested rep)
145    ),
146    /// A large list array
147    LargeList(
148        Box<LevelInfoBuilder>, // Child Values
149        LevelContext,          // Context
150        OffsetBuffer<i64>,     // Offsets
151        Option<NullBuffer>,    // Nulls
152        bool,                  // is_last_level (child has no nested rep)
153    ),
154    /// A fixed size list array
155    FixedSizeList(
156        Box<LevelInfoBuilder>, // Values
157        LevelContext,          // Context
158        usize,                 // List Size
159        Option<NullBuffer>,    // Nulls
160    ),
161    /// A list view array
162    ListView(
163        Box<LevelInfoBuilder>, // Child Values
164        LevelContext,          // Context
165        ScalarBuffer<i32>,     // Offsets
166        ScalarBuffer<i32>,     // Sizes
167        Option<NullBuffer>,    // Nulls
168    ),
169    /// A large list view array
170    LargeListView(
171        Box<LevelInfoBuilder>, // Child Values
172        LevelContext,          // Context
173        ScalarBuffer<i64>,     // Offsets
174        ScalarBuffer<i64>,     // Sizes
175        Option<NullBuffer>,    // Nulls
176    ),
177    /// A struct array
178    Struct(Vec<LevelInfoBuilder>, LevelContext, Option<NullBuffer>),
179}
180
181/// Minimum sub-range length before the bulk-fill fast path in `write_leaf`
182/// becomes profitable for null-heavy leaf columns. Below this, per-call
183/// slice + popcount overhead regresses list/struct paths that call
184/// `write_leaf` many times with tiny ranges. Picked via threshold sweep;
185/// see <https://github.com/apache/arrow-rs/pull/9967> for the rationale.
186const BULK_FILL_MIN_LEN: usize = 64;
187
188impl LevelInfoBuilder {
189    /// Create a new [`LevelInfoBuilder`] for the given [`Field`] and parent [`LevelContext`]
190    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    /// Finish this [`LevelInfoBuilder`] returning the [`ArrayLevels`] for the leaf columns
316    /// as enumerated by a depth-first search
317    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    /// Given an `array`, write the level data for the elements in `range`
330    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    /// Returns `true` if the child contains no nested repetition levels, meaning
355    /// each child element produces exactly one rep_level entry in the leaf.
356    /// This is true for `Primitive` children and `Struct` trees with no list descendants.
357    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    /// Write `range` elements from ListArray `array`
368    ///
369    /// Note: MapArrays are `ListArray<i32>` under the hood and so are dispatched to this method
370    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        // Fast path: entire list array is null; emit bulk null rep/def levels
379        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        // Dispatch to separate functions so the compiler can optimize each
388        // hot loop independently (function body size affects codegen quality).
389        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    /// Batch write for lists whose child has no nested repetition.
397    ///
398    /// "direct" means writing the child rep levels using offsets without scanning.
399    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            // The first element of each list slot needs rep_level =
417            // list_start_rep to mark a new list boundary. Because there's a 1:1
418            // mapping between child elements and rep_level entries, the position
419            // of each slot's first element is directly computable from offsets.
420            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    /// Batch write for lists whose child has nested repetition.
436    ///
437    /// After batch-writing child elements, scans backward through rep_levels
438    /// counting child-element starts to find and stamp slot boundaries.
439    ///
440    /// Scan backward because we don't know start offset before writing.
441    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                    // This algorithm is the same as write_list_direct.
463                    // Use a separate function because the branch code size would affect codegen
464                    // quality of the hot loop of write_list_direct.
465                    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                    // Backward scan: count child-element starts (rep <= ctx.rep_level)
473                    // and stamp list_start_rep at slot boundaries.
474                    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                        // Asserting rep >= ctx.rep_level to ensure low level depth haven't
480                        // being written.
481                        debug_assert!(*rep >= ctx.rep_level);
482                        // Count element starts by skipping nested reps (rep > ctx.rep_level).
483                        //
484                        // `==` would also work here: the child is written before the
485                        // parent, so no entry within the batch has rep < ctx.rep_level.
486                        // Benchmarks show no difference, so keep the more defensive `<=`.
487                        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    /// Shared run-classification loop for write_list_direct and write_list_scan.
506    /// Monomorphized per `emit_non_empty_run` closure type, giving the compiler
507    /// separate optimization contexts for each backfill strategy.
508    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        // Classify each slot, detect run boundaries, flush on transition.
535        #[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            // A null buffer without any null can skip the per-slot validity
571            // checks and use the null-free classification loop below.
572            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    /// Write `range` elements from ListViewArray `array`
610    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                            // Nested element - ignore
633                            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                // TODO: Faster bitmask iteration (#1757)
663                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    /// Write `range` elements from StructArray `array`
693    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        // Fast path: entire struct array is null; emit bulk null def/rep levels
709        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                // TODO: Faster bitmask iteration (#1757)
726                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    /// Write `range` elements from FixedSizeListArray with child data `values` and null bitmap `nulls`.
756    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        // Fast path: entire fixed-size list array is null
764        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                // Step backward over the child rep levels and mark the start of each list
786                rep_levels
787                    .iter_mut()
788                    .rev()
789                    // Filter out reps from nested children
790                    .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        // If list size is 0, ignore values and just write rep/def levels.
801        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 a run of valid rows if not already inside of one
823                        start_idx.get_or_insert(idx);
824                    } else {
825                        // Write out any pending valid rows
826                        if let Some(start) = start_idx.take() {
827                            write_rows(child, start, idx);
828                        }
829                        // Add null row
830                        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                // Write out any remaining valid rows
837                if let Some(start) = start_idx.take() {
838                    write_rows(child, start, range.end);
839                }
840            }
841            // If all rows are valid then write the whole array
842            None => write_rows(child, range.start, range.end),
843        }
844    }
845
846    /// Write a primitive array, as defined by [`is_leaf`]
847    fn write_leaf(info: &mut ArrayLevels, range: Range<usize>) {
848        let len = range.end - range.start;
849
850        // Fast path: entire leaf array is null
851        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                    // Bulk-fill is profitable only on null-heavy ranges long enough to
867                    // amortize the slice/popcount cost; see `BULK_FILL_MIN_LEN` and the
868                    // PR description for the threshold sweep. The gate uses the cached
869                    // buffer-wide `null_count` (O(1)) to stay cheap on the cold path.
870                    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                            // Safety: range.end was asserted to be in bounds earlier
890                            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    /// Visits all children of this node in depth first order
914    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    /// Determine if the fields are compatible for purposes of constructing `LevelBuilderInfo`.
931    ///
932    /// Fields are compatible if they're the same type. Otherwise if one of them is a dictionary
933    /// and the other is a native array, the dictionary values must have the same type as the
934    /// native array
935    fn types_compatible(a: &DataType, b: &DataType) -> bool {
936        // if the Arrow data types are equal, the types are deemed compatible
937        if a.equals_datatype(b) {
938            return true;
939        }
940
941        // get the values out of the dictionaries
942        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        // now that we've got the values from one/both dictionaries, if the values
952        // have the same Arrow data type, they're compatible
953        if a == b {
954            return true;
955        }
956
957        // here we have different Arrow data types, but if the array contains the same type of data
958        // then we consider the type compatible
959        match a {
960            // String, StringView and LargeString are compatible
961            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            // Binary, BinaryView and LargeBinary are compatible
966            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            // otherwise we have incompatible types
971            _ => false,
972        }
973    }
974}
975
976/// The data necessary to write a primitive Arrow array to parquet, taking into account
977/// any non-primitive parents it may have in the arrow representation
978#[derive(Debug, Clone)]
979pub(crate) enum LevelData {
980    Absent,
981    Materialized(Vec<i16>),
982    Uniform { value: i16, count: usize },
983}
984
985// Compare logical level contents rather than physical representation, so a
986// uniform run compares equal to the equivalent materialized buffer.
987impl 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            // No physical level stream exists for this schema. Higher-level
1050            // traversal may still append implicit levels, so this remains a no-op.
1051            Self::Absent => {}
1052            // Start compact: the first appended run can be represented without
1053            // allocating a level buffer.
1054            Self::Materialized(values) if values.is_empty() => {
1055                *self = Self::Uniform { value, count };
1056            }
1057            // Already materialized, so preserve the buffer representation and append.
1058            Self::Materialized(values) => values.extend(std::iter::repeat_n(value, count)),
1059            // Preserve the compact representation while the appended run has
1060            // the same value.
1061            Self::Uniform {
1062                value: uniform_value,
1063                count: uniform_count,
1064            } if *uniform_value == value => {
1065                *uniform_count += count;
1066            }
1067            // A different value breaks the uniform representation. Materialize
1068            // the existing run, then append the new run to the buffer.
1069            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    /// Convert a uniform run into a materialized buffer if needed, then return
1086    /// the mutable level buffer. Returns `None` when no physical level stream exists.
1087    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    /// Array's definition levels
1106    ///
1107    /// Present if `max_def_level != 0`
1108    def_levels: LevelData,
1109
1110    /// Array's optional repetition levels
1111    ///
1112    /// Present if `max_rep_level != 0`
1113    rep_levels: LevelData,
1114
1115    /// The corresponding array identifying non-null slices of data
1116    /// from the primitive array
1117    non_null_indices: Vec<usize>,
1118
1119    /// The maximum definition level for this leaf column
1120    max_def_level: i16,
1121
1122    /// The maximum repetition for this leaf column
1123    max_rep_level: i16,
1124
1125    /// The arrow array
1126    array: ArrayRef,
1127
1128    /// cached logical nulls of the array.
1129    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    /// Create a sliced view of this `ArrayLevels` for a CDC chunk.
1183    ///
1184    /// The chunk's `value_offset`/`num_values` select the relevant slice of
1185    /// `non_null_indices`. The array is sliced to the range covered by
1186    /// those indices, and they are shifted to be relative to the slice.
1187    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        // Select the non-null indices for this chunk.
1192        let nni = &self.non_null_indices[chunk.value_offset..chunk.value_offset + chunk.num_values];
1193        // Compute the array range spanned by the non-null indices.
1194        // When nni is empty (all-null chunk), start=0, end=0 → zero-length
1195        // array slice; write_batch_internal will process only the def/rep
1196        // levels and write no values.
1197        let start = nni.first().copied().unwrap_or(0);
1198        let end = nni.last().map_or(0, |&i| i + 1);
1199        // Shift indices to be relative to the sliced array.
1200        let non_null_indices = nni.iter().map(|&idx| idx - start).collect();
1201        // Slice the array to the computed range.
1202        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    /// Bulk-emit `count` uniform def/rep levels.
1217    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        // based on the example at https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html
1247        // [[a, b, c], [d, e, f, g]], [[h], [i,j]]
1248
1249        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        // Cannot use from_iter_primitive as always infers nullable
1258        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        // This test calculates the levels for a non-null primitive array
1293        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        // This test calculates the levels for a nullable primitive array
1314        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        // if all array values are defined (e.g. batch<list<_>>)
1345        // [[0], [1], [2], [3], [4]]
1346
1347        let leaf_array = Int32Array::from_iter(0..5);
1348        // Cannot use from_iter_primitive as always infers nullable
1349        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        // array: [[0, 0], NULL, [2, 2], [3, 3, 3, 3], [4, 4, 4]]
1374        // all values are defined as we do not have nulls on the root (batch)
1375        // repetition:
1376        //   0: 0, 1
1377        //   1: 0
1378        //   2: 0, 1
1379        //   3: 0, 1, 1, 1
1380        //   4: 0, 1, 1
1381        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        // If some values are null
1411        //
1412        // This emulates an array in the form: <struct<list<?>>
1413        // with values:
1414        // - 0: [0, 1], but is null because of the struct
1415        // - 1: []
1416        // - 2: [2, 3], but is null because of the struct
1417        // - 3: [4, 5, 6, 7]
1418        // - 4: [8, 9, 10]
1419        //
1420        // If the first values of a list are null due to a parent, we have to still account for them
1421        // while indexing, because they would affect the way the child is indexed
1422        // i.e. in the above example, we have to know that [0, 1] has to be skipped
1423        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        // nested lists
1459
1460        // 0: [[100, 101], [102, 103]]
1461        // 1: []
1462        // 2: [[104, 105], [106, 107]]
1463        // 3: [[108, 109], [110, 111], [112, 113], [114, 115]]
1464        // 4: [[116, 117], [118, 119], [120, 121]]
1465
1466        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        // if all array values are defined (e.g. batch<list<_>>)
1516        // The array at this level looks like:
1517        // 0: [a]
1518        // 1: [a]
1519        // 2: [a]
1520        // 3: [a]
1521
1522        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        // 0: null
1547        // 1: [1, 2, 3]
1548        // 2: [4, 5]
1549        // 3: [6, 7]
1550        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        // nested lists
1580        // In a JSON syntax with the schema: <struct<list<list<primitive>>>>, this translates into:
1581        // 0: {"struct": null }
1582        // 1: {"struct": [ [201], [202, 203], [] ]}
1583        // 2: {"struct": [ [204, 205, 206], [207, 208, 209, 210] ]}
1584        // 3: {"struct": [ [], [211, 212, 213, 214, 215] ]}
1585
1586        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        // tests a <struct[a]<struct[b]<int[c]>>
1636        // array:
1637        //  - {a: {b: {c: 1}}}
1638        //  - {a: {b: {c: null}}}
1639        //  - {a: {b: {c: 3}}}
1640        //  - {a: {b: null}}
1641        //  - {a: null}}
1642        //  - {a: {b: {c: 6}}}
1643
1644        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        // this tests the level generation from the arrow_writer equivalent test
1677
1678        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        // this tests the level generation from the equivalent arrow_writer_complex test
1717
1718        // define schema
1719        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, // https://github.com/apache/arrow-rs/issues/245
1738            ),
1739        ]);
1740
1741        // create some data
1742        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        // Construct a buffer for value offsets, for the nested array:
1750        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
1751        let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
1752
1753        // Construct a list array from the above two
1754        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        // build a record batch
1773        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        //////////////////////////////////////////////
1780        // calculate the list's level
1781        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        // test "a" levels
1793        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        // test "b" levels
1807        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        // test "d" levels
1822        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        // test "f" levels
1837        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        // define schema
1855        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        // create some data
1863        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        // build a record batch
1869        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        // create second batch
1876        // define schema
1877        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        // create some data
1885        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        // build a record batch
1891        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        // The 2 levels should not be the same
1898        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        // Note: we are using the JSON Arrow reader for brevity
1906        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            // not nullable, so the keys have max level = 1
1926            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        // calculate the map's level
1935        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        // test key levels
1950        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        // test values levels
1964        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        // define schema
1982        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        // [{a: 1}], [], null, [null, null], [{a: null}], [{a: 2}]
1992        //
1993        // [{a: 1}]
1994        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        // []
2003        list_builder.append(true);
2004
2005        // null
2006        list_builder.append(false);
2007
2008        // [null, null]
2009        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        // [{a: null}]
2023        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        // [{a: 2}]
2032        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        // Test the null mask of a struct array masking out non-empty slices of a child ListArray
2070        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]), // Masked by struct array
2075            Some(vec![Some(4), Some(5)]),
2076            None, // Masked by struct array
2077            None,
2078        ]);
2079        let values = inner.values().clone();
2080
2081        // This test assumes that nulls don't take up space
2082        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        // Test the null mask of a struct array and the null mask of a list array
2112        // masking out non-null elements of their children
2113
2114        let a1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2115            Some(vec![None]), // Masked by list array
2116            Some(vec![]),     // Masked by list array
2117            Some(vec![Some(3), None]),
2118            Some(vec![Some(4), Some(5), None, Some(6)]), // Masked by struct array
2119            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), // Masked by list array
2127            Some(2), // Masked by list array
2128            None,
2129            Some(4), // Masked by struct array
2130            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        // [[1, 2], null, null, [7, 8], null]
2212        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        // define schema
2251        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        // [
2268        //   [{a: 1, b: 2}, null],
2269        //   null,
2270        //   [null, null],
2271        //   [{a: null, b: 3}, {a: 2, b: 4}]
2272        // ]
2273
2274        // [{a: 1, b: 2}, null]
2275        let values = list_builder.values();
2276        // {a: 1, b: 2}
2277        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        // null
2287        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        // null
2299        let values = list_builder.values();
2300        // null
2301        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        // null
2311        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        // [null, null]
2323        let values = list_builder.values();
2324        // null
2325        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        // null
2335        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        // [{a: null, b: 3}, {a: 2, b: 4}]
2347        let values = list_builder.values();
2348        // {a: null, b: 3}
2349        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        // {a: 2, b: 4}
2359        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        // [[{a: 1}, null], null, [null, null], [{a: null}, {a: 2}]]
2387        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        // [[{b: 2}, null], null, [null, null], [{b: 3}, {b: 4}]]
2398        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        // [[[1, null, 3], null], [[4], []], [[5, 6], [null, null]], null]
2447        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        // [NULL, NULL, 3, 0]
2493        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        // Case 1: required field (max_def_level=0, no def/rep levels stored).
2537        // Array has 6 values; all are non-null so non_null_indices covers every position.
2538        // value_offset=2, num_values=3 → non_null_indices[2..5] = [2,3,4].
2539        // Array is sliced (no def_levels → write_batch_internal uses values.len()).
2540        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        // Case 2: optional field (max_def_level=1, def levels present, no rep levels).
2563        // Array: [Some(1), None, Some(3), None, Some(5), Some(6)]
2564        // non_null_indices: [0, 2, 4, 5]
2565        // value_offset=1, num_values=1 → non_null_indices[1..2] = [2].
2566        // Array is not sliced (def_levels present → num_levels from def_levels.len()).
2567        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]); // [2] shifted by -2 (nni[0])
2594        assert_eq!(sliced.array.len(), 1);
2595    }
2596
2597    #[test]
2598    fn test_slice_for_chunk_nested_with_nulls() {
2599        // Regression test for https://github.com/apache/arrow-rs/issues/9637
2600        //
2601        // Simulates a List<Int32?> where null list entries have non-zero child
2602        // ranges (valid per Arrow spec: "a null value may correspond to a
2603        // non-empty segment in the child array"). This creates gaps in the
2604        // leaf array that don't correspond to any levels.
2605        //
2606        // 5 rows with 2 null list entries owning non-empty child ranges:
2607        //   row 0: [1]       → leaf[0]
2608        //   row 1: null list → owns leaf[1..3] (gap of 2)
2609        //   row 2: [2, null] → leaf[3], leaf[4]=null element
2610        //   row 3: null list → owns leaf[5..8] (gap of 3)
2611        //   row 4: [4, 5]   → leaf[8], leaf[9]
2612        //
2613        // def_levels: [3,  0,  3, 2,  0,  3, 3]
2614        // rep_levels: [0,  0,  0, 1,  0,  0, 1]
2615        // non_null_indices: [0, 3, 8, 9]
2616        //   gaps in array: 0→3 (skip 1,2), 3→8 (skip 5,6,7)
2617        let array: ArrayRef = Arc::new(Int32Array::from(vec![
2618            Some(1), // 0: row 0
2619            None,    // 1: gap (null list row 1)
2620            None,    // 2: gap (null list row 1)
2621            Some(2), // 3: row 2
2622            None,    // 4: row 2, null element
2623            None,    // 5: gap (null list row 3)
2624            None,    // 6: gap (null list row 3)
2625            None,    // 7: gap (null list row 3)
2626            Some(4), // 8: row 4
2627            Some(5), // 9: row 4
2628        ]));
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        // Chunk 0: rows 0-1, nni=[0] → array sliced to [0..1]
2641        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        // Chunk 1: rows 2-3, nni=[3] → array sliced to [3..4]
2651        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        // Chunk 2: row 4, nni=[8, 9] → array sliced to [8..10]
2661        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        // All-null chunk: num_values=0 → empty nni slice → zero-length array.
2674        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        // Chunk covering only the two null rows (levels 1..3), zero non-null values.
2686        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        // List<Int32> where every list slot is null.
2700        // Schema: list (nullable) -> item (int32, nullable)
2701        // Data: [null, null, null, null]
2702        //
2703        // Expected: max_def=3, max_rep=1, def/rep levels all 0.
2704        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        // FixedSizeList<Int32; 2> where every list slot is null.
2729        // Schema: list (nullable) -> item (int32, nullable)
2730        // Data: [null, null, null]
2731        //
2732        // Expected: max_def=3, max_rep=1, def/rep levels all 0.
2733        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        // Struct<Int32> where every struct slot is null.
2758        // Schema: a (struct, nullable) -> c (int32, nullable)
2759        // Data: [null, null, null, null]
2760        //
2761        // Expected: max_def=2, def_levels all 0 (struct is null → child never reached),
2762        // leaf values are empty.
2763        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        // Struct<Struct<Int32>> where the outer struct is entirely null.
2788        // Schema: a (struct, nullable) -> b (struct, nullable) -> c (int32, nullable)
2789        // Data: [null, null, null]
2790        //
2791        // Expected: max_def=3, def_levels all 0.
2792        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        // Struct with two leaf children, entirely null.
2822        // Schema: a (struct, nullable) -> { c1 (int32, nullable), c2 (int32, nullable) }
2823        // Data: [null, null]
2824        //
2825        // Both leaf columns should get uniform def_levels=0.
2826        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}