Skip to main content

lance_encoding/
repdef.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Utilities for rep-def levels
5//!
6//! Repetition and definition levels are a way to encode multipile validity / offsets arrays
7//! into a single buffer.  They are a form of "zipping" buffers together that takes advantage
8//! of the fact that, if the outermost array is invalid, then the validity of the inner items
9//! is irrelevant.
10//!
11//! Note: the concept of repetition & definition levels comes from the Dremel paper and has
12//! been implemented in Apache Parquet.  However, the implementation here is not necessarily
13//! compatible with Parquet.  For example, we use 0 to represent the "inner-most" item and
14//! Parquet uses 0 to represent the "outer-most" item.
15//!
16//! # Repetition Levels
17//!
18//! With repetition levels we convert a sparse array of offsets into a dense array of levels.
19//! These levels are marked non-zero whenever a new list begins.  In other words, given the
20//! list array with 3 rows [{<0,1>, <>, <2>}, {<3>}, {}], [], [{<4>}] we would have three
21//! offsets arrays:
22//!
23//! Outer-most ([]): [0, 3, 3, 4]
24//! Middle     ({}): [0, 3, 4, 4, 5]
25//! Inner      (<>): [0, 2, 2, 3, 4, 5]
26//! Values         : [0, 1, 2, 3, 4]
27//!
28//! We can convert these into repetition levels as follows:
29//!
30//! | Values | Repetition |
31//! | ------ | ---------- |
32//! |      0 |          3 | // Start of outer-most list
33//! |      1 |          0 | // Continues inner-most list (no new lists)
34//! |      - |          1 | // Start of new inner-most list (empty list)
35//! |      2 |          1 | // Start of new inner-most list
36//! |      3 |          2 | // Start of new middle list
37//! |      - |          2 | // Start of new inner-most list (empty list)
38//! |      - |          3 | // Start of new outer-most list (empty list)
39//! |      4 |          0 | // Start of new outer-most list
40//!
41//! Note: We actually have MORE repetition levels than values.  This is because the repetition
42//! levels need to be able to represent empty lists.
43//!
44//! # Definition Levels
45//!
46//! Definition levels are simpler.  We can think of them as zipping together various validity bitmaps
47//! (from different levels of nesting) into a single buffer.  For example, we could zip the arrays
48//! [1, 1, 0, 0] and [1, 0, 1, 0] into [11, 10, 01, 00].  However, 00 and 01 are redundant.  If the
49//! outer level is null then the validity of the inner levels is irrelevant.  To save space we instead
50//! encode a "level" which is the "depth" of the null.  Let's look at a more complete example:
51//!
52//! Array: [{"middle": {"inner": 1]}}, NULL, {"middle": NULL}, {"middle": {"inner": NULL}}]
53//!
54//! In Arrow we would have the following validity arrays:
55//! Outer validity : 1, 0, 1, 1
56//! Middle validity: 1, ?, 0, 1
57//! Inner validity : 1, ?, ?, 0
58//! Values         : 1, ?, ?, ?
59//!
60//! The ? values are undefined in the Arrow format.  We can convert these into definition levels as follows:
61//!
62//! | Values | Definition |
63//! | ------ | ---------- |
64//! |      1 |          0 | // Valid at all levels
65//! |      - |          3 | // Null at outer level
66//! |      - |          2 | // Null at middle level
67//! |      - |          1 | // Null at inner level
68//!
69//! # Compression
70//!
71//! Note that we only need 2 bits of definition levels to represent 3 levels of nesting.  Definition
72//! levels are always more compact than the input validity arrays.  However, compressed levels are not
73//! necessarily more compact than the compressed validity arrays.
74//!
75//! Repetition levels are more complex.  If there are very large lists then a sparse array of offsets
76//! (which has one element per list) might be more compact than a dense array of repetition levels
77//! (which has one element per list value, possibly even more if there are empty lists).
78//!
79//! However, both repetition levels and definition levels are typically very compressible with RLE.
80//!
81//! However, in Lance we don't always take advantage of that compression because we want to be able
82//! to zip rep-def levels together with our values.  This gives us fewer IOPS when accessing row values.
83//!
84//! # Utilities in this Module
85//!
86//! - `RepDefBuilder` - Extracts validity and offset information from Arrow arrays.  We use this as we
87//!   shred the incoming data into primitive leaf arrays.  We don't immediately convert into rep-def because
88//!   we need to share and cheaply clone the builder when we have structs (each struct child shares some parent
89//!   validity / offset information)  The `serialize` method is called once all data has been received to create
90//!   the final rep-def levels.
91//!
92//! - `SerializerContext` - This is an internal utility that helps with serializing rep-def levels.
93//!
94//! - `CompositeRepDefUnraveler` - This structure is used to reverse the process.  It starts with a set of
95//!   rep-def levels and then uses the `unravel_validity` and `unravel_offsets` methods to produce validity
96//!   buffers and offset buffers.  It is "composite" because we may be combining sets of rep-def buffers from
97//!   multiple locations (e.g. multiple blocks in a mini-block encoded file).
98//!
99//! - `RepDefSlicer` - This is a utility that helps with slicing rep-def buffers.  These buffers are "kind of"
100//!   transparent (maps 1:1 with the values in the array) but not exactly because of the special (empty/null) lists.
101//!   The slicer helps with this issue and is used when slicing a rep-def buffer into mini-blocks.
102//!
103//! - `build_control_word_iterator` - This takes in rep-def levels and returns an iterator that returns byte-padded
104//!   "control words" which are used when creating full-zip encoded data.
105//!
106//! - `ControlWordParser` - This parser can parse the control words returned by `build_control_word_iterator` and is
107//!   used when decoding full-zip encoded data.
108
109use std::{
110    iter::{Copied, Zip},
111    ops::Range,
112    sync::Arc,
113};
114
115use arrow_array::OffsetSizeTrait;
116use arrow_buffer::{
117    ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, NullBuffer, OffsetBuffer, ScalarBuffer,
118};
119use lance_core::{Error, Result, utils::bit::log_2_ceil};
120
121use crate::{
122    buffer::LanceBuffer,
123    encodings::logical::primitive::sparse::{SparseStructuralPlan, SparseStructuralUnraveler},
124};
125
126pub type LevelBuffer = Vec<u16>;
127
128/// A top-level-row range whose dense rep/def stream fits one mini-block page.
129#[derive(Debug, Clone, PartialEq, Eq)]
130pub(crate) struct MiniBlockRepDefSplit {
131    /// Top-level row offset, relative to the original unsplit page.
132    pub(crate) row_start: u64,
133    /// Number of top-level rows in this split.
134    pub(crate) num_rows: u64,
135    /// Rep/def level range, relative to the original unsplit page.
136    pub(crate) level_range: Range<usize>,
137    /// Visible value offset, relative to the original unsplit page.
138    pub(crate) value_start: u64,
139    /// Number of visible values in this split.
140    pub(crate) num_values: u64,
141}
142
143/// Dense mini-block rep/def budget result for one accumulated page.
144#[derive(Debug, Clone, PartialEq, Eq)]
145pub(crate) enum MiniBlockRepDefBudget {
146    /// The dense rep/def stream fits one mini-block structural page.
147    WithinBudget,
148    /// The dense rep/def stream fits after splitting on top-level row boundaries.
149    RequiresPageSplit(Vec<MiniBlockRepDefSplit>),
150    /// A single top-level row has this many rep/def levels and exceeds the budget.
151    SingleRowOverBudget(u64),
152}
153
154// As we build def levels we add this to special values to indicate that they
155// are special so that we can skip over them when processing lower levels.
156//
157// We assume 16 bits is good enough for rep-def levels.  This _would_ give
158// us 65536 levels of struct nesting and list nesting.  However, we cut that
159// in half for SPECIAL_THRESHOLD because we use the top bit to indicate if an
160// item is a special value (null list / empty list) during construction.
161//
162// We subtract this off at the end of construction to get the actual definition
163// levels.
164const SPECIAL_THRESHOLD: u16 = u16::MAX / 2;
165
166/// Represents information that we extract from a list array as we are
167/// encoding
168#[derive(Clone, Debug)]
169struct OffsetDesc {
170    offsets: Arc<[i64]>,
171    validity: Option<BooleanBuffer>,
172    has_empty_lists: bool,
173    num_values: usize,
174    num_specials: usize,
175}
176
177/// Represents validity information that we extract from non-list arrays (that
178/// have nulls) as we are encoding
179#[derive(Clone, Debug)]
180struct ValidityDesc {
181    validity: Option<BooleanBuffer>,
182    num_values: usize,
183}
184
185/// Represents validity information that we extract from FSL arrays.  This is
186/// just validity (no offsets) but we also record the dimension of the FSL array
187/// as that will impact the next layer
188#[derive(Clone, Debug)]
189struct FslDesc {
190    validity: Option<BooleanBuffer>,
191    dimension: usize,
192    num_values: usize,
193}
194
195// As we build up rep/def from arrow arrays we record a
196// series of RawRepDef objects.  Each one corresponds to layer
197// in the array structure
198#[derive(Clone, Debug)]
199enum RawRepDef {
200    Offsets(OffsetDesc),
201    Validity(ValidityDesc),
202    Fsl(FslDesc),
203}
204
205/// A normalized Arrow structural layer shared by dense and sparse serializers.
206#[derive(Clone, Copy, Debug)]
207pub(crate) enum NormalizedStructuralLayer<'a> {
208    List {
209        offsets: &'a [i64],
210        validity: Option<&'a BooleanBuffer>,
211        num_slots: usize,
212    },
213    Validity {
214        validity: Option<&'a BooleanBuffer>,
215        num_slots: usize,
216    },
217    FixedSizeList {
218        validity: Option<&'a BooleanBuffer>,
219        dimension: usize,
220        num_slots: usize,
221    },
222}
223
224/// Structural layers concatenated across input batches exactly once.
225///
226/// Dense rep/def serialization and sparse metadata planning both consume this
227/// representation so the Arrow nesting is not independently reconstructed.
228#[derive(Debug)]
229pub(crate) struct NormalizedStructuralPlan {
230    layers: Vec<RawRepDef>,
231    dense_all_valid: bool,
232}
233
234impl NormalizedStructuralPlan {
235    pub(crate) fn layers(&self) -> impl ExactSizeIterator<Item = NormalizedStructuralLayer<'_>> {
236        self.layers.iter().map(|layer| match layer {
237            RawRepDef::Offsets(OffsetDesc {
238                offsets,
239                validity,
240                num_values,
241                ..
242            }) => NormalizedStructuralLayer::List {
243                offsets,
244                validity: validity.as_ref(),
245                num_slots: *num_values,
246            },
247            RawRepDef::Validity(ValidityDesc {
248                validity,
249                num_values,
250            }) => NormalizedStructuralLayer::Validity {
251                validity: validity.as_ref(),
252                num_slots: *num_values,
253            },
254            RawRepDef::Fsl(FslDesc {
255                validity,
256                dimension,
257                num_values,
258            }) => NormalizedStructuralLayer::FixedSizeList {
259                validity: validity.as_ref(),
260                dimension: *dimension,
261                num_slots: *num_values,
262            },
263        })
264    }
265
266    fn to_serializer(&self) -> (SerializerContext, Option<u64>) {
267        if self.dense_all_valid {
268            let def_meaning = self
269                .layers
270                .iter()
271                .map(|_| DefinitionInterpretation::AllValidItem)
272                .collect::<Vec<_>>();
273            return (
274                SerializerContext {
275                    def_meaning,
276                    rep_levels: LevelBuffer::default(),
277                    spare_rep: LevelBuffer::default(),
278                    def_levels: LevelBuffer::default(),
279                    spare_def: LevelBuffer::default(),
280                    current_rep: 0,
281                    current_def: 0,
282                    current_len: 0,
283                    current_num_specials: 0,
284                    has_fsl: false,
285                },
286                None,
287            );
288        }
289
290        let total_len = self.layers.last().map_or(0, RawRepDef::num_values)
291            + self
292                .layers
293                .iter()
294                .map(RawRepDef::num_specials)
295                .sum::<usize>();
296        let max_rep = self.layers.iter().map(RawRepDef::max_rep).sum::<u16>();
297        let max_def = self.layers.iter().map(RawRepDef::max_def).sum::<u16>();
298        let bits_per_rep = if max_rep > 0 {
299            u64::from(u16::BITS - max_rep.leading_zeros())
300        } else {
301            0
302        };
303        let bits_per_def = if max_def > 0 {
304            u64::from(u16::BITS - max_def.leading_zeros())
305        } else {
306            0
307        };
308        let bits_per_level =
309            (bits_per_rep + bits_per_def > 0).then_some(bits_per_rep + bits_per_def);
310
311        let num_layers = self.layers.len();
312        let mut context = SerializerContext::new(total_len, num_layers, max_rep, max_def);
313        for layer in &self.layers {
314            match layer {
315                RawRepDef::Validity(def) => context.record_validity(def),
316                RawRepDef::Offsets(rep) => context.record_offsets(rep),
317                RawRepDef::Fsl(fsl) => context.record_fsl(fsl),
318            }
319        }
320        (context, bits_per_level)
321    }
322
323    pub(crate) fn serialize(&self) -> SerializedRepDefs {
324        self.to_serializer().0.build()
325    }
326
327    pub(crate) fn serialize_with_miniblock_repdef_budget(
328        &self,
329        max_levels_for_bits: impl FnOnce(u64) -> u64,
330        num_rows: u64,
331        num_values: u64,
332    ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> {
333        let (context, bits_per_level) = self.to_serializer();
334        context.build_with_miniblock_repdef_budget(
335            bits_per_level.map(max_levels_for_bits),
336            num_rows,
337            num_values,
338        )
339    }
340}
341
342impl RawRepDef {
343    // Are there any nulls in this layer
344    fn has_nulls(&self) -> bool {
345        match self {
346            Self::Offsets(OffsetDesc { validity, .. }) => validity.is_some(),
347            Self::Validity(ValidityDesc { validity, .. }) => validity.is_some(),
348            Self::Fsl(FslDesc { validity, .. }) => validity.is_some(),
349        }
350    }
351
352    // How many values are in this layer
353    fn num_values(&self) -> usize {
354        match self {
355            Self::Offsets(OffsetDesc { num_values, .. }) => *num_values,
356            Self::Validity(ValidityDesc { num_values, .. }) => *num_values,
357            Self::Fsl(FslDesc { num_values, .. }) => *num_values,
358        }
359    }
360
361    /// How many empty/null lists are in this layer
362    fn num_specials(&self) -> usize {
363        match self {
364            Self::Offsets(OffsetDesc { num_specials, .. }) => *num_specials,
365            _ => 0,
366        }
367    }
368
369    /// How many definition levels do we need for this layer
370    fn max_def(&self) -> u16 {
371        match self {
372            Self::Offsets(OffsetDesc {
373                has_empty_lists,
374                validity,
375                ..
376            }) => {
377                let mut max_def = 0;
378                if *has_empty_lists {
379                    max_def += 1;
380                }
381                if validity.is_some() {
382                    max_def += 1;
383                }
384                max_def
385            }
386            Self::Validity(ValidityDesc { validity: None, .. }) => 0,
387            Self::Validity(ValidityDesc { .. }) => 1,
388            Self::Fsl(FslDesc { validity: None, .. }) => 0,
389            Self::Fsl(FslDesc { .. }) => 1,
390        }
391    }
392
393    /// How many repetition levels do we need for this layer
394    fn max_rep(&self) -> u16 {
395        match self {
396            Self::Offsets(_) => 1,
397            _ => 0,
398        }
399    }
400}
401
402/// Represents repetition and definition levels that have been
403/// serialized into a pair of (optional) level buffers
404#[derive(Debug)]
405pub struct SerializedRepDefs {
406    /// The repetition levels, one per item
407    ///
408    /// If None, there are no lists
409    pub repetition_levels: Option<Arc<[u16]>>,
410    /// The definition levels, one per item
411    ///
412    /// If None, there are no nulls
413    pub definition_levels: Option<Arc<[u16]>>,
414    /// The meaning of each definition level
415    pub def_meaning: Vec<DefinitionInterpretation>,
416    /// The maximum level that is "visible" from the lowest level
417    ///
418    /// This is the last level before we encounter a list level of some kind.  Once we've
419    /// hit a list level then nulls in any level beyond do not map to actual items.
420    ///
421    /// This is None if there are no lists
422    pub max_visible_level: Option<u16>,
423    has_fsl: bool,
424}
425
426impl SerializedRepDefs {
427    fn max_visible_level(def_meaning: &[DefinitionInterpretation]) -> Option<u16> {
428        let first_list = def_meaning.iter().position(|level| level.is_list());
429        first_list.map(|first_list| {
430            def_meaning
431                .iter()
432                .map(|level| level.num_def_levels())
433                .take(first_list)
434                .sum::<u16>()
435        })
436    }
437
438    pub fn new(
439        repetition_levels: Option<LevelBuffer>,
440        definition_levels: Option<LevelBuffer>,
441        def_meaning: Vec<DefinitionInterpretation>,
442    ) -> Self {
443        Self::new_with_fixed_size_list_levels(
444            repetition_levels,
445            definition_levels,
446            def_meaning,
447            false,
448        )
449    }
450
451    pub(crate) fn new_with_fixed_size_list_levels(
452        repetition_levels: Option<LevelBuffer>,
453        definition_levels: Option<LevelBuffer>,
454        def_meaning: Vec<DefinitionInterpretation>,
455        has_fsl: bool,
456    ) -> Self {
457        let max_visible_level = Self::max_visible_level(&def_meaning);
458        Self {
459            repetition_levels: repetition_levels.map(Arc::from),
460            definition_levels: definition_levels.map(Arc::from),
461            def_meaning,
462            max_visible_level,
463            has_fsl,
464        }
465    }
466
467    /// Creates an empty SerializedRepDefs (no repetition, all valid)
468    pub fn empty(def_meaning: Vec<DefinitionInterpretation>) -> Self {
469        Self {
470            repetition_levels: None,
471            definition_levels: None,
472            def_meaning,
473            max_visible_level: None,
474            has_fsl: false,
475        }
476    }
477
478    pub fn rep_slicer(&self) -> Option<RepDefSlicer<'_>> {
479        self.repetition_levels
480            .as_ref()
481            .map(|rep| RepDefSlicer::new(self, rep.clone()))
482    }
483
484    pub fn def_slicer(&self) -> Option<RepDefSlicer<'_>> {
485        self.definition_levels
486            .as_ref()
487            .map(|def| RepDefSlicer::new(self, def.clone()))
488    }
489
490    pub(crate) fn has_fixed_size_list_levels(&self) -> bool {
491        self.has_fsl
492    }
493}
494
495/// Slices a level buffer into pieces
496///
497/// This is needed to handle the fact that a level buffer may have more
498/// levels than values due to special (empty/null) lists.
499///
500/// As a result, a call to `slice_next(10)` may return 10 levels or it may
501/// return more than 10 levels if any special values are encountered.
502#[derive(Debug)]
503pub struct RepDefSlicer<'a> {
504    repdef: &'a SerializedRepDefs,
505    to_slice: LanceBuffer,
506    current: usize,
507}
508
509// TODO: All of this logic will need some changing when we compress rep/def levels.
510impl<'a> RepDefSlicer<'a> {
511    fn new(repdef: &'a SerializedRepDefs, levels: Arc<[u16]>) -> Self {
512        Self {
513            repdef,
514            to_slice: LanceBuffer::reinterpret_slice(levels),
515            current: 0,
516        }
517    }
518
519    pub fn num_levels(&self) -> usize {
520        self.to_slice.len() / 2
521    }
522
523    pub fn num_levels_remaining(&self) -> usize {
524        self.num_levels() - self.current
525    }
526
527    pub fn all_levels(&self) -> &LanceBuffer {
528        &self.to_slice
529    }
530
531    /// Returns the rest of the levels not yet sliced
532    ///
533    /// This must be called instead of `slice_next` on the final iteration.
534    /// This is because anytime we slice there may be empty/null lists on the
535    /// boundary that are "free" and the current behavior in `slice_next` is to
536    /// leave them for the next call.
537    ///
538    /// `slice_rest` will slice all remaining levels and return them.
539    pub fn slice_rest(&mut self) -> LanceBuffer {
540        let start = self.current;
541        let remaining = self.num_levels_remaining();
542        self.current = self.num_levels();
543        self.to_slice.slice_with_length(start * 2, remaining * 2)
544    }
545
546    /// Returns enough levels to satisfy the next `num_values` values
547    pub fn slice_next(&mut self, num_values: usize) -> LanceBuffer {
548        let start = self.current;
549        let Some(max_visible_level) = self.repdef.max_visible_level else {
550            // No lists, should be 1:1 mapping from levels to values
551            self.current = start + num_values;
552            return self.to_slice.slice_with_length(start * 2, num_values * 2);
553        };
554        if let Some(def) = self.repdef.definition_levels.as_ref() {
555            // There are lists and there are def levels.  That means there may be
556            // more rep/def levels than values.  We need to scan the def levels to figure
557            // out which items are "invisible" and skip over them
558            let mut def_itr = def[start..].iter();
559            let mut num_taken = 0;
560            let mut num_passed = 0;
561            while num_taken < num_values {
562                let def_level = *def_itr.next().unwrap();
563                if def_level <= max_visible_level {
564                    num_taken += 1;
565                }
566                num_passed += 1;
567            }
568            self.current = start + num_passed;
569            self.to_slice.slice_with_length(start * 2, num_passed * 2)
570        } else {
571            // No def levels, should be 1:1 mapping from levels to values
572            self.current = start + num_values;
573            self.to_slice.slice_with_length(start * 2, num_values * 2)
574        }
575    }
576}
577
578/// This tells us how an array handles definition.  Given a stack of
579/// these and a nested array and a set of definition levels we can calculate
580/// how we should interpret the definition levels.
581///
582/// For example, if the interpretation is [AllValidItem, NullableItem] then
583/// a 0 means "valid item" and a 1 means "null struct".  If the interpretation
584/// is [NullableItem, NullableItem] then a 0 means "valid item" and a 1 means
585/// "null item" and a 2 means "null struct".
586///
587/// Lists are tricky because we might use up to two definition levels for a
588/// single layer of list nesting because we need one value to indicate "empty list"
589/// and another value to indicate "null list".
590#[derive(Debug, Copy, Clone, PartialEq, Eq)]
591pub enum DefinitionInterpretation {
592    AllValidItem,
593    AllValidList,
594    NullableItem,
595    NullableList,
596    EmptyableList,
597    NullableAndEmptyableList,
598}
599
600impl DefinitionInterpretation {
601    /// How many definition levels do we need for this layer
602    pub fn num_def_levels(&self) -> u16 {
603        match self {
604            Self::AllValidItem => 0,
605            Self::AllValidList => 0,
606            Self::NullableItem => 1,
607            Self::NullableList => 1,
608            Self::EmptyableList => 1,
609            Self::NullableAndEmptyableList => 2,
610        }
611    }
612
613    /// Does this layer have nulls?
614    pub fn is_all_valid(&self) -> bool {
615        matches!(
616            self,
617            Self::AllValidItem | Self::AllValidList | Self::EmptyableList
618        )
619    }
620
621    /// Does this layer represent a list?
622    pub fn is_list(&self) -> bool {
623        matches!(
624            self,
625            Self::AllValidList
626                | Self::NullableList
627                | Self::EmptyableList
628                | Self::NullableAndEmptyableList
629        )
630    }
631}
632
633/// The RepDefBuilder is used to collect offsets & validity buffers
634/// from arrow structures.  Once we have those we use the SerializerContext
635/// to build the actual repetition and definition levels.
636///
637/// We know ahead of time how many rep/def levels we will need (number of items
638/// in inner-most array + the number of empty/null lists in any parent arrays).
639///
640/// As a result we try and avoid any re-allocations by pre-allocating the buffers
641/// up front.  We allocate two copies of each buffer which allows us to avoid unsafe
642/// code caused by reading and writing to the same buffer (also, it's unavoidable
643/// because there are times we need to write 'faster' than we read)
644#[derive(Debug)]
645struct SerializerContext {
646    // This is built from outer-to-inner and then reversed at the end
647    def_meaning: Vec<DefinitionInterpretation>,
648    rep_levels: LevelBuffer,
649    spare_rep: LevelBuffer,
650    def_levels: LevelBuffer,
651    spare_def: LevelBuffer,
652    current_rep: u16,
653    current_def: u16,
654    current_len: usize,
655    current_num_specials: usize,
656    has_fsl: bool,
657}
658
659impl SerializerContext {
660    fn new(len: usize, num_layers: usize, max_rep: u16, max_def: u16) -> Self {
661        let def_meaning = Vec::with_capacity(num_layers);
662        Self {
663            rep_levels: if max_rep > 0 {
664                vec![0; len]
665            } else {
666                LevelBuffer::default()
667            },
668            spare_rep: if max_rep > 0 {
669                vec![0; len]
670            } else {
671                LevelBuffer::default()
672            },
673            def_levels: if max_def > 0 {
674                vec![0; len]
675            } else {
676                LevelBuffer::default()
677            },
678            spare_def: if max_def > 0 {
679                vec![0; len]
680            } else {
681                LevelBuffer::default()
682            },
683            def_meaning,
684            current_rep: max_rep,
685            current_def: max_def,
686            current_len: 0,
687            current_num_specials: 0,
688            has_fsl: false,
689        }
690    }
691
692    fn checkout_def(&mut self, meaning: DefinitionInterpretation) -> u16 {
693        let def = self.current_def;
694        self.current_def -= meaning.num_def_levels();
695        self.def_meaning.push(meaning);
696        def
697    }
698
699    fn record_offsets(&mut self, offset_desc: &OffsetDesc) {
700        let rep_level = self.current_rep;
701        let (null_list_level, empty_list_level) =
702            match (offset_desc.validity.is_some(), offset_desc.has_empty_lists) {
703                (true, true) => {
704                    let level =
705                        self.checkout_def(DefinitionInterpretation::NullableAndEmptyableList);
706                    (level - 1, level)
707                }
708                (true, false) => (self.checkout_def(DefinitionInterpretation::NullableList), 0),
709                (false, true) => (
710                    0,
711                    self.checkout_def(DefinitionInterpretation::EmptyableList),
712                ),
713                (false, false) => {
714                    self.checkout_def(DefinitionInterpretation::AllValidList);
715                    (0, 0)
716                }
717            };
718        self.current_rep -= 1;
719
720        if let Some(validity) = &offset_desc.validity {
721            self.do_record_validity(validity, null_list_level);
722        }
723
724        // We write into the spare buffers and read from the active buffers
725        // and then swap at the end.  This way we don't write over what we
726        // are reading.
727
728        let mut new_len = 0;
729        let expected_len = offset_desc.num_values + self.current_num_specials;
730        if expected_len == 0 {
731            // Offsets [0] mean no list values, so no levels.
732            self.current_len = 0;
733            return;
734        }
735        assert!(self.rep_levels.len() >= expected_len - 1);
736        if self.def_levels.is_empty() {
737            let mut write_itr = self.spare_rep.iter_mut();
738            let mut read_iter = self.rep_levels.iter().copied();
739            for w in offset_desc.offsets.windows(2) {
740                let len = w[1] - w[0];
741                // len can't be 0 because then we'd have def levels
742                assert!(len > 0);
743                let rep = read_iter.next().unwrap();
744                let list_level = if rep == 0 { rep_level } else { rep };
745                *write_itr.next().unwrap() = list_level;
746
747                for _ in 1..len {
748                    *write_itr.next().unwrap() = 0;
749                }
750                new_len += len as usize;
751            }
752            std::mem::swap(&mut self.rep_levels, &mut self.spare_rep);
753        } else {
754            assert!(self.def_levels.len() >= expected_len - 1);
755            let mut def_write_itr = self.spare_def.iter_mut();
756            let mut rep_write_itr = self.spare_rep.iter_mut();
757            let mut rep_read_itr = self.rep_levels.iter().copied();
758            let mut def_read_itr = self.def_levels.iter().copied();
759            let specials_to_pass = self.current_num_specials;
760            let mut specials_passed = 0;
761
762            for w in offset_desc.offsets.windows(2) {
763                let mut def = def_read_itr.next().unwrap();
764                // Copy over any higher-level special values in place
765                while def > SPECIAL_THRESHOLD {
766                    *def_write_itr.next().unwrap() = def;
767                    *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap();
768                    def = def_read_itr.next().unwrap();
769                    new_len += 1;
770                    specials_passed += 1;
771                }
772
773                let len = w[1] - w[0];
774                let rep = rep_read_itr.next().unwrap();
775
776                // If the rep_level is 0 then we are the first list level
777                // otherwise we are starting a higher level list so keep
778                // existing rep level
779                let list_level = if rep == 0 { rep_level } else { rep };
780
781                if def == 0 && len > 0 {
782                    // New valid list, write a rep level and then add new 0/0 items
783                    *def_write_itr.next().unwrap() = 0;
784                    *rep_write_itr.next().unwrap() = list_level;
785
786                    for _ in 1..len {
787                        *def_write_itr.next().unwrap() = 0;
788                        *rep_write_itr.next().unwrap() = 0;
789                    }
790
791                    new_len += len as usize;
792                } else if def == 0 {
793                    // Empty list, insert new special
794                    *def_write_itr.next().unwrap() = empty_list_level + SPECIAL_THRESHOLD;
795                    *rep_write_itr.next().unwrap() = list_level;
796                    new_len += 1;
797                } else {
798                    // Either the list is null or one of its struct parents
799                    // is null.  Promote it to a special value.
800                    *def_write_itr.next().unwrap() = def + SPECIAL_THRESHOLD;
801                    *rep_write_itr.next().unwrap() = list_level;
802                    new_len += 1;
803                }
804            }
805
806            // If we have any special values at the end, we need to copy them over
807            while specials_passed < specials_to_pass {
808                *def_write_itr.next().unwrap() = def_read_itr.next().unwrap();
809                *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap();
810                new_len += 1;
811                specials_passed += 1;
812            }
813            std::mem::swap(&mut self.def_levels, &mut self.spare_def);
814            std::mem::swap(&mut self.rep_levels, &mut self.spare_rep);
815        }
816
817        self.current_len = new_len;
818        self.current_num_specials += offset_desc.num_specials;
819    }
820
821    fn do_record_validity(&mut self, validity: &BooleanBuffer, null_level: u16) {
822        assert!(self.def_levels.len() >= validity.len() + self.current_num_specials);
823        debug_assert!(
824            self.current_len == 0 || self.current_len == validity.len() + self.current_num_specials
825        );
826        self.current_len = validity.len();
827
828        let mut def_read_itr = self.def_levels.iter().copied();
829        let mut def_write_itr = self.spare_def.iter_mut();
830
831        let specials_to_pass = self.current_num_specials;
832        let mut specials_passed = 0;
833
834        for incoming_validity in validity.iter() {
835            let mut def = def_read_itr.next().unwrap();
836            while def > SPECIAL_THRESHOLD {
837                *def_write_itr.next().unwrap() = def;
838                def = def_read_itr.next().unwrap();
839                specials_passed += 1;
840            }
841            if def == 0 && !incoming_validity {
842                *def_write_itr.next().unwrap() = null_level;
843            } else {
844                *def_write_itr.next().unwrap() = def;
845            }
846        }
847
848        while specials_passed < specials_to_pass {
849            *def_write_itr.next().unwrap() = def_read_itr.next().unwrap();
850            specials_passed += 1;
851        }
852
853        std::mem::swap(&mut self.def_levels, &mut self.spare_def);
854    }
855
856    fn multiply_levels(&mut self, multiplier: usize) {
857        let old_len = self.current_len;
858        // All non-special values will be broadcasted by the multiplier.  Special values are copied as-is.
859        self.current_len =
860            (self.current_len - self.current_num_specials) * multiplier + self.current_num_specials;
861
862        if self.rep_levels.is_empty() && self.def_levels.is_empty() {
863            // All valid with no rep/def levels, nothing to do
864            return;
865        } else if self.rep_levels.is_empty() {
866            assert!(self.def_levels.len() >= self.current_len);
867            // No rep levels, just multiply the def levels
868            let mut def_read_itr = self.def_levels.iter().copied();
869            let mut def_write_itr = self.spare_def.iter_mut();
870            for _ in 0..old_len {
871                let mut def = def_read_itr.next().unwrap();
872                while def > SPECIAL_THRESHOLD {
873                    *def_write_itr.next().unwrap() = def;
874                    def = def_read_itr.next().unwrap();
875                }
876                for _ in 0..multiplier {
877                    *def_write_itr.next().unwrap() = def;
878                }
879            }
880        } else if self.def_levels.is_empty() {
881            assert!(self.rep_levels.len() >= self.current_len);
882            // No def levels, just multiply the rep levels
883            let mut rep_read_itr = self.rep_levels.iter().copied();
884            let mut rep_write_itr = self.spare_rep.iter_mut();
885            for _ in 0..old_len {
886                let rep = rep_read_itr.next().unwrap();
887                for _ in 0..multiplier {
888                    *rep_write_itr.next().unwrap() = rep;
889                }
890            }
891        } else {
892            assert!(self.rep_levels.len() >= self.current_len);
893            assert!(self.def_levels.len() >= self.current_len);
894            let mut rep_read_itr = self.rep_levels.iter().copied();
895            let mut def_read_itr = self.def_levels.iter().copied();
896            let mut rep_write_itr = self.spare_rep.iter_mut();
897            let mut def_write_itr = self.spare_def.iter_mut();
898            for _ in 0..old_len {
899                let mut def = def_read_itr.next().unwrap();
900                while def > SPECIAL_THRESHOLD {
901                    *def_write_itr.next().unwrap() = def;
902                    *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap();
903                    def = def_read_itr.next().unwrap();
904                }
905                let rep = rep_read_itr.next().unwrap();
906                for _ in 0..multiplier {
907                    *def_write_itr.next().unwrap() = def;
908                    *rep_write_itr.next().unwrap() = rep;
909                }
910            }
911        }
912        std::mem::swap(&mut self.def_levels, &mut self.spare_def);
913        std::mem::swap(&mut self.rep_levels, &mut self.spare_rep);
914    }
915
916    fn record_validity_buf(&mut self, validity: &Option<BooleanBuffer>) {
917        if let Some(validity) = validity {
918            let def_level = self.checkout_def(DefinitionInterpretation::NullableItem);
919            self.do_record_validity(validity, def_level);
920        } else {
921            self.checkout_def(DefinitionInterpretation::AllValidItem);
922        }
923    }
924
925    fn record_validity(&mut self, validity_desc: &ValidityDesc) {
926        self.record_validity_buf(&validity_desc.validity)
927    }
928
929    fn record_fsl(&mut self, fsl_desc: &FslDesc) {
930        self.has_fsl = true;
931        self.record_validity_buf(&fsl_desc.validity);
932        self.multiply_levels(fsl_desc.dimension);
933    }
934
935    fn normalize_specials(&mut self) {
936        for def in self.def_levels.iter_mut() {
937            if *def > SPECIAL_THRESHOLD {
938                *def -= SPECIAL_THRESHOLD;
939            }
940        }
941    }
942
943    fn normalize_specials_and_plan_splits(
944        &mut self,
945        def_meaning: &[DefinitionInterpretation],
946        max_levels_per_page: Option<u64>,
947        num_rows: u64,
948        num_values: u64,
949    ) -> Result<MiniBlockRepDefBudget> {
950        // Extremely sparse lists can have many rep/def levels for very few
951        // visible leaf values.  If this ratio becomes too skewed then a
952        // mini-block rep/def chunk can exceed its packed metadata budget even
953        // though the value buffers are small.  We detect that case while
954        // normalizing special def levels and split on top-level row boundaries
955        // so each emitted dense mini-block page stays within the budget.
956        if self.def_levels.is_empty() {
957            return Ok(MiniBlockRepDefBudget::WithinBudget);
958        }
959
960        if self.rep_levels.is_empty() {
961            self.normalize_specials();
962            return Ok(MiniBlockRepDefBudget::WithinBudget);
963        }
964
965        if self.rep_levels.len() != self.def_levels.len() {
966            return Err(Error::internal(format!(
967                "Cannot plan structural page splits with mismatched rep/def lengths: rep={}, def={}",
968                self.rep_levels.len(),
969                self.def_levels.len()
970            )));
971        }
972
973        let Some(max_levels_per_page) = max_levels_per_page else {
974            self.normalize_specials();
975            return Ok(MiniBlockRepDefBudget::WithinBudget);
976        };
977
978        if num_values == 0 {
979            self.normalize_specials();
980            return Ok(MiniBlockRepDefBudget::WithinBudget);
981        }
982
983        let max_schema_rep = def_meaning.iter().filter(|level| level.is_list()).count() as u16;
984        let max_visible_level = SerializedRepDefs::max_visible_level(def_meaning);
985        let should_plan = !self.has_fsl && max_schema_rep > 0 && max_visible_level.is_some();
986
987        if !should_plan {
988            self.normalize_specials();
989            return Ok(MiniBlockRepDefBudget::WithinBudget);
990        }
991
992        let max_visible_level = max_visible_level.unwrap();
993        let mut splits = Vec::new();
994        let mut counted_rows = 0u64;
995        let mut counted_values = 0u64;
996        let mut saw_structural_overhead = false;
997        let mut single_row_over_budget_levels = None;
998
999        let mut current_row_level_start = None;
1000        let mut current_row_num_values = 0u64;
1001
1002        let mut current_page_row_start = 0u64;
1003        let mut current_page_num_rows = 0u64;
1004        let mut current_page_level_start = 0usize;
1005        let mut current_page_level_end = 0usize;
1006        let mut current_page_value_start = 0u64;
1007        let mut current_page_num_values = 0u64;
1008        let mut current_page_num_levels = 0u64;
1009        let mut current_page_has_structural_overhead = false;
1010
1011        let mut finish_row =
1012            |row_level_start: usize, row_level_end: usize, row_num_values: u64| -> Result<()> {
1013                let row_num_levels = (row_level_end - row_level_start) as u64;
1014                let row_has_structural_overhead = row_num_levels > row_num_values;
1015                saw_structural_overhead |= row_has_structural_overhead;
1016
1017                if row_has_structural_overhead && row_num_levels > max_levels_per_page {
1018                    single_row_over_budget_levels = Some(row_num_levels);
1019                }
1020
1021                if current_page_num_rows > 0
1022                    && (current_page_has_structural_overhead || row_has_structural_overhead)
1023                    && current_page_num_levels + row_num_levels > max_levels_per_page
1024                {
1025                    splits.push(MiniBlockRepDefSplit {
1026                        row_start: current_page_row_start,
1027                        num_rows: current_page_num_rows,
1028                        level_range: current_page_level_start..current_page_level_end,
1029                        value_start: current_page_value_start,
1030                        num_values: current_page_num_values,
1031                    });
1032                    current_page_row_start = counted_rows;
1033                    current_page_num_rows = 0;
1034                    current_page_level_start = row_level_start;
1035                    current_page_value_start = counted_values;
1036                    current_page_num_values = 0;
1037                    current_page_num_levels = 0;
1038                    current_page_has_structural_overhead = false;
1039                }
1040
1041                if current_page_num_rows == 0 {
1042                    current_page_level_start = row_level_start;
1043                }
1044                current_page_num_rows += 1;
1045                current_page_level_end = row_level_end;
1046                current_page_num_values += row_num_values;
1047                current_page_num_levels += row_num_levels;
1048                current_page_has_structural_overhead |= row_has_structural_overhead;
1049                counted_rows += 1;
1050                counted_values += row_num_values;
1051                Ok(())
1052            };
1053
1054        for (idx, (rep_level, def_level)) in self
1055            .rep_levels
1056            .iter()
1057            .copied()
1058            .zip(self.def_levels.iter_mut())
1059            .enumerate()
1060        {
1061            if *def_level > SPECIAL_THRESHOLD {
1062                *def_level -= SPECIAL_THRESHOLD;
1063            }
1064
1065            if rep_level == max_schema_rep {
1066                if let Some(level_start) = current_row_level_start {
1067                    finish_row(level_start, idx, current_row_num_values)?;
1068                    current_row_num_values = 0;
1069                } else if idx != 0 {
1070                    return Err(Error::internal(format!(
1071                        "Cannot plan structural page splits: first top-level row starts at level {}, expected 0",
1072                        idx
1073                    )));
1074                }
1075                current_row_level_start = Some(idx);
1076            }
1077
1078            if current_row_level_start.is_none() {
1079                return Err(Error::internal(
1080                    "Cannot plan structural page splits: found levels before the first top-level row start",
1081                ));
1082            }
1083            if *def_level <= max_visible_level {
1084                current_row_num_values += 1;
1085            }
1086        }
1087
1088        let Some(level_start) = current_row_level_start else {
1089            return Err(Error::internal(
1090                "Cannot plan structural page splits: found no top-level row starts",
1091            ));
1092        };
1093        finish_row(level_start, self.rep_levels.len(), current_row_num_values)?;
1094
1095        if counted_rows != num_rows {
1096            return Err(Error::internal(format!(
1097                "Cannot plan structural page splits: expected {} top-level row starts, found {}",
1098                num_rows, counted_rows
1099            )));
1100        }
1101        if counted_values != num_values {
1102            return Err(Error::internal(format!(
1103                "Cannot plan structural page splits: counted {} visible values, expected {}",
1104                counted_values, num_values
1105            )));
1106        }
1107        if !saw_structural_overhead {
1108            return Ok(MiniBlockRepDefBudget::WithinBudget);
1109        }
1110        if let Some(row_num_levels) = single_row_over_budget_levels {
1111            return Ok(MiniBlockRepDefBudget::SingleRowOverBudget(row_num_levels));
1112        }
1113
1114        if current_page_num_rows > 0 {
1115            splits.push(MiniBlockRepDefSplit {
1116                row_start: current_page_row_start,
1117                num_rows: current_page_num_rows,
1118                level_range: current_page_level_start..current_page_level_end,
1119                value_start: current_page_value_start,
1120                num_values: current_page_num_values,
1121            });
1122        }
1123
1124        if splits.len() > 1 {
1125            Ok(MiniBlockRepDefBudget::RequiresPageSplit(splits))
1126        } else {
1127            Ok(MiniBlockRepDefBudget::WithinBudget)
1128        }
1129    }
1130
1131    fn build(mut self) -> SerializedRepDefs {
1132        if self.current_len == 0 {
1133            return SerializedRepDefs::new_with_fixed_size_list_levels(
1134                None,
1135                None,
1136                self.def_meaning,
1137                self.has_fsl,
1138            );
1139        }
1140
1141        self.normalize_specials();
1142
1143        let definition_levels = if self.def_levels.is_empty() {
1144            None
1145        } else {
1146            Some(self.def_levels)
1147        };
1148        let repetition_levels = if self.rep_levels.is_empty() {
1149            None
1150        } else {
1151            Some(self.rep_levels)
1152        };
1153
1154        // Need to reverse the def meaning since we build rep / def levels in reverse
1155        let def_meaning = self.def_meaning.into_iter().rev().collect::<Vec<_>>();
1156
1157        SerializedRepDefs::new_with_fixed_size_list_levels(
1158            repetition_levels,
1159            definition_levels,
1160            def_meaning,
1161            self.has_fsl,
1162        )
1163    }
1164
1165    fn build_with_miniblock_repdef_budget(
1166        mut self,
1167        max_levels_per_page: Option<u64>,
1168        num_rows: u64,
1169        num_values: u64,
1170    ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> {
1171        if self.current_len == 0 {
1172            return Ok((
1173                SerializedRepDefs::new_with_fixed_size_list_levels(
1174                    None,
1175                    None,
1176                    self.def_meaning,
1177                    self.has_fsl,
1178                ),
1179                MiniBlockRepDefBudget::WithinBudget,
1180            ));
1181        }
1182
1183        // Need to reverse the def meaning since we build rep / def levels in reverse
1184        let def_meaning = std::mem::take(&mut self.def_meaning)
1185            .into_iter()
1186            .rev()
1187            .collect::<Vec<_>>();
1188        let budget = self.normalize_specials_and_plan_splits(
1189            &def_meaning,
1190            max_levels_per_page,
1191            num_rows,
1192            num_values,
1193        )?;
1194
1195        let definition_levels = if self.def_levels.is_empty() {
1196            None
1197        } else {
1198            Some(self.def_levels)
1199        };
1200        let repetition_levels = if self.rep_levels.is_empty() {
1201            None
1202        } else {
1203            Some(self.rep_levels)
1204        };
1205
1206        Ok((
1207            SerializedRepDefs::new_with_fixed_size_list_levels(
1208                repetition_levels,
1209                definition_levels,
1210                def_meaning,
1211                self.has_fsl,
1212            ),
1213            budget,
1214        ))
1215    }
1216}
1217
1218/// A structure used to collect validity buffers and offsets from arrow
1219/// arrays and eventually create repetition and definition levels
1220///
1221/// As we are encoding the structural encoders are given this struct and
1222/// will record the arrow information into it.  Once we hit a leaf node we
1223/// serialize the data into rep/def levels and write these into the page.
1224#[derive(Clone, Default, Debug)]
1225pub struct RepDefBuilder {
1226    // The rep/def info we have collected so far
1227    repdefs: Vec<RawRepDef>,
1228    // The current length, can get larger as we traverse lists (e.g. an
1229    // array might have 5 lists which results in 50 items)
1230    //
1231    // Starts uninitialized until we see the first rep/def item
1232    len: Option<usize>,
1233}
1234
1235impl RepDefBuilder {
1236    fn check_validity_len(&mut self, incoming_len: usize) {
1237        if let Some(len) = self.len {
1238            assert_eq!(incoming_len, len);
1239        } else {
1240            // First validity buffer we've seen
1241            self.len = Some(incoming_len);
1242        }
1243    }
1244
1245    fn num_layers(&self) -> usize {
1246        self.repdefs.len()
1247    }
1248
1249    /// The builder is "empty" if there is no repetition and no nulls.  In this case we don't need
1250    /// to store anything to disk (except the description)
1251    pub fn is_empty(&self) -> bool {
1252        self.repdefs
1253            .iter()
1254            .all(|r| matches!(r, RawRepDef::Validity(ValidityDesc { validity: None, .. })))
1255    }
1256
1257    /// Returns true if there is only a single layer of definition
1258    pub fn is_simple_validity(&self) -> bool {
1259        self.repdefs.len() == 1 && matches!(self.repdefs[0], RawRepDef::Validity(_))
1260    }
1261
1262    /// Registers a nullable validity bitmap
1263    pub fn add_validity_bitmap(&mut self, validity: NullBuffer) {
1264        self.check_validity_len(validity.len());
1265        if validity.null_count() == 0 {
1266            self.add_no_null(validity.len());
1267            return;
1268        }
1269        self.repdefs.push(RawRepDef::Validity(ValidityDesc {
1270            num_values: validity.len(),
1271            validity: Some(validity.into_inner()),
1272        }));
1273    }
1274
1275    /// Registers an all-valid validity layer
1276    pub fn add_no_null(&mut self, len: usize) {
1277        self.check_validity_len(len);
1278        self.repdefs.push(RawRepDef::Validity(ValidityDesc {
1279            validity: None,
1280            num_values: len,
1281        }));
1282    }
1283
1284    pub fn add_fsl(&mut self, validity: Option<NullBuffer>, dimension: usize, num_values: usize) {
1285        if let Some(len) = self.len {
1286            assert_eq!(num_values, len);
1287        }
1288        self.len = Some(num_values * dimension);
1289        debug_assert!(validity.is_none() || validity.as_ref().unwrap().len() == num_values);
1290        self.repdefs.push(RawRepDef::Fsl(FslDesc {
1291            num_values,
1292            validity: validity.map(|v| v.into_inner()),
1293            dimension,
1294        }))
1295    }
1296
1297    fn check_offset_len(&mut self, offsets: &[i64]) {
1298        if let Some(len) = self.len {
1299            assert!(offsets.len() == len + 1);
1300        }
1301        self.len = Some(offsets[offsets.len() - 1] as usize);
1302    }
1303
1304    fn do_add_offsets(
1305        &mut self,
1306        lengths: impl Iterator<Item = i64>,
1307        validity: Option<NullBuffer>,
1308        capacity: usize,
1309    ) -> bool {
1310        let mut num_specials = 0;
1311        let mut has_empty_lists = false;
1312        let mut has_garbage_values = false;
1313        let mut last_off: i64 = 0;
1314
1315        let mut normalized_offsets = Vec::with_capacity(capacity);
1316        normalized_offsets.push(0);
1317
1318        if let Some(ref validity) = validity {
1319            for (len, is_valid) in lengths.zip(validity.iter()) {
1320                match (is_valid, len == 0) {
1321                    (false, is_empty) => {
1322                        num_specials += 1;
1323                        has_garbage_values |= !is_empty;
1324                    }
1325                    (true, true) => {
1326                        num_specials += 1;
1327                        has_empty_lists = true;
1328                    }
1329                    _ => {
1330                        last_off += len;
1331                    }
1332                }
1333                normalized_offsets.push(last_off);
1334            }
1335        } else {
1336            for len in lengths {
1337                if len == 0 {
1338                    num_specials += 1;
1339                    has_empty_lists = true;
1340                }
1341                last_off += len;
1342                normalized_offsets.push(last_off);
1343            }
1344        }
1345
1346        self.check_offset_len(&normalized_offsets);
1347        self.repdefs.push(RawRepDef::Offsets(OffsetDesc {
1348            num_values: normalized_offsets.len() - 1,
1349            offsets: normalized_offsets.into(),
1350            validity: validity.map(|v| v.into_inner()),
1351            has_empty_lists,
1352            num_specials: num_specials as usize,
1353        }));
1354
1355        has_garbage_values
1356    }
1357
1358    /// Adds a layer of offsets
1359    ///
1360    /// Offsets are casted to a common type (i64) and also normalized.  Null lists are
1361    /// always represented by a zero-length (identical) pair of offsets and so the caller
1362    /// should filter out any garbage items before encoding them.  To assist with this the
1363    /// method will return true if any non-empty null lists were found.
1364    pub fn add_offsets<O: OffsetSizeTrait>(
1365        &mut self,
1366        offsets: OffsetBuffer<O>,
1367        validity: Option<NullBuffer>,
1368    ) -> bool {
1369        let inner = offsets.into_inner();
1370        let buffer_len = inner.len();
1371
1372        if O::IS_LARGE {
1373            let i64_buff = ScalarBuffer::<i64>::new(inner.into_inner(), 0, buffer_len);
1374            let lengths = i64_buff.windows(2).map(|off| off[1] - off[0]);
1375            self.do_add_offsets(lengths, validity, buffer_len)
1376        } else {
1377            let i32_buff = ScalarBuffer::<i32>::new(inner.into_inner(), 0, buffer_len);
1378            let lengths = i32_buff.windows(2).map(|off| (off[1] - off[0]) as i64);
1379            self.do_add_offsets(lengths, validity, buffer_len)
1380        }
1381    }
1382
1383    // When we are encoding data it arrives in batches.  For each batch we create a RepDefBuilder and collect the
1384    // various validity buffers and offset buffers from that batch.  Once we have enough batches to write a page we
1385    // need to take this collection of RepDefBuilders and concatenate them and then serialize them into rep/def levels.
1386    //
1387    // TODO: In the future, we may concatenate and serialize at the same time?
1388    //
1389    // This method takes care of the concatenation part.  First we collect all of layer 0 from each builder, then we
1390    // call this method.  Then we collect all of layer 1 from each builder and call this method.  And so on.
1391    //
1392    // That means this method should get a collection of `RawRepDef` where each item is the same kind (all validity or
1393    // all offsets) though the nullability / lengths may be different in each layer.
1394    fn concat_layers<'a>(
1395        layers: impl Iterator<Item = &'a RawRepDef>,
1396        num_layers: usize,
1397    ) -> RawRepDef {
1398        enum LayerKind {
1399            Validity,
1400            Fsl,
1401            Offsets,
1402        }
1403
1404        // We make two passes through the layers.  The first determines if we need to pay the cost of allocating
1405        // buffers.  The second pass actually adds the values.
1406        let mut collected = Vec::with_capacity(num_layers);
1407        let mut has_nulls = false;
1408        let mut layer_kind = LayerKind::Validity;
1409        let mut total_num_specials = 0;
1410        let mut all_dimension = 0;
1411        let mut all_has_empty_lists = false;
1412        let mut all_num_values = 0;
1413        for layer in layers {
1414            has_nulls |= layer.has_nulls();
1415            match layer {
1416                RawRepDef::Validity(_) => {
1417                    layer_kind = LayerKind::Validity;
1418                }
1419                RawRepDef::Offsets(OffsetDesc {
1420                    num_specials,
1421                    has_empty_lists,
1422                    ..
1423                }) => {
1424                    all_has_empty_lists |= *has_empty_lists;
1425                    layer_kind = LayerKind::Offsets;
1426                    total_num_specials += num_specials;
1427                }
1428                RawRepDef::Fsl(FslDesc { dimension, .. }) => {
1429                    layer_kind = LayerKind::Fsl;
1430                    all_dimension = *dimension;
1431                }
1432            }
1433            collected.push(layer);
1434            all_num_values += layer.num_values();
1435        }
1436
1437        // Shortcut if there are no nulls
1438        if !has_nulls {
1439            match layer_kind {
1440                LayerKind::Validity => {
1441                    return RawRepDef::Validity(ValidityDesc {
1442                        validity: None,
1443                        num_values: all_num_values,
1444                    });
1445                }
1446                LayerKind::Fsl => {
1447                    return RawRepDef::Fsl(FslDesc {
1448                        validity: None,
1449                        num_values: all_num_values,
1450                        dimension: all_dimension,
1451                    });
1452                }
1453                LayerKind::Offsets => {}
1454            }
1455        }
1456
1457        // Only allocate if needed
1458        let mut validity_builder = if has_nulls {
1459            BooleanBufferBuilder::new(all_num_values)
1460        } else {
1461            BooleanBufferBuilder::new(0)
1462        };
1463        let mut all_offsets = if matches!(layer_kind, LayerKind::Offsets) {
1464            let mut all_offsets = Vec::with_capacity(all_num_values);
1465            all_offsets.push(0);
1466            all_offsets
1467        } else {
1468            Vec::new()
1469        };
1470
1471        for layer in collected {
1472            match layer {
1473                RawRepDef::Validity(ValidityDesc {
1474                    validity: Some(validity),
1475                    ..
1476                }) => {
1477                    validity_builder.append_buffer(validity);
1478                }
1479                RawRepDef::Validity(ValidityDesc {
1480                    validity: None,
1481                    num_values,
1482                }) => {
1483                    validity_builder.append_n(*num_values, true);
1484                }
1485                RawRepDef::Fsl(FslDesc {
1486                    validity,
1487                    num_values,
1488                    ..
1489                }) => {
1490                    if let Some(validity) = validity {
1491                        validity_builder.append_buffer(validity);
1492                    } else {
1493                        validity_builder.append_n(*num_values, true);
1494                    }
1495                }
1496                RawRepDef::Offsets(OffsetDesc {
1497                    offsets,
1498                    validity: Some(validity),
1499                    has_empty_lists,
1500                    ..
1501                }) => {
1502                    all_has_empty_lists |= has_empty_lists;
1503                    validity_builder.append_buffer(validity);
1504                    let last = *all_offsets.last().unwrap();
1505                    all_offsets.extend(offsets.iter().skip(1).map(|off| *off + last));
1506                }
1507                RawRepDef::Offsets(OffsetDesc {
1508                    offsets,
1509                    validity: None,
1510                    has_empty_lists,
1511                    num_values,
1512                    ..
1513                }) => {
1514                    all_has_empty_lists |= has_empty_lists;
1515                    if has_nulls {
1516                        validity_builder.append_n(*num_values, true);
1517                    }
1518                    let last = *all_offsets.last().unwrap();
1519                    all_offsets.extend(offsets.iter().skip(1).map(|off| *off + last));
1520                }
1521            }
1522        }
1523        let validity = if has_nulls {
1524            Some(validity_builder.finish())
1525        } else {
1526            None
1527        };
1528        match layer_kind {
1529            LayerKind::Fsl => RawRepDef::Fsl(FslDesc {
1530                validity,
1531                num_values: all_num_values,
1532                dimension: all_dimension,
1533            }),
1534            LayerKind::Validity => RawRepDef::Validity(ValidityDesc {
1535                validity,
1536                num_values: all_num_values,
1537            }),
1538            LayerKind::Offsets => RawRepDef::Offsets(OffsetDesc {
1539                offsets: all_offsets.into(),
1540                validity,
1541                has_empty_lists: all_has_empty_lists,
1542                num_values: all_num_values,
1543                num_specials: total_num_specials,
1544            }),
1545        }
1546    }
1547
1548    /// Converts the validity / offsets buffers that have been gathered so far
1549    /// into repetition and definition levels
1550    pub fn serialize(builders: Vec<Self>) -> SerializedRepDefs {
1551        Self::normalize(builders).serialize()
1552    }
1553
1554    pub(crate) fn normalize(builders: Vec<Self>) -> NormalizedStructuralPlan {
1555        assert!(!builders.is_empty());
1556        let num_layers = builders[0].num_layers();
1557        debug_assert!(
1558            builders
1559                .iter()
1560                .all(|builder| builder.num_layers() == num_layers)
1561        );
1562        let layers = (0..num_layers)
1563            .map(|layer_index| {
1564                Self::concat_layers(
1565                    builders.iter().map(|b| &b.repdefs[layer_index]),
1566                    builders.len(),
1567                )
1568            })
1569            .collect::<Vec<_>>();
1570        NormalizedStructuralPlan {
1571            layers,
1572            dense_all_valid: builders.iter().all(Self::is_empty),
1573        }
1574    }
1575}
1576
1577/// Starts with serialized repetition and definition levels and unravels
1578/// them into validity buffers and offsets buffers
1579///
1580/// This is used during decoding to create the necessary arrow structures
1581#[derive(Debug)]
1582pub struct RepDefUnraveler {
1583    sparse: Option<SparseStructuralUnraveler>,
1584    rep_levels: Option<LevelBuffer>,
1585    def_levels: Option<LevelBuffer>,
1586    // Maps from definition level to the rep level at which that definition level is visible
1587    levels_to_rep: Vec<u16>,
1588    def_meaning: Arc<[DefinitionInterpretation]>,
1589    // Current definition level to compare to.
1590    current_def_cmp: u16,
1591    // Current rep level, determines which specials we can see
1592    current_rep_cmp: u16,
1593    // Current layer index, 0 means inner-most layer and it counts up from there.  Used to index
1594    // into special_defs
1595    current_layer: usize,
1596    // Number of items in the inner-most layer (needed if the definition levels are not present)
1597    num_items: u64,
1598}
1599
1600impl RepDefUnraveler {
1601    /// Creates a new unraveler from serialized repetition and definition information
1602    pub fn new(
1603        rep_levels: Option<LevelBuffer>,
1604        def_levels: Option<LevelBuffer>,
1605        def_meaning: Arc<[DefinitionInterpretation]>,
1606        num_items: u64,
1607    ) -> Self {
1608        let mut levels_to_rep = Vec::with_capacity(def_meaning.len());
1609        let mut rep_counter = 0;
1610        // Level=0 is always visible and means valid item
1611        levels_to_rep.push(0);
1612        for meaning in def_meaning.as_ref() {
1613            match meaning {
1614                DefinitionInterpretation::AllValidItem | DefinitionInterpretation::AllValidList => {
1615                    // There is no corresponding level, so nothing to put in levels_to_rep
1616                }
1617                DefinitionInterpretation::NullableItem => {
1618                    // Some null structs are not visible at inner rep levels in cases like LIST<STRUCT<LIST<...>>>
1619                    levels_to_rep.push(rep_counter);
1620                }
1621                DefinitionInterpretation::NullableList => {
1622                    rep_counter += 1;
1623                    levels_to_rep.push(rep_counter);
1624                }
1625                DefinitionInterpretation::EmptyableList => {
1626                    rep_counter += 1;
1627                    levels_to_rep.push(rep_counter);
1628                }
1629                DefinitionInterpretation::NullableAndEmptyableList => {
1630                    rep_counter += 1;
1631                    levels_to_rep.push(rep_counter);
1632                    levels_to_rep.push(rep_counter);
1633                }
1634            }
1635        }
1636        Self {
1637            sparse: None,
1638            rep_levels,
1639            def_levels,
1640            current_def_cmp: 0,
1641            current_rep_cmp: 0,
1642            levels_to_rep,
1643            current_layer: 0,
1644            def_meaning,
1645            num_items,
1646        }
1647    }
1648
1649    pub(crate) fn new_sparse(plan: SparseStructuralPlan) -> Self {
1650        Self {
1651            sparse: Some(SparseStructuralUnraveler::new(plan)),
1652            rep_levels: None,
1653            def_levels: None,
1654            levels_to_rep: Vec::new(),
1655            def_meaning: Arc::new([]),
1656            current_def_cmp: 0,
1657            current_rep_cmp: 0,
1658            current_layer: 0,
1659            num_items: 0,
1660        }
1661    }
1662
1663    fn ensure_exhausted(&self) -> Result<()> {
1664        if let Some(sparse) = &self.sparse {
1665            sparse.ensure_exhausted()?;
1666        }
1667        Ok(())
1668    }
1669
1670    fn is_sparse(&self) -> bool {
1671        self.sparse.is_some()
1672    }
1673
1674    pub fn is_all_valid(&self) -> bool {
1675        if let Some(sparse) = &self.sparse {
1676            return sparse.is_all_valid();
1677        }
1678        self.def_levels.is_none() || self.def_meaning[self.current_layer].is_all_valid()
1679    }
1680
1681    /// If the current level is a repetition layer then this returns the number of lists
1682    /// at this level.
1683    ///
1684    /// This is not valid to call when the current level is a struct/primitive layer because
1685    /// in some cases there may be no rep or def information to know this.
1686    pub fn max_lists(&self) -> Result<usize> {
1687        if let Some(sparse) = &self.sparse {
1688            return sparse.max_lists();
1689        }
1690        debug_assert!(
1691            self.def_meaning[self.current_layer] != DefinitionInterpretation::NullableItem
1692        );
1693        Ok(self
1694            .rep_levels
1695            .as_ref()
1696            // Worst case every rep item is max_rep and a new list
1697            .map(|levels| levels.len())
1698            .unwrap_or(0))
1699    }
1700
1701    /// Unravels a layer of offsets from the unraveler into the given offset width
1702    ///
1703    /// When decoding a list the caller should first unravel the offsets and then
1704    /// unravel the validity (this is the opposite order used during encoding)
1705    pub fn unravel_offsets<T: ArrowNativeType>(
1706        &mut self,
1707        offsets: &mut Vec<T>,
1708        validity: Option<&mut BooleanBufferBuilder>,
1709    ) -> Result<()> {
1710        if let Some(sparse) = self.sparse.as_mut() {
1711            return sparse.unravel_offsets(offsets, validity);
1712        }
1713        let rep_levels = self
1714            .rep_levels
1715            .as_mut()
1716            .expect("Expected repetition level but data didn't contain repetition");
1717        let valid_level = self.current_def_cmp;
1718        let (null_level, empty_level) = match self.def_meaning[self.current_layer] {
1719            DefinitionInterpretation::NullableList => {
1720                self.current_def_cmp += 1;
1721                (valid_level + 1, 0)
1722            }
1723            DefinitionInterpretation::EmptyableList => {
1724                self.current_def_cmp += 1;
1725                (0, valid_level + 1)
1726            }
1727            DefinitionInterpretation::NullableAndEmptyableList => {
1728                self.current_def_cmp += 2;
1729                (valid_level + 1, valid_level + 2)
1730            }
1731            DefinitionInterpretation::AllValidList => (0, 0),
1732            _ => unreachable!(),
1733        };
1734        self.current_layer += 1;
1735
1736        // This is the highest def level that is still visible.  Once we hit a list then
1737        // we stop looking because any null / empty list (or list masked by a higher level
1738        // null) will not be visible
1739        let mut max_level = null_level.max(empty_level).max(valid_level);
1740        // Anything higher than this (but less than max_level) is a null struct masking our
1741        // list.  We will materialize this is a null list.
1742        let upper_null = max_level;
1743        for level in self.def_meaning[self.current_layer..].iter() {
1744            match level {
1745                DefinitionInterpretation::NullableItem => {
1746                    max_level += 1;
1747                }
1748                DefinitionInterpretation::AllValidItem => {}
1749                _ => {
1750                    break;
1751                }
1752            }
1753        }
1754
1755        let mut curlen: usize = offsets.last().map(|o| o.as_usize()).unwrap_or(0);
1756
1757        // If offsets is empty this is a no-op.  If offsets is not empty that means we already
1758        // added a set of offsets.  For example, we might have added [0, 3, 5] (2 lists).  Now
1759        // say we want to add [0, 1, 4] (2 lists).  We should get [0, 3, 5, 6, 9] (4 lists).  If
1760        // we don't pop here we get [0, 3, 5, 5, 6, 9] which is wrong.
1761        //
1762        // Or, to think about it another way, if every unraveler adds the starting 0 and the trailing
1763        // length then we have N + unravelers.len() values instead of N + 1.
1764        offsets.pop();
1765
1766        let to_offset = |val: usize| {
1767            T::from_usize(val)
1768            .ok_or_else(|| Error::invalid_input("A single batch had more than i32::MAX values and so a large container type is required"))
1769        };
1770        self.current_rep_cmp += 1;
1771        if let Some(def_levels) = &mut self.def_levels {
1772            assert!(rep_levels.len() == def_levels.len());
1773            // It's possible validity is None even if we have def levels.  For example, we might have
1774            // empty lists (which require def levels) but no nulls.
1775            let mut push_validity: Box<dyn FnMut(bool)> = if let Some(validity) = validity {
1776                Box::new(|is_valid| validity.append(is_valid))
1777            } else {
1778                Box::new(|_| {})
1779            };
1780            // This is a strange access pattern.  We are iterating over the rep/def levels and
1781            // at the same time writing the rep/def levels.  This means we need both a mutable
1782            // and immutable reference to the rep/def levels.
1783            let mut read_idx = 0;
1784            let mut write_idx = 0;
1785            while read_idx < rep_levels.len() {
1786                // SAFETY: We assert that rep_levels and def_levels have the same
1787                // len and read_idx and write_idx can never go past the end.
1788                unsafe {
1789                    let rep_val = *rep_levels.get_unchecked(read_idx);
1790                    if rep_val != 0 {
1791                        let def_val = *def_levels.get_unchecked(read_idx);
1792                        // Copy over
1793                        *rep_levels.get_unchecked_mut(write_idx) = rep_val - 1;
1794                        *def_levels.get_unchecked_mut(write_idx) = def_val;
1795                        write_idx += 1;
1796
1797                        if def_val == 0 {
1798                            // This is a valid list
1799                            offsets.push(to_offset(curlen)?);
1800                            curlen += 1;
1801                            push_validity(true);
1802                        } else if def_val > max_level {
1803                            // This is not visible at this rep level, do not add to offsets, but keep in repdef
1804                        } else if def_val == null_level || def_val > upper_null {
1805                            // This is a null list (or a list masked by a null struct)
1806                            offsets.push(to_offset(curlen)?);
1807                            push_validity(false);
1808                        } else if def_val == empty_level {
1809                            // This is an empty list
1810                            offsets.push(to_offset(curlen)?);
1811                            push_validity(true);
1812                        } else {
1813                            // New valid list starting with null item
1814                            offsets.push(to_offset(curlen)?);
1815                            curlen += 1;
1816                            push_validity(true);
1817                        }
1818                    } else {
1819                        curlen += 1;
1820                    }
1821                    read_idx += 1;
1822                }
1823            }
1824            offsets.push(to_offset(curlen)?);
1825            rep_levels.truncate(write_idx);
1826            def_levels.truncate(write_idx);
1827            Ok(())
1828        } else {
1829            // SAFETY: See above loop
1830            let mut read_idx = 0;
1831            let mut write_idx = 0;
1832            let old_offsets_len = offsets.len();
1833            while read_idx < rep_levels.len() {
1834                // SAFETY: read_idx / write_idx cannot go past rep_levels.len()
1835                unsafe {
1836                    let rep_val = *rep_levels.get_unchecked(read_idx);
1837                    if rep_val != 0 {
1838                        // Finish the current list
1839                        offsets.push(to_offset(curlen)?);
1840                        *rep_levels.get_unchecked_mut(write_idx) = rep_val - 1;
1841                        write_idx += 1;
1842                    }
1843                    curlen += 1;
1844                    read_idx += 1;
1845                }
1846            }
1847            let num_new_lists = offsets.len() - old_offsets_len;
1848            offsets.push(to_offset(curlen)?);
1849            // Truncate to the number of lists THIS unraveler produced (write_idx),
1850            // not `offsets.len() - 1` — the latter includes offsets contributed by
1851            // earlier unravelers in a multi-page read, which would leave too many
1852            // rep levels for the next (outer) layer and over-count its lists.
1853            rep_levels.truncate(write_idx);
1854            if let Some(validity) = validity {
1855                // Even though we don't have validity it is possible another unraveler did and so we need
1856                // to push all valids
1857                validity.append_n(num_new_lists, true);
1858            }
1859            Ok(())
1860        }
1861    }
1862
1863    pub fn skip_validity(&mut self) -> Result<()> {
1864        if let Some(sparse) = self.sparse.as_mut() {
1865            return sparse.skip_validity();
1866        }
1867        debug_assert!(self.is_all_valid());
1868        self.current_layer += 1;
1869        Ok(())
1870    }
1871
1872    /// Unravels a layer of validity from the definition levels
1873    pub fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) -> Result<()> {
1874        if let Some(sparse) = self.sparse.as_mut() {
1875            return sparse.unravel_validity(validity);
1876        }
1877        let meaning = self.def_meaning[self.current_layer];
1878        if meaning == DefinitionInterpretation::AllValidItem || self.def_levels.is_none() {
1879            self.current_layer += 1;
1880            validity.append_n(self.num_items as usize, true);
1881            return Ok(());
1882        }
1883
1884        self.current_layer += 1;
1885        let def_levels = &self.def_levels.as_ref().unwrap();
1886
1887        let current_def_cmp = self.current_def_cmp;
1888        self.current_def_cmp += 1;
1889
1890        for is_valid in def_levels.iter().filter_map(|&level| {
1891            if self.levels_to_rep[level as usize] <= self.current_rep_cmp {
1892                Some(level <= current_def_cmp)
1893            } else {
1894                None
1895            }
1896        }) {
1897            validity.append(is_valid);
1898        }
1899        Ok(())
1900    }
1901
1902    /// Removes all but the first definition level of each fixed-size-list slot
1903    ///
1904    /// The definition levels arrive with one entry per item.  A fixed-size-list
1905    /// layer has a single definition level per slot (all `dimension` items in a
1906    /// slot share it) so we keep every `dimension`-th level and drop the rest.
1907    ///
1908    /// `dimension` must be non-zero.  A zero dimension can only come from a
1909    /// malformed schema (writers reject it, see
1910    /// [`lance_core::datatypes::validate_fixed_size_list_dimensions`]) and is
1911    /// rejected here rather than allowed to run off the end of the buffer.
1912    pub fn decimate(&mut self, dimension: usize) -> Result<()> {
1913        if dimension == 0 {
1914            return Err(Error::invalid_input(
1915                "Cannot decimate repetition/definition levels with a fixed-size-list dimension of 0; dimension must be a positive integer",
1916            ));
1917        }
1918        if let Some(sparse) = self.sparse.as_mut() {
1919            return sparse.decimate(dimension);
1920        }
1921        if self.rep_levels.is_some() {
1922            // If we need to support this then I think we need to walk through the rep def levels to find
1923            // the spots at which we keep.  E.g. if we have:
1924            //  rep: 1 0 0 1 0 1 0 0 0 1 0 0
1925            //  def: 1 1 1 0 1 0 1 1 0 1 1 0
1926            //  dimension: 2
1927            //
1928            // The output should be:
1929            //  rep: 1 0 0 1 0 0 0
1930            //  def: 1 1 1 0 1 1 0
1931            //
1932            // Maybe there's some special logic for empty/null lists?  I'll save the headache for future me.
1933            todo!("Not yet supported FSL<...List<...>>");
1934        }
1935        let Some(def_levels) = self.def_levels.as_mut() else {
1936            return Ok(());
1937        };
1938        let mut read_idx = 0;
1939        let mut write_idx = 0;
1940        while read_idx < def_levels.len() {
1941            // SAFETY: `read_idx` is checked against the length by the loop condition and
1942            // `dimension >= 1` (checked above) means `write_idx <= read_idx`, so both
1943            // indices are in bounds.
1944            unsafe {
1945                *def_levels.get_unchecked_mut(write_idx) = *def_levels.get_unchecked(read_idx);
1946            }
1947            write_idx += 1;
1948            read_idx += dimension;
1949        }
1950        def_levels.truncate(write_idx);
1951        Ok(())
1952    }
1953}
1954
1955/// As we decode we may extract rep/def information from multiple pages (or multiple
1956/// chunks within a page).
1957///
1958/// For each chunk we create an unraveler.  Each unraveler can have a completely different
1959/// interpretation (e.g. one page might contain null items but no null structs and the next
1960/// page might have null structs but no null items).
1961///
1962/// Concatenating these unravelers would be tricky and expensive so instead we have a
1963/// composite unraveler which unravels across multiple unravelers.
1964///
1965/// Note: this class should be used even if there is only one page / unraveler.  This is
1966/// because the `RepDefUnraveler`'s API is more complex (it's meant to be called by this
1967/// class)
1968#[derive(Debug)]
1969pub struct CompositeRepDefUnraveler {
1970    unravelers: Vec<RepDefUnraveler>,
1971    comparisons: Vec<Self>,
1972}
1973
1974impl CompositeRepDefUnraveler {
1975    pub fn new(unravelers: Vec<RepDefUnraveler>) -> Self {
1976        Self {
1977            unravelers,
1978            comparisons: Vec::new(),
1979        }
1980    }
1981
1982    pub(crate) fn add_compatibility_check(&mut self, other: Self) {
1983        self.comparisons.push(other);
1984    }
1985
1986    pub(crate) fn has_sparse(&self) -> bool {
1987        self.unravelers.iter().any(RepDefUnraveler::is_sparse)
1988            || self.comparisons.iter().any(Self::has_sparse)
1989    }
1990
1991    pub(crate) fn ensure_exhausted(&self) -> Result<()> {
1992        for unraveler in &self.unravelers {
1993            unraveler.ensure_exhausted()?;
1994        }
1995        for comparison in &self.comparisons {
1996            comparison.ensure_exhausted()?;
1997        }
1998        Ok(())
1999    }
2000
2001    fn null_buffers_equal(
2002        left: &Option<NullBuffer>,
2003        right: &Option<NullBuffer>,
2004        expected_len: usize,
2005    ) -> bool {
2006        match (left, right) {
2007            (None, None) => true,
2008            (Some(left), Some(right)) => {
2009                left.len() == expected_len
2010                    && right.len() == expected_len
2011                    && left.iter().eq(right.iter())
2012            }
2013            (None, Some(right)) => right.len() == expected_len && right.null_count() == 0,
2014            (Some(left), None) => left.len() == expected_len && left.null_count() == 0,
2015        }
2016    }
2017
2018    fn decimate(&mut self, dimension: usize) -> Result<()> {
2019        for unraveler in &mut self.unravelers {
2020            unraveler.decimate(dimension)?;
2021        }
2022        for comparison in &mut self.comparisons {
2023            comparison.decimate(dimension)?;
2024        }
2025        Ok(())
2026    }
2027
2028    /// Unravels a layer of validity
2029    ///
2030    /// Returns None if there are no null items in this layer
2031    pub fn unravel_validity(&mut self, num_values: usize) -> Result<Option<NullBuffer>> {
2032        let is_all_valid = self
2033            .unravelers
2034            .iter()
2035            .all(|unraveler| unraveler.is_all_valid());
2036
2037        let validity = if is_all_valid {
2038            for unraveler in self.unravelers.iter_mut() {
2039                unraveler.skip_validity()?;
2040            }
2041            None
2042        } else {
2043            let mut validity = BooleanBufferBuilder::new(num_values);
2044            for unraveler in self.unravelers.iter_mut() {
2045                unraveler.unravel_validity(&mut validity)?;
2046            }
2047            Some(NullBuffer::new(validity.finish()))
2048        };
2049        for comparison in &mut self.comparisons {
2050            let other = comparison.unravel_validity(num_values)?;
2051            if !Self::null_buffers_equal(&validity, &other, num_values) {
2052                return Err(Error::invalid_input_source(
2053                    format!(
2054                        "Structural sibling fields have incompatible validity metadata for {num_values} values"
2055                    )
2056                    .into(),
2057                ));
2058            }
2059        }
2060        Ok(validity)
2061    }
2062
2063    pub fn unravel_fsl_validity(
2064        &mut self,
2065        num_values: usize,
2066        dimension: usize,
2067    ) -> Result<Option<NullBuffer>> {
2068        self.decimate(dimension)?;
2069        self.unravel_validity(num_values)
2070    }
2071
2072    /// Unravels a layer of offsets (and the validity for that layer)
2073    pub fn unravel_offsets<T: ArrowNativeType>(
2074        &mut self,
2075    ) -> Result<(OffsetBuffer<T>, Option<NullBuffer>)> {
2076        let mut is_all_valid = true;
2077        let mut max_num_lists: usize = 0;
2078        for unraveler in self.unravelers.iter() {
2079            is_all_valid &= unraveler.is_all_valid();
2080            max_num_lists = max_num_lists
2081                .checked_add(unraveler.max_lists()?)
2082                .ok_or_else(|| {
2083                    Error::invalid_input_source(
2084                        "Combined repetition/definition list count exceeds usize::MAX".into(),
2085                    )
2086                })?;
2087        }
2088
2089        let mut validity = if is_all_valid {
2090            None
2091        } else {
2092            // Note: This is probably an over-estimate and potentially even an under-estimate.  We only know
2093            // right now how many items we have and not how many rows.  (TODO: Shouldn't we know the # of rows?)
2094            Some(BooleanBufferBuilder::new(max_num_lists))
2095        };
2096
2097        let mut offsets = Vec::with_capacity(max_num_lists + 1);
2098
2099        for unraveler in self.unravelers.iter_mut() {
2100            unraveler.unravel_offsets(&mut offsets, validity.as_mut())?;
2101        }
2102
2103        let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
2104        let validity = validity.map(|mut v| NullBuffer::new(v.finish()));
2105        for comparison in &mut self.comparisons {
2106            let (other_offsets, other_validity) = comparison.unravel_offsets::<T>()?;
2107            if offsets.as_ref() != other_offsets.as_ref()
2108                || !Self::null_buffers_equal(
2109                    &validity,
2110                    &other_validity,
2111                    offsets.len().saturating_sub(1),
2112                )
2113            {
2114                return Err(Error::invalid_input_source(
2115                    format!(
2116                        "Structural sibling fields have incompatible list metadata for {} slots",
2117                        offsets.len().saturating_sub(1)
2118                    )
2119                    .into(),
2120                ));
2121            }
2122        }
2123
2124        Ok((offsets, validity))
2125    }
2126}
2127
2128/// A [`ControlWordIterator`] when there are both repetition and definition levels
2129///
2130/// The iterator will put the repetition level in the upper bits and the definition
2131/// level in the lower bits.  The number of bits used for each level is determined
2132/// by the width of the repetition and definition levels.
2133#[derive(Debug)]
2134pub struct BinaryControlWordIterator<I: Iterator<Item = (u16, u16)>, W> {
2135    repdef: I,
2136    def_width: usize,
2137    max_rep: u16,
2138    max_visible_def: u16,
2139    rep_mask: u16,
2140    def_mask: u16,
2141    bits_rep: u8,
2142    bits_def: u8,
2143    phantom: std::marker::PhantomData<W>,
2144}
2145
2146impl<I: Iterator<Item = (u16, u16)>> BinaryControlWordIterator<I, u8> {
2147    fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2148        let next = self.repdef.next()?;
2149        let control_word: u8 =
2150            (((next.0 & self.rep_mask) as u8) << self.def_width) + ((next.1 & self.def_mask) as u8);
2151        buf.push(control_word);
2152        let is_new_row = next.0 == self.max_rep;
2153        let is_visible = next.1 <= self.max_visible_def;
2154        let is_valid_item = next.1 == 0;
2155        Some(ControlWordDesc {
2156            is_new_row,
2157            is_visible,
2158            is_valid_item,
2159        })
2160    }
2161}
2162
2163impl<I: Iterator<Item = (u16, u16)>> BinaryControlWordIterator<I, u16> {
2164    fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2165        let next = self.repdef.next()?;
2166        let control_word: u16 =
2167            ((next.0 & self.rep_mask) << self.def_width) + (next.1 & self.def_mask);
2168        let control_word = control_word.to_le_bytes();
2169        buf.push(control_word[0]);
2170        buf.push(control_word[1]);
2171        let is_new_row = next.0 == self.max_rep;
2172        let is_visible = next.1 <= self.max_visible_def;
2173        let is_valid_item = next.1 == 0;
2174        Some(ControlWordDesc {
2175            is_new_row,
2176            is_visible,
2177            is_valid_item,
2178        })
2179    }
2180}
2181
2182impl<I: Iterator<Item = (u16, u16)>> BinaryControlWordIterator<I, u32> {
2183    fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2184        let next = self.repdef.next()?;
2185        let control_word: u32 = (((next.0 & self.rep_mask) as u32) << self.def_width)
2186            + ((next.1 & self.def_mask) as u32);
2187        let control_word = control_word.to_le_bytes();
2188        buf.push(control_word[0]);
2189        buf.push(control_word[1]);
2190        buf.push(control_word[2]);
2191        buf.push(control_word[3]);
2192        let is_new_row = next.0 == self.max_rep;
2193        let is_visible = next.1 <= self.max_visible_def;
2194        let is_valid_item = next.1 == 0;
2195        Some(ControlWordDesc {
2196            is_new_row,
2197            is_visible,
2198            is_valid_item,
2199        })
2200    }
2201}
2202
2203/// A [`ControlWordIterator`] when there are only definition levels or only repetition levels
2204#[derive(Debug)]
2205pub struct UnaryControlWordIterator<I: Iterator<Item = u16>, W> {
2206    repdef: I,
2207    level_mask: u16,
2208    bits_rep: u8,
2209    bits_def: u8,
2210    max_rep: u16,
2211    phantom: std::marker::PhantomData<W>,
2212}
2213
2214impl<I: Iterator<Item = u16>> UnaryControlWordIterator<I, u8> {
2215    fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2216        let next = self.repdef.next()?;
2217        buf.push((next & self.level_mask) as u8);
2218        let is_new_row = self.max_rep == 0 || next == self.max_rep;
2219        let is_valid_item = next == 0 || self.bits_def == 0;
2220        Some(ControlWordDesc {
2221            is_new_row,
2222            // Either there is no rep, in which case there are no invisible items
2223            // or there is no def, in which case there are no invisible items
2224            is_visible: true,
2225            is_valid_item,
2226        })
2227    }
2228}
2229
2230impl<I: Iterator<Item = u16>> UnaryControlWordIterator<I, u16> {
2231    fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2232        let next = self.repdef.next().unwrap() & self.level_mask;
2233        let control_word = next.to_le_bytes();
2234        buf.push(control_word[0]);
2235        buf.push(control_word[1]);
2236        let is_new_row = self.max_rep == 0 || next == self.max_rep;
2237        let is_valid_item = next == 0 || self.bits_def == 0;
2238        Some(ControlWordDesc {
2239            is_new_row,
2240            is_visible: true,
2241            is_valid_item,
2242        })
2243    }
2244}
2245
2246impl<I: Iterator<Item = u16>> UnaryControlWordIterator<I, u32> {
2247    fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2248        let next = self.repdef.next()?;
2249        let next = (next & self.level_mask) as u32;
2250        let control_word = next.to_le_bytes();
2251        buf.push(control_word[0]);
2252        buf.push(control_word[1]);
2253        buf.push(control_word[2]);
2254        buf.push(control_word[3]);
2255        let is_new_row = self.max_rep == 0 || next as u16 == self.max_rep;
2256        let is_valid_item = next == 0 || self.bits_def == 0;
2257        Some(ControlWordDesc {
2258            is_new_row,
2259            is_visible: true,
2260            is_valid_item,
2261        })
2262    }
2263}
2264
2265/// A [`ControlWordIterator`] when there are no repetition or definition levels
2266#[derive(Debug)]
2267pub struct NilaryControlWordIterator {
2268    len: usize,
2269    idx: usize,
2270}
2271
2272impl NilaryControlWordIterator {
2273    fn append_next(&mut self) -> Option<ControlWordDesc> {
2274        if self.idx == self.len {
2275            None
2276        } else {
2277            self.idx += 1;
2278            Some(ControlWordDesc {
2279                is_new_row: true,
2280                is_visible: true,
2281                is_valid_item: true,
2282            })
2283        }
2284    }
2285}
2286
2287/// Helper function to get a bit mask of the given width
2288fn get_mask(width: u16) -> u16 {
2289    (1 << width) - 1
2290}
2291
2292// We're really going out of our way to avoid boxing here but this will be called on a per-value basis
2293// so it is in the critical path.
2294type SpecificBinaryControlWordIterator<'a, T> = BinaryControlWordIterator<
2295    Zip<Copied<std::slice::Iter<'a, u16>>, Copied<std::slice::Iter<'a, u16>>>,
2296    T,
2297>;
2298
2299/// An iterator that generates control words from repetition and definition levels
2300///
2301/// "Control word" is just a fancy term for a single u8/u16/u32 that contains both
2302/// the repetition and definition in it.
2303///
2304/// In the large majority of case we only need a single byte to represent both the
2305/// repetition and definition levels.  However, if there is deep nesting then we may
2306/// need two bytes.  In the worst case we need 4 bytes though this suggests hundreds of
2307/// levels of nesting which seems unlikely to encounter in practice.
2308#[derive(Debug)]
2309pub enum ControlWordIterator<'a> {
2310    Binary8(SpecificBinaryControlWordIterator<'a, u8>),
2311    Binary16(SpecificBinaryControlWordIterator<'a, u16>),
2312    Binary32(SpecificBinaryControlWordIterator<'a, u32>),
2313    Unary8(UnaryControlWordIterator<Copied<std::slice::Iter<'a, u16>>, u8>),
2314    Unary16(UnaryControlWordIterator<Copied<std::slice::Iter<'a, u16>>, u16>),
2315    Unary32(UnaryControlWordIterator<Copied<std::slice::Iter<'a, u16>>, u32>),
2316    Nilary(NilaryControlWordIterator),
2317}
2318
2319/// Describes the properties of a control word
2320#[derive(Debug)]
2321pub struct ControlWordDesc {
2322    pub is_new_row: bool,
2323    pub is_visible: bool,
2324    pub is_valid_item: bool,
2325}
2326
2327impl ControlWordIterator<'_> {
2328    /// Appends the next control word to the buffer
2329    ///
2330    /// Returns true if this is the start of a new item (i.e. the repetition level is maxed out)
2331    pub fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2332        match self {
2333            Self::Binary8(iter) => iter.append_next(buf),
2334            Self::Binary16(iter) => iter.append_next(buf),
2335            Self::Binary32(iter) => iter.append_next(buf),
2336            Self::Unary8(iter) => iter.append_next(buf),
2337            Self::Unary16(iter) => iter.append_next(buf),
2338            Self::Unary32(iter) => iter.append_next(buf),
2339            Self::Nilary(iter) => iter.append_next(),
2340        }
2341    }
2342
2343    /// Return true if the control word iterator has repetition levels
2344    pub fn has_repetition(&self) -> bool {
2345        match self {
2346            Self::Binary8(_) | Self::Binary16(_) | Self::Binary32(_) => true,
2347            Self::Unary8(iter) => iter.bits_rep > 0,
2348            Self::Unary16(iter) => iter.bits_rep > 0,
2349            Self::Unary32(iter) => iter.bits_rep > 0,
2350            Self::Nilary(_) => false,
2351        }
2352    }
2353
2354    /// Returns the number of bytes per control word
2355    pub fn bytes_per_word(&self) -> usize {
2356        match self {
2357            Self::Binary8(_) => 1,
2358            Self::Binary16(_) => 2,
2359            Self::Binary32(_) => 4,
2360            Self::Unary8(_) => 1,
2361            Self::Unary16(_) => 2,
2362            Self::Unary32(_) => 4,
2363            Self::Nilary(_) => 0,
2364        }
2365    }
2366
2367    /// Returns the number of bits used for the repetition level
2368    pub fn bits_rep(&self) -> u8 {
2369        match self {
2370            Self::Binary8(iter) => iter.bits_rep,
2371            Self::Binary16(iter) => iter.bits_rep,
2372            Self::Binary32(iter) => iter.bits_rep,
2373            Self::Unary8(iter) => iter.bits_rep,
2374            Self::Unary16(iter) => iter.bits_rep,
2375            Self::Unary32(iter) => iter.bits_rep,
2376            Self::Nilary(_) => 0,
2377        }
2378    }
2379
2380    /// Returns the number of bits used for the definition level
2381    pub fn bits_def(&self) -> u8 {
2382        match self {
2383            Self::Binary8(iter) => iter.bits_def,
2384            Self::Binary16(iter) => iter.bits_def,
2385            Self::Binary32(iter) => iter.bits_def,
2386            Self::Unary8(iter) => iter.bits_def,
2387            Self::Unary16(iter) => iter.bits_def,
2388            Self::Unary32(iter) => iter.bits_def,
2389            Self::Nilary(_) => 0,
2390        }
2391    }
2392}
2393
2394/// Builds a [`ControlWordIterator`] from repetition and definition levels
2395/// by first calculating the width needed and then creating the iterator
2396/// with the appropriate width
2397pub fn build_control_word_iterator<'a>(
2398    rep: Option<&'a [u16]>,
2399    max_rep: u16,
2400    def: Option<&'a [u16]>,
2401    max_def: u16,
2402    max_visible_def: u16,
2403    len: usize,
2404) -> ControlWordIterator<'a> {
2405    let rep_width = if max_rep == 0 {
2406        0
2407    } else {
2408        log_2_ceil(max_rep as u32) as u16
2409    };
2410    let rep_mask = if max_rep == 0 { 0 } else { get_mask(rep_width) };
2411    let def_width = if max_def == 0 {
2412        0
2413    } else {
2414        log_2_ceil(max_def as u32) as u16
2415    };
2416    let def_mask = if max_def == 0 { 0 } else { get_mask(def_width) };
2417    let total_width = rep_width + def_width;
2418    match (rep, def) {
2419        (Some(rep), Some(def)) => {
2420            let iter = rep.iter().copied().zip(def.iter().copied());
2421            let def_width = def_width as usize;
2422            if total_width <= 8 {
2423                ControlWordIterator::Binary8(BinaryControlWordIterator {
2424                    repdef: iter,
2425                    rep_mask,
2426                    def_mask,
2427                    def_width,
2428                    max_rep,
2429                    max_visible_def,
2430                    bits_rep: rep_width as u8,
2431                    bits_def: def_width as u8,
2432                    phantom: std::marker::PhantomData,
2433                })
2434            } else if total_width <= 16 {
2435                ControlWordIterator::Binary16(BinaryControlWordIterator {
2436                    repdef: iter,
2437                    rep_mask,
2438                    def_mask,
2439                    def_width,
2440                    max_rep,
2441                    max_visible_def,
2442                    bits_rep: rep_width as u8,
2443                    bits_def: def_width as u8,
2444                    phantom: std::marker::PhantomData,
2445                })
2446            } else {
2447                ControlWordIterator::Binary32(BinaryControlWordIterator {
2448                    repdef: iter,
2449                    rep_mask,
2450                    def_mask,
2451                    def_width,
2452                    max_rep,
2453                    max_visible_def,
2454                    bits_rep: rep_width as u8,
2455                    bits_def: def_width as u8,
2456                    phantom: std::marker::PhantomData,
2457                })
2458            }
2459        }
2460        (Some(lev), None) => {
2461            let iter = lev.iter().copied();
2462            if total_width <= 8 {
2463                ControlWordIterator::Unary8(UnaryControlWordIterator {
2464                    repdef: iter,
2465                    level_mask: rep_mask,
2466                    bits_rep: total_width as u8,
2467                    bits_def: 0,
2468                    max_rep,
2469                    phantom: std::marker::PhantomData,
2470                })
2471            } else if total_width <= 16 {
2472                ControlWordIterator::Unary16(UnaryControlWordIterator {
2473                    repdef: iter,
2474                    level_mask: rep_mask,
2475                    bits_rep: total_width as u8,
2476                    bits_def: 0,
2477                    max_rep,
2478                    phantom: std::marker::PhantomData,
2479                })
2480            } else {
2481                ControlWordIterator::Unary32(UnaryControlWordIterator {
2482                    repdef: iter,
2483                    level_mask: rep_mask,
2484                    bits_rep: total_width as u8,
2485                    bits_def: 0,
2486                    max_rep,
2487                    phantom: std::marker::PhantomData,
2488                })
2489            }
2490        }
2491        (None, Some(lev)) => {
2492            let iter = lev.iter().copied();
2493            if total_width <= 8 {
2494                ControlWordIterator::Unary8(UnaryControlWordIterator {
2495                    repdef: iter,
2496                    level_mask: def_mask,
2497                    bits_rep: 0,
2498                    bits_def: total_width as u8,
2499                    max_rep: 0,
2500                    phantom: std::marker::PhantomData,
2501                })
2502            } else if total_width <= 16 {
2503                ControlWordIterator::Unary16(UnaryControlWordIterator {
2504                    repdef: iter,
2505                    level_mask: def_mask,
2506                    bits_rep: 0,
2507                    bits_def: total_width as u8,
2508                    max_rep: 0,
2509                    phantom: std::marker::PhantomData,
2510                })
2511            } else {
2512                ControlWordIterator::Unary32(UnaryControlWordIterator {
2513                    repdef: iter,
2514                    level_mask: def_mask,
2515                    bits_rep: 0,
2516                    bits_def: total_width as u8,
2517                    max_rep: 0,
2518                    phantom: std::marker::PhantomData,
2519                })
2520            }
2521        }
2522        (None, None) => ControlWordIterator::Nilary(NilaryControlWordIterator { len, idx: 0 }),
2523    }
2524}
2525
2526/// A parser to unwrap control words into repetition and definition levels
2527///
2528/// This is the inverse of the [`ControlWordIterator`].
2529#[derive(Copy, Clone, Debug)]
2530pub enum ControlWordParser {
2531    // First item is the bits to shift, second is the mask to apply (the mask can be
2532    // calculated from the bits to shift but we don't want to calculate it each time)
2533    BOTH8(u8, u32),
2534    BOTH16(u8, u32),
2535    BOTH32(u8, u32),
2536    REP8,
2537    REP16,
2538    REP32,
2539    DEF8,
2540    DEF16,
2541    DEF32,
2542    NIL,
2543}
2544
2545impl ControlWordParser {
2546    fn parse_both<const WORD_SIZE: u8>(
2547        src: &[u8],
2548        dst_rep: &mut Vec<u16>,
2549        dst_def: &mut Vec<u16>,
2550        bits_to_shift: u8,
2551        mask_to_apply: u32,
2552    ) {
2553        match WORD_SIZE {
2554            1 => {
2555                let word = src[0];
2556                let rep = word >> bits_to_shift;
2557                let def = word & (mask_to_apply as u8);
2558                dst_rep.push(rep as u16);
2559                dst_def.push(def as u16);
2560            }
2561            2 => {
2562                let word = u16::from_le_bytes([src[0], src[1]]);
2563                let rep = word >> bits_to_shift;
2564                let def = word & mask_to_apply as u16;
2565                dst_rep.push(rep);
2566                dst_def.push(def);
2567            }
2568            4 => {
2569                let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]);
2570                let rep = word >> bits_to_shift;
2571                let def = word & mask_to_apply;
2572                dst_rep.push(rep as u16);
2573                dst_def.push(def as u16);
2574            }
2575            _ => unreachable!(),
2576        }
2577    }
2578
2579    fn parse_desc_both<const WORD_SIZE: u8>(
2580        src: &[u8],
2581        bits_to_shift: u8,
2582        mask_to_apply: u32,
2583        max_rep: u16,
2584        max_visible_def: u16,
2585    ) -> ControlWordDesc {
2586        match WORD_SIZE {
2587            1 => {
2588                let word = src[0];
2589                let rep = word >> bits_to_shift;
2590                let def = word & (mask_to_apply as u8);
2591                let is_visible = def as u16 <= max_visible_def;
2592                let is_new_row = rep as u16 == max_rep;
2593                let is_valid_item = def == 0;
2594                ControlWordDesc {
2595                    is_visible,
2596                    is_new_row,
2597                    is_valid_item,
2598                }
2599            }
2600            2 => {
2601                let word = u16::from_le_bytes([src[0], src[1]]);
2602                let rep = word >> bits_to_shift;
2603                let def = word & mask_to_apply as u16;
2604                let is_visible = def <= max_visible_def;
2605                let is_new_row = rep == max_rep;
2606                let is_valid_item = def == 0;
2607                ControlWordDesc {
2608                    is_visible,
2609                    is_new_row,
2610                    is_valid_item,
2611                }
2612            }
2613            4 => {
2614                let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]);
2615                let rep = word >> bits_to_shift;
2616                let def = word & mask_to_apply;
2617                let is_visible = def as u16 <= max_visible_def;
2618                let is_new_row = rep as u16 == max_rep;
2619                let is_valid_item = def == 0;
2620                ControlWordDesc {
2621                    is_visible,
2622                    is_new_row,
2623                    is_valid_item,
2624                }
2625            }
2626            _ => unreachable!(),
2627        }
2628    }
2629
2630    fn parse_one<const WORD_SIZE: u8>(src: &[u8], dst: &mut Vec<u16>) {
2631        match WORD_SIZE {
2632            1 => {
2633                let word = src[0];
2634                dst.push(word as u16);
2635            }
2636            2 => {
2637                let word = u16::from_le_bytes([src[0], src[1]]);
2638                dst.push(word);
2639            }
2640            4 => {
2641                let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]);
2642                dst.push(word as u16);
2643            }
2644            _ => unreachable!(),
2645        }
2646    }
2647
2648    fn parse_rep_desc_one<const WORD_SIZE: u8>(src: &[u8], max_rep: u16) -> ControlWordDesc {
2649        match WORD_SIZE {
2650            1 => ControlWordDesc {
2651                is_new_row: src[0] as u16 == max_rep,
2652                is_visible: true,
2653                is_valid_item: true,
2654            },
2655            2 => ControlWordDesc {
2656                is_new_row: u16::from_le_bytes([src[0], src[1]]) == max_rep,
2657                is_visible: true,
2658                is_valid_item: true,
2659            },
2660            4 => ControlWordDesc {
2661                is_new_row: u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as u16 == max_rep,
2662                is_visible: true,
2663                is_valid_item: true,
2664            },
2665            _ => unreachable!(),
2666        }
2667    }
2668
2669    fn parse_def_desc_one<const WORD_SIZE: u8>(src: &[u8]) -> ControlWordDesc {
2670        match WORD_SIZE {
2671            1 => ControlWordDesc {
2672                is_new_row: true,
2673                is_visible: true,
2674                is_valid_item: src[0] == 0,
2675            },
2676            2 => ControlWordDesc {
2677                is_new_row: true,
2678                is_visible: true,
2679                is_valid_item: u16::from_le_bytes([src[0], src[1]]) == 0,
2680            },
2681            4 => ControlWordDesc {
2682                is_new_row: true,
2683                is_visible: true,
2684                is_valid_item: u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as u16 == 0,
2685            },
2686            _ => unreachable!(),
2687        }
2688    }
2689
2690    /// Returns the number of bytes per control word
2691    pub fn bytes_per_word(&self) -> usize {
2692        match self {
2693            Self::BOTH8(..) => 1,
2694            Self::BOTH16(..) => 2,
2695            Self::BOTH32(..) => 4,
2696            Self::REP8 => 1,
2697            Self::REP16 => 2,
2698            Self::REP32 => 4,
2699            Self::DEF8 => 1,
2700            Self::DEF16 => 2,
2701            Self::DEF32 => 4,
2702            Self::NIL => 0,
2703        }
2704    }
2705
2706    /// Appends the next control word to the rep & def buffers
2707    ///
2708    /// `src` should be pointing at the first byte (little endian) of the control word
2709    ///
2710    /// `dst_rep` and `dst_def` are the buffers to append the rep and def levels to.
2711    /// They will not be appended to if not needed.
2712    pub fn parse(&self, src: &[u8], dst_rep: &mut Vec<u16>, dst_def: &mut Vec<u16>) {
2713        match self {
2714            Self::BOTH8(bits_to_shift, mask_to_apply) => {
2715                Self::parse_both::<1>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply)
2716            }
2717            Self::BOTH16(bits_to_shift, mask_to_apply) => {
2718                Self::parse_both::<2>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply)
2719            }
2720            Self::BOTH32(bits_to_shift, mask_to_apply) => {
2721                Self::parse_both::<4>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply)
2722            }
2723            Self::REP8 => Self::parse_one::<1>(src, dst_rep),
2724            Self::REP16 => Self::parse_one::<2>(src, dst_rep),
2725            Self::REP32 => Self::parse_one::<4>(src, dst_rep),
2726            Self::DEF8 => Self::parse_one::<1>(src, dst_def),
2727            Self::DEF16 => Self::parse_one::<2>(src, dst_def),
2728            Self::DEF32 => Self::parse_one::<4>(src, dst_def),
2729            Self::NIL => {}
2730        }
2731    }
2732
2733    /// Return true if the control words contain repetition information
2734    pub fn has_rep(&self) -> bool {
2735        match self {
2736            Self::BOTH8(..)
2737            | Self::BOTH16(..)
2738            | Self::BOTH32(..)
2739            | Self::REP8
2740            | Self::REP16
2741            | Self::REP32 => true,
2742            Self::DEF8 | Self::DEF16 | Self::DEF32 | Self::NIL => false,
2743        }
2744    }
2745
2746    /// Temporarily parses the control word to inspect its properties but does not append to any buffers
2747    pub fn parse_desc(&self, src: &[u8], max_rep: u16, max_visible_def: u16) -> ControlWordDesc {
2748        match self {
2749            Self::BOTH8(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<1>(
2750                src,
2751                *bits_to_shift,
2752                *mask_to_apply,
2753                max_rep,
2754                max_visible_def,
2755            ),
2756            Self::BOTH16(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<2>(
2757                src,
2758                *bits_to_shift,
2759                *mask_to_apply,
2760                max_rep,
2761                max_visible_def,
2762            ),
2763            Self::BOTH32(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<4>(
2764                src,
2765                *bits_to_shift,
2766                *mask_to_apply,
2767                max_rep,
2768                max_visible_def,
2769            ),
2770            Self::REP8 => Self::parse_rep_desc_one::<1>(src, max_rep),
2771            Self::REP16 => Self::parse_rep_desc_one::<2>(src, max_rep),
2772            Self::REP32 => Self::parse_rep_desc_one::<4>(src, max_rep),
2773            Self::DEF8 => Self::parse_def_desc_one::<1>(src),
2774            Self::DEF16 => Self::parse_def_desc_one::<2>(src),
2775            Self::DEF32 => Self::parse_def_desc_one::<4>(src),
2776            Self::NIL => ControlWordDesc {
2777                is_new_row: true,
2778                is_valid_item: true,
2779                is_visible: true,
2780            },
2781        }
2782    }
2783
2784    /// Creates a new parser from the number of bits used for the repetition and definition levels
2785    pub fn new(bits_rep: u8, bits_def: u8) -> Self {
2786        let total_bits = bits_rep + bits_def;
2787
2788        enum WordSize {
2789            One,
2790            Two,
2791            Four,
2792        }
2793
2794        let word_size = if total_bits <= 8 {
2795            WordSize::One
2796        } else if total_bits <= 16 {
2797            WordSize::Two
2798        } else {
2799            WordSize::Four
2800        };
2801
2802        match (bits_rep > 0, bits_def > 0, word_size) {
2803            (false, false, _) => Self::NIL,
2804            (false, true, WordSize::One) => Self::DEF8,
2805            (false, true, WordSize::Two) => Self::DEF16,
2806            (false, true, WordSize::Four) => Self::DEF32,
2807            (true, false, WordSize::One) => Self::REP8,
2808            (true, false, WordSize::Two) => Self::REP16,
2809            (true, false, WordSize::Four) => Self::REP32,
2810            (true, true, WordSize::One) => Self::BOTH8(bits_def, get_mask(bits_def as u16) as u32),
2811            (true, true, WordSize::Two) => Self::BOTH16(bits_def, get_mask(bits_def as u16) as u32),
2812            (true, true, WordSize::Four) => {
2813                Self::BOTH32(bits_def, get_mask(bits_def as u16) as u32)
2814            }
2815        }
2816    }
2817}
2818
2819#[cfg(test)]
2820mod tests {
2821    use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
2822
2823    use crate::encodings::logical::primitive::sparse::{
2824        SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, SparseValidityMeaning,
2825        SparseValiditySet,
2826    };
2827    use crate::repdef::{
2828        CompositeRepDefUnraveler, DefinitionInterpretation, RepDefUnraveler, SerializedRepDefs,
2829    };
2830
2831    use super::RepDefBuilder;
2832
2833    fn validity(values: &[bool]) -> NullBuffer {
2834        NullBuffer::from_iter(values.iter().copied())
2835    }
2836
2837    fn offsets_32(values: &[i32]) -> OffsetBuffer<i32> {
2838        OffsetBuffer::<i32>::new(ScalarBuffer::from_iter(values.iter().copied()))
2839    }
2840
2841    fn offsets_64(values: &[i64]) -> OffsetBuffer<i64> {
2842        OffsetBuffer::<i64>::new(ScalarBuffer::from_iter(values.iter().copied()))
2843    }
2844
2845    #[test]
2846    fn sparse_sibling_validity_mismatch_is_invalid_input() {
2847        let sparse = |positions| {
2848            RepDefUnraveler::new_sparse(SparseStructuralPlan {
2849                layers: vec![SparseStructuralLayerPlan::Validity {
2850                    num_slots: 2,
2851                    validity: SparseValiditySet {
2852                        meaning: SparseValidityMeaning::NullPositions,
2853                        positions,
2854                    },
2855                }],
2856                num_items: 2,
2857                num_visible_items: 2,
2858            })
2859        };
2860        let mut repdef = CompositeRepDefUnraveler::new(vec![sparse(SparsePositionSet::Empty)]);
2861        repdef.add_compatibility_check(CompositeRepDefUnraveler::new(vec![sparse(
2862            SparsePositionSet::Explicit(vec![0]),
2863        )]));
2864
2865        let err = repdef.unravel_validity(2).unwrap_err();
2866        assert!(matches!(err, lance_core::Error::InvalidInput { .. }));
2867        assert!(err.to_string().contains("incompatible validity metadata"));
2868    }
2869
2870    #[test]
2871    fn test_repdef_empty_offsets() {
2872        // Empty offsets should serialize without panicking.
2873        let mut builder = RepDefBuilder::default();
2874        builder.add_offsets(offsets_32(&[0]), None);
2875        let repdefs = RepDefBuilder::serialize(vec![builder]);
2876        assert!(repdefs.repetition_levels.is_none());
2877        assert!(repdefs.definition_levels.is_none());
2878    }
2879
2880    #[test]
2881    fn test_repdef_basic() {
2882        // Basic case, rep & def
2883        let mut builder = RepDefBuilder::default();
2884        builder.add_offsets(
2885            offsets_64(&[0, 2, 2, 5]),
2886            Some(validity(&[true, false, true])),
2887        );
2888        builder.add_offsets(
2889            offsets_64(&[0, 1, 3, 5, 5, 9]),
2890            Some(validity(&[true, true, true, false, true])),
2891        );
2892        builder.add_validity_bitmap(validity(&[
2893            true, true, true, false, false, false, true, true, false,
2894        ]));
2895
2896        let repdefs = RepDefBuilder::serialize(vec![builder]);
2897        let rep = repdefs.repetition_levels.unwrap();
2898        let def = repdefs.definition_levels.unwrap();
2899
2900        assert_eq!(vec![0, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1], *def);
2901        assert_eq!(vec![2, 1, 0, 2, 2, 0, 1, 1, 0, 0, 0], *rep);
2902
2903        // [[I], [I, I]], NULL, [[NULL, NULL], NULL, [NULL, I, I, NULL]]
2904
2905        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
2906            Some(rep.as_ref().to_vec()),
2907            Some(def.as_ref().to_vec()),
2908            repdefs.def_meaning.into(),
2909            9,
2910        )]);
2911
2912        // Note: validity doesn't exactly round-trip because repdef normalizes some of the
2913        // redundant validity values
2914        assert_eq!(
2915            unraveler.unravel_validity(9).unwrap(),
2916            Some(validity(&[
2917                true, true, true, false, false, false, true, true, false
2918            ]))
2919        );
2920        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
2921        assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 5, 9]).inner());
2922        assert_eq!(val, Some(validity(&[true, true, true, false, true])));
2923        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
2924        assert_eq!(off.inner(), offsets_32(&[0, 2, 2, 5]).inner());
2925        assert_eq!(val, Some(validity(&[true, false, true])));
2926    }
2927
2928    #[test]
2929    fn test_repdef_simple_null_empty_list() {
2930        let check = |repdefs: SerializedRepDefs, last_def: DefinitionInterpretation| {
2931            let rep = repdefs.repetition_levels.unwrap();
2932            let def = repdefs.definition_levels.unwrap();
2933
2934            assert_eq!([1, 0, 1, 1, 0, 0], *rep);
2935            assert_eq!([0, 0, 2, 0, 1, 0], *def);
2936            assert_eq!(
2937                vec![DefinitionInterpretation::NullableItem, last_def,],
2938                repdefs.def_meaning
2939            );
2940        };
2941
2942        // Null list and empty list should be serialized mostly the same
2943
2944        // Null case
2945        let mut builder = RepDefBuilder::default();
2946        builder.add_offsets(
2947            offsets_32(&[0, 2, 2, 5]),
2948            Some(validity(&[true, false, true])),
2949        );
2950        builder.add_validity_bitmap(validity(&[true, true, true, false, true]));
2951
2952        let repdefs = RepDefBuilder::serialize(vec![builder]);
2953
2954        check(repdefs, DefinitionInterpretation::NullableList);
2955
2956        // Empty case
2957        let mut builder = RepDefBuilder::default();
2958        builder.add_offsets(offsets_32(&[0, 2, 2, 5]), None);
2959        builder.add_validity_bitmap(validity(&[true, true, true, false, true]));
2960
2961        let repdefs = RepDefBuilder::serialize(vec![builder]);
2962
2963        check(repdefs, DefinitionInterpretation::EmptyableList);
2964    }
2965
2966    #[test]
2967    fn test_repdef_empty_list_at_end() {
2968        // Regresses a failure we encountered when the last item was an empty list
2969        let mut builder = RepDefBuilder::default();
2970        builder.add_offsets(offsets_32(&[0, 2, 5, 5]), None);
2971        builder.add_validity_bitmap(validity(&[true, true, true, false, true]));
2972
2973        let repdefs = RepDefBuilder::serialize(vec![builder]);
2974
2975        let rep = repdefs.repetition_levels.unwrap();
2976        let def = repdefs.definition_levels.unwrap();
2977
2978        assert_eq!([1, 0, 1, 0, 0, 1], *rep);
2979        assert_eq!([0, 0, 0, 1, 0, 2], *def);
2980        assert_eq!(
2981            vec![
2982                DefinitionInterpretation::NullableItem,
2983                DefinitionInterpretation::EmptyableList,
2984            ],
2985            repdefs.def_meaning
2986        );
2987    }
2988
2989    #[test]
2990    fn test_repdef_abnormal_nulls() {
2991        // List nulls are allowed to have non-empty offsets and garbage values
2992        // and the add_offsets call should normalize this
2993        let mut builder = RepDefBuilder::default();
2994        builder.add_offsets(
2995            offsets_32(&[0, 2, 5, 8]),
2996            Some(validity(&[true, false, true])),
2997        );
2998        // Note: we pass 5 here and not 8.  If add_offsets tells us there is garbage nulls they
2999        // should be removed before continuing
3000        builder.add_no_null(5);
3001
3002        let repdefs = RepDefBuilder::serialize(vec![builder]);
3003
3004        let rep = repdefs.repetition_levels.unwrap();
3005        let def = repdefs.definition_levels.unwrap();
3006
3007        assert_eq!([1, 0, 1, 1, 0, 0], *rep);
3008        assert_eq!([0, 0, 1, 0, 0, 0], *def);
3009
3010        assert_eq!(
3011            vec![
3012                DefinitionInterpretation::AllValidItem,
3013                DefinitionInterpretation::NullableList,
3014            ],
3015            repdefs.def_meaning
3016        );
3017    }
3018
3019    #[test]
3020    fn test_repdef_fsl() {
3021        let mut builder = RepDefBuilder::default();
3022        builder.add_fsl(Some(validity(&[true, false])), 2, 2);
3023        builder.add_fsl(None, 2, 4);
3024        builder.add_validity_bitmap(validity(&[
3025            true, false, true, false, true, false, true, false,
3026        ]));
3027
3028        let repdefs = RepDefBuilder::serialize(vec![builder]);
3029
3030        assert_eq!(
3031            vec![
3032                DefinitionInterpretation::NullableItem,
3033                DefinitionInterpretation::AllValidItem,
3034                DefinitionInterpretation::NullableItem
3035            ],
3036            repdefs.def_meaning
3037        );
3038
3039        assert!(repdefs.repetition_levels.is_none());
3040
3041        let def = repdefs.definition_levels.unwrap();
3042
3043        assert_eq!([0, 1, 0, 1, 2, 2, 2, 2], *def);
3044
3045        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3046            None,
3047            Some(def.as_ref().to_vec()),
3048            repdefs.def_meaning.into(),
3049            8,
3050        )]);
3051
3052        assert_eq!(
3053            unraveler.unravel_validity(8).unwrap(),
3054            Some(validity(&[
3055                true, false, true, false, false, false, false, false
3056            ]))
3057        );
3058        assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None);
3059        assert_eq!(
3060            unraveler.unravel_fsl_validity(2, 2).unwrap(),
3061            Some(validity(&[true, false]))
3062        );
3063    }
3064
3065    #[test]
3066    fn test_repdef_fsl_zero_dimension_is_invalid_input() {
3067        // A zero dimension can only reach us from a malformed schema.  Decimating with it
3068        // used to loop forever, writing past the end of the definition levels buffer.
3069        let mut builder = RepDefBuilder::default();
3070        builder.add_fsl(Some(validity(&[true, false])), 2, 2);
3071        builder.add_validity_bitmap(validity(&[true, false, true, false]));
3072
3073        let repdefs = RepDefBuilder::serialize(vec![builder]);
3074        let def = repdefs.definition_levels.unwrap();
3075
3076        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3077            None,
3078            Some(def.as_ref().to_vec()),
3079            repdefs.def_meaning.into(),
3080            4,
3081        )]);
3082        // Consume the item layer so the fixed-size-list layer is next
3083        unraveler.unravel_validity(4).unwrap();
3084
3085        let err = unraveler.unravel_fsl_validity(2, 0).unwrap_err();
3086        assert!(matches!(err, lance_core::Error::InvalidInput { .. }));
3087        assert!(
3088            err.to_string()
3089                .contains("dimension must be a positive integer"),
3090            "unexpected error: {}",
3091            err
3092        );
3093    }
3094
3095    #[test]
3096    fn test_repdef_fsl_allvalid_item() {
3097        let mut builder = RepDefBuilder::default();
3098        builder.add_fsl(Some(validity(&[true, false])), 2, 2);
3099        builder.add_fsl(None, 2, 4);
3100        builder.add_no_null(8);
3101
3102        let repdefs = RepDefBuilder::serialize(vec![builder]);
3103
3104        assert_eq!(
3105            vec![
3106                DefinitionInterpretation::AllValidItem,
3107                DefinitionInterpretation::AllValidItem,
3108                DefinitionInterpretation::NullableItem
3109            ],
3110            repdefs.def_meaning
3111        );
3112
3113        assert!(repdefs.repetition_levels.is_none());
3114
3115        let def = repdefs.definition_levels.unwrap();
3116
3117        assert_eq!([0, 0, 0, 0, 1, 1, 1, 1], *def);
3118
3119        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3120            None,
3121            Some(def.as_ref().to_vec()),
3122            repdefs.def_meaning.into(),
3123            8,
3124        )]);
3125
3126        assert_eq!(unraveler.unravel_validity(8).unwrap(), None);
3127        assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None);
3128        assert_eq!(
3129            unraveler.unravel_fsl_validity(2, 2).unwrap(),
3130            Some(validity(&[true, false]))
3131        );
3132    }
3133
3134    #[test]
3135    fn test_repdef_sliced_offsets() {
3136        // Sliced lists may have offsets that don't start with zero.  The
3137        // add_offsets call needs to normalize these to operate correctly.
3138        let mut builder = RepDefBuilder::default();
3139        builder.add_offsets(
3140            offsets_32(&[5, 7, 7, 10]),
3141            Some(validity(&[true, false, true])),
3142        );
3143        builder.add_no_null(5);
3144
3145        let repdefs = RepDefBuilder::serialize(vec![builder]);
3146
3147        let rep = repdefs.repetition_levels.unwrap();
3148        let def = repdefs.definition_levels.unwrap();
3149
3150        assert_eq!([1, 0, 1, 1, 0, 0], *rep);
3151        assert_eq!([0, 0, 1, 0, 0, 0], *def);
3152
3153        assert_eq!(
3154            vec![
3155                DefinitionInterpretation::AllValidItem,
3156                DefinitionInterpretation::NullableList,
3157            ],
3158            repdefs.def_meaning
3159        );
3160    }
3161
3162    #[test]
3163    fn test_repdef_complex_null_empty() {
3164        let mut builder = RepDefBuilder::default();
3165        builder.add_offsets(
3166            offsets_32(&[0, 4, 4, 4, 6]),
3167            Some(validity(&[true, false, true, true])),
3168        );
3169        builder.add_offsets(
3170            offsets_32(&[0, 1, 1, 2, 2, 2, 3]),
3171            Some(validity(&[true, false, true, false, true, true])),
3172        );
3173        builder.add_no_null(3);
3174
3175        let repdefs = RepDefBuilder::serialize(vec![builder]);
3176
3177        let rep = repdefs.repetition_levels.unwrap();
3178        let def = repdefs.definition_levels.unwrap();
3179
3180        assert_eq!([2, 1, 1, 1, 2, 2, 2, 1], *rep);
3181        assert_eq!([0, 1, 0, 1, 3, 4, 2, 0], *def);
3182    }
3183
3184    #[test]
3185    fn test_repdef_empty_list_no_null() {
3186        // Tests when we have some empty lists but no null lists.  This case
3187        // caused some bugs because we have definition but no nulls
3188        let mut builder = RepDefBuilder::default();
3189        builder.add_offsets(offsets_32(&[0, 4, 4, 4, 6]), None);
3190        builder.add_no_null(6);
3191
3192        let repdefs = RepDefBuilder::serialize(vec![builder]);
3193
3194        let rep = repdefs.repetition_levels.unwrap();
3195        let def = repdefs.definition_levels.unwrap();
3196
3197        assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3198        assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def);
3199
3200        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3201            Some(rep.as_ref().to_vec()),
3202            Some(def.as_ref().to_vec()),
3203            repdefs.def_meaning.into(),
3204            8,
3205        )]);
3206
3207        assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3208        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3209        assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3210        assert_eq!(val, None);
3211    }
3212
3213    #[test]
3214    fn test_repdef_all_valid() {
3215        let mut builder = RepDefBuilder::default();
3216        builder.add_offsets(offsets_64(&[0, 2, 3, 5]), None);
3217        builder.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None);
3218        builder.add_no_null(9);
3219
3220        let repdefs = RepDefBuilder::serialize(vec![builder]);
3221        let rep = repdefs.repetition_levels.unwrap();
3222        assert!(repdefs.definition_levels.is_none());
3223
3224        assert_eq!([2, 1, 0, 2, 0, 2, 0, 1, 0], *rep);
3225
3226        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3227            Some(rep.as_ref().to_vec()),
3228            None,
3229            repdefs.def_meaning.into(),
3230            9,
3231        )]);
3232
3233        assert_eq!(unraveler.unravel_validity(9).unwrap(), None);
3234        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3235        assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 7, 9]).inner());
3236        assert_eq!(val, None);
3237        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3238        assert_eq!(off.inner(), offsets_32(&[0, 2, 3, 5]).inner());
3239        assert_eq!(val, None);
3240    }
3241
3242    #[test]
3243    fn test_repdef_nested_list_multibatch_matches_single() {
3244        // Single builder: List<List<i32>>, 3 rows.
3245        //   outer [0,2,3,5] -> rows have 2,1,2 inner lists
3246        //   inner [0,1,3,5,7,9] -> 5 inner lists, lengths 1,2,2,2,2 (9 leaf)
3247        let mut single = RepDefBuilder::default();
3248        single.add_offsets(offsets_64(&[0, 2, 3, 5]), None);
3249        single.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None);
3250        single.add_no_null(9);
3251        let single_rep = RepDefBuilder::serialize(vec![single])
3252            .repetition_levels
3253            .unwrap();
3254
3255        // Same logical data split into two batches:
3256        //   batch0 = rows 0,1 : outer [0,2,3], inner [0,1,3,5] (3 inner, 5 leaf)
3257        //   batch1 = row 2    : outer [0,2],   inner [0,2,4]   (2 inner, 4 leaf)
3258        let mut b0 = RepDefBuilder::default();
3259        b0.add_offsets(offsets_64(&[0, 2, 3]), None);
3260        b0.add_offsets(offsets_64(&[0, 1, 3, 5]), None);
3261        b0.add_no_null(5);
3262        let mut b1 = RepDefBuilder::default();
3263        b1.add_offsets(offsets_64(&[0, 2]), None);
3264        b1.add_offsets(offsets_64(&[0, 2, 4]), None);
3265        b1.add_no_null(4);
3266        let multi_rep = RepDefBuilder::serialize(vec![b0, b1])
3267            .repetition_levels
3268            .unwrap();
3269
3270        assert_eq!(
3271            *single_rep, *multi_rep,
3272            "multi-batch nested-list rep levels must equal single-batch"
3273        );
3274    }
3275
3276    #[test]
3277    fn test_only_empty_lists() {
3278        let mut builder = RepDefBuilder::default();
3279        builder.add_offsets(offsets_32(&[0, 4, 4, 4, 6]), None);
3280        builder.add_no_null(6);
3281
3282        let repdefs = RepDefBuilder::serialize(vec![builder]);
3283
3284        let rep = repdefs.repetition_levels.unwrap();
3285        let def = repdefs.definition_levels.unwrap();
3286
3287        assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3288        assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def);
3289
3290        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3291            Some(rep.as_ref().to_vec()),
3292            Some(def.as_ref().to_vec()),
3293            repdefs.def_meaning.into(),
3294            8,
3295        )]);
3296
3297        assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3298        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3299        assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3300        assert_eq!(val, None);
3301    }
3302
3303    #[test]
3304    fn test_only_null_lists() {
3305        let mut builder = RepDefBuilder::default();
3306        builder.add_offsets(
3307            offsets_32(&[0, 4, 4, 4, 6]),
3308            Some(validity(&[true, false, false, true])),
3309        );
3310        builder.add_no_null(6);
3311
3312        let repdefs = RepDefBuilder::serialize(vec![builder]);
3313
3314        let rep = repdefs.repetition_levels.unwrap();
3315        let def = repdefs.definition_levels.unwrap();
3316
3317        assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3318        assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def);
3319
3320        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3321            Some(rep.as_ref().to_vec()),
3322            Some(def.as_ref().to_vec()),
3323            repdefs.def_meaning.into(),
3324            8,
3325        )]);
3326
3327        assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3328        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3329        assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3330        assert_eq!(val, Some(validity(&[true, false, false, true])));
3331    }
3332
3333    #[test]
3334    fn test_null_and_empty_lists() {
3335        let mut builder = RepDefBuilder::default();
3336        builder.add_offsets(
3337            offsets_32(&[0, 4, 4, 4, 6]),
3338            Some(validity(&[true, false, true, true])),
3339        );
3340        builder.add_no_null(6);
3341
3342        let repdefs = RepDefBuilder::serialize(vec![builder]);
3343
3344        let rep = repdefs.repetition_levels.unwrap();
3345        let def = repdefs.definition_levels.unwrap();
3346
3347        assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3348        assert_eq!([0, 0, 0, 0, 1, 2, 0, 0], *def);
3349
3350        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3351            Some(rep.as_ref().to_vec()),
3352            Some(def.as_ref().to_vec()),
3353            repdefs.def_meaning.into(),
3354            8,
3355        )]);
3356
3357        assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3358        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3359        assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3360        assert_eq!(val, Some(validity(&[true, false, true, true])));
3361    }
3362
3363    #[test]
3364    fn test_repdef_null_struct_valid_list() {
3365        // This regresses a bug
3366
3367        let rep = vec![1, 0, 0, 0];
3368        let def = vec![2, 0, 2, 2];
3369        // AllValidList<NullableStruct<NullableItem>>
3370        let def_meaning = vec![
3371            DefinitionInterpretation::NullableItem,
3372            DefinitionInterpretation::NullableItem,
3373            DefinitionInterpretation::AllValidList,
3374        ];
3375        let num_items = 4;
3376
3377        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3378            Some(rep),
3379            Some(def),
3380            def_meaning.into(),
3381            num_items,
3382        )]);
3383
3384        assert_eq!(
3385            unraveler.unravel_validity(4).unwrap(),
3386            Some(validity(&[false, true, false, false]))
3387        );
3388        assert_eq!(
3389            unraveler.unravel_validity(4).unwrap(),
3390            Some(validity(&[false, true, false, false]))
3391        );
3392        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3393        assert_eq!(off.inner(), offsets_32(&[0, 4]).inner());
3394        assert_eq!(val, None);
3395    }
3396
3397    #[test]
3398    fn test_repdef_no_rep() {
3399        let mut builder = RepDefBuilder::default();
3400        builder.add_no_null(5);
3401        builder.add_validity_bitmap(validity(&[false, false, true, true, true]));
3402        builder.add_validity_bitmap(validity(&[false, true, true, true, false]));
3403
3404        let repdefs = RepDefBuilder::serialize(vec![builder]);
3405        assert!(repdefs.repetition_levels.is_none());
3406        let def = repdefs.definition_levels.unwrap();
3407
3408        assert_eq!([2, 2, 0, 0, 1], *def);
3409
3410        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3411            None,
3412            Some(def.as_ref().to_vec()),
3413            repdefs.def_meaning.into(),
3414            5,
3415        )]);
3416
3417        assert_eq!(
3418            unraveler.unravel_validity(5).unwrap(),
3419            Some(validity(&[false, false, true, true, false]))
3420        );
3421        assert_eq!(
3422            unraveler.unravel_validity(5).unwrap(),
3423            Some(validity(&[false, false, true, true, true]))
3424        );
3425        assert_eq!(unraveler.unravel_validity(5).unwrap(), None);
3426    }
3427
3428    #[test]
3429    fn test_composite_unravel() {
3430        let mut builder = RepDefBuilder::default();
3431        builder.add_offsets(
3432            offsets_64(&[0, 2, 2, 5]),
3433            Some(validity(&[true, false, true])),
3434        );
3435        builder.add_no_null(5);
3436        let repdef1 = RepDefBuilder::serialize(vec![builder]);
3437
3438        let mut builder = RepDefBuilder::default();
3439        builder.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None);
3440        builder.add_no_null(9);
3441        let repdef2 = RepDefBuilder::serialize(vec![builder]);
3442
3443        let rep1 = repdef1.repetition_levels.clone().unwrap();
3444        let def1 = repdef1.definition_levels.clone().unwrap();
3445        let rep2 = repdef2.repetition_levels.clone().unwrap();
3446        assert!(repdef2.definition_levels.is_none());
3447
3448        assert_eq!([1, 0, 1, 1, 0, 0], *rep1);
3449        assert_eq!([0, 0, 1, 0, 0, 0], *def1);
3450        assert_eq!([1, 1, 0, 1, 0, 1, 0, 1, 0], *rep2);
3451
3452        let unravel1 = RepDefUnraveler::new(
3453            repdef1.repetition_levels.map(|l| l.to_vec()),
3454            repdef1.definition_levels.map(|l| l.to_vec()),
3455            repdef1.def_meaning.into(),
3456            5,
3457        );
3458        let unravel2 = RepDefUnraveler::new(
3459            repdef2.repetition_levels.map(|l| l.to_vec()),
3460            repdef2.definition_levels.map(|l| l.to_vec()),
3461            repdef2.def_meaning.into(),
3462            9,
3463        );
3464
3465        let mut unraveler = CompositeRepDefUnraveler::new(vec![unravel1, unravel2]);
3466
3467        assert!(unraveler.unravel_validity(9).unwrap().is_none());
3468        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3469        assert_eq!(
3470            off.inner(),
3471            offsets_32(&[0, 2, 2, 5, 6, 8, 10, 12, 14]).inner()
3472        );
3473        assert_eq!(
3474            val,
3475            Some(validity(&[true, false, true, true, true, true, true, true]))
3476        );
3477    }
3478
3479    #[test]
3480    fn test_repdef_multiple_builders() {
3481        // Basic case, rep & def
3482        let mut builder1 = RepDefBuilder::default();
3483        builder1.add_offsets(offsets_64(&[0, 2]), None);
3484        builder1.add_offsets(offsets_64(&[0, 1, 3]), None);
3485        builder1.add_validity_bitmap(validity(&[true, true, true]));
3486
3487        let mut builder2 = RepDefBuilder::default();
3488        builder2.add_offsets(offsets_64(&[0, 0, 3]), Some(validity(&[false, true])));
3489        builder2.add_offsets(
3490            offsets_64(&[0, 2, 2, 6]),
3491            Some(validity(&[true, false, true])),
3492        );
3493        builder2.add_validity_bitmap(validity(&[false, false, false, true, true, false]));
3494
3495        let repdefs = RepDefBuilder::serialize(vec![builder1, builder2]);
3496
3497        let rep = repdefs.repetition_levels.unwrap();
3498        let def = repdefs.definition_levels.unwrap();
3499
3500        assert_eq!([2, 1, 0, 2, 2, 0, 1, 1, 0, 0, 0], *rep);
3501        assert_eq!([0, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1], *def);
3502    }
3503
3504    #[test]
3505    fn test_all_valid_validity_bitmap_serializes_as_no_null() {
3506        let mut from_bitmap = RepDefBuilder::default();
3507        from_bitmap.add_validity_bitmap(validity(&[true, true, true, true]));
3508
3509        let mut from_no_null = RepDefBuilder::default();
3510        from_no_null.add_no_null(4);
3511
3512        let from_bitmap = RepDefBuilder::serialize(vec![from_bitmap]);
3513        let from_no_null = RepDefBuilder::serialize(vec![from_no_null]);
3514
3515        assert!(from_bitmap.repetition_levels.is_none());
3516        assert!(from_bitmap.definition_levels.is_none());
3517        assert_eq!(from_bitmap.def_meaning, from_no_null.def_meaning);
3518        assert_eq!(
3519            from_bitmap.max_visible_level,
3520            from_no_null.max_visible_level
3521        );
3522    }
3523
3524    #[test]
3525    fn test_slicer() {
3526        let mut builder = RepDefBuilder::default();
3527        builder.add_offsets(
3528            offsets_64(&[0, 2, 2, 30, 30]),
3529            Some(validity(&[true, false, true, true])),
3530        );
3531        builder.add_no_null(30);
3532
3533        let repdefs = RepDefBuilder::serialize(vec![builder]);
3534
3535        let mut rep_slicer = repdefs.rep_slicer().unwrap();
3536
3537        // First 5 items include a null list so we get 6 levels (12 bytes)
3538        assert_eq!(rep_slicer.slice_next(5).len(), 12);
3539        // Next 20 are all plain
3540        assert_eq!(rep_slicer.slice_next(20).len(), 40);
3541        // Last 5 include an empty list so we get 6 levels (12 bytes)
3542        assert_eq!(rep_slicer.slice_rest().len(), 12);
3543
3544        let mut def_slicer = repdefs.rep_slicer().unwrap();
3545
3546        // First 5 items include a null list so we get 6 levels (12 bytes)
3547        assert_eq!(def_slicer.slice_next(5).len(), 12);
3548        // Next 20 are all plain
3549        assert_eq!(def_slicer.slice_next(20).len(), 40);
3550        // Last 5 include an empty list so we get 6 levels (12 bytes)
3551        assert_eq!(def_slicer.slice_rest().len(), 12);
3552    }
3553
3554    #[test]
3555    fn test_control_words() {
3556        // Convert to control words, verify expected, convert back, verify same as original
3557        fn check(
3558            rep: &[u16],
3559            def: &[u16],
3560            expected_values: Vec<u8>,
3561            expected_bytes_per_word: usize,
3562            expected_bits_rep: u8,
3563            expected_bits_def: u8,
3564        ) {
3565            let num_vals = rep.len().max(def.len());
3566            let max_rep = rep.iter().max().copied().unwrap_or(0);
3567            let max_def = def.iter().max().copied().unwrap_or(0);
3568
3569            let in_rep = if rep.is_empty() { None } else { Some(rep) };
3570            let in_def = if def.is_empty() { None } else { Some(def) };
3571
3572            let mut iter = super::build_control_word_iterator(
3573                in_rep,
3574                max_rep,
3575                in_def,
3576                max_def,
3577                max_def + 1,
3578                expected_values.len(),
3579            );
3580            assert_eq!(iter.bytes_per_word(), expected_bytes_per_word);
3581            assert_eq!(iter.bits_rep(), expected_bits_rep);
3582            assert_eq!(iter.bits_def(), expected_bits_def);
3583            let mut cw_vec = Vec::with_capacity(num_vals * iter.bytes_per_word());
3584
3585            for _ in 0..num_vals {
3586                iter.append_next(&mut cw_vec);
3587            }
3588            assert!(iter.append_next(&mut cw_vec).is_none());
3589
3590            assert_eq!(expected_values, cw_vec);
3591
3592            let parser = super::ControlWordParser::new(expected_bits_rep, expected_bits_def);
3593
3594            let mut rep_out = Vec::with_capacity(num_vals);
3595            let mut def_out = Vec::with_capacity(num_vals);
3596
3597            if expected_bytes_per_word > 0 {
3598                for slice in cw_vec.chunks_exact(expected_bytes_per_word) {
3599                    parser.parse(slice, &mut rep_out, &mut def_out);
3600                }
3601            }
3602
3603            assert_eq!(rep, rep_out.as_slice());
3604            assert_eq!(def, def_out.as_slice());
3605        }
3606
3607        // Each will need 4 bits and so we should get 1-byte control words
3608        let rep = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3609        let def = &[5_u16, 3, 1, 2, 12, 15, 0, 2];
3610        let expected = vec![
3611            0b00000101, // 0, 5
3612            0b01110011, // 7, 3
3613            0b00110001, // 3, 1
3614            0b00100010, // 2, 2
3615            0b10011100, // 9, 12
3616            0b10001111, // 8, 15
3617            0b11000000, // 12, 0
3618            0b01010010, // 5, 2
3619        ];
3620        check(rep, def, expected, 1, 4, 4);
3621
3622        // Now we need 5 bits for def so we get 2-byte control words
3623        let rep = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3624        let def = &[5_u16, 3, 1, 2, 12, 22, 0, 2];
3625        let expected = vec![
3626            0b00000101, 0b00000000, // 0, 5
3627            0b11100011, 0b00000000, // 7, 3
3628            0b01100001, 0b00000000, // 3, 1
3629            0b01000010, 0b00000000, // 2, 2
3630            0b00101100, 0b00000001, // 9, 12
3631            0b00010110, 0b00000001, // 8, 22
3632            0b10000000, 0b00000001, // 12, 0
3633            0b10100010, 0b00000000, // 5, 2
3634        ];
3635        check(rep, def, expected, 2, 4, 5);
3636
3637        // Just rep, 4 bits so 1 byte each
3638        let levels = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3639        let expected = vec![
3640            0b00000000, // 0
3641            0b00000111, // 7
3642            0b00000011, // 3
3643            0b00000010, // 2
3644            0b00001001, // 9
3645            0b00001000, // 8
3646            0b00001100, // 12
3647            0b00000101, // 5
3648        ];
3649        check(levels, &[], expected.clone(), 1, 4, 0);
3650
3651        // Just def
3652        check(&[], levels, expected, 1, 0, 4);
3653
3654        // No rep, no def, no bytes
3655        check(&[], &[], Vec::default(), 0, 0, 0);
3656    }
3657
3658    #[test]
3659    fn test_control_words_rep_index() {
3660        fn check(
3661            rep: &[u16],
3662            def: &[u16],
3663            expected_new_rows: Vec<bool>,
3664            expected_is_visible: Vec<bool>,
3665        ) {
3666            let num_vals = rep.len().max(def.len());
3667            let max_rep = rep.iter().max().copied().unwrap_or(0);
3668            let max_def = def.iter().max().copied().unwrap_or(0);
3669
3670            let in_rep = if rep.is_empty() { None } else { Some(rep) };
3671            let in_def = if def.is_empty() { None } else { Some(def) };
3672
3673            let mut iter = super::build_control_word_iterator(
3674                in_rep,
3675                max_rep,
3676                in_def,
3677                max_def,
3678                /*max_visible_def=*/ 2,
3679                expected_new_rows.len(),
3680            );
3681
3682            let mut cw_vec = Vec::with_capacity(num_vals * iter.bytes_per_word());
3683            let mut expected_new_rows = expected_new_rows.iter().copied();
3684            let mut expected_is_visible = expected_is_visible.iter().copied();
3685            for _ in 0..expected_new_rows.len() {
3686                let word_desc = iter.append_next(&mut cw_vec).unwrap();
3687                assert_eq!(word_desc.is_new_row, expected_new_rows.next().unwrap());
3688                assert_eq!(word_desc.is_visible, expected_is_visible.next().unwrap());
3689            }
3690            assert!(iter.append_next(&mut cw_vec).is_none());
3691        }
3692
3693        // 2 means new list
3694        let rep = &[2_u16, 1, 0, 2, 2, 0, 1, 1, 0, 2, 0];
3695        // These values don't matter for this test
3696        let def = &[0_u16, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1];
3697
3698        // Rep & def
3699        check(
3700            rep,
3701            def,
3702            vec![
3703                true, false, false, true, true, false, false, false, false, true, false,
3704            ],
3705            vec![
3706                true, true, true, false, true, true, true, true, true, true, true,
3707            ],
3708        );
3709        // Rep only
3710        check(
3711            rep,
3712            &[],
3713            vec![
3714                true, false, false, true, true, false, false, false, false, true, false,
3715            ],
3716            vec![true; 11],
3717        );
3718        // No repetition
3719        check(
3720            &[],
3721            def,
3722            vec![
3723                true, true, true, true, true, true, true, true, true, true, true,
3724            ],
3725            vec![true; 11],
3726        );
3727        // No repetition, no definition
3728        check(
3729            &[],
3730            &[],
3731            vec![
3732                true, true, true, true, true, true, true, true, true, true, true,
3733            ],
3734            vec![true; 11],
3735        );
3736    }
3737
3738    #[test]
3739    fn regress_empty_list_case() {
3740        // This regresses a case where we had 3 null lists inside a struct
3741        let mut builder = RepDefBuilder::default();
3742        builder.add_validity_bitmap(validity(&[true, false, true]));
3743        builder.add_offsets(
3744            offsets_32(&[0, 0, 0, 0]),
3745            Some(validity(&[false, false, false])),
3746        );
3747        builder.add_no_null(0);
3748
3749        let repdefs = RepDefBuilder::serialize(vec![builder]);
3750        let rep = repdefs.repetition_levels.unwrap();
3751        let def = repdefs.definition_levels.unwrap();
3752
3753        assert_eq!([1, 1, 1], *rep);
3754        assert_eq!([1, 2, 1], *def);
3755
3756        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3757            Some(rep.as_ref().to_vec()),
3758            Some(def.as_ref().to_vec()),
3759            repdefs.def_meaning.into(),
3760            0,
3761        )]);
3762
3763        assert_eq!(unraveler.unravel_validity(0).unwrap(), None);
3764        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3765        assert_eq!(off.inner(), offsets_32(&[0, 0, 0, 0]).inner());
3766        assert_eq!(val, Some(validity(&[false, false, false])));
3767        let val = unraveler.unravel_validity(3).unwrap().unwrap();
3768        assert_eq!(val.inner(), validity(&[true, false, true]).inner());
3769    }
3770
3771    #[test]
3772    fn regress_list_ends_null_case() {
3773        let mut builder = RepDefBuilder::default();
3774        builder.add_offsets(
3775            offsets_64(&[0, 1, 2, 2]),
3776            Some(validity(&[true, true, false])),
3777        );
3778        builder.add_offsets(offsets_64(&[0, 1, 1]), Some(validity(&[true, false])));
3779        builder.add_no_null(1);
3780
3781        let repdefs = RepDefBuilder::serialize(vec![builder]);
3782        let rep = repdefs.repetition_levels.unwrap();
3783        let def = repdefs.definition_levels.unwrap();
3784
3785        assert_eq!([2, 2, 2], *rep);
3786        assert_eq!([0, 1, 2], *def);
3787
3788        let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3789            Some(rep.as_ref().to_vec()),
3790            Some(def.as_ref().to_vec()),
3791            repdefs.def_meaning.into(),
3792            1,
3793        )]);
3794
3795        assert_eq!(unraveler.unravel_validity(1).unwrap(), None);
3796        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3797        assert_eq!(off.inner(), offsets_32(&[0, 1, 1]).inner());
3798        assert_eq!(val, Some(validity(&[true, false])));
3799        let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3800        assert_eq!(off.inner(), offsets_32(&[0, 1, 2, 2]).inner());
3801        assert_eq!(val, Some(validity(&[true, true, false])));
3802    }
3803
3804    #[test]
3805    fn test_mixed_unraveler() {
3806        // This tests cases where the validity is different between two different pages
3807        // because one page has nulls and the other doesn't.
3808
3809        // Simple case with one layer of validity and no repetition
3810        let mut unraveler = CompositeRepDefUnraveler::new(vec![
3811            RepDefUnraveler::new(
3812                None,
3813                Some(vec![0, 1, 0, 1]),
3814                vec![DefinitionInterpretation::NullableItem].into(),
3815                4,
3816            ),
3817            RepDefUnraveler::new(
3818                None,
3819                None,
3820                vec![DefinitionInterpretation::AllValidItem].into(),
3821                4,
3822            ),
3823        ]);
3824
3825        assert_eq!(
3826            unraveler.unravel_validity(8).unwrap(),
3827            Some(validity(&[
3828                true, false, true, false, true, true, true, true
3829            ]))
3830        );
3831
3832        // More complex case with two layers of validity and repetition
3833        let def1 = Some(vec![0, 1, 2]);
3834        let rep1 = Some(vec![1, 0, 1]);
3835
3836        let def2 = Some(vec![1, 0, 0]);
3837        let rep2 = Some(vec![1, 1, 0]);
3838
3839        let mut unraveler = CompositeRepDefUnraveler::new(vec![
3840            RepDefUnraveler::new(
3841                rep1,
3842                def1,
3843                vec![
3844                    DefinitionInterpretation::NullableItem,
3845                    DefinitionInterpretation::EmptyableList,
3846                ]
3847                .into(),
3848                2,
3849            ),
3850            RepDefUnraveler::new(
3851                rep2,
3852                def2,
3853                vec![
3854                    DefinitionInterpretation::AllValidItem,
3855                    DefinitionInterpretation::NullableList,
3856                ]
3857                .into(),
3858                2,
3859            ),
3860        ]);
3861
3862        assert_eq!(
3863            unraveler.unravel_validity(4).unwrap(),
3864            Some(validity(&[true, false, true, true]))
3865        );
3866        assert_eq!(
3867            unraveler.unravel_offsets::<i32>().unwrap(),
3868            (
3869                offsets_32(&[0, 2, 2, 2, 4]),
3870                Some(validity(&[true, true, false, true]))
3871            )
3872        );
3873    }
3874
3875    #[test]
3876    fn test_mixed_unraveler_nullable_without_def_levels() {
3877        // A page can keep nullable layer metadata even when all definition levels are 0
3878        // and no definition buffer needs to be materialized. This should decode as all-valid.
3879        let mut unraveler = CompositeRepDefUnraveler::new(vec![
3880            RepDefUnraveler::new(
3881                None,
3882                Some(vec![0, 1, 0, 1]),
3883                vec![DefinitionInterpretation::NullableItem].into(),
3884                4,
3885            ),
3886            RepDefUnraveler::new(
3887                None,
3888                None,
3889                vec![DefinitionInterpretation::NullableItem].into(),
3890                4,
3891            ),
3892        ]);
3893
3894        assert_eq!(
3895            unraveler.unravel_validity(8).unwrap(),
3896            Some(validity(&[
3897                true, false, true, false, true, true, true, true
3898            ]))
3899        );
3900    }
3901}