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
426pub fn max_visible_level(def_meaning: &[DefinitionInterpretation]) -> u16 {
436 def_meaning
437 .iter()
438 .take_while(|level| !level.is_list())
439 .map(|level| level.num_def_levels())
440 .sum()
441}
442
443impl SerializedRepDefs {
444 fn max_visible_level(def_meaning: &[DefinitionInterpretation]) -> Option<u16> {
447 def_meaning
448 .iter()
449 .any(|level| level.is_list())
450 .then(|| max_visible_level(def_meaning))
451 }
452
453 pub fn new(
454 repetition_levels: Option<LevelBuffer>,
455 definition_levels: Option<LevelBuffer>,
456 def_meaning: Vec<DefinitionInterpretation>,
457 ) -> Self {
458 Self::new_with_fixed_size_list_levels(
459 repetition_levels,
460 definition_levels,
461 def_meaning,
462 false,
463 )
464 }
465
466 pub(crate) fn new_with_fixed_size_list_levels(
467 repetition_levels: Option<LevelBuffer>,
468 definition_levels: Option<LevelBuffer>,
469 def_meaning: Vec<DefinitionInterpretation>,
470 has_fsl: bool,
471 ) -> Self {
472 let max_visible_level = Self::max_visible_level(&def_meaning);
473 Self {
474 repetition_levels: repetition_levels.map(Arc::from),
475 definition_levels: definition_levels.map(Arc::from),
476 def_meaning,
477 max_visible_level,
478 has_fsl,
479 }
480 }
481
482 pub fn empty(def_meaning: Vec<DefinitionInterpretation>) -> Self {
484 Self {
485 repetition_levels: None,
486 definition_levels: None,
487 def_meaning,
488 max_visible_level: None,
489 has_fsl: false,
490 }
491 }
492
493 pub fn rep_slicer(&self) -> Option<RepDefSlicer<'_>> {
494 self.repetition_levels
495 .as_ref()
496 .map(|rep| RepDefSlicer::new(self, rep.clone()))
497 }
498
499 pub fn def_slicer(&self) -> Option<RepDefSlicer<'_>> {
500 self.definition_levels
501 .as_ref()
502 .map(|def| RepDefSlicer::new(self, def.clone()))
503 }
504
505 pub(crate) fn has_fixed_size_list_levels(&self) -> bool {
506 self.has_fsl
507 }
508}
509
510#[derive(Debug)]
518pub struct RepDefSlicer<'a> {
519 repdef: &'a SerializedRepDefs,
520 to_slice: LanceBuffer,
521 current: usize,
522}
523
524impl<'a> RepDefSlicer<'a> {
526 fn new(repdef: &'a SerializedRepDefs, levels: Arc<[u16]>) -> Self {
527 Self {
528 repdef,
529 to_slice: LanceBuffer::reinterpret_slice(levels),
530 current: 0,
531 }
532 }
533
534 pub fn num_levels(&self) -> usize {
535 self.to_slice.len() / 2
536 }
537
538 pub fn num_levels_remaining(&self) -> usize {
539 self.num_levels() - self.current
540 }
541
542 pub fn all_levels(&self) -> &LanceBuffer {
543 &self.to_slice
544 }
545
546 pub fn slice_rest(&mut self) -> LanceBuffer {
555 let start = self.current;
556 let remaining = self.num_levels_remaining();
557 self.current = self.num_levels();
558 self.to_slice.slice_with_length(start * 2, remaining * 2)
559 }
560
561 pub fn slice_next(&mut self, num_values: usize) -> LanceBuffer {
563 let start = self.current;
564 let Some(max_visible_level) = self.repdef.max_visible_level else {
565 self.current = start + num_values;
567 return self.to_slice.slice_with_length(start * 2, num_values * 2);
568 };
569 if let Some(def) = self.repdef.definition_levels.as_ref() {
570 let mut def_itr = def[start..].iter();
574 let mut num_taken = 0;
575 let mut num_passed = 0;
576 while num_taken < num_values {
577 let def_level = *def_itr.next().unwrap();
578 if def_level <= max_visible_level {
579 num_taken += 1;
580 }
581 num_passed += 1;
582 }
583 self.current = start + num_passed;
584 self.to_slice.slice_with_length(start * 2, num_passed * 2)
585 } else {
586 self.current = start + num_values;
588 self.to_slice.slice_with_length(start * 2, num_values * 2)
589 }
590 }
591}
592
593#[derive(Debug, Copy, Clone, PartialEq, Eq)]
606pub enum DefinitionInterpretation {
607 AllValidItem,
608 AllValidList,
609 NullableItem,
610 NullableList,
611 EmptyableList,
612 NullableAndEmptyableList,
613}
614
615impl DefinitionInterpretation {
616 pub fn num_def_levels(&self) -> u16 {
618 match self {
619 Self::AllValidItem => 0,
620 Self::AllValidList => 0,
621 Self::NullableItem => 1,
622 Self::NullableList => 1,
623 Self::EmptyableList => 1,
624 Self::NullableAndEmptyableList => 2,
625 }
626 }
627
628 pub fn is_all_valid(&self) -> bool {
630 matches!(
631 self,
632 Self::AllValidItem | Self::AllValidList | Self::EmptyableList
633 )
634 }
635
636 pub fn is_list(&self) -> bool {
638 matches!(
639 self,
640 Self::AllValidList
641 | Self::NullableList
642 | Self::EmptyableList
643 | Self::NullableAndEmptyableList
644 )
645 }
646}
647
648#[derive(Debug)]
660struct SerializerContext {
661 def_meaning: Vec<DefinitionInterpretation>,
663 rep_levels: LevelBuffer,
664 spare_rep: LevelBuffer,
665 def_levels: LevelBuffer,
666 spare_def: LevelBuffer,
667 current_rep: u16,
668 current_def: u16,
669 current_len: usize,
670 current_num_specials: usize,
671 has_fsl: bool,
672}
673
674impl SerializerContext {
675 fn new(len: usize, num_layers: usize, max_rep: u16, max_def: u16) -> Self {
676 let def_meaning = Vec::with_capacity(num_layers);
677 Self {
678 rep_levels: if max_rep > 0 {
679 vec![0; len]
680 } else {
681 LevelBuffer::default()
682 },
683 spare_rep: if max_rep > 0 {
684 vec![0; len]
685 } else {
686 LevelBuffer::default()
687 },
688 def_levels: if max_def > 0 {
689 vec![0; len]
690 } else {
691 LevelBuffer::default()
692 },
693 spare_def: if max_def > 0 {
694 vec![0; len]
695 } else {
696 LevelBuffer::default()
697 },
698 def_meaning,
699 current_rep: max_rep,
700 current_def: max_def,
701 current_len: 0,
702 current_num_specials: 0,
703 has_fsl: false,
704 }
705 }
706
707 fn checkout_def(&mut self, meaning: DefinitionInterpretation) -> u16 {
708 let def = self.current_def;
709 self.current_def -= meaning.num_def_levels();
710 self.def_meaning.push(meaning);
711 def
712 }
713
714 fn record_offsets(&mut self, offset_desc: &OffsetDesc) {
715 let rep_level = self.current_rep;
716 let (null_list_level, empty_list_level) =
717 match (offset_desc.validity.is_some(), offset_desc.has_empty_lists) {
718 (true, true) => {
719 let level =
720 self.checkout_def(DefinitionInterpretation::NullableAndEmptyableList);
721 (level - 1, level)
722 }
723 (true, false) => (self.checkout_def(DefinitionInterpretation::NullableList), 0),
724 (false, true) => (
725 0,
726 self.checkout_def(DefinitionInterpretation::EmptyableList),
727 ),
728 (false, false) => {
729 self.checkout_def(DefinitionInterpretation::AllValidList);
730 (0, 0)
731 }
732 };
733 self.current_rep -= 1;
734
735 if let Some(validity) = &offset_desc.validity {
736 self.do_record_validity(validity, null_list_level);
737 }
738
739 let mut new_len = 0;
744 let expected_len = offset_desc.num_values + self.current_num_specials;
745 if expected_len == 0 {
746 self.current_len = 0;
748 return;
749 }
750 assert!(self.rep_levels.len() >= expected_len - 1);
751 if self.def_levels.is_empty() {
752 let mut write_itr = self.spare_rep.iter_mut();
753 let mut read_iter = self.rep_levels.iter().copied();
754 for w in offset_desc.offsets.windows(2) {
755 let len = w[1] - w[0];
756 assert!(len > 0);
758 let rep = read_iter.next().unwrap();
759 let list_level = if rep == 0 { rep_level } else { rep };
760 *write_itr.next().unwrap() = list_level;
761
762 for _ in 1..len {
763 *write_itr.next().unwrap() = 0;
764 }
765 new_len += len as usize;
766 }
767 std::mem::swap(&mut self.rep_levels, &mut self.spare_rep);
768 } else {
769 assert!(self.def_levels.len() >= expected_len - 1);
770 let mut def_write_itr = self.spare_def.iter_mut();
771 let mut rep_write_itr = self.spare_rep.iter_mut();
772 let mut rep_read_itr = self.rep_levels.iter().copied();
773 let mut def_read_itr = self.def_levels.iter().copied();
774 let specials_to_pass = self.current_num_specials;
775 let mut specials_passed = 0;
776
777 for w in offset_desc.offsets.windows(2) {
778 let mut def = def_read_itr.next().unwrap();
779 while def > SPECIAL_THRESHOLD {
781 *def_write_itr.next().unwrap() = def;
782 *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap();
783 def = def_read_itr.next().unwrap();
784 new_len += 1;
785 specials_passed += 1;
786 }
787
788 let len = w[1] - w[0];
789 let rep = rep_read_itr.next().unwrap();
790
791 let list_level = if rep == 0 { rep_level } else { rep };
795
796 if def == 0 && len > 0 {
797 *def_write_itr.next().unwrap() = 0;
799 *rep_write_itr.next().unwrap() = list_level;
800
801 for _ in 1..len {
802 *def_write_itr.next().unwrap() = 0;
803 *rep_write_itr.next().unwrap() = 0;
804 }
805
806 new_len += len as usize;
807 } else if def == 0 {
808 *def_write_itr.next().unwrap() = empty_list_level + SPECIAL_THRESHOLD;
810 *rep_write_itr.next().unwrap() = list_level;
811 new_len += 1;
812 } else {
813 *def_write_itr.next().unwrap() = def + SPECIAL_THRESHOLD;
816 *rep_write_itr.next().unwrap() = list_level;
817 new_len += 1;
818 }
819 }
820
821 while specials_passed < specials_to_pass {
823 *def_write_itr.next().unwrap() = def_read_itr.next().unwrap();
824 *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap();
825 new_len += 1;
826 specials_passed += 1;
827 }
828 std::mem::swap(&mut self.def_levels, &mut self.spare_def);
829 std::mem::swap(&mut self.rep_levels, &mut self.spare_rep);
830 }
831
832 self.current_len = new_len;
833 self.current_num_specials += offset_desc.num_specials;
834 }
835
836 fn do_record_validity(&mut self, validity: &BooleanBuffer, null_level: u16) {
837 assert!(self.def_levels.len() >= validity.len() + self.current_num_specials);
838 debug_assert!(
839 self.current_len == 0 || self.current_len == validity.len() + self.current_num_specials
840 );
841 self.current_len = validity.len() + self.current_num_specials;
842
843 let mut def_read_itr = self.def_levels.iter().copied();
844 let mut def_write_itr = self.spare_def.iter_mut();
845
846 let specials_to_pass = self.current_num_specials;
847 let mut specials_passed = 0;
848
849 for incoming_validity in validity.iter() {
850 let mut def = def_read_itr.next().unwrap();
851 while def > SPECIAL_THRESHOLD {
852 *def_write_itr.next().unwrap() = def;
853 def = def_read_itr.next().unwrap();
854 specials_passed += 1;
855 }
856 if def == 0 && !incoming_validity {
857 *def_write_itr.next().unwrap() = null_level;
858 } else {
859 *def_write_itr.next().unwrap() = def;
860 }
861 }
862
863 while specials_passed < specials_to_pass {
864 *def_write_itr.next().unwrap() = def_read_itr.next().unwrap();
865 specials_passed += 1;
866 }
867
868 std::mem::swap(&mut self.def_levels, &mut self.spare_def);
869 }
870
871 fn multiply_levels(&mut self, multiplier: usize) {
872 let old_len = self.current_len;
873 self.current_len =
875 (self.current_len - self.current_num_specials) * multiplier + self.current_num_specials;
876
877 if self.rep_levels.is_empty() && self.def_levels.is_empty() {
878 return;
880 } else if self.rep_levels.is_empty() {
881 assert!(self.def_levels.len() >= self.current_len);
882 let mut def_read_itr = self.def_levels.iter().copied();
884 let mut def_write_itr = self.spare_def.iter_mut();
885 for _ in 0..old_len {
886 let mut def = def_read_itr.next().unwrap();
887 while def > SPECIAL_THRESHOLD {
888 *def_write_itr.next().unwrap() = def;
889 def = def_read_itr.next().unwrap();
890 }
891 for _ in 0..multiplier {
892 *def_write_itr.next().unwrap() = def;
893 }
894 }
895 } else if self.def_levels.is_empty() {
896 assert!(self.rep_levels.len() >= self.current_len);
897 let mut rep_read_itr = self.rep_levels.iter().copied();
899 let mut rep_write_itr = self.spare_rep.iter_mut();
900 for _ in 0..old_len {
901 let rep = rep_read_itr.next().unwrap();
902 for _ in 0..multiplier {
903 *rep_write_itr.next().unwrap() = rep;
904 }
905 }
906 } else {
907 assert!(self.rep_levels.len() >= self.current_len);
908 assert!(self.def_levels.len() >= self.current_len);
909 let mut rep_read_itr = self.rep_levels.iter().copied();
910 let mut def_read_itr = self.def_levels.iter().copied();
911 let mut rep_write_itr = self.spare_rep.iter_mut();
912 let mut def_write_itr = self.spare_def.iter_mut();
913 for _ in 0..old_len {
914 let mut def = def_read_itr.next().unwrap();
915 while def > SPECIAL_THRESHOLD {
916 *def_write_itr.next().unwrap() = def;
917 *rep_write_itr.next().unwrap() = rep_read_itr.next().unwrap();
918 def = def_read_itr.next().unwrap();
919 }
920 let rep = rep_read_itr.next().unwrap();
921 for _ in 0..multiplier {
922 *def_write_itr.next().unwrap() = def;
923 *rep_write_itr.next().unwrap() = rep;
924 }
925 }
926 }
927 std::mem::swap(&mut self.def_levels, &mut self.spare_def);
928 std::mem::swap(&mut self.rep_levels, &mut self.spare_rep);
929 }
930
931 fn record_validity_buf(&mut self, validity: &Option<BooleanBuffer>) {
932 if let Some(validity) = validity {
933 let def_level = self.checkout_def(DefinitionInterpretation::NullableItem);
934 self.do_record_validity(validity, def_level);
935 } else {
936 self.checkout_def(DefinitionInterpretation::AllValidItem);
937 }
938 }
939
940 fn record_validity(&mut self, validity_desc: &ValidityDesc) {
941 self.record_validity_buf(&validity_desc.validity)
942 }
943
944 fn record_fsl(&mut self, fsl_desc: &FslDesc) {
945 self.has_fsl = true;
946 self.record_validity_buf(&fsl_desc.validity);
947 self.multiply_levels(fsl_desc.dimension);
948 }
949
950 fn normalize_specials(&mut self) {
951 for def in self.def_levels.iter_mut() {
952 if *def > SPECIAL_THRESHOLD {
953 *def -= SPECIAL_THRESHOLD;
954 }
955 }
956 }
957
958 fn normalize_specials_and_plan_splits(
959 &mut self,
960 def_meaning: &[DefinitionInterpretation],
961 max_levels_per_page: Option<u64>,
962 num_rows: u64,
963 num_values: u64,
964 ) -> Result<MiniBlockRepDefBudget> {
965 if self.def_levels.is_empty() {
972 return Ok(MiniBlockRepDefBudget::WithinBudget);
973 }
974
975 if self.rep_levels.is_empty() {
976 self.normalize_specials();
977 return Ok(MiniBlockRepDefBudget::WithinBudget);
978 }
979
980 if self.rep_levels.len() != self.def_levels.len() {
981 return Err(Error::internal(format!(
982 "Cannot plan structural page splits with mismatched rep/def lengths: rep={}, def={}",
983 self.rep_levels.len(),
984 self.def_levels.len()
985 )));
986 }
987
988 let Some(max_levels_per_page) = max_levels_per_page else {
989 self.normalize_specials();
990 return Ok(MiniBlockRepDefBudget::WithinBudget);
991 };
992
993 if num_values == 0 {
994 self.normalize_specials();
995 return Ok(MiniBlockRepDefBudget::WithinBudget);
996 }
997
998 let max_schema_rep = def_meaning.iter().filter(|level| level.is_list()).count() as u16;
999 let max_visible_level = SerializedRepDefs::max_visible_level(def_meaning);
1000 let should_plan = !self.has_fsl && max_schema_rep > 0 && max_visible_level.is_some();
1001
1002 if !should_plan {
1003 self.normalize_specials();
1004 return Ok(MiniBlockRepDefBudget::WithinBudget);
1005 }
1006
1007 let max_visible_level = max_visible_level.unwrap();
1008 let mut splits = Vec::new();
1009 let mut counted_rows = 0u64;
1010 let mut counted_values = 0u64;
1011 let mut saw_structural_overhead = false;
1012 let mut single_row_over_budget_levels = None;
1013
1014 let mut current_row_level_start = None;
1015 let mut current_row_num_values = 0u64;
1016
1017 let mut current_page_row_start = 0u64;
1018 let mut current_page_num_rows = 0u64;
1019 let mut current_page_level_start = 0usize;
1020 let mut current_page_level_end = 0usize;
1021 let mut current_page_value_start = 0u64;
1022 let mut current_page_num_values = 0u64;
1023 let mut current_page_num_levels = 0u64;
1024 let mut current_page_has_structural_overhead = false;
1025
1026 let mut finish_row =
1027 |row_level_start: usize, row_level_end: usize, row_num_values: u64| -> Result<()> {
1028 let row_num_levels = (row_level_end - row_level_start) as u64;
1029 let row_has_structural_overhead = row_num_levels > row_num_values;
1030 saw_structural_overhead |= row_has_structural_overhead;
1031
1032 if row_has_structural_overhead && row_num_levels > max_levels_per_page {
1033 single_row_over_budget_levels = Some(row_num_levels);
1034 }
1035
1036 if current_page_num_rows > 0
1037 && (current_page_has_structural_overhead || row_has_structural_overhead)
1038 && current_page_num_levels + row_num_levels > max_levels_per_page
1039 {
1040 splits.push(MiniBlockRepDefSplit {
1041 row_start: current_page_row_start,
1042 num_rows: current_page_num_rows,
1043 level_range: current_page_level_start..current_page_level_end,
1044 value_start: current_page_value_start,
1045 num_values: current_page_num_values,
1046 });
1047 current_page_row_start = counted_rows;
1048 current_page_num_rows = 0;
1049 current_page_level_start = row_level_start;
1050 current_page_value_start = counted_values;
1051 current_page_num_values = 0;
1052 current_page_num_levels = 0;
1053 current_page_has_structural_overhead = false;
1054 }
1055
1056 if current_page_num_rows == 0 {
1057 current_page_level_start = row_level_start;
1058 }
1059 current_page_num_rows += 1;
1060 current_page_level_end = row_level_end;
1061 current_page_num_values += row_num_values;
1062 current_page_num_levels += row_num_levels;
1063 current_page_has_structural_overhead |= row_has_structural_overhead;
1064 counted_rows += 1;
1065 counted_values += row_num_values;
1066 Ok(())
1067 };
1068
1069 for (idx, (rep_level, def_level)) in self
1070 .rep_levels
1071 .iter()
1072 .copied()
1073 .zip(self.def_levels.iter_mut())
1074 .enumerate()
1075 {
1076 if *def_level > SPECIAL_THRESHOLD {
1077 *def_level -= SPECIAL_THRESHOLD;
1078 }
1079
1080 if rep_level == max_schema_rep {
1081 if let Some(level_start) = current_row_level_start {
1082 finish_row(level_start, idx, current_row_num_values)?;
1083 current_row_num_values = 0;
1084 } else if idx != 0 {
1085 return Err(Error::internal(format!(
1086 "Cannot plan structural page splits: first top-level row starts at level {}, expected 0",
1087 idx
1088 )));
1089 }
1090 current_row_level_start = Some(idx);
1091 }
1092
1093 if current_row_level_start.is_none() {
1094 return Err(Error::internal(
1095 "Cannot plan structural page splits: found levels before the first top-level row start",
1096 ));
1097 }
1098 if *def_level <= max_visible_level {
1099 current_row_num_values += 1;
1100 }
1101 }
1102
1103 let Some(level_start) = current_row_level_start else {
1104 return Err(Error::internal(
1105 "Cannot plan structural page splits: found no top-level row starts",
1106 ));
1107 };
1108 finish_row(level_start, self.rep_levels.len(), current_row_num_values)?;
1109
1110 if counted_rows != num_rows {
1111 return Err(Error::internal(format!(
1112 "Cannot plan structural page splits: expected {} top-level row starts, found {}",
1113 num_rows, counted_rows
1114 )));
1115 }
1116 if counted_values != num_values {
1117 return Err(Error::internal(format!(
1118 "Cannot plan structural page splits: counted {} visible values, expected {}",
1119 counted_values, num_values
1120 )));
1121 }
1122 if !saw_structural_overhead {
1123 return Ok(MiniBlockRepDefBudget::WithinBudget);
1124 }
1125 if let Some(row_num_levels) = single_row_over_budget_levels {
1126 return Ok(MiniBlockRepDefBudget::SingleRowOverBudget(row_num_levels));
1127 }
1128
1129 if current_page_num_rows > 0 {
1130 splits.push(MiniBlockRepDefSplit {
1131 row_start: current_page_row_start,
1132 num_rows: current_page_num_rows,
1133 level_range: current_page_level_start..current_page_level_end,
1134 value_start: current_page_value_start,
1135 num_values: current_page_num_values,
1136 });
1137 }
1138
1139 if splits.len() > 1 {
1140 Ok(MiniBlockRepDefBudget::RequiresPageSplit(splits))
1141 } else {
1142 Ok(MiniBlockRepDefBudget::WithinBudget)
1143 }
1144 }
1145
1146 fn build(mut self) -> SerializedRepDefs {
1147 if self.current_len == 0 {
1148 return SerializedRepDefs::new_with_fixed_size_list_levels(
1149 None,
1150 None,
1151 self.def_meaning,
1152 self.has_fsl,
1153 );
1154 }
1155
1156 self.normalize_specials();
1157
1158 let definition_levels = if self.def_levels.is_empty() {
1159 None
1160 } else {
1161 Some(self.def_levels)
1162 };
1163 let repetition_levels = if self.rep_levels.is_empty() {
1164 None
1165 } else {
1166 Some(self.rep_levels)
1167 };
1168
1169 let def_meaning = self.def_meaning.into_iter().rev().collect::<Vec<_>>();
1171
1172 SerializedRepDefs::new_with_fixed_size_list_levels(
1173 repetition_levels,
1174 definition_levels,
1175 def_meaning,
1176 self.has_fsl,
1177 )
1178 }
1179
1180 fn build_with_miniblock_repdef_budget(
1181 mut self,
1182 max_levels_per_page: Option<u64>,
1183 num_rows: u64,
1184 num_values: u64,
1185 ) -> Result<(SerializedRepDefs, MiniBlockRepDefBudget)> {
1186 if self.current_len == 0 {
1187 return Ok((
1188 SerializedRepDefs::new_with_fixed_size_list_levels(
1189 None,
1190 None,
1191 self.def_meaning,
1192 self.has_fsl,
1193 ),
1194 MiniBlockRepDefBudget::WithinBudget,
1195 ));
1196 }
1197
1198 let def_meaning = std::mem::take(&mut self.def_meaning)
1200 .into_iter()
1201 .rev()
1202 .collect::<Vec<_>>();
1203 let budget = self.normalize_specials_and_plan_splits(
1204 &def_meaning,
1205 max_levels_per_page,
1206 num_rows,
1207 num_values,
1208 )?;
1209
1210 let definition_levels = if self.def_levels.is_empty() {
1211 None
1212 } else {
1213 Some(self.def_levels)
1214 };
1215 let repetition_levels = if self.rep_levels.is_empty() {
1216 None
1217 } else {
1218 Some(self.rep_levels)
1219 };
1220
1221 Ok((
1222 SerializedRepDefs::new_with_fixed_size_list_levels(
1223 repetition_levels,
1224 definition_levels,
1225 def_meaning,
1226 self.has_fsl,
1227 ),
1228 budget,
1229 ))
1230 }
1231}
1232
1233#[derive(Clone, Default, Debug)]
1240pub struct RepDefBuilder {
1241 repdefs: Vec<RawRepDef>,
1243 len: Option<usize>,
1248}
1249
1250impl RepDefBuilder {
1251 fn check_validity_len(&mut self, incoming_len: usize) {
1252 if let Some(len) = self.len {
1253 assert_eq!(incoming_len, len);
1254 } else {
1255 self.len = Some(incoming_len);
1257 }
1258 }
1259
1260 fn num_layers(&self) -> usize {
1261 self.repdefs.len()
1262 }
1263
1264 pub fn is_empty(&self) -> bool {
1267 self.repdefs
1268 .iter()
1269 .all(|r| matches!(r, RawRepDef::Validity(ValidityDesc { validity: None, .. })))
1270 }
1271
1272 pub fn is_simple_validity(&self) -> bool {
1274 self.repdefs.len() == 1 && matches!(self.repdefs[0], RawRepDef::Validity(_))
1275 }
1276
1277 pub fn add_validity_bitmap(&mut self, validity: NullBuffer) {
1279 self.check_validity_len(validity.len());
1280 if validity.null_count() == 0 {
1281 self.add_no_null(validity.len());
1282 return;
1283 }
1284 self.repdefs.push(RawRepDef::Validity(ValidityDesc {
1285 num_values: validity.len(),
1286 validity: Some(validity.into_inner()),
1287 }));
1288 }
1289
1290 pub fn add_no_null(&mut self, len: usize) {
1292 self.check_validity_len(len);
1293 self.repdefs.push(RawRepDef::Validity(ValidityDesc {
1294 validity: None,
1295 num_values: len,
1296 }));
1297 }
1298
1299 pub fn add_fsl(&mut self, validity: Option<NullBuffer>, dimension: usize, num_values: usize) {
1300 if let Some(len) = self.len {
1301 assert_eq!(num_values, len);
1302 }
1303 self.len = Some(num_values * dimension);
1304 debug_assert!(validity.is_none() || validity.as_ref().unwrap().len() == num_values);
1305 self.repdefs.push(RawRepDef::Fsl(FslDesc {
1306 num_values,
1307 validity: validity.map(|v| v.into_inner()),
1308 dimension,
1309 }))
1310 }
1311
1312 fn check_offset_len(&mut self, offsets: &[i64]) {
1313 if let Some(len) = self.len {
1314 assert!(offsets.len() == len + 1);
1315 }
1316 self.len = Some(offsets[offsets.len() - 1] as usize);
1317 }
1318
1319 fn do_add_offsets(
1320 &mut self,
1321 lengths: impl Iterator<Item = i64>,
1322 validity: Option<NullBuffer>,
1323 capacity: usize,
1324 ) -> bool {
1325 let mut num_specials = 0;
1326 let mut has_empty_lists = false;
1327 let mut has_garbage_values = false;
1328 let mut last_off: i64 = 0;
1329
1330 let mut normalized_offsets = Vec::with_capacity(capacity);
1331 normalized_offsets.push(0);
1332
1333 if let Some(ref validity) = validity {
1334 for (len, is_valid) in lengths.zip(validity.iter()) {
1335 match (is_valid, len == 0) {
1336 (false, is_empty) => {
1337 num_specials += 1;
1338 has_garbage_values |= !is_empty;
1339 }
1340 (true, true) => {
1341 num_specials += 1;
1342 has_empty_lists = true;
1343 }
1344 _ => {
1345 last_off += len;
1346 }
1347 }
1348 normalized_offsets.push(last_off);
1349 }
1350 } else {
1351 for len in lengths {
1352 if len == 0 {
1353 num_specials += 1;
1354 has_empty_lists = true;
1355 }
1356 last_off += len;
1357 normalized_offsets.push(last_off);
1358 }
1359 }
1360
1361 self.check_offset_len(&normalized_offsets);
1362 self.repdefs.push(RawRepDef::Offsets(OffsetDesc {
1363 num_values: normalized_offsets.len() - 1,
1364 offsets: normalized_offsets.into(),
1365 validity: validity.map(|v| v.into_inner()),
1366 has_empty_lists,
1367 num_specials: num_specials as usize,
1368 }));
1369
1370 has_garbage_values
1371 }
1372
1373 pub fn add_offsets<O: OffsetSizeTrait>(
1380 &mut self,
1381 offsets: OffsetBuffer<O>,
1382 validity: Option<NullBuffer>,
1383 ) -> bool {
1384 let inner = offsets.into_inner();
1385 let buffer_len = inner.len();
1386
1387 if O::IS_LARGE {
1388 let i64_buff = ScalarBuffer::<i64>::new(inner.into_inner(), 0, buffer_len);
1389 let lengths = i64_buff.windows(2).map(|off| off[1] - off[0]);
1390 self.do_add_offsets(lengths, validity, buffer_len)
1391 } else {
1392 let i32_buff = ScalarBuffer::<i32>::new(inner.into_inner(), 0, buffer_len);
1393 let lengths = i32_buff.windows(2).map(|off| (off[1] - off[0]) as i64);
1394 self.do_add_offsets(lengths, validity, buffer_len)
1395 }
1396 }
1397
1398 fn concat_layers<'a>(
1410 layers: impl Iterator<Item = &'a RawRepDef>,
1411 num_layers: usize,
1412 ) -> RawRepDef {
1413 enum LayerKind {
1414 Validity,
1415 Fsl,
1416 Offsets,
1417 }
1418
1419 let mut collected = Vec::with_capacity(num_layers);
1422 let mut has_nulls = false;
1423 let mut layer_kind = LayerKind::Validity;
1424 let mut total_num_specials = 0;
1425 let mut all_dimension = 0;
1426 let mut all_has_empty_lists = false;
1427 let mut all_num_values = 0;
1428 for layer in layers {
1429 has_nulls |= layer.has_nulls();
1430 match layer {
1431 RawRepDef::Validity(_) => {
1432 layer_kind = LayerKind::Validity;
1433 }
1434 RawRepDef::Offsets(OffsetDesc {
1435 num_specials,
1436 has_empty_lists,
1437 ..
1438 }) => {
1439 all_has_empty_lists |= *has_empty_lists;
1440 layer_kind = LayerKind::Offsets;
1441 total_num_specials += num_specials;
1442 }
1443 RawRepDef::Fsl(FslDesc { dimension, .. }) => {
1444 layer_kind = LayerKind::Fsl;
1445 all_dimension = *dimension;
1446 }
1447 }
1448 collected.push(layer);
1449 all_num_values += layer.num_values();
1450 }
1451
1452 if !has_nulls {
1454 match layer_kind {
1455 LayerKind::Validity => {
1456 return RawRepDef::Validity(ValidityDesc {
1457 validity: None,
1458 num_values: all_num_values,
1459 });
1460 }
1461 LayerKind::Fsl => {
1462 return RawRepDef::Fsl(FslDesc {
1463 validity: None,
1464 num_values: all_num_values,
1465 dimension: all_dimension,
1466 });
1467 }
1468 LayerKind::Offsets => {}
1469 }
1470 }
1471
1472 let mut validity_builder = if has_nulls {
1474 BooleanBufferBuilder::new(all_num_values)
1475 } else {
1476 BooleanBufferBuilder::new(0)
1477 };
1478 let mut all_offsets = if matches!(layer_kind, LayerKind::Offsets) {
1479 let mut all_offsets = Vec::with_capacity(all_num_values);
1480 all_offsets.push(0);
1481 all_offsets
1482 } else {
1483 Vec::new()
1484 };
1485
1486 for layer in collected {
1487 match layer {
1488 RawRepDef::Validity(ValidityDesc {
1489 validity: Some(validity),
1490 ..
1491 }) => {
1492 validity_builder.append_buffer(validity);
1493 }
1494 RawRepDef::Validity(ValidityDesc {
1495 validity: None,
1496 num_values,
1497 }) => {
1498 validity_builder.append_n(*num_values, true);
1499 }
1500 RawRepDef::Fsl(FslDesc {
1501 validity,
1502 num_values,
1503 ..
1504 }) => {
1505 if let Some(validity) = validity {
1506 validity_builder.append_buffer(validity);
1507 } else {
1508 validity_builder.append_n(*num_values, true);
1509 }
1510 }
1511 RawRepDef::Offsets(OffsetDesc {
1512 offsets,
1513 validity: Some(validity),
1514 has_empty_lists,
1515 ..
1516 }) => {
1517 all_has_empty_lists |= has_empty_lists;
1518 validity_builder.append_buffer(validity);
1519 let last = *all_offsets.last().unwrap();
1520 all_offsets.extend(offsets.iter().skip(1).map(|off| *off + last));
1521 }
1522 RawRepDef::Offsets(OffsetDesc {
1523 offsets,
1524 validity: None,
1525 has_empty_lists,
1526 num_values,
1527 ..
1528 }) => {
1529 all_has_empty_lists |= has_empty_lists;
1530 if has_nulls {
1531 validity_builder.append_n(*num_values, true);
1532 }
1533 let last = *all_offsets.last().unwrap();
1534 all_offsets.extend(offsets.iter().skip(1).map(|off| *off + last));
1535 }
1536 }
1537 }
1538 let validity = if has_nulls {
1539 Some(validity_builder.finish())
1540 } else {
1541 None
1542 };
1543 match layer_kind {
1544 LayerKind::Fsl => RawRepDef::Fsl(FslDesc {
1545 validity,
1546 num_values: all_num_values,
1547 dimension: all_dimension,
1548 }),
1549 LayerKind::Validity => RawRepDef::Validity(ValidityDesc {
1550 validity,
1551 num_values: all_num_values,
1552 }),
1553 LayerKind::Offsets => RawRepDef::Offsets(OffsetDesc {
1554 offsets: all_offsets.into(),
1555 validity,
1556 has_empty_lists: all_has_empty_lists,
1557 num_values: all_num_values,
1558 num_specials: total_num_specials,
1559 }),
1560 }
1561 }
1562
1563 pub fn serialize(builders: Vec<Self>) -> SerializedRepDefs {
1566 Self::normalize(builders).serialize()
1567 }
1568
1569 pub(crate) fn normalize(builders: Vec<Self>) -> NormalizedStructuralPlan {
1570 assert!(!builders.is_empty());
1571 let num_layers = builders[0].num_layers();
1572 debug_assert!(
1573 builders
1574 .iter()
1575 .all(|builder| builder.num_layers() == num_layers)
1576 );
1577 let layers = (0..num_layers)
1578 .map(|layer_index| {
1579 Self::concat_layers(
1580 builders.iter().map(|b| &b.repdefs[layer_index]),
1581 builders.len(),
1582 )
1583 })
1584 .collect::<Vec<_>>();
1585 NormalizedStructuralPlan {
1586 layers,
1587 dense_all_valid: builders.iter().all(Self::is_empty),
1588 }
1589 }
1590}
1591
1592#[derive(Debug)]
1597pub struct RepDefUnraveler {
1598 sparse: Option<SparseStructuralUnraveler>,
1599 rep_levels: Option<LevelBuffer>,
1600 def_levels: Option<LevelBuffer>,
1601 levels_to_rep: Vec<u16>,
1603 def_meaning: Arc<[DefinitionInterpretation]>,
1604 current_def_cmp: u16,
1606 current_rep_cmp: u16,
1608 current_layer: usize,
1611 num_items: u64,
1613}
1614
1615impl RepDefUnraveler {
1616 pub fn new(
1618 rep_levels: Option<LevelBuffer>,
1619 def_levels: Option<LevelBuffer>,
1620 def_meaning: Arc<[DefinitionInterpretation]>,
1621 num_items: u64,
1622 ) -> Self {
1623 let mut levels_to_rep = Vec::with_capacity(def_meaning.len());
1624 let mut rep_counter = 0;
1625 levels_to_rep.push(0);
1627 for meaning in def_meaning.as_ref() {
1628 match meaning {
1629 DefinitionInterpretation::AllValidItem | DefinitionInterpretation::AllValidList => {
1630 }
1632 DefinitionInterpretation::NullableItem => {
1633 levels_to_rep.push(rep_counter);
1635 }
1636 DefinitionInterpretation::NullableList => {
1637 rep_counter += 1;
1638 levels_to_rep.push(rep_counter);
1639 }
1640 DefinitionInterpretation::EmptyableList => {
1641 rep_counter += 1;
1642 levels_to_rep.push(rep_counter);
1643 }
1644 DefinitionInterpretation::NullableAndEmptyableList => {
1645 rep_counter += 1;
1646 levels_to_rep.push(rep_counter);
1647 levels_to_rep.push(rep_counter);
1648 }
1649 }
1650 }
1651 Self {
1652 sparse: None,
1653 rep_levels,
1654 def_levels,
1655 current_def_cmp: 0,
1656 current_rep_cmp: 0,
1657 levels_to_rep,
1658 current_layer: 0,
1659 def_meaning,
1660 num_items,
1661 }
1662 }
1663
1664 pub(crate) fn new_sparse(plan: SparseStructuralPlan) -> Self {
1665 Self {
1666 sparse: Some(SparseStructuralUnraveler::new(plan)),
1667 rep_levels: None,
1668 def_levels: None,
1669 levels_to_rep: Vec::new(),
1670 def_meaning: Arc::new([]),
1671 current_def_cmp: 0,
1672 current_rep_cmp: 0,
1673 current_layer: 0,
1674 num_items: 0,
1675 }
1676 }
1677
1678 fn ensure_exhausted(&self) -> Result<()> {
1679 if let Some(sparse) = &self.sparse {
1680 sparse.ensure_exhausted()?;
1681 }
1682 Ok(())
1683 }
1684
1685 fn is_sparse(&self) -> bool {
1686 self.sparse.is_some()
1687 }
1688
1689 pub fn is_all_valid(&self) -> bool {
1690 if let Some(sparse) = &self.sparse {
1691 return sparse.is_all_valid();
1692 }
1693 self.def_levels.is_none() || self.def_meaning[self.current_layer].is_all_valid()
1694 }
1695
1696 pub fn max_lists(&self) -> Result<usize> {
1702 if let Some(sparse) = &self.sparse {
1703 return sparse.max_lists();
1704 }
1705 debug_assert!(
1706 self.def_meaning[self.current_layer] != DefinitionInterpretation::NullableItem
1707 );
1708 Ok(self
1709 .rep_levels
1710 .as_ref()
1711 .map(|levels| levels.len())
1713 .unwrap_or(0))
1714 }
1715
1716 pub fn unravel_offsets<T: ArrowNativeType>(
1721 &mut self,
1722 offsets: &mut Vec<T>,
1723 validity: Option<&mut BooleanBufferBuilder>,
1724 ) -> Result<()> {
1725 if let Some(sparse) = self.sparse.as_mut() {
1726 return sparse.unravel_offsets(offsets, validity);
1727 }
1728 let rep_levels = self
1729 .rep_levels
1730 .as_mut()
1731 .expect("Expected repetition level but data didn't contain repetition");
1732 let valid_level = self.current_def_cmp;
1733 let (null_level, empty_level) = match self.def_meaning[self.current_layer] {
1734 DefinitionInterpretation::NullableList => {
1735 self.current_def_cmp += 1;
1736 (valid_level + 1, 0)
1737 }
1738 DefinitionInterpretation::EmptyableList => {
1739 self.current_def_cmp += 1;
1740 (0, valid_level + 1)
1741 }
1742 DefinitionInterpretation::NullableAndEmptyableList => {
1743 self.current_def_cmp += 2;
1744 (valid_level + 1, valid_level + 2)
1745 }
1746 DefinitionInterpretation::AllValidList => (0, 0),
1747 _ => unreachable!(),
1748 };
1749 self.current_layer += 1;
1750
1751 let mut max_level = null_level.max(empty_level).max(valid_level);
1755 let upper_null = max_level;
1758 for level in self.def_meaning[self.current_layer..].iter() {
1759 match level {
1760 DefinitionInterpretation::NullableItem => {
1761 max_level += 1;
1762 }
1763 DefinitionInterpretation::AllValidItem => {}
1764 _ => {
1765 break;
1766 }
1767 }
1768 }
1769
1770 let mut curlen: usize = offsets.last().map(|o| o.as_usize()).unwrap_or(0);
1771
1772 offsets.pop();
1780
1781 let to_offset = |val: usize| {
1782 T::from_usize(val)
1783 .ok_or_else(|| Error::invalid_input("A single batch had more than i32::MAX values and so a large container type is required"))
1784 };
1785 self.current_rep_cmp += 1;
1786 if let Some(def_levels) = &mut self.def_levels {
1787 assert!(rep_levels.len() == def_levels.len());
1788 let mut push_validity: Box<dyn FnMut(bool)> = if let Some(validity) = validity {
1791 Box::new(|is_valid| validity.append(is_valid))
1792 } else {
1793 Box::new(|_| {})
1794 };
1795 let mut read_idx = 0;
1799 let mut write_idx = 0;
1800 while read_idx < rep_levels.len() {
1801 unsafe {
1804 let rep_val = *rep_levels.get_unchecked(read_idx);
1805 if rep_val != 0 {
1806 let def_val = *def_levels.get_unchecked(read_idx);
1807 *rep_levels.get_unchecked_mut(write_idx) = rep_val - 1;
1809 *def_levels.get_unchecked_mut(write_idx) = def_val;
1810 write_idx += 1;
1811
1812 if def_val == 0 {
1813 offsets.push(to_offset(curlen)?);
1815 curlen += 1;
1816 push_validity(true);
1817 } else if def_val > max_level {
1818 } else if def_val == null_level || def_val > upper_null {
1820 offsets.push(to_offset(curlen)?);
1822 push_validity(false);
1823 } else if def_val == empty_level {
1824 offsets.push(to_offset(curlen)?);
1826 push_validity(true);
1827 } else {
1828 offsets.push(to_offset(curlen)?);
1830 curlen += 1;
1831 push_validity(true);
1832 }
1833 } else {
1834 curlen += 1;
1835 }
1836 read_idx += 1;
1837 }
1838 }
1839 offsets.push(to_offset(curlen)?);
1840 rep_levels.truncate(write_idx);
1841 def_levels.truncate(write_idx);
1842 Ok(())
1843 } else {
1844 let mut read_idx = 0;
1846 let mut write_idx = 0;
1847 let old_offsets_len = offsets.len();
1848 while read_idx < rep_levels.len() {
1849 unsafe {
1851 let rep_val = *rep_levels.get_unchecked(read_idx);
1852 if rep_val != 0 {
1853 offsets.push(to_offset(curlen)?);
1855 *rep_levels.get_unchecked_mut(write_idx) = rep_val - 1;
1856 write_idx += 1;
1857 }
1858 curlen += 1;
1859 read_idx += 1;
1860 }
1861 }
1862 let num_new_lists = offsets.len() - old_offsets_len;
1863 offsets.push(to_offset(curlen)?);
1864 rep_levels.truncate(write_idx);
1869 if let Some(validity) = validity {
1870 validity.append_n(num_new_lists, true);
1873 }
1874 Ok(())
1875 }
1876 }
1877
1878 pub fn skip_validity(&mut self) -> Result<()> {
1879 if let Some(sparse) = self.sparse.as_mut() {
1880 return sparse.skip_validity();
1881 }
1882 debug_assert!(self.is_all_valid());
1883 self.current_layer += 1;
1884 Ok(())
1885 }
1886
1887 pub fn unravel_validity(&mut self, validity: &mut BooleanBufferBuilder) -> Result<()> {
1889 if let Some(sparse) = self.sparse.as_mut() {
1890 return sparse.unravel_validity(validity);
1891 }
1892 let meaning = self.def_meaning[self.current_layer];
1893 if meaning == DefinitionInterpretation::AllValidItem || self.def_levels.is_none() {
1894 self.current_layer += 1;
1895 validity.append_n(self.num_items as usize, true);
1896 return Ok(());
1897 }
1898
1899 self.current_layer += 1;
1900 let def_levels = &self.def_levels.as_ref().unwrap();
1901
1902 let current_def_cmp = self.current_def_cmp;
1903 self.current_def_cmp += 1;
1904
1905 for is_valid in def_levels.iter().filter_map(|&level| {
1906 if self.levels_to_rep[level as usize] <= self.current_rep_cmp {
1907 Some(level <= current_def_cmp)
1908 } else {
1909 None
1910 }
1911 }) {
1912 validity.append(is_valid);
1913 }
1914 Ok(())
1915 }
1916
1917 pub fn decimate(&mut self, dimension: usize) -> Result<()> {
1928 if dimension == 0 {
1929 return Err(Error::invalid_input(
1930 "Cannot decimate repetition/definition levels with a fixed-size-list dimension of 0; dimension must be a positive integer",
1931 ));
1932 }
1933 if let Some(sparse) = self.sparse.as_mut() {
1934 return sparse.decimate(dimension);
1935 }
1936 if self.rep_levels.is_some() {
1937 todo!("Not yet supported FSL<...List<...>>");
1949 }
1950 let Some(def_levels) = self.def_levels.as_mut() else {
1951 return Ok(());
1952 };
1953 let mut read_idx = 0;
1954 let mut write_idx = 0;
1955 while read_idx < def_levels.len() {
1956 unsafe {
1960 *def_levels.get_unchecked_mut(write_idx) = *def_levels.get_unchecked(read_idx);
1961 }
1962 write_idx += 1;
1963 read_idx += dimension;
1964 }
1965 def_levels.truncate(write_idx);
1966 Ok(())
1967 }
1968}
1969
1970#[derive(Debug)]
1984pub struct CompositeRepDefUnraveler {
1985 unravelers: Vec<RepDefUnraveler>,
1986 comparisons: Vec<Self>,
1987}
1988
1989impl CompositeRepDefUnraveler {
1990 pub fn new(unravelers: Vec<RepDefUnraveler>) -> Self {
1991 Self {
1992 unravelers,
1993 comparisons: Vec::new(),
1994 }
1995 }
1996
1997 pub(crate) fn add_compatibility_check(&mut self, other: Self) {
1998 self.comparisons.push(other);
1999 }
2000
2001 pub(crate) fn has_sparse(&self) -> bool {
2002 self.unravelers.iter().any(RepDefUnraveler::is_sparse)
2003 || self.comparisons.iter().any(Self::has_sparse)
2004 }
2005
2006 pub(crate) fn ensure_exhausted(&self) -> Result<()> {
2007 for unraveler in &self.unravelers {
2008 unraveler.ensure_exhausted()?;
2009 }
2010 for comparison in &self.comparisons {
2011 comparison.ensure_exhausted()?;
2012 }
2013 Ok(())
2014 }
2015
2016 fn null_buffers_equal(
2017 left: &Option<NullBuffer>,
2018 right: &Option<NullBuffer>,
2019 expected_len: usize,
2020 ) -> bool {
2021 match (left, right) {
2022 (None, None) => true,
2023 (Some(left), Some(right)) => {
2024 left.len() == expected_len
2025 && right.len() == expected_len
2026 && left.iter().eq(right.iter())
2027 }
2028 (None, Some(right)) => right.len() == expected_len && right.null_count() == 0,
2029 (Some(left), None) => left.len() == expected_len && left.null_count() == 0,
2030 }
2031 }
2032
2033 fn decimate(&mut self, dimension: usize) -> Result<()> {
2034 for unraveler in &mut self.unravelers {
2035 unraveler.decimate(dimension)?;
2036 }
2037 for comparison in &mut self.comparisons {
2038 comparison.decimate(dimension)?;
2039 }
2040 Ok(())
2041 }
2042
2043 pub fn unravel_validity(&mut self, num_values: usize) -> Result<Option<NullBuffer>> {
2047 let is_all_valid = self
2048 .unravelers
2049 .iter()
2050 .all(|unraveler| unraveler.is_all_valid());
2051
2052 let validity = if is_all_valid {
2053 for unraveler in self.unravelers.iter_mut() {
2054 unraveler.skip_validity()?;
2055 }
2056 None
2057 } else {
2058 let mut validity = BooleanBufferBuilder::new(num_values);
2059 for unraveler in self.unravelers.iter_mut() {
2060 unraveler.unravel_validity(&mut validity)?;
2061 }
2062 Some(NullBuffer::new(validity.finish()))
2063 };
2064 for comparison in &mut self.comparisons {
2065 let other = comparison.unravel_validity(num_values)?;
2066 if !Self::null_buffers_equal(&validity, &other, num_values) {
2067 return Err(Error::invalid_input_source(
2068 format!(
2069 "Structural sibling fields have incompatible validity metadata for {num_values} values"
2070 )
2071 .into(),
2072 ));
2073 }
2074 }
2075 Ok(validity)
2076 }
2077
2078 pub fn unravel_fsl_validity(
2079 &mut self,
2080 num_values: usize,
2081 dimension: usize,
2082 ) -> Result<Option<NullBuffer>> {
2083 self.decimate(dimension)?;
2084 self.unravel_validity(num_values)
2085 }
2086
2087 pub fn unravel_offsets<T: ArrowNativeType>(
2089 &mut self,
2090 ) -> Result<(OffsetBuffer<T>, Option<NullBuffer>)> {
2091 let mut is_all_valid = true;
2092 let mut max_num_lists: usize = 0;
2093 for unraveler in self.unravelers.iter() {
2094 is_all_valid &= unraveler.is_all_valid();
2095 max_num_lists = max_num_lists
2096 .checked_add(unraveler.max_lists()?)
2097 .ok_or_else(|| {
2098 Error::invalid_input_source(
2099 "Combined repetition/definition list count exceeds usize::MAX".into(),
2100 )
2101 })?;
2102 }
2103
2104 let mut validity = if is_all_valid {
2105 None
2106 } else {
2107 Some(BooleanBufferBuilder::new(max_num_lists))
2110 };
2111
2112 let mut offsets = Vec::with_capacity(max_num_lists + 1);
2113
2114 for unraveler in self.unravelers.iter_mut() {
2115 unraveler.unravel_offsets(&mut offsets, validity.as_mut())?;
2116 }
2117
2118 let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
2119 let validity = validity.map(|mut v| NullBuffer::new(v.finish()));
2120 for comparison in &mut self.comparisons {
2121 let (other_offsets, other_validity) = comparison.unravel_offsets::<T>()?;
2122 if offsets.as_ref() != other_offsets.as_ref()
2123 || !Self::null_buffers_equal(
2124 &validity,
2125 &other_validity,
2126 offsets.len().saturating_sub(1),
2127 )
2128 {
2129 return Err(Error::invalid_input_source(
2130 format!(
2131 "Structural sibling fields have incompatible list metadata for {} slots",
2132 offsets.len().saturating_sub(1)
2133 )
2134 .into(),
2135 ));
2136 }
2137 }
2138
2139 Ok((offsets, validity))
2140 }
2141}
2142
2143#[derive(Debug)]
2149pub struct BinaryControlWordIterator<I: Iterator<Item = (u16, u16)>, W> {
2150 repdef: I,
2151 def_width: usize,
2152 max_rep: u16,
2153 max_visible_def: u16,
2154 rep_mask: u16,
2155 def_mask: u16,
2156 bits_rep: u8,
2157 bits_def: u8,
2158 phantom: std::marker::PhantomData<W>,
2159}
2160
2161impl<I: Iterator<Item = (u16, u16)>> BinaryControlWordIterator<I, u8> {
2162 fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2163 let next = self.repdef.next()?;
2164 let control_word: u8 =
2165 (((next.0 & self.rep_mask) as u8) << self.def_width) + ((next.1 & self.def_mask) as u8);
2166 buf.push(control_word);
2167 let is_new_row = next.0 == self.max_rep;
2168 let is_visible = next.1 <= self.max_visible_def;
2169 let is_valid_item = next.1 == 0;
2170 Some(ControlWordDesc {
2171 is_new_row,
2172 is_visible,
2173 is_valid_item,
2174 })
2175 }
2176}
2177
2178impl<I: Iterator<Item = (u16, u16)>> BinaryControlWordIterator<I, u16> {
2179 fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2180 let next = self.repdef.next()?;
2181 let control_word: u16 =
2182 ((next.0 & self.rep_mask) << self.def_width) + (next.1 & self.def_mask);
2183 let control_word = control_word.to_le_bytes();
2184 buf.push(control_word[0]);
2185 buf.push(control_word[1]);
2186 let is_new_row = next.0 == self.max_rep;
2187 let is_visible = next.1 <= self.max_visible_def;
2188 let is_valid_item = next.1 == 0;
2189 Some(ControlWordDesc {
2190 is_new_row,
2191 is_visible,
2192 is_valid_item,
2193 })
2194 }
2195}
2196
2197impl<I: Iterator<Item = (u16, u16)>> BinaryControlWordIterator<I, u32> {
2198 fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2199 let next = self.repdef.next()?;
2200 let control_word: u32 = (((next.0 & self.rep_mask) as u32) << self.def_width)
2201 + ((next.1 & self.def_mask) as u32);
2202 let control_word = control_word.to_le_bytes();
2203 buf.push(control_word[0]);
2204 buf.push(control_word[1]);
2205 buf.push(control_word[2]);
2206 buf.push(control_word[3]);
2207 let is_new_row = next.0 == self.max_rep;
2208 let is_visible = next.1 <= self.max_visible_def;
2209 let is_valid_item = next.1 == 0;
2210 Some(ControlWordDesc {
2211 is_new_row,
2212 is_visible,
2213 is_valid_item,
2214 })
2215 }
2216}
2217
2218#[derive(Debug)]
2220pub struct UnaryControlWordIterator<I: Iterator<Item = u16>, W> {
2221 repdef: I,
2222 level_mask: u16,
2223 bits_rep: u8,
2224 bits_def: u8,
2225 max_rep: u16,
2226 phantom: std::marker::PhantomData<W>,
2227}
2228
2229impl<I: Iterator<Item = u16>> UnaryControlWordIterator<I, u8> {
2230 fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2231 let next = self.repdef.next()?;
2232 buf.push((next & self.level_mask) as u8);
2233 let is_new_row = self.max_rep == 0 || next == self.max_rep;
2234 let is_valid_item = next == 0 || self.bits_def == 0;
2235 Some(ControlWordDesc {
2236 is_new_row,
2237 is_visible: true,
2240 is_valid_item,
2241 })
2242 }
2243}
2244
2245impl<I: Iterator<Item = u16>> UnaryControlWordIterator<I, u16> {
2246 fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2247 let next = self.repdef.next().unwrap() & self.level_mask;
2248 let control_word = next.to_le_bytes();
2249 buf.push(control_word[0]);
2250 buf.push(control_word[1]);
2251 let is_new_row = self.max_rep == 0 || next == self.max_rep;
2252 let is_valid_item = next == 0 || self.bits_def == 0;
2253 Some(ControlWordDesc {
2254 is_new_row,
2255 is_visible: true,
2256 is_valid_item,
2257 })
2258 }
2259}
2260
2261impl<I: Iterator<Item = u16>> UnaryControlWordIterator<I, u32> {
2262 fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2263 let next = self.repdef.next()?;
2264 let next = (next & self.level_mask) as u32;
2265 let control_word = next.to_le_bytes();
2266 buf.push(control_word[0]);
2267 buf.push(control_word[1]);
2268 buf.push(control_word[2]);
2269 buf.push(control_word[3]);
2270 let is_new_row = self.max_rep == 0 || next as u16 == self.max_rep;
2271 let is_valid_item = next == 0 || self.bits_def == 0;
2272 Some(ControlWordDesc {
2273 is_new_row,
2274 is_visible: true,
2275 is_valid_item,
2276 })
2277 }
2278}
2279
2280#[derive(Debug)]
2282pub struct NilaryControlWordIterator {
2283 len: usize,
2284 idx: usize,
2285}
2286
2287impl NilaryControlWordIterator {
2288 fn append_next(&mut self) -> Option<ControlWordDesc> {
2289 if self.idx == self.len {
2290 None
2291 } else {
2292 self.idx += 1;
2293 Some(ControlWordDesc {
2294 is_new_row: true,
2295 is_visible: true,
2296 is_valid_item: true,
2297 })
2298 }
2299 }
2300}
2301
2302fn get_mask(width: u16) -> u16 {
2304 (1 << width) - 1
2305}
2306
2307type SpecificBinaryControlWordIterator<'a, T> = BinaryControlWordIterator<
2310 Zip<Copied<std::slice::Iter<'a, u16>>, Copied<std::slice::Iter<'a, u16>>>,
2311 T,
2312>;
2313
2314#[derive(Debug)]
2324pub enum ControlWordIterator<'a> {
2325 Binary8(SpecificBinaryControlWordIterator<'a, u8>),
2326 Binary16(SpecificBinaryControlWordIterator<'a, u16>),
2327 Binary32(SpecificBinaryControlWordIterator<'a, u32>),
2328 Unary8(UnaryControlWordIterator<Copied<std::slice::Iter<'a, u16>>, u8>),
2329 Unary16(UnaryControlWordIterator<Copied<std::slice::Iter<'a, u16>>, u16>),
2330 Unary32(UnaryControlWordIterator<Copied<std::slice::Iter<'a, u16>>, u32>),
2331 Nilary(NilaryControlWordIterator),
2332}
2333
2334#[derive(Debug)]
2336pub struct ControlWordDesc {
2337 pub is_new_row: bool,
2338 pub is_visible: bool,
2339 pub is_valid_item: bool,
2340}
2341
2342impl ControlWordIterator<'_> {
2343 pub fn append_next(&mut self, buf: &mut Vec<u8>) -> Option<ControlWordDesc> {
2347 match self {
2348 Self::Binary8(iter) => iter.append_next(buf),
2349 Self::Binary16(iter) => iter.append_next(buf),
2350 Self::Binary32(iter) => iter.append_next(buf),
2351 Self::Unary8(iter) => iter.append_next(buf),
2352 Self::Unary16(iter) => iter.append_next(buf),
2353 Self::Unary32(iter) => iter.append_next(buf),
2354 Self::Nilary(iter) => iter.append_next(),
2355 }
2356 }
2357
2358 pub fn has_repetition(&self) -> bool {
2360 match self {
2361 Self::Binary8(_) | Self::Binary16(_) | Self::Binary32(_) => true,
2362 Self::Unary8(iter) => iter.bits_rep > 0,
2363 Self::Unary16(iter) => iter.bits_rep > 0,
2364 Self::Unary32(iter) => iter.bits_rep > 0,
2365 Self::Nilary(_) => false,
2366 }
2367 }
2368
2369 pub fn bytes_per_word(&self) -> usize {
2371 match self {
2372 Self::Binary8(_) => 1,
2373 Self::Binary16(_) => 2,
2374 Self::Binary32(_) => 4,
2375 Self::Unary8(_) => 1,
2376 Self::Unary16(_) => 2,
2377 Self::Unary32(_) => 4,
2378 Self::Nilary(_) => 0,
2379 }
2380 }
2381
2382 pub fn bits_rep(&self) -> u8 {
2384 match self {
2385 Self::Binary8(iter) => iter.bits_rep,
2386 Self::Binary16(iter) => iter.bits_rep,
2387 Self::Binary32(iter) => iter.bits_rep,
2388 Self::Unary8(iter) => iter.bits_rep,
2389 Self::Unary16(iter) => iter.bits_rep,
2390 Self::Unary32(iter) => iter.bits_rep,
2391 Self::Nilary(_) => 0,
2392 }
2393 }
2394
2395 pub fn bits_def(&self) -> u8 {
2397 match self {
2398 Self::Binary8(iter) => iter.bits_def,
2399 Self::Binary16(iter) => iter.bits_def,
2400 Self::Binary32(iter) => iter.bits_def,
2401 Self::Unary8(iter) => iter.bits_def,
2402 Self::Unary16(iter) => iter.bits_def,
2403 Self::Unary32(iter) => iter.bits_def,
2404 Self::Nilary(_) => 0,
2405 }
2406 }
2407}
2408
2409pub fn build_control_word_iterator<'a>(
2413 rep: Option<&'a [u16]>,
2414 max_rep: u16,
2415 def: Option<&'a [u16]>,
2416 max_def: u16,
2417 max_visible_def: u16,
2418 len: usize,
2419) -> ControlWordIterator<'a> {
2420 let rep_width = if max_rep == 0 {
2421 0
2422 } else {
2423 log_2_ceil(max_rep as u32) as u16
2424 };
2425 let rep_mask = if max_rep == 0 { 0 } else { get_mask(rep_width) };
2426 let def_width = if max_def == 0 {
2427 0
2428 } else {
2429 log_2_ceil(max_def as u32) as u16
2430 };
2431 let def_mask = if max_def == 0 { 0 } else { get_mask(def_width) };
2432 let total_width = rep_width + def_width;
2433 if total_width == 0 {
2440 return ControlWordIterator::Nilary(NilaryControlWordIterator { len, idx: 0 });
2441 }
2442 match (rep, def) {
2443 (Some(rep), Some(def)) => {
2444 let iter = rep.iter().copied().zip(def.iter().copied());
2445 let def_width = def_width as usize;
2446 if total_width <= 8 {
2447 ControlWordIterator::Binary8(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 } else if total_width <= 16 {
2459 ControlWordIterator::Binary16(BinaryControlWordIterator {
2460 repdef: iter,
2461 rep_mask,
2462 def_mask,
2463 def_width,
2464 max_rep,
2465 max_visible_def,
2466 bits_rep: rep_width as u8,
2467 bits_def: def_width as u8,
2468 phantom: std::marker::PhantomData,
2469 })
2470 } else {
2471 ControlWordIterator::Binary32(BinaryControlWordIterator {
2472 repdef: iter,
2473 rep_mask,
2474 def_mask,
2475 def_width,
2476 max_rep,
2477 max_visible_def,
2478 bits_rep: rep_width as u8,
2479 bits_def: def_width as u8,
2480 phantom: std::marker::PhantomData,
2481 })
2482 }
2483 }
2484 (Some(lev), None) => {
2485 let iter = lev.iter().copied();
2486 if total_width <= 8 {
2487 ControlWordIterator::Unary8(UnaryControlWordIterator {
2488 repdef: iter,
2489 level_mask: rep_mask,
2490 bits_rep: total_width as u8,
2491 bits_def: 0,
2492 max_rep,
2493 phantom: std::marker::PhantomData,
2494 })
2495 } else if total_width <= 16 {
2496 ControlWordIterator::Unary16(UnaryControlWordIterator {
2497 repdef: iter,
2498 level_mask: rep_mask,
2499 bits_rep: total_width as u8,
2500 bits_def: 0,
2501 max_rep,
2502 phantom: std::marker::PhantomData,
2503 })
2504 } else {
2505 ControlWordIterator::Unary32(UnaryControlWordIterator {
2506 repdef: iter,
2507 level_mask: rep_mask,
2508 bits_rep: total_width as u8,
2509 bits_def: 0,
2510 max_rep,
2511 phantom: std::marker::PhantomData,
2512 })
2513 }
2514 }
2515 (None, Some(lev)) => {
2516 let iter = lev.iter().copied();
2517 if total_width <= 8 {
2518 ControlWordIterator::Unary8(UnaryControlWordIterator {
2519 repdef: iter,
2520 level_mask: def_mask,
2521 bits_rep: 0,
2522 bits_def: total_width as u8,
2523 max_rep: 0,
2524 phantom: std::marker::PhantomData,
2525 })
2526 } else if total_width <= 16 {
2527 ControlWordIterator::Unary16(UnaryControlWordIterator {
2528 repdef: iter,
2529 level_mask: def_mask,
2530 bits_rep: 0,
2531 bits_def: total_width as u8,
2532 max_rep: 0,
2533 phantom: std::marker::PhantomData,
2534 })
2535 } else {
2536 ControlWordIterator::Unary32(UnaryControlWordIterator {
2537 repdef: iter,
2538 level_mask: def_mask,
2539 bits_rep: 0,
2540 bits_def: total_width as u8,
2541 max_rep: 0,
2542 phantom: std::marker::PhantomData,
2543 })
2544 }
2545 }
2546 (None, None) => ControlWordIterator::Nilary(NilaryControlWordIterator { len, idx: 0 }),
2547 }
2548}
2549
2550#[derive(Copy, Clone, Debug)]
2554pub enum ControlWordParser {
2555 BOTH8(u8, u32),
2558 BOTH16(u8, u32),
2559 BOTH32(u8, u32),
2560 REP8,
2561 REP16,
2562 REP32,
2563 DEF8,
2564 DEF16,
2565 DEF32,
2566 NIL,
2567}
2568
2569impl ControlWordParser {
2570 fn parse_both<const WORD_SIZE: u8>(
2571 src: &[u8],
2572 dst_rep: &mut Vec<u16>,
2573 dst_def: &mut Vec<u16>,
2574 bits_to_shift: u8,
2575 mask_to_apply: u32,
2576 ) {
2577 match WORD_SIZE {
2578 1 => {
2579 let word = src[0];
2580 let rep = word >> bits_to_shift;
2581 let def = word & (mask_to_apply as u8);
2582 dst_rep.push(rep as u16);
2583 dst_def.push(def as u16);
2584 }
2585 2 => {
2586 let word = u16::from_le_bytes([src[0], src[1]]);
2587 let rep = word >> bits_to_shift;
2588 let def = word & mask_to_apply as u16;
2589 dst_rep.push(rep);
2590 dst_def.push(def);
2591 }
2592 4 => {
2593 let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]);
2594 let rep = word >> bits_to_shift;
2595 let def = word & mask_to_apply;
2596 dst_rep.push(rep as u16);
2597 dst_def.push(def as u16);
2598 }
2599 _ => unreachable!(),
2600 }
2601 }
2602
2603 fn parse_desc_both<const WORD_SIZE: u8>(
2604 src: &[u8],
2605 bits_to_shift: u8,
2606 mask_to_apply: u32,
2607 max_rep: u16,
2608 max_visible_def: u16,
2609 ) -> ControlWordDesc {
2610 match WORD_SIZE {
2611 1 => {
2612 let word = src[0];
2613 let rep = word >> bits_to_shift;
2614 let def = word & (mask_to_apply as u8);
2615 let is_visible = def as u16 <= max_visible_def;
2616 let is_new_row = rep as u16 == max_rep;
2617 let is_valid_item = def == 0;
2618 ControlWordDesc {
2619 is_visible,
2620 is_new_row,
2621 is_valid_item,
2622 }
2623 }
2624 2 => {
2625 let word = u16::from_le_bytes([src[0], src[1]]);
2626 let rep = word >> bits_to_shift;
2627 let def = word & mask_to_apply as u16;
2628 let is_visible = def <= max_visible_def;
2629 let is_new_row = rep == max_rep;
2630 let is_valid_item = def == 0;
2631 ControlWordDesc {
2632 is_visible,
2633 is_new_row,
2634 is_valid_item,
2635 }
2636 }
2637 4 => {
2638 let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]);
2639 let rep = word >> bits_to_shift;
2640 let def = word & mask_to_apply;
2641 let is_visible = def as u16 <= max_visible_def;
2642 let is_new_row = rep as u16 == max_rep;
2643 let is_valid_item = def == 0;
2644 ControlWordDesc {
2645 is_visible,
2646 is_new_row,
2647 is_valid_item,
2648 }
2649 }
2650 _ => unreachable!(),
2651 }
2652 }
2653
2654 fn parse_one<const WORD_SIZE: u8>(src: &[u8], dst: &mut Vec<u16>) {
2655 match WORD_SIZE {
2656 1 => {
2657 let word = src[0];
2658 dst.push(word as u16);
2659 }
2660 2 => {
2661 let word = u16::from_le_bytes([src[0], src[1]]);
2662 dst.push(word);
2663 }
2664 4 => {
2665 let word = u32::from_le_bytes([src[0], src[1], src[2], src[3]]);
2666 dst.push(word as u16);
2667 }
2668 _ => unreachable!(),
2669 }
2670 }
2671
2672 fn parse_rep_desc_one<const WORD_SIZE: u8>(src: &[u8], max_rep: u16) -> ControlWordDesc {
2673 match WORD_SIZE {
2674 1 => ControlWordDesc {
2675 is_new_row: src[0] as u16 == max_rep,
2676 is_visible: true,
2677 is_valid_item: true,
2678 },
2679 2 => ControlWordDesc {
2680 is_new_row: u16::from_le_bytes([src[0], src[1]]) == max_rep,
2681 is_visible: true,
2682 is_valid_item: true,
2683 },
2684 4 => ControlWordDesc {
2685 is_new_row: u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as u16 == max_rep,
2686 is_visible: true,
2687 is_valid_item: true,
2688 },
2689 _ => unreachable!(),
2690 }
2691 }
2692
2693 fn parse_def_desc_one<const WORD_SIZE: u8>(src: &[u8]) -> ControlWordDesc {
2694 match WORD_SIZE {
2695 1 => ControlWordDesc {
2696 is_new_row: true,
2697 is_visible: true,
2698 is_valid_item: src[0] == 0,
2699 },
2700 2 => ControlWordDesc {
2701 is_new_row: true,
2702 is_visible: true,
2703 is_valid_item: u16::from_le_bytes([src[0], src[1]]) == 0,
2704 },
2705 4 => ControlWordDesc {
2706 is_new_row: true,
2707 is_visible: true,
2708 is_valid_item: u32::from_le_bytes([src[0], src[1], src[2], src[3]]) as u16 == 0,
2709 },
2710 _ => unreachable!(),
2711 }
2712 }
2713
2714 pub fn bytes_per_word(&self) -> usize {
2716 match self {
2717 Self::BOTH8(..) => 1,
2718 Self::BOTH16(..) => 2,
2719 Self::BOTH32(..) => 4,
2720 Self::REP8 => 1,
2721 Self::REP16 => 2,
2722 Self::REP32 => 4,
2723 Self::DEF8 => 1,
2724 Self::DEF16 => 2,
2725 Self::DEF32 => 4,
2726 Self::NIL => 0,
2727 }
2728 }
2729
2730 pub fn parse(&self, src: &[u8], dst_rep: &mut Vec<u16>, dst_def: &mut Vec<u16>) {
2737 match self {
2738 Self::BOTH8(bits_to_shift, mask_to_apply) => {
2739 Self::parse_both::<1>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply)
2740 }
2741 Self::BOTH16(bits_to_shift, mask_to_apply) => {
2742 Self::parse_both::<2>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply)
2743 }
2744 Self::BOTH32(bits_to_shift, mask_to_apply) => {
2745 Self::parse_both::<4>(src, dst_rep, dst_def, *bits_to_shift, *mask_to_apply)
2746 }
2747 Self::REP8 => Self::parse_one::<1>(src, dst_rep),
2748 Self::REP16 => Self::parse_one::<2>(src, dst_rep),
2749 Self::REP32 => Self::parse_one::<4>(src, dst_rep),
2750 Self::DEF8 => Self::parse_one::<1>(src, dst_def),
2751 Self::DEF16 => Self::parse_one::<2>(src, dst_def),
2752 Self::DEF32 => Self::parse_one::<4>(src, dst_def),
2753 Self::NIL => {}
2754 }
2755 }
2756
2757 pub fn has_rep(&self) -> bool {
2759 match self {
2760 Self::BOTH8(..)
2761 | Self::BOTH16(..)
2762 | Self::BOTH32(..)
2763 | Self::REP8
2764 | Self::REP16
2765 | Self::REP32 => true,
2766 Self::DEF8 | Self::DEF16 | Self::DEF32 | Self::NIL => false,
2767 }
2768 }
2769
2770 pub fn parse_desc(&self, src: &[u8], max_rep: u16, max_visible_def: u16) -> ControlWordDesc {
2772 match self {
2773 Self::BOTH8(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<1>(
2774 src,
2775 *bits_to_shift,
2776 *mask_to_apply,
2777 max_rep,
2778 max_visible_def,
2779 ),
2780 Self::BOTH16(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<2>(
2781 src,
2782 *bits_to_shift,
2783 *mask_to_apply,
2784 max_rep,
2785 max_visible_def,
2786 ),
2787 Self::BOTH32(bits_to_shift, mask_to_apply) => Self::parse_desc_both::<4>(
2788 src,
2789 *bits_to_shift,
2790 *mask_to_apply,
2791 max_rep,
2792 max_visible_def,
2793 ),
2794 Self::REP8 => Self::parse_rep_desc_one::<1>(src, max_rep),
2795 Self::REP16 => Self::parse_rep_desc_one::<2>(src, max_rep),
2796 Self::REP32 => Self::parse_rep_desc_one::<4>(src, max_rep),
2797 Self::DEF8 => Self::parse_def_desc_one::<1>(src),
2798 Self::DEF16 => Self::parse_def_desc_one::<2>(src),
2799 Self::DEF32 => Self::parse_def_desc_one::<4>(src),
2800 Self::NIL => ControlWordDesc {
2801 is_new_row: true,
2802 is_valid_item: true,
2803 is_visible: true,
2804 },
2805 }
2806 }
2807
2808 pub fn new(bits_rep: u8, bits_def: u8) -> Self {
2810 let total_bits = bits_rep + bits_def;
2811
2812 enum WordSize {
2813 One,
2814 Two,
2815 Four,
2816 }
2817
2818 let word_size = if total_bits <= 8 {
2819 WordSize::One
2820 } else if total_bits <= 16 {
2821 WordSize::Two
2822 } else {
2823 WordSize::Four
2824 };
2825
2826 match (bits_rep > 0, bits_def > 0, word_size) {
2827 (false, false, _) => Self::NIL,
2828 (false, true, WordSize::One) => Self::DEF8,
2829 (false, true, WordSize::Two) => Self::DEF16,
2830 (false, true, WordSize::Four) => Self::DEF32,
2831 (true, false, WordSize::One) => Self::REP8,
2832 (true, false, WordSize::Two) => Self::REP16,
2833 (true, false, WordSize::Four) => Self::REP32,
2834 (true, true, WordSize::One) => Self::BOTH8(bits_def, get_mask(bits_def as u16) as u32),
2835 (true, true, WordSize::Two) => Self::BOTH16(bits_def, get_mask(bits_def as u16) as u32),
2836 (true, true, WordSize::Four) => {
2837 Self::BOTH32(bits_def, get_mask(bits_def as u16) as u32)
2838 }
2839 }
2840 }
2841}
2842
2843#[cfg(test)]
2844mod tests {
2845 use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
2846
2847 use crate::encodings::logical::primitive::sparse::{
2848 SparsePositionPlan, SparseStructuralLayerPlan, SparseStructuralPlan, SparseValidityMeaning,
2849 SparseValiditySet,
2850 };
2851 use crate::repdef::{
2852 CompositeRepDefUnraveler, DefinitionInterpretation, RepDefUnraveler, SerializedRepDefs,
2853 max_visible_level,
2854 };
2855
2856 use super::RepDefBuilder;
2857
2858 #[rstest::rstest]
2862 #[case::no_lists(
2864 &[DefinitionInterpretation::NullableItem, DefinitionInterpretation::NullableItem],
2865 2
2866 )]
2867 #[case::nullable_leaf_under_list(
2869 &[DefinitionInterpretation::NullableItem, DefinitionInterpretation::NullableList],
2870 1
2871 )]
2872 #[case::nullable_struct_above_list(
2874 &[
2875 DefinitionInterpretation::AllValidItem,
2876 DefinitionInterpretation::NullableList,
2877 DefinitionInterpretation::NullableItem,
2878 ],
2879 0
2880 )]
2881 fn test_max_visible_level(
2882 #[case] def_meaning: &[DefinitionInterpretation],
2883 #[case] expected: u16,
2884 ) {
2885 assert_eq!(max_visible_level(def_meaning), expected);
2886 }
2887
2888 fn validity(values: &[bool]) -> NullBuffer {
2889 NullBuffer::from_iter(values.iter().copied())
2890 }
2891
2892 fn offsets_32(values: &[i32]) -> OffsetBuffer<i32> {
2893 OffsetBuffer::<i32>::new(ScalarBuffer::from_iter(values.iter().copied()))
2894 }
2895
2896 fn offsets_64(values: &[i64]) -> OffsetBuffer<i64> {
2897 OffsetBuffer::<i64>::new(ScalarBuffer::from_iter(values.iter().copied()))
2898 }
2899
2900 #[test]
2901 fn sparse_sibling_validity_mismatch_is_invalid_input() {
2902 let sparse = |positions| {
2903 RepDefUnraveler::new_sparse(SparseStructuralPlan {
2904 layers: vec![SparseStructuralLayerPlan::Validity {
2905 num_slots: 2,
2906 validity: SparseValiditySet {
2907 meaning: SparseValidityMeaning::NullPositions,
2908 positions,
2909 },
2910 }],
2911 num_visible_items: 2,
2912 })
2913 };
2914 let mut repdef = CompositeRepDefUnraveler::new(vec![sparse(SparsePositionPlan::Empty)]);
2915 repdef.add_compatibility_check(CompositeRepDefUnraveler::new(vec![sparse(
2916 SparsePositionPlan::Explicit(vec![0]),
2917 )]));
2918
2919 let err = repdef.unravel_validity(2).unwrap_err();
2920 assert!(matches!(err, lance_core::Error::InvalidInput { .. }));
2921 assert!(err.to_string().contains("incompatible validity metadata"));
2922 }
2923
2924 #[test]
2925 fn test_repdef_empty_offsets() {
2926 let mut builder = RepDefBuilder::default();
2928 builder.add_offsets(offsets_32(&[0]), None);
2929 let repdefs = RepDefBuilder::serialize(vec![builder]);
2930 assert!(repdefs.repetition_levels.is_none());
2931 assert!(repdefs.definition_levels.is_none());
2932 }
2933
2934 #[test]
2935 fn test_repdef_basic() {
2936 let mut builder = RepDefBuilder::default();
2938 builder.add_offsets(
2939 offsets_64(&[0, 2, 2, 5]),
2940 Some(validity(&[true, false, true])),
2941 );
2942 builder.add_offsets(
2943 offsets_64(&[0, 1, 3, 5, 5, 9]),
2944 Some(validity(&[true, true, true, false, true])),
2945 );
2946 builder.add_validity_bitmap(validity(&[
2947 true, true, true, false, false, false, true, true, false,
2948 ]));
2949
2950 let repdefs = RepDefBuilder::serialize(vec![builder]);
2951 let rep = repdefs.repetition_levels.unwrap();
2952 let def = repdefs.definition_levels.unwrap();
2953
2954 assert_eq!(vec![0, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1], *def);
2955 assert_eq!(vec![2, 1, 0, 2, 2, 0, 1, 1, 0, 0, 0], *rep);
2956
2957 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
2960 Some(rep.as_ref().to_vec()),
2961 Some(def.as_ref().to_vec()),
2962 repdefs.def_meaning.into(),
2963 9,
2964 )]);
2965
2966 assert_eq!(
2969 unraveler.unravel_validity(9).unwrap(),
2970 Some(validity(&[
2971 true, true, true, false, false, false, true, true, false
2972 ]))
2973 );
2974 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
2975 assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 5, 9]).inner());
2976 assert_eq!(val, Some(validity(&[true, true, true, false, true])));
2977 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
2978 assert_eq!(off.inner(), offsets_32(&[0, 2, 2, 5]).inner());
2979 assert_eq!(val, Some(validity(&[true, false, true])));
2980 }
2981
2982 #[test]
2983 fn test_repdef_simple_null_empty_list() {
2984 let check = |repdefs: SerializedRepDefs, last_def: DefinitionInterpretation| {
2985 let rep = repdefs.repetition_levels.unwrap();
2986 let def = repdefs.definition_levels.unwrap();
2987
2988 assert_eq!([1, 0, 1, 1, 0, 0], *rep);
2989 assert_eq!([0, 0, 2, 0, 1, 0], *def);
2990 assert_eq!(
2991 vec![DefinitionInterpretation::NullableItem, last_def,],
2992 repdefs.def_meaning
2993 );
2994 };
2995
2996 let mut builder = RepDefBuilder::default();
3000 builder.add_offsets(
3001 offsets_32(&[0, 2, 2, 5]),
3002 Some(validity(&[true, false, true])),
3003 );
3004 builder.add_validity_bitmap(validity(&[true, true, true, false, true]));
3005
3006 let repdefs = RepDefBuilder::serialize(vec![builder]);
3007
3008 check(repdefs, DefinitionInterpretation::NullableList);
3009
3010 let mut builder = RepDefBuilder::default();
3012 builder.add_offsets(offsets_32(&[0, 2, 2, 5]), None);
3013 builder.add_validity_bitmap(validity(&[true, true, true, false, true]));
3014
3015 let repdefs = RepDefBuilder::serialize(vec![builder]);
3016
3017 check(repdefs, DefinitionInterpretation::EmptyableList);
3018 }
3019
3020 #[test]
3021 fn test_repdef_empty_list_at_end() {
3022 let mut builder = RepDefBuilder::default();
3024 builder.add_offsets(offsets_32(&[0, 2, 5, 5]), None);
3025 builder.add_validity_bitmap(validity(&[true, true, true, false, true]));
3026
3027 let repdefs = RepDefBuilder::serialize(vec![builder]);
3028
3029 let rep = repdefs.repetition_levels.unwrap();
3030 let def = repdefs.definition_levels.unwrap();
3031
3032 assert_eq!([1, 0, 1, 0, 0, 1], *rep);
3033 assert_eq!([0, 0, 0, 1, 0, 2], *def);
3034 assert_eq!(
3035 vec![
3036 DefinitionInterpretation::NullableItem,
3037 DefinitionInterpretation::EmptyableList,
3038 ],
3039 repdefs.def_meaning
3040 );
3041 }
3042
3043 #[test]
3044 fn test_repdef_abnormal_nulls() {
3045 let mut builder = RepDefBuilder::default();
3048 builder.add_offsets(
3049 offsets_32(&[0, 2, 5, 8]),
3050 Some(validity(&[true, false, true])),
3051 );
3052 builder.add_no_null(5);
3055
3056 let repdefs = RepDefBuilder::serialize(vec![builder]);
3057
3058 let rep = repdefs.repetition_levels.unwrap();
3059 let def = repdefs.definition_levels.unwrap();
3060
3061 assert_eq!([1, 0, 1, 1, 0, 0], *rep);
3062 assert_eq!([0, 0, 1, 0, 0, 0], *def);
3063
3064 assert_eq!(
3065 vec![
3066 DefinitionInterpretation::AllValidItem,
3067 DefinitionInterpretation::NullableList,
3068 ],
3069 repdefs.def_meaning
3070 );
3071 }
3072
3073 #[test]
3074 fn test_repdef_fsl() {
3075 let mut builder = RepDefBuilder::default();
3076 builder.add_fsl(Some(validity(&[true, false])), 2, 2);
3077 builder.add_fsl(None, 2, 4);
3078 builder.add_validity_bitmap(validity(&[
3079 true, false, true, false, true, false, true, false,
3080 ]));
3081
3082 let repdefs = RepDefBuilder::serialize(vec![builder]);
3083
3084 assert_eq!(
3085 vec![
3086 DefinitionInterpretation::NullableItem,
3087 DefinitionInterpretation::AllValidItem,
3088 DefinitionInterpretation::NullableItem
3089 ],
3090 repdefs.def_meaning
3091 );
3092
3093 assert!(repdefs.repetition_levels.is_none());
3094
3095 let def = repdefs.definition_levels.unwrap();
3096
3097 assert_eq!([0, 1, 0, 1, 2, 2, 2, 2], *def);
3098
3099 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3100 None,
3101 Some(def.as_ref().to_vec()),
3102 repdefs.def_meaning.into(),
3103 8,
3104 )]);
3105
3106 assert_eq!(
3107 unraveler.unravel_validity(8).unwrap(),
3108 Some(validity(&[
3109 true, false, true, false, false, false, false, false
3110 ]))
3111 );
3112 assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None);
3113 assert_eq!(
3114 unraveler.unravel_fsl_validity(2, 2).unwrap(),
3115 Some(validity(&[true, false]))
3116 );
3117 }
3118
3119 #[test]
3120 fn test_repdef_fsl_zero_dimension_is_invalid_input() {
3121 let mut builder = RepDefBuilder::default();
3124 builder.add_fsl(Some(validity(&[true, false])), 2, 2);
3125 builder.add_validity_bitmap(validity(&[true, false, true, false]));
3126
3127 let repdefs = RepDefBuilder::serialize(vec![builder]);
3128 let def = repdefs.definition_levels.unwrap();
3129
3130 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3131 None,
3132 Some(def.as_ref().to_vec()),
3133 repdefs.def_meaning.into(),
3134 4,
3135 )]);
3136 unraveler.unravel_validity(4).unwrap();
3138
3139 let err = unraveler.unravel_fsl_validity(2, 0).unwrap_err();
3140 assert!(matches!(err, lance_core::Error::InvalidInput { .. }));
3141 assert!(
3142 err.to_string()
3143 .contains("dimension must be a positive integer"),
3144 "unexpected error: {}",
3145 err
3146 );
3147 }
3148
3149 #[test]
3150 fn test_repdef_fsl_allvalid_item() {
3151 let mut builder = RepDefBuilder::default();
3152 builder.add_fsl(Some(validity(&[true, false])), 2, 2);
3153 builder.add_fsl(None, 2, 4);
3154 builder.add_no_null(8);
3155
3156 let repdefs = RepDefBuilder::serialize(vec![builder]);
3157
3158 assert_eq!(
3159 vec![
3160 DefinitionInterpretation::AllValidItem,
3161 DefinitionInterpretation::AllValidItem,
3162 DefinitionInterpretation::NullableItem
3163 ],
3164 repdefs.def_meaning
3165 );
3166
3167 assert!(repdefs.repetition_levels.is_none());
3168
3169 let def = repdefs.definition_levels.unwrap();
3170
3171 assert_eq!([0, 0, 0, 0, 1, 1, 1, 1], *def);
3172
3173 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3174 None,
3175 Some(def.as_ref().to_vec()),
3176 repdefs.def_meaning.into(),
3177 8,
3178 )]);
3179
3180 assert_eq!(unraveler.unravel_validity(8).unwrap(), None);
3181 assert_eq!(unraveler.unravel_fsl_validity(4, 2).unwrap(), None);
3182 assert_eq!(
3183 unraveler.unravel_fsl_validity(2, 2).unwrap(),
3184 Some(validity(&[true, false]))
3185 );
3186 }
3187
3188 #[test]
3189 fn test_repdef_sliced_offsets() {
3190 let mut builder = RepDefBuilder::default();
3193 builder.add_offsets(
3194 offsets_32(&[5, 7, 7, 10]),
3195 Some(validity(&[true, false, true])),
3196 );
3197 builder.add_no_null(5);
3198
3199 let repdefs = RepDefBuilder::serialize(vec![builder]);
3200
3201 let rep = repdefs.repetition_levels.unwrap();
3202 let def = repdefs.definition_levels.unwrap();
3203
3204 assert_eq!([1, 0, 1, 1, 0, 0], *rep);
3205 assert_eq!([0, 0, 1, 0, 0, 0], *def);
3206
3207 assert_eq!(
3208 vec![
3209 DefinitionInterpretation::AllValidItem,
3210 DefinitionInterpretation::NullableList,
3211 ],
3212 repdefs.def_meaning
3213 );
3214 }
3215
3216 #[test]
3217 fn test_repdef_complex_null_empty() {
3218 let mut builder = RepDefBuilder::default();
3219 builder.add_offsets(
3220 offsets_32(&[0, 4, 4, 4, 6]),
3221 Some(validity(&[true, false, true, true])),
3222 );
3223 builder.add_offsets(
3224 offsets_32(&[0, 1, 1, 2, 2, 2, 3]),
3225 Some(validity(&[true, false, true, false, true, true])),
3226 );
3227 builder.add_no_null(3);
3228
3229 let repdefs = RepDefBuilder::serialize(vec![builder]);
3230
3231 let rep = repdefs.repetition_levels.unwrap();
3232 let def = repdefs.definition_levels.unwrap();
3233
3234 assert_eq!([2, 1, 1, 1, 2, 2, 2, 1], *rep);
3235 assert_eq!([0, 1, 0, 1, 3, 4, 2, 0], *def);
3236 }
3237
3238 #[test]
3239 fn test_repdef_empty_list_no_null() {
3240 let mut builder = RepDefBuilder::default();
3243 builder.add_offsets(offsets_32(&[0, 4, 4, 4, 6]), None);
3244 builder.add_no_null(6);
3245
3246 let repdefs = RepDefBuilder::serialize(vec![builder]);
3247
3248 let rep = repdefs.repetition_levels.unwrap();
3249 let def = repdefs.definition_levels.unwrap();
3250
3251 assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3252 assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def);
3253
3254 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3255 Some(rep.as_ref().to_vec()),
3256 Some(def.as_ref().to_vec()),
3257 repdefs.def_meaning.into(),
3258 8,
3259 )]);
3260
3261 assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3262 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3263 assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3264 assert_eq!(val, None);
3265 }
3266
3267 #[test]
3268 fn test_repdef_all_valid() {
3269 let mut builder = RepDefBuilder::default();
3270 builder.add_offsets(offsets_64(&[0, 2, 3, 5]), None);
3271 builder.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None);
3272 builder.add_no_null(9);
3273
3274 let repdefs = RepDefBuilder::serialize(vec![builder]);
3275 let rep = repdefs.repetition_levels.unwrap();
3276 assert!(repdefs.definition_levels.is_none());
3277
3278 assert_eq!([2, 1, 0, 2, 0, 2, 0, 1, 0], *rep);
3279
3280 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3281 Some(rep.as_ref().to_vec()),
3282 None,
3283 repdefs.def_meaning.into(),
3284 9,
3285 )]);
3286
3287 assert_eq!(unraveler.unravel_validity(9).unwrap(), None);
3288 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3289 assert_eq!(off.inner(), offsets_32(&[0, 1, 3, 5, 7, 9]).inner());
3290 assert_eq!(val, None);
3291 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3292 assert_eq!(off.inner(), offsets_32(&[0, 2, 3, 5]).inner());
3293 assert_eq!(val, None);
3294 }
3295
3296 #[test]
3297 fn test_repdef_nested_list_multibatch_matches_single() {
3298 let mut single = RepDefBuilder::default();
3302 single.add_offsets(offsets_64(&[0, 2, 3, 5]), None);
3303 single.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None);
3304 single.add_no_null(9);
3305 let single_rep = RepDefBuilder::serialize(vec![single])
3306 .repetition_levels
3307 .unwrap();
3308
3309 let mut b0 = RepDefBuilder::default();
3313 b0.add_offsets(offsets_64(&[0, 2, 3]), None);
3314 b0.add_offsets(offsets_64(&[0, 1, 3, 5]), None);
3315 b0.add_no_null(5);
3316 let mut b1 = RepDefBuilder::default();
3317 b1.add_offsets(offsets_64(&[0, 2]), None);
3318 b1.add_offsets(offsets_64(&[0, 2, 4]), None);
3319 b1.add_no_null(4);
3320 let multi_rep = RepDefBuilder::serialize(vec![b0, b1])
3321 .repetition_levels
3322 .unwrap();
3323
3324 assert_eq!(
3325 *single_rep, *multi_rep,
3326 "multi-batch nested-list rep levels must equal single-batch"
3327 );
3328 }
3329
3330 #[test]
3331 fn test_only_empty_lists() {
3332 let mut builder = RepDefBuilder::default();
3333 builder.add_offsets(offsets_32(&[0, 4, 4, 4, 6]), None);
3334 builder.add_no_null(6);
3335
3336 let repdefs = RepDefBuilder::serialize(vec![builder]);
3337
3338 let rep = repdefs.repetition_levels.unwrap();
3339 let def = repdefs.definition_levels.unwrap();
3340
3341 assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3342 assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def);
3343
3344 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3345 Some(rep.as_ref().to_vec()),
3346 Some(def.as_ref().to_vec()),
3347 repdefs.def_meaning.into(),
3348 8,
3349 )]);
3350
3351 assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3352 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3353 assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3354 assert_eq!(val, None);
3355 }
3356
3357 #[test]
3358 fn test_only_null_lists() {
3359 let mut builder = RepDefBuilder::default();
3360 builder.add_offsets(
3361 offsets_32(&[0, 4, 4, 4, 6]),
3362 Some(validity(&[true, false, false, true])),
3363 );
3364 builder.add_no_null(6);
3365
3366 let repdefs = RepDefBuilder::serialize(vec![builder]);
3367
3368 let rep = repdefs.repetition_levels.unwrap();
3369 let def = repdefs.definition_levels.unwrap();
3370
3371 assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3372 assert_eq!([0, 0, 0, 0, 1, 1, 0, 0], *def);
3373
3374 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3375 Some(rep.as_ref().to_vec()),
3376 Some(def.as_ref().to_vec()),
3377 repdefs.def_meaning.into(),
3378 8,
3379 )]);
3380
3381 assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3382 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3383 assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3384 assert_eq!(val, Some(validity(&[true, false, false, true])));
3385 }
3386
3387 #[test]
3388 fn test_null_and_empty_lists() {
3389 let mut builder = RepDefBuilder::default();
3390 builder.add_offsets(
3391 offsets_32(&[0, 4, 4, 4, 6]),
3392 Some(validity(&[true, false, true, true])),
3393 );
3394 builder.add_no_null(6);
3395
3396 let repdefs = RepDefBuilder::serialize(vec![builder]);
3397
3398 let rep = repdefs.repetition_levels.unwrap();
3399 let def = repdefs.definition_levels.unwrap();
3400
3401 assert_eq!([1, 0, 0, 0, 1, 1, 1, 0], *rep);
3402 assert_eq!([0, 0, 0, 0, 1, 2, 0, 0], *def);
3403
3404 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3405 Some(rep.as_ref().to_vec()),
3406 Some(def.as_ref().to_vec()),
3407 repdefs.def_meaning.into(),
3408 8,
3409 )]);
3410
3411 assert_eq!(unraveler.unravel_validity(6).unwrap(), None);
3412 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3413 assert_eq!(off.inner(), offsets_32(&[0, 4, 4, 4, 6]).inner());
3414 assert_eq!(val, Some(validity(&[true, false, true, true])));
3415 }
3416
3417 #[test]
3418 fn test_repdef_nullable_struct_in_null_and_empty_lists() {
3419 let mut builder = RepDefBuilder::default();
3420 builder.add_offsets(
3421 offsets_32(&[0, 0, 3, 3]),
3422 Some(validity(&[false, true, true])),
3423 );
3424 builder.add_validity_bitmap(validity(&[false, true, true]));
3425 builder.add_validity_bitmap(validity(&[true, false, true]));
3426
3427 let repdefs = RepDefBuilder::serialize(vec![builder]);
3428 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3429 repdefs.repetition_levels.map(|levels| levels.to_vec()),
3430 repdefs.definition_levels.map(|levels| levels.to_vec()),
3431 repdefs.def_meaning.into(),
3432 3,
3433 )]);
3434
3435 assert_eq!(
3436 unraveler.unravel_validity(3).unwrap(),
3437 Some(validity(&[false, false, true]))
3438 );
3439 assert_eq!(
3440 unraveler.unravel_validity(3).unwrap(),
3441 Some(validity(&[false, true, true]))
3442 );
3443 let (offsets, nulls) = unraveler.unravel_offsets::<i32>().unwrap();
3444 assert_eq!(offsets.inner(), offsets_32(&[0, 0, 3, 3]).inner());
3445 assert_eq!(nulls, Some(validity(&[false, true, true])));
3446 }
3447
3448 #[test]
3449 fn test_repdef_null_struct_valid_list() {
3450 let rep = vec![1, 0, 0, 0];
3453 let def = vec![2, 0, 2, 2];
3454 let def_meaning = vec![
3456 DefinitionInterpretation::NullableItem,
3457 DefinitionInterpretation::NullableItem,
3458 DefinitionInterpretation::AllValidList,
3459 ];
3460 let num_items = 4;
3461
3462 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3463 Some(rep),
3464 Some(def),
3465 def_meaning.into(),
3466 num_items,
3467 )]);
3468
3469 assert_eq!(
3470 unraveler.unravel_validity(4).unwrap(),
3471 Some(validity(&[false, true, false, false]))
3472 );
3473 assert_eq!(
3474 unraveler.unravel_validity(4).unwrap(),
3475 Some(validity(&[false, true, false, false]))
3476 );
3477 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3478 assert_eq!(off.inner(), offsets_32(&[0, 4]).inner());
3479 assert_eq!(val, None);
3480 }
3481
3482 #[test]
3483 fn test_repdef_no_rep() {
3484 let mut builder = RepDefBuilder::default();
3485 builder.add_no_null(5);
3486 builder.add_validity_bitmap(validity(&[false, false, true, true, true]));
3487 builder.add_validity_bitmap(validity(&[false, true, true, true, false]));
3488
3489 let repdefs = RepDefBuilder::serialize(vec![builder]);
3490 assert!(repdefs.repetition_levels.is_none());
3491 let def = repdefs.definition_levels.unwrap();
3492
3493 assert_eq!([2, 2, 0, 0, 1], *def);
3494
3495 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3496 None,
3497 Some(def.as_ref().to_vec()),
3498 repdefs.def_meaning.into(),
3499 5,
3500 )]);
3501
3502 assert_eq!(
3503 unraveler.unravel_validity(5).unwrap(),
3504 Some(validity(&[false, false, true, true, false]))
3505 );
3506 assert_eq!(
3507 unraveler.unravel_validity(5).unwrap(),
3508 Some(validity(&[false, false, true, true, true]))
3509 );
3510 assert_eq!(unraveler.unravel_validity(5).unwrap(), None);
3511 }
3512
3513 #[test]
3514 fn test_composite_unravel() {
3515 let mut builder = RepDefBuilder::default();
3516 builder.add_offsets(
3517 offsets_64(&[0, 2, 2, 5]),
3518 Some(validity(&[true, false, true])),
3519 );
3520 builder.add_no_null(5);
3521 let repdef1 = RepDefBuilder::serialize(vec![builder]);
3522
3523 let mut builder = RepDefBuilder::default();
3524 builder.add_offsets(offsets_64(&[0, 1, 3, 5, 7, 9]), None);
3525 builder.add_no_null(9);
3526 let repdef2 = RepDefBuilder::serialize(vec![builder]);
3527
3528 let rep1 = repdef1.repetition_levels.clone().unwrap();
3529 let def1 = repdef1.definition_levels.clone().unwrap();
3530 let rep2 = repdef2.repetition_levels.clone().unwrap();
3531 assert!(repdef2.definition_levels.is_none());
3532
3533 assert_eq!([1, 0, 1, 1, 0, 0], *rep1);
3534 assert_eq!([0, 0, 1, 0, 0, 0], *def1);
3535 assert_eq!([1, 1, 0, 1, 0, 1, 0, 1, 0], *rep2);
3536
3537 let unravel1 = RepDefUnraveler::new(
3538 repdef1.repetition_levels.map(|l| l.to_vec()),
3539 repdef1.definition_levels.map(|l| l.to_vec()),
3540 repdef1.def_meaning.into(),
3541 5,
3542 );
3543 let unravel2 = RepDefUnraveler::new(
3544 repdef2.repetition_levels.map(|l| l.to_vec()),
3545 repdef2.definition_levels.map(|l| l.to_vec()),
3546 repdef2.def_meaning.into(),
3547 9,
3548 );
3549
3550 let mut unraveler = CompositeRepDefUnraveler::new(vec![unravel1, unravel2]);
3551
3552 assert!(unraveler.unravel_validity(9).unwrap().is_none());
3553 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3554 assert_eq!(
3555 off.inner(),
3556 offsets_32(&[0, 2, 2, 5, 6, 8, 10, 12, 14]).inner()
3557 );
3558 assert_eq!(
3559 val,
3560 Some(validity(&[true, false, true, true, true, true, true, true]))
3561 );
3562 }
3563
3564 #[test]
3565 fn test_repdef_multiple_builders() {
3566 let mut builder1 = RepDefBuilder::default();
3568 builder1.add_offsets(offsets_64(&[0, 2]), None);
3569 builder1.add_offsets(offsets_64(&[0, 1, 3]), None);
3570 builder1.add_validity_bitmap(validity(&[true, true, true]));
3571
3572 let mut builder2 = RepDefBuilder::default();
3573 builder2.add_offsets(offsets_64(&[0, 0, 3]), Some(validity(&[false, true])));
3574 builder2.add_offsets(
3575 offsets_64(&[0, 2, 2, 6]),
3576 Some(validity(&[true, false, true])),
3577 );
3578 builder2.add_validity_bitmap(validity(&[false, false, false, true, true, false]));
3579
3580 let repdefs = RepDefBuilder::serialize(vec![builder1, builder2]);
3581
3582 let rep = repdefs.repetition_levels.unwrap();
3583 let def = repdefs.definition_levels.unwrap();
3584
3585 assert_eq!([2, 1, 0, 2, 2, 0, 1, 1, 0, 0, 0], *rep);
3586 assert_eq!([0, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1], *def);
3587 }
3588
3589 #[test]
3590 fn test_all_valid_validity_bitmap_serializes_as_no_null() {
3591 let mut from_bitmap = RepDefBuilder::default();
3592 from_bitmap.add_validity_bitmap(validity(&[true, true, true, true]));
3593
3594 let mut from_no_null = RepDefBuilder::default();
3595 from_no_null.add_no_null(4);
3596
3597 let from_bitmap = RepDefBuilder::serialize(vec![from_bitmap]);
3598 let from_no_null = RepDefBuilder::serialize(vec![from_no_null]);
3599
3600 assert!(from_bitmap.repetition_levels.is_none());
3601 assert!(from_bitmap.definition_levels.is_none());
3602 assert_eq!(from_bitmap.def_meaning, from_no_null.def_meaning);
3603 assert_eq!(
3604 from_bitmap.max_visible_level,
3605 from_no_null.max_visible_level
3606 );
3607 }
3608
3609 #[test]
3610 fn test_slicer() {
3611 let mut builder = RepDefBuilder::default();
3612 builder.add_offsets(
3613 offsets_64(&[0, 2, 2, 30, 30]),
3614 Some(validity(&[true, false, true, true])),
3615 );
3616 builder.add_no_null(30);
3617
3618 let repdefs = RepDefBuilder::serialize(vec![builder]);
3619
3620 let mut rep_slicer = repdefs.rep_slicer().unwrap();
3621
3622 assert_eq!(rep_slicer.slice_next(5).len(), 12);
3624 assert_eq!(rep_slicer.slice_next(20).len(), 40);
3626 assert_eq!(rep_slicer.slice_rest().len(), 12);
3628
3629 let mut def_slicer = repdefs.rep_slicer().unwrap();
3630
3631 assert_eq!(def_slicer.slice_next(5).len(), 12);
3633 assert_eq!(def_slicer.slice_next(20).len(), 40);
3635 assert_eq!(def_slicer.slice_rest().len(), 12);
3637 }
3638
3639 #[test]
3640 fn test_control_words() {
3641 fn check(
3643 rep: &[u16],
3644 def: &[u16],
3645 expected_values: Vec<u8>,
3646 expected_bytes_per_word: usize,
3647 expected_bits_rep: u8,
3648 expected_bits_def: u8,
3649 ) {
3650 let num_vals = rep.len().max(def.len());
3651 let max_rep = rep.iter().max().copied().unwrap_or(0);
3652 let max_def = def.iter().max().copied().unwrap_or(0);
3653
3654 let in_rep = if rep.is_empty() { None } else { Some(rep) };
3655 let in_def = if def.is_empty() { None } else { Some(def) };
3656
3657 let mut iter = super::build_control_word_iterator(
3658 in_rep,
3659 max_rep,
3660 in_def,
3661 max_def,
3662 max_def + 1,
3663 expected_values.len(),
3664 );
3665 assert_eq!(iter.bytes_per_word(), expected_bytes_per_word);
3666 assert_eq!(iter.bits_rep(), expected_bits_rep);
3667 assert_eq!(iter.bits_def(), expected_bits_def);
3668 let mut cw_vec = Vec::with_capacity(num_vals * iter.bytes_per_word());
3669
3670 for _ in 0..num_vals {
3671 iter.append_next(&mut cw_vec);
3672 }
3673 assert!(iter.append_next(&mut cw_vec).is_none());
3674
3675 assert_eq!(expected_values, cw_vec);
3676
3677 let parser = super::ControlWordParser::new(expected_bits_rep, expected_bits_def);
3678
3679 let mut rep_out = Vec::with_capacity(num_vals);
3680 let mut def_out = Vec::with_capacity(num_vals);
3681
3682 if expected_bytes_per_word > 0 {
3683 for slice in cw_vec.chunks_exact(expected_bytes_per_word) {
3684 parser.parse(slice, &mut rep_out, &mut def_out);
3685 }
3686 }
3687
3688 assert_eq!(rep, rep_out.as_slice());
3689 assert_eq!(def, def_out.as_slice());
3690 }
3691
3692 let rep = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3694 let def = &[5_u16, 3, 1, 2, 12, 15, 0, 2];
3695 let expected = vec![
3696 0b00000101, 0b01110011, 0b00110001, 0b00100010, 0b10011100, 0b10001111, 0b11000000, 0b01010010, ];
3705 check(rep, def, expected, 1, 4, 4);
3706
3707 let rep = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3709 let def = &[5_u16, 3, 1, 2, 12, 22, 0, 2];
3710 let expected = vec![
3711 0b00000101, 0b00000000, 0b11100011, 0b00000000, 0b01100001, 0b00000000, 0b01000010, 0b00000000, 0b00101100, 0b00000001, 0b00010110, 0b00000001, 0b10000000, 0b00000001, 0b10100010, 0b00000000, ];
3720 check(rep, def, expected, 2, 4, 5);
3721
3722 let levels = &[0_u16, 7, 3, 2, 9, 8, 12, 5];
3724 let expected = vec![
3725 0b00000000, 0b00000111, 0b00000011, 0b00000010, 0b00001001, 0b00001000, 0b00001100, 0b00000101, ];
3734 check(levels, &[], expected.clone(), 1, 4, 0);
3735
3736 check(&[], levels, expected, 1, 0, 4);
3738
3739 check(&[], &[], Vec::default(), 0, 0, 0);
3741 }
3742
3743 #[test]
3744 fn test_control_words_def_only_all_zero_levels() {
3745 let def = [0_u16; 5];
3753
3754 let mut iter = super::build_control_word_iterator(
3755 None,
3756 0,
3757 Some(&def),
3758 0,
3759 u16::MAX,
3760 def.len(),
3761 );
3762 assert_eq!(iter.bytes_per_word(), 0);
3763 assert_eq!(iter.bits_rep(), 0);
3764 assert_eq!(iter.bits_def(), 0);
3765
3766 let mut cw_vec = Vec::new();
3767 for _ in 0..def.len() {
3768 let word_desc = iter.append_next(&mut cw_vec).unwrap();
3769 assert!(word_desc.is_new_row);
3770 assert!(word_desc.is_visible);
3771 }
3772 assert!(iter.append_next(&mut cw_vec).is_none());
3773 assert!(cw_vec.is_empty());
3776
3777 let parser = super::ControlWordParser::new(0, 0);
3780 let mut rep_out = Vec::new();
3781 let mut def_out = Vec::new();
3782 parser.parse(&[], &mut rep_out, &mut def_out);
3783 assert!(rep_out.is_empty());
3784 assert!(def_out.is_empty());
3785 }
3786
3787 #[test]
3788 fn test_control_words_rep_only_all_zero_levels() {
3789 let rep = [0_u16; 5];
3795 let mut iter = super::build_control_word_iterator(
3796 Some(&rep),
3797 0,
3798 None,
3799 0,
3800 u16::MAX,
3801 rep.len(),
3802 );
3803 assert_eq!(iter.bytes_per_word(), 0);
3804 let mut cw_vec = Vec::new();
3805 for _ in 0..rep.len() {
3806 assert!(iter.append_next(&mut cw_vec).unwrap().is_new_row);
3807 }
3808 assert!(iter.append_next(&mut cw_vec).is_none());
3809 assert!(cw_vec.is_empty());
3810 }
3811
3812 #[test]
3813 fn test_control_words_both_levels_all_zero() {
3814 let rep = [0_u16; 5];
3819 let def = [0_u16; 5];
3820 let mut iter = super::build_control_word_iterator(
3821 Some(&rep),
3822 0,
3823 Some(&def),
3824 0,
3825 u16::MAX,
3826 rep.len(),
3827 );
3828 assert_eq!(iter.bytes_per_word(), 0);
3829 assert_eq!(iter.bits_rep(), 0);
3830 assert_eq!(iter.bits_def(), 0);
3831 let mut cw_vec = Vec::new();
3832 for _ in 0..rep.len() {
3833 assert!(iter.append_next(&mut cw_vec).unwrap().is_new_row);
3834 }
3835 assert!(iter.append_next(&mut cw_vec).is_none());
3836 assert!(cw_vec.is_empty());
3837 }
3838
3839 #[test]
3840 fn test_control_words_rep_index() {
3841 fn check(
3842 rep: &[u16],
3843 def: &[u16],
3844 expected_new_rows: Vec<bool>,
3845 expected_is_visible: Vec<bool>,
3846 ) {
3847 let num_vals = rep.len().max(def.len());
3848 let max_rep = rep.iter().max().copied().unwrap_or(0);
3849 let max_def = def.iter().max().copied().unwrap_or(0);
3850
3851 let in_rep = if rep.is_empty() { None } else { Some(rep) };
3852 let in_def = if def.is_empty() { None } else { Some(def) };
3853
3854 let mut iter = super::build_control_word_iterator(
3855 in_rep,
3856 max_rep,
3857 in_def,
3858 max_def,
3859 2,
3860 expected_new_rows.len(),
3861 );
3862
3863 let mut cw_vec = Vec::with_capacity(num_vals * iter.bytes_per_word());
3864 let mut expected_new_rows = expected_new_rows.iter().copied();
3865 let mut expected_is_visible = expected_is_visible.iter().copied();
3866 for _ in 0..expected_new_rows.len() {
3867 let word_desc = iter.append_next(&mut cw_vec).unwrap();
3868 assert_eq!(word_desc.is_new_row, expected_new_rows.next().unwrap());
3869 assert_eq!(word_desc.is_visible, expected_is_visible.next().unwrap());
3870 }
3871 assert!(iter.append_next(&mut cw_vec).is_none());
3872 }
3873
3874 let rep = &[2_u16, 1, 0, 2, 2, 0, 1, 1, 0, 2, 0];
3876 let def = &[0_u16, 0, 0, 3, 1, 1, 2, 1, 0, 0, 1];
3878
3879 check(
3881 rep,
3882 def,
3883 vec![
3884 true, false, false, true, true, false, false, false, false, true, false,
3885 ],
3886 vec![
3887 true, true, true, false, true, true, true, true, true, true, true,
3888 ],
3889 );
3890 check(
3892 rep,
3893 &[],
3894 vec![
3895 true, false, false, true, true, false, false, false, false, true, false,
3896 ],
3897 vec![true; 11],
3898 );
3899 check(
3901 &[],
3902 def,
3903 vec![
3904 true, true, true, true, true, true, true, true, true, true, true,
3905 ],
3906 vec![true; 11],
3907 );
3908 check(
3910 &[],
3911 &[],
3912 vec![
3913 true, true, true, true, true, true, true, true, true, true, true,
3914 ],
3915 vec![true; 11],
3916 );
3917 }
3918
3919 #[test]
3920 fn regress_empty_list_case() {
3921 let mut builder = RepDefBuilder::default();
3923 builder.add_validity_bitmap(validity(&[true, false, true]));
3924 builder.add_offsets(
3925 offsets_32(&[0, 0, 0, 0]),
3926 Some(validity(&[false, false, false])),
3927 );
3928 builder.add_no_null(0);
3929
3930 let repdefs = RepDefBuilder::serialize(vec![builder]);
3931 let rep = repdefs.repetition_levels.unwrap();
3932 let def = repdefs.definition_levels.unwrap();
3933
3934 assert_eq!([1, 1, 1], *rep);
3935 assert_eq!([1, 2, 1], *def);
3936
3937 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3938 Some(rep.as_ref().to_vec()),
3939 Some(def.as_ref().to_vec()),
3940 repdefs.def_meaning.into(),
3941 0,
3942 )]);
3943
3944 assert_eq!(unraveler.unravel_validity(0).unwrap(), None);
3945 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3946 assert_eq!(off.inner(), offsets_32(&[0, 0, 0, 0]).inner());
3947 assert_eq!(val, Some(validity(&[false, false, false])));
3948 let val = unraveler.unravel_validity(3).unwrap().unwrap();
3949 assert_eq!(val.inner(), validity(&[true, false, true]).inner());
3950 }
3951
3952 #[test]
3953 fn regress_list_ends_null_case() {
3954 let mut builder = RepDefBuilder::default();
3955 builder.add_offsets(
3956 offsets_64(&[0, 1, 2, 2]),
3957 Some(validity(&[true, true, false])),
3958 );
3959 builder.add_offsets(offsets_64(&[0, 1, 1]), Some(validity(&[true, false])));
3960 builder.add_no_null(1);
3961
3962 let repdefs = RepDefBuilder::serialize(vec![builder]);
3963 let rep = repdefs.repetition_levels.unwrap();
3964 let def = repdefs.definition_levels.unwrap();
3965
3966 assert_eq!([2, 2, 2], *rep);
3967 assert_eq!([0, 1, 2], *def);
3968
3969 let mut unraveler = CompositeRepDefUnraveler::new(vec![RepDefUnraveler::new(
3970 Some(rep.as_ref().to_vec()),
3971 Some(def.as_ref().to_vec()),
3972 repdefs.def_meaning.into(),
3973 1,
3974 )]);
3975
3976 assert_eq!(unraveler.unravel_validity(1).unwrap(), None);
3977 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3978 assert_eq!(off.inner(), offsets_32(&[0, 1, 1]).inner());
3979 assert_eq!(val, Some(validity(&[true, false])));
3980 let (off, val) = unraveler.unravel_offsets::<i32>().unwrap();
3981 assert_eq!(off.inner(), offsets_32(&[0, 1, 2, 2]).inner());
3982 assert_eq!(val, Some(validity(&[true, true, false])));
3983 }
3984
3985 #[test]
3986 fn test_mixed_unraveler() {
3987 let mut unraveler = CompositeRepDefUnraveler::new(vec![
3992 RepDefUnraveler::new(
3993 None,
3994 Some(vec![0, 1, 0, 1]),
3995 vec![DefinitionInterpretation::NullableItem].into(),
3996 4,
3997 ),
3998 RepDefUnraveler::new(
3999 None,
4000 None,
4001 vec![DefinitionInterpretation::AllValidItem].into(),
4002 4,
4003 ),
4004 ]);
4005
4006 assert_eq!(
4007 unraveler.unravel_validity(8).unwrap(),
4008 Some(validity(&[
4009 true, false, true, false, true, true, true, true
4010 ]))
4011 );
4012
4013 let def1 = Some(vec![0, 1, 2]);
4015 let rep1 = Some(vec![1, 0, 1]);
4016
4017 let def2 = Some(vec![1, 0, 0]);
4018 let rep2 = Some(vec![1, 1, 0]);
4019
4020 let mut unraveler = CompositeRepDefUnraveler::new(vec![
4021 RepDefUnraveler::new(
4022 rep1,
4023 def1,
4024 vec![
4025 DefinitionInterpretation::NullableItem,
4026 DefinitionInterpretation::EmptyableList,
4027 ]
4028 .into(),
4029 2,
4030 ),
4031 RepDefUnraveler::new(
4032 rep2,
4033 def2,
4034 vec![
4035 DefinitionInterpretation::AllValidItem,
4036 DefinitionInterpretation::NullableList,
4037 ]
4038 .into(),
4039 2,
4040 ),
4041 ]);
4042
4043 assert_eq!(
4044 unraveler.unravel_validity(4).unwrap(),
4045 Some(validity(&[true, false, true, true]))
4046 );
4047 assert_eq!(
4048 unraveler.unravel_offsets::<i32>().unwrap(),
4049 (
4050 offsets_32(&[0, 2, 2, 2, 4]),
4051 Some(validity(&[true, true, false, true]))
4052 )
4053 );
4054 }
4055
4056 #[test]
4057 fn test_mixed_unraveler_nullable_without_def_levels() {
4058 let mut unraveler = CompositeRepDefUnraveler::new(vec![
4061 RepDefUnraveler::new(
4062 None,
4063 Some(vec![0, 1, 0, 1]),
4064 vec![DefinitionInterpretation::NullableItem].into(),
4065 4,
4066 ),
4067 RepDefUnraveler::new(
4068 None,
4069 None,
4070 vec![DefinitionInterpretation::NullableItem].into(),
4071 4,
4072 ),
4073 ]);
4074
4075 assert_eq!(
4076 unraveler.unravel_validity(8).unwrap(),
4077 Some(validity(&[
4078 true, false, true, false, true, true, true, true
4079 ]))
4080 );
4081 }
4082}