1use 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#[derive(Debug, Clone, PartialEq, Eq)]
130pub(crate) struct MiniBlockRepDefSplit {
131 pub(crate) row_start: u64,
133 pub(crate) num_rows: u64,
135 pub(crate) level_range: Range<usize>,
137 pub(crate) value_start: u64,
139 pub(crate) num_values: u64,
141}
142
143#[derive(Debug, Clone, PartialEq, Eq)]
145pub(crate) enum MiniBlockRepDefBudget {
146 WithinBudget,
148 RequiresPageSplit(Vec<MiniBlockRepDefSplit>),
150 SingleRowOverBudget(u64),
152}
153
154const SPECIAL_THRESHOLD: u16 = u16::MAX / 2;
165
166#[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#[derive(Clone, Debug)]
180struct ValidityDesc {
181 validity: Option<BooleanBuffer>,
182 num_values: usize,
183}
184
185#[derive(Clone, Debug)]
189struct FslDesc {
190 validity: Option<BooleanBuffer>,
191 dimension: usize,
192 num_values: usize,
193}
194
195#[derive(Clone, Debug)]
199enum RawRepDef {
200 Offsets(OffsetDesc),
201 Validity(ValidityDesc),
202 Fsl(FslDesc),
203}
204
205#[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#[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 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 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 fn num_specials(&self) -> usize {
363 match self {
364 Self::Offsets(OffsetDesc { num_specials, .. }) => *num_specials,
365 _ => 0,
366 }
367 }
368
369 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 fn max_rep(&self) -> u16 {
395 match self {
396 Self::Offsets(_) => 1,
397 _ => 0,
398 }
399 }
400}
401
402#[derive(Debug)]
405pub struct SerializedRepDefs {
406 pub repetition_levels: Option<Arc<[u16]>>,
410 pub definition_levels: Option<Arc<[u16]>>,
414 pub def_meaning: Vec<DefinitionInterpretation>,
416 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 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#[derive(Debug)]
503pub struct RepDefSlicer<'a> {
504 repdef: &'a SerializedRepDefs,
505 to_slice: LanceBuffer,
506 current: usize,
507}
508
509impl<'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 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 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 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 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 self.current = start + num_values;
573 self.to_slice.slice_with_length(start * 2, num_values * 2)
574 }
575 }
576}
577
578#[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 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 pub fn is_all_valid(&self) -> bool {
615 matches!(
616 self,
617 Self::AllValidItem | Self::AllValidList | Self::EmptyableList
618 )
619 }
620
621 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#[derive(Debug)]
645struct SerializerContext {
646 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 let mut new_len = 0;
729 let expected_len = offset_desc.num_values + self.current_num_specials;
730 if expected_len == 0 {
731 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 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 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 let list_level = if rep == 0 { rep_level } else { rep };
780
781 if def == 0 && len > 0 {
782 *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 *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 *def_write_itr.next().unwrap() = def + SPECIAL_THRESHOLD;
801 *rep_write_itr.next().unwrap() = list_level;
802 new_len += 1;
803 }
804 }
805
806 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 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 return;
865 } else if self.rep_levels.is_empty() {
866 assert!(self.def_levels.len() >= self.current_len);
867 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 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 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 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 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#[derive(Clone, Default, Debug)]
1225pub struct RepDefBuilder {
1226 repdefs: Vec<RawRepDef>,
1228 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 self.len = Some(incoming_len);
1242 }
1243 }
1244
1245 fn num_layers(&self) -> usize {
1246 self.repdefs.len()
1247 }
1248
1249 pub fn is_empty(&self) -> bool {
1252 self.repdefs
1253 .iter()
1254 .all(|r| matches!(r, RawRepDef::Validity(ValidityDesc { validity: None, .. })))
1255 }
1256
1257 pub fn is_simple_validity(&self) -> bool {
1259 self.repdefs.len() == 1 && matches!(self.repdefs[0], RawRepDef::Validity(_))
1260 }
1261
1262 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 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 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 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 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 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 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 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#[derive(Debug)]
1582pub struct RepDefUnraveler {
1583 sparse: Option<SparseStructuralUnraveler>,
1584 rep_levels: Option<LevelBuffer>,
1585 def_levels: Option<LevelBuffer>,
1586 levels_to_rep: Vec<u16>,
1588 def_meaning: Arc<[DefinitionInterpretation]>,
1589 current_def_cmp: u16,
1591 current_rep_cmp: u16,
1593 current_layer: usize,
1596 num_items: u64,
1598}
1599
1600impl RepDefUnraveler {
1601 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 levels_to_rep.push(0);
1612 for meaning in def_meaning.as_ref() {
1613 match meaning {
1614 DefinitionInterpretation::AllValidItem | DefinitionInterpretation::AllValidList => {
1615 }
1617 DefinitionInterpretation::NullableItem => {
1618 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 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 .map(|levels| levels.len())
1698 .unwrap_or(0))
1699 }
1700
1701 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 let mut max_level = null_level.max(empty_level).max(valid_level);
1740 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 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 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 let mut read_idx = 0;
1784 let mut write_idx = 0;
1785 while read_idx < rep_levels.len() {
1786 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 *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 offsets.push(to_offset(curlen)?);
1800 curlen += 1;
1801 push_validity(true);
1802 } else if def_val > max_level {
1803 } else if def_val == null_level || def_val > upper_null {
1805 offsets.push(to_offset(curlen)?);
1807 push_validity(false);
1808 } else if def_val == empty_level {
1809 offsets.push(to_offset(curlen)?);
1811 push_validity(true);
1812 } else {
1813 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 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 unsafe {
1836 let rep_val = *rep_levels.get_unchecked(read_idx);
1837 if rep_val != 0 {
1838 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 rep_levels.truncate(write_idx);
1854 if let Some(validity) = validity {
1855 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 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 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 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 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#[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 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 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 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#[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#[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 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#[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
2287fn get_mask(width: u16) -> u16 {
2289 (1 << width) - 1
2290}
2291
2292type SpecificBinaryControlWordIterator<'a, T> = BinaryControlWordIterator<
2295 Zip<Copied<std::slice::Iter<'a, u16>>, Copied<std::slice::Iter<'a, u16>>>,
2296 T,
2297>;
2298
2299#[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#[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 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 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 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 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 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
2394pub 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#[derive(Copy, Clone, Debug)]
2530pub enum ControlWordParser {
2531 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 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 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 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 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 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 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 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 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 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 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 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 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 let mut builder = RepDefBuilder::default();
2994 builder.add_offsets(
2995 offsets_32(&[0, 2, 5, 8]),
2996 Some(validity(&[true, false, true])),
2997 );
2998 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 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 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 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 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 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 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 let rep = vec![1, 0, 0, 0];
3368 let def = vec![2, 0, 2, 2];
3369 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 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 assert_eq!(rep_slicer.slice_next(5).len(), 12);
3539 assert_eq!(rep_slicer.slice_next(20).len(), 40);
3541 assert_eq!(rep_slicer.slice_rest().len(), 12);
3543
3544 let mut def_slicer = repdefs.rep_slicer().unwrap();
3545
3546 assert_eq!(def_slicer.slice_next(5).len(), 12);
3548 assert_eq!(def_slicer.slice_next(20).len(), 40);
3550 assert_eq!(def_slicer.slice_rest().len(), 12);
3552 }
3553
3554 #[test]
3555 fn test_control_words() {
3556 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 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, 0b01110011, 0b00110001, 0b00100010, 0b10011100, 0b10001111, 0b11000000, 0b01010010, ];
3620 check(rep, def, expected, 1, 4, 4);
3621
3622 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, 0b11100011, 0b00000000, 0b01100001, 0b00000000, 0b01000010, 0b00000000, 0b00101100, 0b00000001, 0b00010110, 0b00000001, 0b10000000, 0b00000001, 0b10100010, 0b00000000, ];
3635 check(rep, def, expected, 2, 4, 5);
3636
3637 let levels = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3639 let expected = vec![
3640 0b00000000, 0b00000111, 0b00000011, 0b00000010, 0b00001001, 0b00001000, 0b00001100, 0b00000101, ];
3649 check(levels, &[], expected.clone(), 1, 4, 0);
3650
3651 check(&[], levels, expected, 1, 0, 4);
3653
3654 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 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 let rep = &[2_u16, 1, 0, 2, 2, 0, 1, 1, 0, 2, 0];
3695 let def = &[0_u16, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1];
3697
3698 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 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 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 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 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 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 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 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}