1#[cfg(feature = "arrow")]
19use crate::column::writer::LevelDataRef;
20use crate::errors::{ParquetError, Result};
21use crate::file::properties::CdcOptions;
22use crate::schema::types::ColumnDescriptor;
23
24use super::CdcChunk;
25use super::cdc_generated::{GEARHASH_TABLE, NUM_GEARHASH_TABLES};
26
27#[derive(Debug)]
88pub(crate) struct ContentDefinedChunker {
89 max_def_level: i16,
91 max_rep_level: i16,
93 repeated_ancestor_def_level: i16,
95
96 min_chunk_size: i64,
101 max_chunk_size: i64,
109 rolling_hash_mask: u64,
111
112 rolling_hash: u64,
114 has_matched: bool,
116 nth_run: usize,
118 chunk_size: i64,
120}
121
122impl ContentDefinedChunker {
123 pub fn new(desc: &ColumnDescriptor, options: &CdcOptions) -> Result<Self> {
124 let rolling_hash_mask = Self::calculate_mask(
125 options.min_chunk_size as i64,
126 options.max_chunk_size as i64,
127 options.norm_level,
128 )?;
129 Ok(Self {
130 max_def_level: desc.max_def_level(),
131 max_rep_level: desc.max_rep_level(),
132 repeated_ancestor_def_level: desc.repeated_ancestor_def_level(),
133 min_chunk_size: options.min_chunk_size as i64,
134 max_chunk_size: options.max_chunk_size as i64,
135 rolling_hash_mask,
136 rolling_hash: 0,
137 has_matched: false,
138 nth_run: 0,
139 chunk_size: 0,
140 })
141 }
142
143 fn calculate_mask(min_chunk_size: i64, max_chunk_size: i64, norm_level: i32) -> Result<u64> {
148 if min_chunk_size < 0 {
149 return Err(ParquetError::General(
150 "min_chunk_size must be non-negative".to_string(),
151 ));
152 }
153 if max_chunk_size <= min_chunk_size {
154 return Err(ParquetError::General(
155 "max_chunk_size must be greater than min_chunk_size".to_string(),
156 ));
157 }
158
159 let avg_chunk_size = min_chunk_size.midpoint(max_chunk_size);
160 let target_size = (avg_chunk_size - min_chunk_size) / NUM_GEARHASH_TABLES as i64;
163
164 let mask_bits = if target_size > 0 {
166 63 - target_size.leading_zeros() as i32
167 } else {
168 0
169 };
170
171 let effective_bits = mask_bits - norm_level;
172
173 if !(1..=63).contains(&effective_bits) {
174 return Err(ParquetError::General(format!(
175 "The number of bits in the CDC mask must be between 1 and 63, got {effective_bits}"
176 )));
177 }
178
179 Ok(u64::MAX << (64 - effective_bits))
181 }
182
183 #[inline]
190 fn roll(&mut self, bytes: &[u8]) {
191 self.chunk_size += bytes.len() as i64;
192 if self.chunk_size < self.min_chunk_size {
193 return;
194 }
195 for &b in bytes {
196 self.rolling_hash = self
197 .rolling_hash
198 .wrapping_shl(1)
199 .wrapping_add(GEARHASH_TABLE[self.nth_run][b as usize]);
200 self.has_matched =
201 self.has_matched || ((self.rolling_hash & self.rolling_hash_mask) == 0);
202 }
203 }
204
205 #[inline(always)]
210 fn roll_fixed<const N: usize>(&mut self, bytes: &[u8; N]) {
211 self.chunk_size += N as i64;
212 if self.chunk_size < self.min_chunk_size {
213 return;
214 }
215 for j in 0..N {
216 self.rolling_hash = self
217 .rolling_hash
218 .wrapping_shl(1)
219 .wrapping_add(GEARHASH_TABLE[self.nth_run][bytes[j] as usize]);
220 self.has_matched =
221 self.has_matched || ((self.rolling_hash & self.rolling_hash_mask) == 0);
222 }
223 }
224
225 #[inline]
227 fn roll_level(&mut self, level: i16) {
228 self.roll_fixed(&level.to_le_bytes());
229 }
230
231 #[inline]
249 fn need_new_chunk(&mut self) -> bool {
250 if self.has_matched {
251 self.has_matched = false;
252 self.nth_run += 1;
253 if self.nth_run >= NUM_GEARHASH_TABLES {
254 self.nth_run = 0;
255 self.chunk_size = 0;
256 return true;
257 }
258 }
259 if self.chunk_size >= self.max_chunk_size {
260 self.chunk_size = 0;
261 return true;
262 }
263 false
264 }
265
266 fn calculate<F>(
279 &mut self,
280 def_levels: LevelDataRef<'_>,
281 rep_levels: LevelDataRef<'_>,
282 num_levels: usize,
283 mut roll_value: F,
284 ) -> Vec<CdcChunk>
285 where
286 F: FnMut(&mut Self, usize),
287 {
288 let has_def_levels = self.max_def_level > 0;
289 let has_rep_levels = self.max_rep_level > 0;
290
291 let mut chunks = Vec::new();
292 let mut prev_offset: usize = 0;
293 let mut prev_value_offset: usize = 0;
294 let mut value_offset: usize = 0;
295
296 if !has_rep_levels && !has_def_levels {
297 for offset in 0..num_levels {
305 roll_value(self, offset);
306 if self.need_new_chunk() {
307 chunks.push(CdcChunk {
308 level_offset: prev_offset,
309 num_levels: offset - prev_offset,
310 value_offset: prev_offset,
311 num_values: offset - prev_offset,
312 });
313 prev_offset = offset;
314 }
315 }
316 prev_value_offset = prev_offset;
317 value_offset = num_levels;
318 } else if !has_rep_levels {
319 for offset in 0..num_levels {
328 let def_level = def_levels
329 .value_at(offset)
330 .expect("def_levels required when max_def_level > 0");
331 self.roll_level(def_level);
332 if def_level == self.max_def_level {
333 roll_value(self, offset);
337 }
338 if self.need_new_chunk() {
341 chunks.push(CdcChunk {
342 level_offset: prev_offset,
343 num_levels: offset - prev_offset,
344 value_offset: prev_value_offset,
345 num_values: value_offset - prev_value_offset,
346 });
347 prev_offset = offset;
348 prev_value_offset = value_offset;
349 }
350 if def_level == self.max_def_level {
351 value_offset += 1;
352 }
353 }
354 } else {
355 let mut leaf_offset: usize = 0;
386
387 for offset in 0..num_levels {
388 let def_level = def_levels
389 .value_at(offset)
390 .expect("def_levels required for nested data");
391 let rep_level = rep_levels
392 .value_at(offset)
393 .expect("rep_levels required for nested data");
394
395 self.roll_level(def_level);
396 self.roll_level(rep_level);
397 if def_level == self.max_def_level {
398 roll_value(self, leaf_offset);
399 }
400
401 if rep_level == 0 && self.need_new_chunk() {
404 let levels_to_write = offset - prev_offset;
405 if levels_to_write > 0 {
406 chunks.push(CdcChunk {
407 level_offset: prev_offset,
408 num_levels: levels_to_write,
409 value_offset: prev_value_offset,
410 num_values: value_offset - prev_value_offset,
411 });
412 prev_offset = offset;
413 prev_value_offset = value_offset;
414 }
415 }
416 if def_level == self.max_def_level {
417 value_offset += 1;
418 }
419 if def_level >= self.repeated_ancestor_def_level {
420 leaf_offset += 1;
421 }
422 }
423 }
424
425 if prev_offset < num_levels {
427 chunks.push(CdcChunk {
428 level_offset: prev_offset,
429 num_levels: num_levels - prev_offset,
430 value_offset: prev_value_offset,
431 num_values: value_offset - prev_value_offset,
432 });
433 }
434
435 #[cfg(debug_assertions)]
436 self.validate_chunks(&chunks, num_levels, value_offset);
437
438 chunks
439 }
440
441 #[cfg(feature = "arrow")]
444 pub(crate) fn get_arrow_chunks(
445 &mut self,
446 def_levels: LevelDataRef<'_>,
447 rep_levels: LevelDataRef<'_>,
448 array: &dyn arrow_array::Array,
449 ) -> Result<Vec<CdcChunk>> {
450 use arrow_array::cast::AsArray;
451 use arrow_schema::DataType;
452
453 let num_levels = match (def_levels.len(), rep_levels.len()) {
458 (0, 0) => array.len(),
459 (d, r) => d.max(r),
460 };
461
462 macro_rules! fixed_width {
463 ($N:literal) => {{
464 let data = array.to_data();
465 let buffer = data.buffers()[0].as_slice();
466 let values = &buffer[data.offset() * $N..];
467 self.calculate(def_levels, rep_levels, num_levels, |c, i| {
468 let offset = i * $N;
469 let slice = &values[offset..offset + $N];
470 c.roll_fixed::<$N>(slice.try_into().unwrap());
471 })
472 }};
473 }
474
475 macro_rules! binary_like {
476 ($a:expr) => {{
477 let a = $a;
478 self.calculate(def_levels, rep_levels, num_levels, |c, i| {
479 c.roll(a.value(i).as_ref());
480 })
481 }};
482 }
483
484 let dtype = array.data_type();
485 let chunks = match dtype {
486 DataType::Null => self.calculate(def_levels, rep_levels, num_levels, |_, _| {}),
487 DataType::Boolean => {
488 let a = array.as_boolean();
489 self.calculate(def_levels, rep_levels, num_levels, |c, i| {
490 c.roll_fixed(&[a.value(i) as u8]);
491 })
492 }
493 DataType::Int8 | DataType::UInt8 => fixed_width!(1),
494 DataType::Int16 | DataType::UInt16 | DataType::Float16 => fixed_width!(2),
495 DataType::Int32
496 | DataType::UInt32
497 | DataType::Float32
498 | DataType::Date32
499 | DataType::Time32(_)
500 | DataType::Interval(arrow_schema::IntervalUnit::YearMonth)
501 | DataType::Decimal32(_, _) => fixed_width!(4),
502 DataType::Int64
503 | DataType::UInt64
504 | DataType::Float64
505 | DataType::Date64
506 | DataType::Time64(_)
507 | DataType::Timestamp(_, _)
508 | DataType::Duration(_)
509 | DataType::Interval(arrow_schema::IntervalUnit::DayTime)
510 | DataType::Decimal64(_, _) => fixed_width!(8),
511 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano)
512 | DataType::Decimal128(_, _) => fixed_width!(16),
513 DataType::Decimal256(_, _) => fixed_width!(32),
514 DataType::FixedSizeBinary(_) => binary_like!(array.as_fixed_size_binary()),
515 DataType::Binary => binary_like!(array.as_binary::<i32>()),
516 DataType::LargeBinary => binary_like!(array.as_binary::<i64>()),
517 DataType::Utf8 => binary_like!(array.as_string::<i32>()),
518 DataType::LargeUtf8 => binary_like!(array.as_string::<i64>()),
519 DataType::BinaryView => binary_like!(array.as_binary_view()),
520 DataType::Utf8View => binary_like!(array.as_string_view()),
521 DataType::Dictionary(_, _) => {
522 let dict = array.as_any_dictionary();
523 self.get_arrow_chunks(def_levels, rep_levels, dict.keys())?
524 }
525 _ => {
526 return Err(ParquetError::General(format!(
527 "content-defined chunking is not supported for data type {dtype:?}",
528 )));
529 }
530 };
531 Ok(chunks)
532 }
533
534 #[cfg(debug_assertions)]
535 fn validate_chunks(&self, chunks: &[CdcChunk], num_levels: usize, total_values: usize) {
536 assert!(!chunks.is_empty(), "chunks must be non-empty");
537
538 let first = &chunks[0];
539 assert_eq!(first.level_offset, 0, "first chunk must start at level 0");
540 assert_eq!(first.value_offset, 0, "first chunk must start at value 0");
541
542 let mut sum_levels = first.num_levels;
543 let mut sum_values = first.num_values;
544 for i in 1..chunks.len() {
545 let chunk = &chunks[i];
546 let prev = &chunks[i - 1];
547 assert!(chunk.num_levels > 0, "chunk must have levels");
548 assert_eq!(
549 chunk.level_offset,
550 prev.level_offset + prev.num_levels,
551 "level offsets must be contiguous"
552 );
553 assert_eq!(
554 chunk.value_offset,
555 prev.value_offset + prev.num_values,
556 "value offsets must be contiguous"
557 );
558 sum_levels += chunk.num_levels;
559 sum_values += chunk.num_values;
560 }
561 assert_eq!(sum_levels, num_levels, "chunks must cover all levels");
562 assert_eq!(sum_values, total_values, "chunks must cover all values");
563
564 let last = chunks.last().unwrap();
565 assert_eq!(
566 last.level_offset + last.num_levels,
567 num_levels,
568 "last chunk must end at num_levels"
569 );
570 }
571}
572
573#[cfg(test)]
574mod tests {
575 use super::*;
576 use crate::basic::Type as PhysicalType;
577 #[cfg(feature = "arrow")]
578 use crate::column::writer::LevelDataRef;
579 use crate::schema::types::{ColumnPath, Type};
580 use std::sync::Arc;
581
582 fn make_desc(max_def_level: i16, max_rep_level: i16) -> ColumnDescriptor {
583 let tp = Type::primitive_type_builder("col", PhysicalType::INT32)
584 .build()
585 .unwrap();
586 ColumnDescriptor::new(
587 Arc::new(tp),
588 max_def_level,
589 max_rep_level,
590 ColumnPath::new(vec![]),
591 )
592 }
593
594 #[test]
595 fn test_calculate_mask_defaults() {
596 let mask = ContentDefinedChunker::calculate_mask(256 * 1024, 1024 * 1024, 0).unwrap();
597 let expected = u64::MAX << (64 - 15);
600 assert_eq!(mask, expected);
601 }
602
603 #[test]
604 fn test_calculate_mask_with_norm_level() {
605 let mask = ContentDefinedChunker::calculate_mask(256 * 1024, 1024 * 1024, 1).unwrap();
606 let expected = u64::MAX << (64 - 14);
607 assert_eq!(mask, expected);
608 }
609
610 #[test]
611 fn test_calculate_mask_invalid() {
612 assert!(ContentDefinedChunker::calculate_mask(-1, 100, 0).is_err());
613 assert!(ContentDefinedChunker::calculate_mask(100, 50, 0).is_err());
614 assert!(ContentDefinedChunker::calculate_mask(100, 100, 0).is_err());
615 }
616
617 #[test]
618 fn test_non_nested_non_null_single_chunk() {
619 let options = CdcOptions {
620 min_chunk_size: 8,
621 max_chunk_size: 1024,
622 norm_level: 0,
623 };
624 let mut chunker = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap();
625
626 let num_values = 4;
628 let chunks = chunker.calculate(
629 LevelDataRef::Absent,
630 LevelDataRef::Absent,
631 num_values,
632 |c, i| {
633 c.roll_fixed::<4>(&(i as i32).to_le_bytes());
634 },
635 );
636 assert_eq!(chunks.len(), 1);
637 assert_eq!(chunks[0].level_offset, 0);
638 assert_eq!(chunks[0].value_offset, 0);
639 assert_eq!(chunks[0].num_levels, 4);
640 }
641
642 #[test]
643 fn test_max_chunk_size_forces_boundary() {
644 let options = CdcOptions {
645 min_chunk_size: 256,
646 max_chunk_size: 1024,
647 norm_level: 0,
648 };
649 let mut chunker = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap();
650
651 let num_values = 2000;
654 let chunks = chunker.calculate(
655 LevelDataRef::Absent,
656 LevelDataRef::Absent,
657 num_values,
658 |c, i| {
659 c.roll_fixed::<4>(&(i as i32).to_le_bytes());
660 },
661 );
662
663 assert!(chunks.len() > 1);
665
666 let mut total_levels = 0;
668 for (i, chunk) in chunks.iter().enumerate() {
669 assert_eq!(chunk.level_offset, total_levels);
670 if i < chunks.len() - 1 {
671 assert!(chunk.num_levels > 0);
672 }
673 total_levels += chunk.num_levels;
674 }
675 assert_eq!(total_levels, num_values);
676 }
677
678 #[test]
679 fn test_deterministic_chunks() {
680 let options = CdcOptions {
681 min_chunk_size: 4,
682 max_chunk_size: 64,
683 norm_level: 0,
684 };
685
686 let roll = |c: &mut ContentDefinedChunker, i: usize| {
687 c.roll_fixed::<8>(&(i as i64).to_le_bytes());
688 };
689
690 let mut chunker1 = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap();
691 let chunks1 = chunker1.calculate(LevelDataRef::Absent, LevelDataRef::Absent, 200, roll);
692
693 let mut chunker2 = ContentDefinedChunker::new(&make_desc(0, 0), &options).unwrap();
694 let chunks2 = chunker2.calculate(LevelDataRef::Absent, LevelDataRef::Absent, 200, roll);
695
696 assert_eq!(chunks1.len(), chunks2.len());
697 for (a, b) in chunks1.iter().zip(chunks2.iter()) {
698 assert_eq!(a.level_offset, b.level_offset);
699 assert_eq!(a.num_levels, b.num_levels);
700 assert_eq!(a.value_offset, b.value_offset);
701 assert_eq!(a.num_values, b.num_values);
702 }
703 }
704
705 #[test]
706 fn test_nullable_non_nested() {
707 let options = CdcOptions {
708 min_chunk_size: 4,
709 max_chunk_size: 64,
710 norm_level: 0,
711 };
712 let mut chunker = ContentDefinedChunker::new(&make_desc(1, 0), &options).unwrap();
713
714 let num_levels = 20;
715 let def_levels: Vec<i16> = (0..num_levels).map(|i| i16::from(i % 3 != 0)).collect();
718 let expected_non_null: usize = def_levels.iter().filter(|&&d| d == 1).count();
719
720 let chunks = chunker.calculate(
721 LevelDataRef::Materialized(&def_levels),
722 LevelDataRef::Absent,
723 num_levels,
724 |c, i| {
725 c.roll_fixed::<4>(&(i as i32).to_le_bytes());
726 },
727 );
728
729 assert!(!chunks.is_empty());
730 let total_levels: usize = chunks.iter().map(|c| c.num_levels).sum();
731 let total_values: usize = chunks.iter().map(|c| c.num_values).sum();
732 assert_eq!(total_levels, num_levels);
733 assert_eq!(total_values, expected_non_null);
734 assert!(total_values < total_levels);
736 }
737}
738
739#[cfg(all(test, feature = "arrow"))]
742mod arrow_tests {
743 use std::borrow::Borrow;
744 use std::cmp::Ordering;
745 use std::sync::Arc;
746
747 use arrow::util::data_gen::create_random_batch;
748 use arrow_array::cast::AsArray;
749 use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, RecordBatch};
750 use arrow_buffer::Buffer;
751 use arrow_data::ArrayData;
752 use arrow_schema::{DataType, Field, Fields, Schema};
753
754 use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
755 use crate::arrow::arrow_writer::ArrowWriter;
756 use crate::column::writer::LevelDataRef;
757 use crate::file::properties::{CdcOptions, WriterProperties};
758 use crate::file::reader::{FileReader, SerializedFileReader};
759
760 const CDC_MIN_CHUNK_SIZE: usize = 4 * 1024;
763 const CDC_MAX_CHUNK_SIZE: usize = 16 * 1024;
764 const CDC_PART_SIZE: usize = 128 * 1024;
765 const CDC_EDIT_SIZE: usize = 128;
766 const CDC_ROW_GROUP_LENGTH: usize = 1024 * 1024;
767
768 fn test_hash(seed: u64, index: u64) -> u64 {
772 let mut h = (index.wrapping_add(seed)).wrapping_mul(0xc4ceb9fe1a85ec53u64);
773 h ^= h >> 33;
774 h = h.wrapping_mul(0xff51afd7ed558ccdu64);
775 h ^= h >> 33;
776 h = h.wrapping_mul(0xc4ceb9fe1a85ec53u64);
777 h ^= h >> 33;
778 h
779 }
780
781 fn generate_array(dtype: &DataType, nullable: bool, length: usize, seed: u64) -> ArrayRef {
783 macro_rules! gen_primitive {
784 ($array_type:ty, $cast:expr) => {{
785 if nullable {
786 let arr: $array_type = (0..length)
787 .map(|i| {
788 let val = test_hash(seed, i as u64);
789 if val % 10 == 0 {
790 None
791 } else {
792 Some($cast(val))
793 }
794 })
795 .collect();
796 Arc::new(arr) as ArrayRef
797 } else {
798 let arr: $array_type = (0..length)
799 .map(|i| Some($cast(test_hash(seed, i as u64))))
800 .collect();
801 Arc::new(arr) as ArrayRef
802 }
803 }};
804 }
805
806 match dtype {
807 DataType::Boolean => {
808 if nullable {
809 let arr: BooleanArray = (0..length)
810 .map(|i| {
811 let val = test_hash(seed, i as u64);
812 if val.is_multiple_of(10) {
813 None
814 } else {
815 Some(val.is_multiple_of(2))
816 }
817 })
818 .collect();
819 Arc::new(arr)
820 } else {
821 let arr: BooleanArray = (0..length)
822 .map(|i| Some(test_hash(seed, i as u64).is_multiple_of(2)))
823 .collect();
824 Arc::new(arr)
825 }
826 }
827 DataType::Int32 => gen_primitive!(Int32Array, |v: u64| v as i32),
828 DataType::Int64 => {
829 gen_primitive!(arrow_array::Int64Array, |v: u64| v as i64)
830 }
831 DataType::Float64 => {
832 gen_primitive!(arrow_array::Float64Array, |v: u64| (v % 100000) as f64
833 / 1000.0)
834 }
835 DataType::Utf8 => {
836 let arr: arrow_array::StringArray = if nullable {
837 (0..length)
838 .map(|i| {
839 let val = test_hash(seed, i as u64);
840 if val.is_multiple_of(10) {
841 None
842 } else {
843 Some(format!("str_{val}"))
844 }
845 })
846 .collect()
847 } else {
848 (0..length)
849 .map(|i| Some(format!("str_{}", test_hash(seed, i as u64))))
850 .collect()
851 };
852 Arc::new(arr)
853 }
854 DataType::Binary => {
855 let arr: arrow_array::BinaryArray = if nullable {
856 (0..length)
857 .map(|i| {
858 let val = test_hash(seed, i as u64);
859 if val.is_multiple_of(10) {
860 None
861 } else {
862 Some(format!("bin_{val}").into_bytes())
863 }
864 })
865 .collect()
866 } else {
867 (0..length)
868 .map(|i| Some(format!("bin_{}", test_hash(seed, i as u64)).into_bytes()))
869 .collect()
870 };
871 Arc::new(arr)
872 }
873 DataType::FixedSizeBinary(size) => {
874 let size = *size;
875 let mut builder = arrow_array::builder::FixedSizeBinaryBuilder::new(size);
876 for i in 0..length {
877 let val = test_hash(seed, i as u64);
878 if nullable && val.is_multiple_of(10) {
879 builder.append_null();
880 } else {
881 let s = format!("bin_{val}");
882 let bytes = s.as_bytes();
883 let mut buf = vec![0u8; size as usize];
884 let copy_len = bytes.len().min(size as usize);
885 buf[..copy_len].copy_from_slice(&bytes[..copy_len]);
886 builder.append_value(&buf).unwrap();
887 }
888 }
889 Arc::new(builder.finish())
890 }
891 DataType::Date32 => {
892 gen_primitive!(arrow_array::Date32Array, |v: u64| v as i32)
893 }
894 DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => {
895 gen_primitive!(arrow_array::TimestampNanosecondArray, |v: u64| v as i64)
896 }
897 _ => panic!("Unsupported test data type: {dtype:?}"),
898 }
899 }
900
901 fn generate_table(schema: &Arc<Schema>, length: usize, seed: u64) -> RecordBatch {
903 let arrays: Vec<ArrayRef> = schema
904 .fields()
905 .iter()
906 .enumerate()
907 .map(|(i, field)| {
908 generate_array(
909 field.data_type(),
910 field.is_nullable(),
911 length,
912 seed + i as u64 * 10,
913 )
914 })
915 .collect();
916 RecordBatch::try_new(schema.clone(), arrays).unwrap()
917 }
918
919 fn cdc_byte_width(dtype: &DataType) -> usize {
922 match dtype {
923 DataType::Boolean => 1,
924 DataType::Int8 | DataType::UInt8 => 1,
925 DataType::Int16 | DataType::UInt16 | DataType::Float16 => 2,
926 DataType::Int32
927 | DataType::UInt32
928 | DataType::Float32
929 | DataType::Date32
930 | DataType::Time32(_) => 4,
931 DataType::Int64
932 | DataType::UInt64
933 | DataType::Float64
934 | DataType::Date64
935 | DataType::Time64(_)
936 | DataType::Timestamp(_, _)
937 | DataType::Duration(_) => 8,
938 DataType::Decimal128(_, _) => 16,
939 DataType::Decimal256(_, _) => 32,
940 DataType::FixedSizeBinary(n) => *n as usize,
941 _ => 0, }
943 }
944
945 fn bytes_per_record(dtype: &DataType, nullable: bool) -> usize {
947 let bw = cdc_byte_width(dtype);
948 if bw > 0 {
949 if nullable { bw + 2 } else { bw }
950 } else {
951 16 }
953 }
954
955 fn calculate_cdc_size(array: &dyn Array, nullable: bool) -> i64 {
957 let dtype = array.data_type();
958 let bw = cdc_byte_width(dtype);
959 let result = if bw > 0 {
960 let valid_count = array.len() - array.null_count();
962 (valid_count * bw) as i64
963 } else {
964 match dtype {
966 DataType::Utf8 => {
967 let a = array.as_string::<i32>();
968 (0..a.len())
969 .filter(|&i| a.is_valid(i))
970 .map(|i| a.value(i).len() as i64)
971 .sum()
972 }
973 DataType::Binary => {
974 let a = array.as_binary::<i32>();
975 (0..a.len())
976 .filter(|&i| a.is_valid(i))
977 .map(|i| a.value(i).len() as i64)
978 .sum()
979 }
980 DataType::LargeBinary => {
981 let a = array.as_binary::<i64>();
982 (0..a.len())
983 .filter(|&i| a.is_valid(i))
984 .map(|i| a.value(i).len() as i64)
985 .sum()
986 }
987 _ => panic!("CDC size calculation not implemented for {dtype:?}"),
988 }
989 };
990
991 if nullable {
992 result + array.len() as i64 * 2
994 } else {
995 result
996 }
997 }
998
999 struct ColumnInfo {
1001 page_lengths: Vec<i64>,
1002 has_dictionary_page: bool,
1003 }
1004
1005 fn get_column_info(data: &[u8], column_index: usize) -> Vec<ColumnInfo> {
1007 let reader = SerializedFileReader::new(bytes::Bytes::from(data.to_vec())).unwrap();
1008 let metadata = reader.metadata();
1009 let mut result = Vec::new();
1010 for rg in 0..metadata.num_row_groups() {
1011 let rg_reader = reader.get_row_group(rg).unwrap();
1012 let col_reader = rg_reader.get_column_page_reader(column_index).unwrap();
1013 let mut info = ColumnInfo {
1014 page_lengths: Vec::new(),
1015 has_dictionary_page: false,
1016 };
1017 for page in col_reader {
1018 let page = page.unwrap();
1019 match page.page_type() {
1020 crate::basic::PageType::DATA_PAGE | crate::basic::PageType::DATA_PAGE_V2 => {
1021 info.page_lengths.push(page.num_values() as i64);
1022 }
1023 crate::basic::PageType::DICTIONARY_PAGE => {
1024 info.has_dictionary_page = true;
1025 }
1026 crate::basic::PageType::INDEX_PAGE => {}
1027 }
1028 }
1029 result.push(info);
1030 }
1031 result
1032 }
1033
1034 fn assert_cdc_chunk_sizes(
1037 array: &ArrayRef,
1038 info: &ColumnInfo,
1039 nullable: bool,
1040 min_chunk_size: usize,
1041 max_chunk_size: usize,
1042 expect_dictionary_page: bool,
1043 ) {
1044 let expect_dict = match array.data_type() {
1046 DataType::Boolean | DataType::FixedSizeBinary(_) => false,
1047 _ => expect_dictionary_page,
1048 };
1049 assert_eq!(
1050 info.has_dictionary_page,
1051 expect_dict,
1052 "dictionary page mismatch for {:?}",
1053 array.data_type()
1054 );
1055
1056 let page_lengths = &info.page_lengths;
1057 assert!(
1058 page_lengths.len() > 1,
1059 "CDC should produce multiple pages, got {page_lengths:?}"
1060 );
1061
1062 let bw = cdc_byte_width(array.data_type());
1063 if bw > 0
1065 || matches!(
1066 array.data_type(),
1067 DataType::Utf8 | DataType::Binary | DataType::LargeBinary
1068 )
1069 {
1070 let mut offset = 0i64;
1071 for (i, &page_len) in page_lengths.iter().enumerate() {
1072 let slice = array.slice(offset as usize, page_len as usize);
1073 let cdc_size = calculate_cdc_size(slice.as_ref(), nullable);
1074 if i < page_lengths.len() - 1 {
1075 assert!(
1076 cdc_size >= min_chunk_size as i64,
1077 "Page {i}: CDC size {cdc_size} < min {min_chunk_size}, pages={page_lengths:?}"
1078 );
1079 }
1080 assert!(
1081 cdc_size <= max_chunk_size as i64,
1082 "Page {i}: CDC size {cdc_size} > max {max_chunk_size}, pages={page_lengths:?}"
1083 );
1084 offset += page_len;
1085 }
1086 assert_eq!(
1087 offset,
1088 array.len() as i64,
1089 "page lengths must sum to array length"
1090 );
1091 }
1092 }
1093
1094 fn write_with_cdc_options(
1097 batches: &[&RecordBatch],
1098 min_chunk_size: usize,
1099 max_chunk_size: usize,
1100 max_row_group_rows: Option<usize>,
1101 enable_dictionary: bool,
1102 ) -> Vec<u8> {
1103 assert!(!batches.is_empty());
1104 let schema = batches[0].schema();
1105 let mut builder = WriterProperties::builder()
1106 .set_dictionary_enabled(enable_dictionary)
1107 .set_content_defined_chunking(Some(CdcOptions {
1108 min_chunk_size,
1109 max_chunk_size,
1110 norm_level: 0,
1111 }));
1112 if let Some(max_rows) = max_row_group_rows {
1113 builder = builder.set_max_row_group_row_count(Some(max_rows));
1114 }
1115 let props = builder.build();
1116 let mut buf = Vec::new();
1117 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap();
1118 for batch in batches {
1119 writer.write(batch).unwrap();
1120 }
1121 writer.close().unwrap();
1122
1123 let readback = read_batches(&buf);
1125 let original_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1126 let readback_rows: usize = readback.iter().map(|b| b.num_rows()).sum();
1127 assert_eq!(original_rows, readback_rows, "Roundtrip row count mismatch");
1128 if original_rows > 0 {
1129 let original = concat_batches(batches.iter().copied());
1130 let roundtrip = concat_batches(&readback);
1131 assert_eq!(original, roundtrip, "Roundtrip validation failed");
1132 }
1133
1134 buf
1135 }
1136
1137 #[test]
1138 fn cdc_all_null_arrow_column_writes_data_pages() {
1139 let array = Arc::new(Int32Array::from(vec![None::<i32>; 4096])) as ArrayRef;
1140 let schema = Arc::new(Schema::new(vec![Field::new("f0", DataType::Int32, true)]));
1141 let batch = RecordBatch::try_new(schema, vec![array.clone()]).unwrap();
1142
1143 let data = write_with_cdc_options(&[&batch], 64, 256, Some(4096), false);
1144 let info = get_column_info(&data, 0);
1145
1146 assert_eq!(info.len(), 1);
1147 assert!(
1148 !info[0].page_lengths.is_empty(),
1149 "all-null CDC write should still emit data pages"
1150 );
1151 assert_eq!(
1152 info[0].page_lengths.iter().sum::<i64>(),
1153 array.len() as i64,
1154 "all-null CDC pages should account for every input row"
1155 );
1156 }
1157
1158 fn read_batches(data: &[u8]) -> Vec<RecordBatch> {
1159 let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(data.to_vec()))
1160 .unwrap()
1161 .build()
1162 .unwrap();
1163 reader.collect::<std::result::Result<Vec<_>, _>>().unwrap()
1164 }
1165
1166 fn concat_batches(batches: impl IntoIterator<Item = impl Borrow<RecordBatch>>) -> RecordBatch {
1167 let batches: Vec<_> = batches.into_iter().collect();
1168 let schema = batches[0].borrow().schema();
1169 let batches = batches.iter().map(|b| b.borrow());
1170 arrow_select::concat::concat_batches(&schema, batches).unwrap()
1171 }
1172
1173 fn find_differences(first: &[i64], second: &[i64]) -> Vec<(Vec<i64>, Vec<i64>)> {
1176 let n = first.len();
1177 let m = second.len();
1178 let mut dp = vec![vec![0usize; m + 1]; n + 1];
1179 for i in 0..n {
1180 for j in 0..m {
1181 if first[i] == second[j] {
1182 dp[i + 1][j + 1] = dp[i][j] + 1;
1183 } else {
1184 dp[i + 1][j + 1] = dp[i + 1][j].max(dp[i][j + 1]);
1185 }
1186 }
1187 }
1188 let mut common = Vec::new();
1189 let (mut i, mut j) = (n, m);
1190 while i > 0 && j > 0 {
1191 if first[i - 1] == second[j - 1] {
1192 common.push((i - 1, j - 1));
1193 i -= 1;
1194 j -= 1;
1195 } else if dp[i - 1][j] >= dp[i][j - 1] {
1196 i -= 1;
1197 } else {
1198 j -= 1;
1199 }
1200 }
1201 common.reverse();
1202
1203 let mut result = Vec::new();
1204 let (mut last_i, mut last_j) = (0usize, 0usize);
1205 for (ci, cj) in &common {
1206 if *ci > last_i || *cj > last_j {
1207 result.push((first[last_i..*ci].to_vec(), second[last_j..*cj].to_vec()));
1208 }
1209 last_i = ci + 1;
1210 last_j = cj + 1;
1211 }
1212 if last_i < n || last_j < m {
1213 result.push((first[last_i..].to_vec(), second[last_j..].to_vec()));
1214 }
1215
1216 let mut merged: Vec<(Vec<i64>, Vec<i64>)> = Vec::new();
1218 for diff in result {
1219 if let Some(prev) = merged.last_mut() {
1220 if prev.0.is_empty() && diff.1.is_empty() {
1221 prev.0 = diff.0;
1222 continue;
1223 } else if prev.1.is_empty() && diff.0.is_empty() {
1224 prev.1 = diff.1;
1225 continue;
1226 }
1227 }
1228 merged.push(diff);
1229 }
1230 merged
1231 }
1232
1233 fn assert_page_length_differences(
1236 original: &ColumnInfo,
1237 modified: &ColumnInfo,
1238 exact_equal_diffs: usize,
1239 exact_larger_diffs: usize,
1240 exact_smaller_diffs: usize,
1241 edit_length: i64,
1242 ) {
1243 let diffs = find_differences(&original.page_lengths, &modified.page_lengths);
1244 let expected = exact_equal_diffs + exact_larger_diffs + exact_smaller_diffs;
1245
1246 if diffs.len() != expected {
1247 eprintln!("Original: {:?}", original.page_lengths);
1248 eprintln!("Modified: {:?}", modified.page_lengths);
1249 for d in &diffs {
1250 eprintln!(" Diff: {:?} vs {:?}", d.0, d.1);
1251 }
1252 }
1253 assert_eq!(
1254 diffs.len(),
1255 expected,
1256 "Expected {expected} diffs, got {}",
1257 diffs.len()
1258 );
1259
1260 let (mut eq, mut larger, mut smaller) = (0usize, 0usize, 0usize);
1261 for (left, right) in &diffs {
1262 let left_sum: i64 = left.iter().sum();
1263 let right_sum: i64 = right.iter().sum();
1264 match left_sum.cmp(&right_sum) {
1265 Ordering::Equal => eq += 1,
1266 Ordering::Less => {
1267 larger += 1;
1268 assert_eq!(
1269 left_sum + edit_length,
1270 right_sum,
1271 "Larger diff mismatch: {left_sum} + {edit_length} != {right_sum}"
1272 );
1273 }
1274 Ordering::Greater => {
1275 smaller += 1;
1276 assert_eq!(
1277 left_sum,
1278 right_sum + edit_length,
1279 "Smaller diff mismatch: {left_sum} != {right_sum} + {edit_length}"
1280 );
1281 }
1282 }
1283 }
1284
1285 assert_eq!(eq, exact_equal_diffs, "equal diffs count");
1286 assert_eq!(larger, exact_larger_diffs, "larger diffs count");
1287 assert_eq!(smaller, exact_smaller_diffs, "smaller diffs count");
1288 }
1289
1290 fn assert_page_length_differences_update(
1293 original: &ColumnInfo,
1294 modified: &ColumnInfo,
1295 max_equal_diffs: usize,
1296 ) {
1297 let diffs = find_differences(&original.page_lengths, &modified.page_lengths);
1298 assert!(
1299 diffs.len() <= max_equal_diffs,
1300 "Expected at most {max_equal_diffs} diffs, got {}",
1301 diffs.len()
1302 );
1303 for (left, right) in &diffs {
1304 let left_sum: i64 = left.iter().sum();
1305 let right_sum: i64 = right.iter().sum();
1306 assert_eq!(
1307 left_sum, right_sum,
1308 "Update diff should not change total row count"
1309 );
1310 }
1311 }
1312
1313 #[test]
1316 fn test_find_differences_basic() {
1317 let diffs = find_differences(&[1, 2, 3, 4, 5], &[1, 7, 8, 4, 5]);
1318 assert_eq!(diffs.len(), 1);
1319 assert_eq!(diffs[0].0, vec![2, 3]);
1320 assert_eq!(diffs[0].1, vec![7, 8]);
1321 }
1322
1323 #[test]
1324 fn test_find_differences_multiple() {
1325 let diffs = find_differences(&[1, 2, 3, 4, 5, 6, 7], &[1, 8, 9, 4, 10, 6, 11]);
1326 assert_eq!(diffs.len(), 3);
1327 assert_eq!(diffs[0].0, vec![2, 3]);
1328 assert_eq!(diffs[0].1, vec![8, 9]);
1329 assert_eq!(diffs[1].0, vec![5]);
1330 assert_eq!(diffs[1].1, vec![10]);
1331 assert_eq!(diffs[2].0, vec![7]);
1332 assert_eq!(diffs[2].1, vec![11]);
1333 }
1334
1335 #[test]
1336 fn test_find_differences_different_lengths() {
1337 let diffs = find_differences(&[1, 2, 3], &[1, 2, 3, 4, 5]);
1338 assert_eq!(diffs.len(), 1);
1339 assert!(diffs[0].0.is_empty());
1340 assert_eq!(diffs[0].1, vec![4, 5]);
1341 }
1342
1343 #[test]
1344 fn test_find_differences_empty() {
1345 let diffs = find_differences(&[], &[]);
1346 assert!(diffs.is_empty());
1347 }
1348
1349 #[test]
1350 fn test_find_differences_changes_at_both_ends() {
1351 let diffs = find_differences(&[1, 2, 3, 4, 5, 6, 7, 8, 9], &[0, 0, 2, 3, 4, 5, 7, 7, 8]);
1352 assert_eq!(diffs.len(), 3);
1353 assert_eq!(diffs[0].0, vec![1]);
1354 assert_eq!(diffs[0].1, vec![0, 0]);
1355 assert_eq!(diffs[1].0, vec![6]);
1356 assert_eq!(diffs[1].1, vec![7]);
1357 assert_eq!(diffs[2].0, vec![9]);
1358 assert!(diffs[2].1.is_empty());
1359 }
1360
1361 #[test]
1362 fn test_find_differences_additional() {
1363 let diffs = find_differences(
1364 &[445, 312, 393, 401, 410, 138, 558, 457],
1365 &[445, 312, 393, 393, 410, 138, 558, 457],
1366 );
1367 assert_eq!(diffs.len(), 1);
1368 assert_eq!(diffs[0].0, vec![401]);
1369 assert_eq!(diffs[0].1, vec![393]);
1370 }
1371
1372 macro_rules! cdc_single_rg_tests {
1375 ($mod_name:ident, $dtype:expr, $nullable:expr) => {
1376 mod $mod_name {
1377 use super::*;
1378
1379 fn config() -> (DataType, bool, usize, usize) {
1380 let dtype: DataType = $dtype;
1381 let nullable: bool = $nullable;
1382 let bpr = bytes_per_record(&dtype, nullable);
1383 let part_length = CDC_PART_SIZE / bpr;
1384 let edit_length = CDC_EDIT_SIZE / bpr;
1385 (dtype, nullable, part_length, edit_length)
1386 }
1387
1388 fn make_schema(dtype: &DataType, nullable: bool) -> Arc<Schema> {
1389 Arc::new(Schema::new(vec![Field::new("f0", dtype.clone(), nullable)]))
1390 }
1391
1392 #[test]
1393 #[cfg_attr(miri, ignore)] fn delete_once() {
1395 let (dtype, nullable, part_length, edit_length) = config();
1396 let schema = make_schema(&dtype, nullable);
1397
1398 let part1 = generate_table(&schema, part_length, 0);
1399 let part2 = generate_table(&schema, edit_length, 1);
1400 let part3 = generate_table(&schema, part_length, part_length as u64);
1401
1402 let base = concat_batches([&part1, &part2, &part3]);
1403 let modified = concat_batches([&part1, &part3]);
1404
1405 for enable_dictionary in [false, true] {
1406 let base_data = write_with_cdc_options(
1407 &[&base],
1408 CDC_MIN_CHUNK_SIZE,
1409 CDC_MAX_CHUNK_SIZE,
1410 Some(CDC_ROW_GROUP_LENGTH),
1411 enable_dictionary,
1412 );
1413 let mod_data = write_with_cdc_options(
1414 &[&modified],
1415 CDC_MIN_CHUNK_SIZE,
1416 CDC_MAX_CHUNK_SIZE,
1417 Some(CDC_ROW_GROUP_LENGTH),
1418 enable_dictionary,
1419 );
1420
1421 let base_info = get_column_info(&base_data, 0);
1422 let mod_info = get_column_info(&mod_data, 0);
1423 assert_eq!(base_info.len(), 1);
1424 assert_eq!(mod_info.len(), 1);
1425
1426 assert_cdc_chunk_sizes(
1427 &base.column(0).clone(),
1428 &base_info[0],
1429 nullable,
1430 CDC_MIN_CHUNK_SIZE,
1431 CDC_MAX_CHUNK_SIZE,
1432 enable_dictionary,
1433 );
1434 assert_cdc_chunk_sizes(
1435 &modified.column(0).clone(),
1436 &mod_info[0],
1437 nullable,
1438 CDC_MIN_CHUNK_SIZE,
1439 CDC_MAX_CHUNK_SIZE,
1440 enable_dictionary,
1441 );
1442
1443 assert_page_length_differences(
1444 &base_info[0],
1445 &mod_info[0],
1446 0,
1447 0,
1448 1,
1449 edit_length as i64,
1450 );
1451 }
1452 }
1453
1454 #[test]
1455 #[cfg_attr(miri, ignore)] fn delete_twice() {
1457 let (dtype, nullable, part_length, edit_length) = config();
1458 let schema = make_schema(&dtype, nullable);
1459
1460 let part1 = generate_table(&schema, part_length, 0);
1461 let part2 = generate_table(&schema, edit_length, 1);
1462 let part3 = generate_table(&schema, part_length, part_length as u64);
1463 let part4 = generate_table(&schema, edit_length, 2);
1464 let part5 = generate_table(&schema, part_length, 2 * part_length as u64);
1465
1466 let base = concat_batches([&part1, &part2, &part3, &part4, &part5]);
1467 let modified = concat_batches([&part1, &part3, &part5]);
1468
1469 for enable_dictionary in [false, true] {
1470 let base_data = write_with_cdc_options(
1471 &[&base],
1472 CDC_MIN_CHUNK_SIZE,
1473 CDC_MAX_CHUNK_SIZE,
1474 Some(CDC_ROW_GROUP_LENGTH),
1475 enable_dictionary,
1476 );
1477 let mod_data = write_with_cdc_options(
1478 &[&modified],
1479 CDC_MIN_CHUNK_SIZE,
1480 CDC_MAX_CHUNK_SIZE,
1481 Some(CDC_ROW_GROUP_LENGTH),
1482 enable_dictionary,
1483 );
1484
1485 let base_info = get_column_info(&base_data, 0);
1486 let mod_info = get_column_info(&mod_data, 0);
1487 assert_eq!(base_info.len(), 1);
1488 assert_eq!(mod_info.len(), 1);
1489
1490 assert_cdc_chunk_sizes(
1491 &base.column(0).clone(),
1492 &base_info[0],
1493 nullable,
1494 CDC_MIN_CHUNK_SIZE,
1495 CDC_MAX_CHUNK_SIZE,
1496 enable_dictionary,
1497 );
1498 assert_cdc_chunk_sizes(
1499 &modified.column(0).clone(),
1500 &mod_info[0],
1501 nullable,
1502 CDC_MIN_CHUNK_SIZE,
1503 CDC_MAX_CHUNK_SIZE,
1504 enable_dictionary,
1505 );
1506
1507 assert_page_length_differences(
1508 &base_info[0],
1509 &mod_info[0],
1510 0,
1511 0,
1512 2,
1513 edit_length as i64,
1514 );
1515 }
1516 }
1517
1518 #[test]
1519 #[cfg_attr(miri, ignore)] fn insert_once() {
1521 let (dtype, nullable, part_length, edit_length) = config();
1522 let schema = make_schema(&dtype, nullable);
1523
1524 let part1 = generate_table(&schema, part_length, 0);
1525 let part2 = generate_table(&schema, edit_length, 1);
1526 let part3 = generate_table(&schema, part_length, part_length as u64);
1527
1528 let base = concat_batches([&part1, &part3]);
1529 let modified = concat_batches([&part1, &part2, &part3]);
1530
1531 for enable_dictionary in [false, true] {
1532 let base_data = write_with_cdc_options(
1533 &[&base],
1534 CDC_MIN_CHUNK_SIZE,
1535 CDC_MAX_CHUNK_SIZE,
1536 Some(CDC_ROW_GROUP_LENGTH),
1537 enable_dictionary,
1538 );
1539 let mod_data = write_with_cdc_options(
1540 &[&modified],
1541 CDC_MIN_CHUNK_SIZE,
1542 CDC_MAX_CHUNK_SIZE,
1543 Some(CDC_ROW_GROUP_LENGTH),
1544 enable_dictionary,
1545 );
1546
1547 let base_info = get_column_info(&base_data, 0);
1548 let mod_info = get_column_info(&mod_data, 0);
1549 assert_eq!(base_info.len(), 1);
1550 assert_eq!(mod_info.len(), 1);
1551
1552 assert_cdc_chunk_sizes(
1553 &base.column(0).clone(),
1554 &base_info[0],
1555 nullable,
1556 CDC_MIN_CHUNK_SIZE,
1557 CDC_MAX_CHUNK_SIZE,
1558 enable_dictionary,
1559 );
1560 assert_cdc_chunk_sizes(
1561 &modified.column(0).clone(),
1562 &mod_info[0],
1563 nullable,
1564 CDC_MIN_CHUNK_SIZE,
1565 CDC_MAX_CHUNK_SIZE,
1566 enable_dictionary,
1567 );
1568
1569 assert_page_length_differences(
1570 &base_info[0],
1571 &mod_info[0],
1572 0,
1573 1,
1574 0,
1575 edit_length as i64,
1576 );
1577 }
1578 }
1579
1580 #[test]
1581 #[cfg_attr(miri, ignore)] fn insert_twice() {
1583 let (dtype, nullable, part_length, edit_length) = config();
1584 let schema = make_schema(&dtype, nullable);
1585
1586 let part1 = generate_table(&schema, part_length, 0);
1587 let part2 = generate_table(&schema, edit_length, 1);
1588 let part3 = generate_table(&schema, part_length, part_length as u64);
1589 let part4 = generate_table(&schema, edit_length, 2);
1590 let part5 = generate_table(&schema, part_length, 2 * part_length as u64);
1591
1592 let base = concat_batches([&part1, &part3, &part5]);
1593 let modified = concat_batches([&part1, &part2, &part3, &part4, &part5]);
1594
1595 for enable_dictionary in [false, true] {
1596 let base_data = write_with_cdc_options(
1597 &[&base],
1598 CDC_MIN_CHUNK_SIZE,
1599 CDC_MAX_CHUNK_SIZE,
1600 Some(CDC_ROW_GROUP_LENGTH),
1601 enable_dictionary,
1602 );
1603 let mod_data = write_with_cdc_options(
1604 &[&modified],
1605 CDC_MIN_CHUNK_SIZE,
1606 CDC_MAX_CHUNK_SIZE,
1607 Some(CDC_ROW_GROUP_LENGTH),
1608 enable_dictionary,
1609 );
1610
1611 let base_info = get_column_info(&base_data, 0);
1612 let mod_info = get_column_info(&mod_data, 0);
1613 assert_eq!(base_info.len(), 1);
1614 assert_eq!(mod_info.len(), 1);
1615
1616 assert_cdc_chunk_sizes(
1617 &base.column(0).clone(),
1618 &base_info[0],
1619 nullable,
1620 CDC_MIN_CHUNK_SIZE,
1621 CDC_MAX_CHUNK_SIZE,
1622 enable_dictionary,
1623 );
1624 assert_cdc_chunk_sizes(
1625 &modified.column(0).clone(),
1626 &mod_info[0],
1627 nullable,
1628 CDC_MIN_CHUNK_SIZE,
1629 CDC_MAX_CHUNK_SIZE,
1630 enable_dictionary,
1631 );
1632
1633 assert_page_length_differences(
1634 &base_info[0],
1635 &mod_info[0],
1636 0,
1637 2,
1638 0,
1639 edit_length as i64,
1640 );
1641 }
1642 }
1643
1644 #[test]
1645 #[cfg_attr(miri, ignore)] fn update_once() {
1647 let (dtype, nullable, part_length, edit_length) = config();
1648 let schema = make_schema(&dtype, nullable);
1649
1650 let part1 = generate_table(&schema, part_length, 0);
1651 let part2 = generate_table(&schema, edit_length, 1);
1652 let part3 = generate_table(&schema, part_length, part_length as u64);
1653 let part4 = generate_table(&schema, edit_length, 2);
1654
1655 let base = concat_batches([&part1, &part2, &part3]);
1656 let modified = concat_batches([&part1, &part4, &part3]);
1657
1658 for enable_dictionary in [false, true] {
1659 let base_data = write_with_cdc_options(
1660 &[&base],
1661 CDC_MIN_CHUNK_SIZE,
1662 CDC_MAX_CHUNK_SIZE,
1663 Some(CDC_ROW_GROUP_LENGTH),
1664 enable_dictionary,
1665 );
1666 let mod_data = write_with_cdc_options(
1667 &[&modified],
1668 CDC_MIN_CHUNK_SIZE,
1669 CDC_MAX_CHUNK_SIZE,
1670 Some(CDC_ROW_GROUP_LENGTH),
1671 enable_dictionary,
1672 );
1673
1674 let base_info = get_column_info(&base_data, 0);
1675 let mod_info = get_column_info(&mod_data, 0);
1676 assert_eq!(base_info.len(), 1);
1677 assert_eq!(mod_info.len(), 1);
1678
1679 assert_cdc_chunk_sizes(
1680 &base.column(0).clone(),
1681 &base_info[0],
1682 nullable,
1683 CDC_MIN_CHUNK_SIZE,
1684 CDC_MAX_CHUNK_SIZE,
1685 enable_dictionary,
1686 );
1687 assert_cdc_chunk_sizes(
1688 &modified.column(0).clone(),
1689 &mod_info[0],
1690 nullable,
1691 CDC_MIN_CHUNK_SIZE,
1692 CDC_MAX_CHUNK_SIZE,
1693 enable_dictionary,
1694 );
1695
1696 assert_page_length_differences_update(&base_info[0], &mod_info[0], 1);
1697 }
1698 }
1699
1700 #[test]
1701 #[cfg_attr(miri, ignore)] fn update_twice() {
1703 let (dtype, nullable, part_length, edit_length) = config();
1704 let schema = make_schema(&dtype, nullable);
1705
1706 let part1 = generate_table(&schema, part_length, 0);
1707 let part2 = generate_table(&schema, edit_length, 1);
1708 let part3 = generate_table(&schema, part_length, part_length as u64);
1709 let part4 = generate_table(&schema, edit_length, 2);
1710 let part5 = generate_table(&schema, part_length, 2 * part_length as u64);
1711 let part6 = generate_table(&schema, edit_length, 3);
1712 let part7 = generate_table(&schema, edit_length, 4);
1713
1714 let base = concat_batches([&part1, &part2, &part3, &part4, &part5]);
1715 let modified = concat_batches([&part1, &part6, &part3, &part7, &part5]);
1716
1717 for enable_dictionary in [false, true] {
1718 let base_data = write_with_cdc_options(
1719 &[&base],
1720 CDC_MIN_CHUNK_SIZE,
1721 CDC_MAX_CHUNK_SIZE,
1722 Some(CDC_ROW_GROUP_LENGTH),
1723 enable_dictionary,
1724 );
1725 let mod_data = write_with_cdc_options(
1726 &[&modified],
1727 CDC_MIN_CHUNK_SIZE,
1728 CDC_MAX_CHUNK_SIZE,
1729 Some(CDC_ROW_GROUP_LENGTH),
1730 enable_dictionary,
1731 );
1732
1733 let base_info = get_column_info(&base_data, 0);
1734 let mod_info = get_column_info(&mod_data, 0);
1735 assert_eq!(base_info.len(), 1);
1736 assert_eq!(mod_info.len(), 1);
1737
1738 assert_cdc_chunk_sizes(
1739 &base.column(0).clone(),
1740 &base_info[0],
1741 nullable,
1742 CDC_MIN_CHUNK_SIZE,
1743 CDC_MAX_CHUNK_SIZE,
1744 enable_dictionary,
1745 );
1746 assert_cdc_chunk_sizes(
1747 &modified.column(0).clone(),
1748 &mod_info[0],
1749 nullable,
1750 CDC_MIN_CHUNK_SIZE,
1751 CDC_MAX_CHUNK_SIZE,
1752 enable_dictionary,
1753 );
1754
1755 assert_page_length_differences_update(&base_info[0], &mod_info[0], 2);
1756 }
1757 }
1758
1759 #[test]
1760 #[cfg_attr(miri, ignore)] fn prepend() {
1762 let (dtype, nullable, part_length, edit_length) = config();
1763 let schema = make_schema(&dtype, nullable);
1764
1765 let part1 = generate_table(&schema, part_length, 0);
1766 let part2 = generate_table(&schema, edit_length, 1);
1767 let part3 = generate_table(&schema, part_length, part_length as u64);
1768 let part4 = generate_table(&schema, edit_length, 2);
1769
1770 let base = concat_batches([&part1, &part2, &part3]);
1771 let modified = concat_batches([&part4, &part1, &part2, &part3]);
1772
1773 for enable_dictionary in [false, true] {
1774 let base_data = write_with_cdc_options(
1775 &[&base],
1776 CDC_MIN_CHUNK_SIZE,
1777 CDC_MAX_CHUNK_SIZE,
1778 Some(CDC_ROW_GROUP_LENGTH),
1779 enable_dictionary,
1780 );
1781 let mod_data = write_with_cdc_options(
1782 &[&modified],
1783 CDC_MIN_CHUNK_SIZE,
1784 CDC_MAX_CHUNK_SIZE,
1785 Some(CDC_ROW_GROUP_LENGTH),
1786 enable_dictionary,
1787 );
1788
1789 let base_info = get_column_info(&base_data, 0);
1790 let mod_info = get_column_info(&mod_data, 0);
1791 assert_eq!(base_info.len(), 1);
1792 assert_eq!(mod_info.len(), 1);
1793
1794 assert_cdc_chunk_sizes(
1795 &base.column(0).clone(),
1796 &base_info[0],
1797 nullable,
1798 CDC_MIN_CHUNK_SIZE,
1799 CDC_MAX_CHUNK_SIZE,
1800 enable_dictionary,
1801 );
1802 assert_cdc_chunk_sizes(
1803 &modified.column(0).clone(),
1804 &mod_info[0],
1805 nullable,
1806 CDC_MIN_CHUNK_SIZE,
1807 CDC_MAX_CHUNK_SIZE,
1808 enable_dictionary,
1809 );
1810
1811 assert!(
1812 mod_info[0].page_lengths.len() >= base_info[0].page_lengths.len(),
1813 "Modified should have same or more pages"
1814 );
1815
1816 assert_page_length_differences(
1817 &base_info[0],
1818 &mod_info[0],
1819 0,
1820 1,
1821 0,
1822 edit_length as i64,
1823 );
1824 }
1825 }
1826
1827 #[test]
1828 #[cfg_attr(miri, ignore)] fn append() {
1830 let (dtype, nullable, part_length, edit_length) = config();
1831 let schema = make_schema(&dtype, nullable);
1832
1833 let part1 = generate_table(&schema, part_length, 0);
1834 let part2 = generate_table(&schema, edit_length, 1);
1835 let part3 = generate_table(&schema, part_length, part_length as u64);
1836 let part4 = generate_table(&schema, edit_length, 2);
1837
1838 let base = concat_batches([&part1, &part2, &part3]);
1839 let modified = concat_batches([&part1, &part2, &part3, &part4]);
1840
1841 for enable_dictionary in [false, true] {
1842 let base_data = write_with_cdc_options(
1843 &[&base],
1844 CDC_MIN_CHUNK_SIZE,
1845 CDC_MAX_CHUNK_SIZE,
1846 Some(CDC_ROW_GROUP_LENGTH),
1847 enable_dictionary,
1848 );
1849 let mod_data = write_with_cdc_options(
1850 &[&modified],
1851 CDC_MIN_CHUNK_SIZE,
1852 CDC_MAX_CHUNK_SIZE,
1853 Some(CDC_ROW_GROUP_LENGTH),
1854 enable_dictionary,
1855 );
1856
1857 let base_info = get_column_info(&base_data, 0);
1858 let mod_info = get_column_info(&mod_data, 0);
1859 assert_eq!(base_info.len(), 1);
1860 assert_eq!(mod_info.len(), 1);
1861
1862 assert_cdc_chunk_sizes(
1863 &base.column(0).clone(),
1864 &base_info[0],
1865 nullable,
1866 CDC_MIN_CHUNK_SIZE,
1867 CDC_MAX_CHUNK_SIZE,
1868 enable_dictionary,
1869 );
1870 assert_cdc_chunk_sizes(
1871 &modified.column(0).clone(),
1872 &mod_info[0],
1873 nullable,
1874 CDC_MIN_CHUNK_SIZE,
1875 CDC_MAX_CHUNK_SIZE,
1876 enable_dictionary,
1877 );
1878
1879 let bp = &base_info[0].page_lengths;
1880 let mp = &mod_info[0].page_lengths;
1881 assert!(mp.len() >= bp.len());
1882 for i in 0..bp.len() - 1 {
1883 assert_eq!(bp[i], mp[i], "Page {i} should be identical");
1884 }
1885 assert!(mp[bp.len() - 1] >= bp[bp.len() - 1]);
1886 }
1887 }
1888
1889 #[test]
1890 fn empty_table() {
1891 let (dtype, nullable, _, _) = config();
1892 let schema = make_schema(&dtype, nullable);
1893
1894 let empty = RecordBatch::new_empty(schema);
1895 for enable_dictionary in [false, true] {
1896 let data = write_with_cdc_options(
1897 &[&empty],
1898 CDC_MIN_CHUNK_SIZE,
1899 CDC_MAX_CHUNK_SIZE,
1900 Some(CDC_ROW_GROUP_LENGTH),
1901 enable_dictionary,
1902 );
1903 let info = get_column_info(&data, 0);
1904 if !info.is_empty() {
1906 assert!(info[0].page_lengths.is_empty());
1907 }
1908 }
1909 }
1910
1911 #[test]
1912 #[cfg_attr(miri, ignore)] fn array_offsets() {
1914 let (dtype, nullable, part_length, edit_length) = config();
1915 let schema = make_schema(&dtype, nullable);
1916
1917 let table = concat_batches([
1918 &generate_table(&schema, part_length, 0),
1919 &generate_table(&schema, edit_length, 1),
1920 &generate_table(&schema, part_length, part_length as u64),
1921 ]);
1922
1923 for offset in [0usize, 512, 1024] {
1924 if offset >= table.num_rows() {
1925 continue;
1926 }
1927 let sliced = table.slice(offset, table.num_rows() - offset);
1928 let data = write_with_cdc_options(
1929 &[&sliced],
1930 CDC_MIN_CHUNK_SIZE,
1931 CDC_MAX_CHUNK_SIZE,
1932 Some(CDC_ROW_GROUP_LENGTH),
1933 true,
1934 );
1935 let info = get_column_info(&data, 0);
1936 assert_eq!(info.len(), 1);
1937
1938 assert_cdc_chunk_sizes(
1940 &sliced.column(0).clone(),
1941 &info[0],
1942 nullable,
1943 CDC_MIN_CHUNK_SIZE,
1944 CDC_MAX_CHUNK_SIZE,
1945 true,
1946 );
1947 }
1948 }
1949 }
1950 };
1951 }
1952
1953 cdc_single_rg_tests!(cdc_bool_non_null, DataType::Boolean, false);
1955 cdc_single_rg_tests!(cdc_i32_non_null, DataType::Int32, false);
1956 cdc_single_rg_tests!(cdc_i64_nullable, DataType::Int64, true);
1957 cdc_single_rg_tests!(cdc_f64_nullable, DataType::Float64, true);
1958 cdc_single_rg_tests!(cdc_utf8_non_null, DataType::Utf8, false);
1959 cdc_single_rg_tests!(cdc_binary_nullable, DataType::Binary, true);
1960 cdc_single_rg_tests!(cdc_fsb16_nullable, DataType::FixedSizeBinary(16), true);
1961 cdc_single_rg_tests!(cdc_date32_non_null, DataType::Date32, false);
1962 cdc_single_rg_tests!(
1963 cdc_timestamp_nullable,
1964 DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, None),
1965 true
1966 );
1967
1968 mod cdc_multiple_row_groups {
1971 use super::*;
1972
1973 const PART_LENGTH: usize = 128 * 1024;
1974 const EDIT_LENGTH: usize = 128;
1975 const ROW_GROUP_LENGTH: usize = 64 * 1024;
1976
1977 fn schema() -> Arc<Schema> {
1978 Arc::new(Schema::new(vec![
1979 Field::new("int32", DataType::Int32, true),
1980 Field::new("float64", DataType::Float64, true),
1981 Field::new("bool", DataType::Boolean, false),
1982 ]))
1983 }
1984
1985 #[test]
1986 #[cfg_attr(miri, ignore)] fn insert_once() {
1988 let s = schema();
1989 let part1 = generate_table(&s, PART_LENGTH, 0);
1990 let part2 = generate_table(&s, PART_LENGTH, 2);
1991 let part3 = generate_table(&s, PART_LENGTH, 4);
1992 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
1993 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
1994
1995 let base = concat_batches([&part1, &edit1, &part2, &part3]);
1996 let modified = concat_batches([&part1, &edit1, &edit2, &part2, &part3]);
1997 assert_eq!(modified.num_rows(), base.num_rows() + EDIT_LENGTH);
1998
1999 let base_data = write_with_cdc_options(
2000 &[&base],
2001 CDC_MIN_CHUNK_SIZE,
2002 CDC_MAX_CHUNK_SIZE,
2003 Some(ROW_GROUP_LENGTH),
2004 false,
2005 );
2006 let mod_data = write_with_cdc_options(
2007 &[&modified],
2008 CDC_MIN_CHUNK_SIZE,
2009 CDC_MAX_CHUNK_SIZE,
2010 Some(ROW_GROUP_LENGTH),
2011 false,
2012 );
2013
2014 for col in 0..s.fields().len() {
2015 let base_info = get_column_info(&base_data, col);
2016 let mod_info = get_column_info(&mod_data, col);
2017
2018 assert_eq!(base_info.len(), 7, "expected 7 row groups for col {col}");
2019 assert_eq!(mod_info.len(), 7);
2020
2021 assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths);
2023 assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths);
2024
2025 for i in 2..mod_info.len() - 1 {
2027 assert_page_length_differences(
2028 &base_info[i],
2029 &mod_info[i],
2030 0,
2031 1,
2032 1,
2033 EDIT_LENGTH as i64,
2034 );
2035 }
2036 assert_page_length_differences(
2038 base_info.last().unwrap(),
2039 mod_info.last().unwrap(),
2040 0,
2041 1,
2042 0,
2043 EDIT_LENGTH as i64,
2044 );
2045 }
2046 }
2047
2048 #[test]
2049 #[cfg_attr(miri, ignore)] fn delete_once() {
2051 let s = schema();
2052 let part1 = generate_table(&s, PART_LENGTH, 0);
2053 let part2 = generate_table(&s, PART_LENGTH, 2);
2054 let part3 = generate_table(&s, PART_LENGTH, 4);
2055 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2056 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2057
2058 let base = concat_batches([&part1, &edit1, &part2, &part3, &edit2]);
2059 let modified = concat_batches([&part1, &part2, &part3, &edit2]);
2060
2061 let base_data = write_with_cdc_options(
2062 &[&base],
2063 CDC_MIN_CHUNK_SIZE,
2064 CDC_MAX_CHUNK_SIZE,
2065 Some(ROW_GROUP_LENGTH),
2066 false,
2067 );
2068 let mod_data = write_with_cdc_options(
2069 &[&modified],
2070 CDC_MIN_CHUNK_SIZE,
2071 CDC_MAX_CHUNK_SIZE,
2072 Some(ROW_GROUP_LENGTH),
2073 false,
2074 );
2075
2076 for col in 0..s.fields().len() {
2077 let base_info = get_column_info(&base_data, col);
2078 let mod_info = get_column_info(&mod_data, col);
2079
2080 assert_eq!(base_info.len(), 7);
2081 assert_eq!(mod_info.len(), 7);
2082
2083 assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths);
2084 assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths);
2085
2086 for i in 2..mod_info.len() - 1 {
2087 assert_page_length_differences(
2088 &base_info[i],
2089 &mod_info[i],
2090 0,
2091 1,
2092 1,
2093 EDIT_LENGTH as i64,
2094 );
2095 }
2096 assert_page_length_differences(
2097 base_info.last().unwrap(),
2098 mod_info.last().unwrap(),
2099 0,
2100 0,
2101 1,
2102 EDIT_LENGTH as i64,
2103 );
2104 }
2105 }
2106
2107 #[test]
2108 #[cfg_attr(miri, ignore)] fn update_once() {
2110 let s = schema();
2111 let part1 = generate_table(&s, PART_LENGTH, 0);
2112 let part2 = generate_table(&s, PART_LENGTH, 2);
2113 let part3 = generate_table(&s, PART_LENGTH, 4);
2114 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2115 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2116 let edit3 = generate_table(&s, EDIT_LENGTH, 5);
2117
2118 let base = concat_batches([&part1, &edit1, &part2, &part3, &edit2]);
2119 let modified = concat_batches([&part1, &edit3, &part2, &part3, &edit2]);
2120
2121 let base_data = write_with_cdc_options(
2122 &[&base],
2123 CDC_MIN_CHUNK_SIZE,
2124 CDC_MAX_CHUNK_SIZE,
2125 Some(ROW_GROUP_LENGTH),
2126 false,
2127 );
2128 let mod_data = write_with_cdc_options(
2129 &[&modified],
2130 CDC_MIN_CHUNK_SIZE,
2131 CDC_MAX_CHUNK_SIZE,
2132 Some(ROW_GROUP_LENGTH),
2133 false,
2134 );
2135
2136 for col in 0..s.fields().len() {
2137 let nullable = s.field(col).is_nullable();
2138 let base_info = get_column_info(&base_data, col);
2139 let mod_info = get_column_info(&mod_data, col);
2140
2141 assert_eq!(base_info.len(), 7);
2142 assert_eq!(mod_info.len(), 7);
2143
2144 assert_cdc_chunk_sizes(
2146 &base.column(col).slice(0, ROW_GROUP_LENGTH),
2147 &base_info[0],
2148 nullable,
2149 CDC_MIN_CHUNK_SIZE,
2150 CDC_MAX_CHUNK_SIZE,
2151 false,
2152 );
2153
2154 assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths);
2155 assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths);
2156
2157 assert_page_length_differences_update(&base_info[2], &mod_info[2], 1);
2159
2160 for i in 3..mod_info.len() {
2162 assert_eq!(base_info[i].page_lengths, mod_info[i].page_lengths);
2163 }
2164 }
2165 }
2166
2167 #[test]
2168 #[cfg_attr(miri, ignore)] fn append() {
2170 let s = schema();
2171 let part1 = generate_table(&s, PART_LENGTH, 0);
2172 let part2 = generate_table(&s, PART_LENGTH, 2);
2173 let part3 = generate_table(&s, PART_LENGTH, 4);
2174 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2175 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2176
2177 let base = concat_batches([&part1, &edit1, &part2, &part3]);
2178 let modified = concat_batches([&part1, &edit1, &part2, &part3, &edit2]);
2179
2180 let base_data = write_with_cdc_options(
2181 &[&base],
2182 CDC_MIN_CHUNK_SIZE,
2183 CDC_MAX_CHUNK_SIZE,
2184 Some(ROW_GROUP_LENGTH),
2185 false,
2186 );
2187 let mod_data = write_with_cdc_options(
2188 &[&modified],
2189 CDC_MIN_CHUNK_SIZE,
2190 CDC_MAX_CHUNK_SIZE,
2191 Some(ROW_GROUP_LENGTH),
2192 false,
2193 );
2194
2195 for col in 0..s.fields().len() {
2196 let nullable = s.field(col).is_nullable();
2197 let base_info = get_column_info(&base_data, col);
2198 let mod_info = get_column_info(&mod_data, col);
2199
2200 assert_eq!(base_info.len(), 7);
2201 assert_eq!(mod_info.len(), 7);
2202
2203 assert_cdc_chunk_sizes(
2205 &base.column(col).slice(0, ROW_GROUP_LENGTH),
2206 &base_info[0],
2207 nullable,
2208 CDC_MIN_CHUNK_SIZE,
2209 CDC_MAX_CHUNK_SIZE,
2210 false,
2211 );
2212
2213 for i in 0..base_info.len() - 1 {
2215 assert_eq!(base_info[i].page_lengths, mod_info[i].page_lengths);
2216 }
2217
2218 let bp = &base_info.last().unwrap().page_lengths;
2220 let mp = &mod_info.last().unwrap().page_lengths;
2221 assert!(mp.len() >= bp.len());
2222 for i in 0..bp.len() - 1 {
2223 assert_eq!(bp[i], mp[i]);
2224 }
2225 }
2226 }
2227 }
2228
2229 #[test]
2232 #[cfg_attr(miri, ignore)] fn test_cdc_array_offsets_direct() {
2234 use crate::basic::Type as PhysicalType;
2235 use crate::schema::types::{ColumnDescriptor, ColumnPath, Type};
2236
2237 let options = CdcOptions {
2238 min_chunk_size: CDC_MIN_CHUNK_SIZE,
2239 max_chunk_size: CDC_MAX_CHUNK_SIZE,
2240 norm_level: 0,
2241 };
2242 let desc = {
2243 let tp = Type::primitive_type_builder("col", PhysicalType::INT32)
2244 .build()
2245 .unwrap();
2246 ColumnDescriptor::new(Arc::new(tp), 0, 0, ColumnPath::new(vec![]))
2247 };
2248
2249 let bpr = bytes_per_record(&DataType::Int32, false);
2250 let n = CDC_PART_SIZE / bpr;
2251 let offset = 10usize;
2252
2253 let array: Int32Array = (0..n).map(|i| test_hash(0, i as u64) as i32).collect();
2254 let mut chunker = super::ContentDefinedChunker::new(&desc, &options).unwrap();
2255 let chunks = chunker
2256 .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &array)
2257 .unwrap();
2258
2259 let sliced = array.slice(offset, n - offset);
2260 let mut chunker2 = super::ContentDefinedChunker::new(&desc, &options).unwrap();
2261 let chunks2 = chunker2
2262 .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &sliced)
2263 .unwrap();
2264
2265 let values: Vec<usize> = chunks.iter().map(|c| c.num_values).collect();
2266 let values2: Vec<usize> = chunks2.iter().map(|c| c.num_values).collect();
2267
2268 assert!(values.len() > 1, "expected multiple chunks, got {values:?}");
2269 assert_eq!(values.len(), values2.len(), "chunk count must match");
2270
2271 assert_eq!(
2272 values[0] - values2[0],
2273 offset,
2274 "offsetted first chunk should be {offset} values shorter"
2275 );
2276 assert_eq!(
2277 &values[1..],
2278 &values2[1..],
2279 "all chunks after the first must be identical"
2280 );
2281 }
2282
2283 #[test]
2288 #[cfg_attr(miri, ignore)] fn test_cdc_list_roundtrip() {
2290 let schema = Arc::new(Schema::new(vec![
2291 Field::new(
2292 "_1",
2293 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
2294 true,
2295 ),
2296 Field::new(
2297 "_2",
2298 DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))),
2299 true,
2300 ),
2301 Field::new(
2302 "_3",
2303 DataType::LargeList(Arc::new(Field::new_list_field(DataType::Utf8, true))),
2304 true,
2305 ),
2306 ]));
2307 let batch = create_random_batch(schema, 10_000, 0.25, 0.75).unwrap();
2308 write_with_cdc_options(
2309 &[&batch],
2310 CDC_MIN_CHUNK_SIZE,
2311 CDC_MAX_CHUNK_SIZE,
2312 None,
2313 true,
2314 );
2315 }
2316
2317 #[test]
2319 #[cfg_attr(miri, ignore)] fn test_cdc_deeply_nested_roundtrip() {
2321 let inner_field = Field::new_list_field(DataType::Int32, true);
2322 let inner_type = DataType::List(Arc::new(inner_field));
2323 let outer_field = Field::new_list_field(inner_type.clone(), true);
2324 let list_list_type = DataType::List(Arc::new(outer_field));
2325
2326 let struct_inner_field = Field::new_list_field(DataType::Int32, true);
2327 let struct_inner_type = DataType::List(Arc::new(struct_inner_field));
2328 let struct_fields = Fields::from(vec![Field::new("a", struct_inner_type, true)]);
2329 let struct_type = DataType::Struct(struct_fields);
2330 let struct_list_field = Field::new_list_field(struct_type, true);
2331 let list_struct_type = DataType::List(Arc::new(struct_list_field));
2332
2333 let schema = Arc::new(Schema::new(vec![
2334 Field::new("list_list", list_list_type, true),
2335 Field::new("list_struct_list", list_struct_type, true),
2336 ]));
2337 let batch = create_random_batch(schema, 10_000, 0.25, 0.75).unwrap();
2338 write_with_cdc_options(
2339 &[&batch],
2340 CDC_MIN_CHUNK_SIZE,
2341 CDC_MAX_CHUNK_SIZE,
2342 None,
2343 true,
2344 );
2345 }
2346
2347 #[test]
2353 fn test_cdc_list_non_empty_null_segments() {
2354 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2361 let offsets = Buffer::from_iter([0_i32, 2, 5, 7, 9, 10]);
2362 let null_bitmap = Buffer::from([0b00010101]); let list_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
2365 let list_data = unsafe {
2366 ArrayData::new_unchecked(
2367 list_type.clone(),
2368 5,
2369 None,
2370 Some(null_bitmap),
2371 0,
2372 vec![offsets],
2373 vec![values.to_data()],
2374 )
2375 };
2376 let list_array = arrow_array::make_array(list_data);
2377
2378 let schema = Arc::new(Schema::new(vec![Field::new("col", list_type, true)]));
2379 let batch = RecordBatch::try_new(schema, vec![list_array]).unwrap();
2380
2381 let buf = write_with_cdc_options(
2382 &[&batch],
2383 CDC_MIN_CHUNK_SIZE,
2384 CDC_MAX_CHUNK_SIZE,
2385 None,
2386 true,
2387 );
2388 let read = concat_batches(read_batches(&buf));
2389 let read_list = read.column(0).as_list::<i32>();
2390 assert_eq!(read_list.len(), 5);
2391 assert!(read_list.is_valid(0));
2392 assert!(read_list.is_null(1));
2393 assert!(read_list.is_valid(2));
2394 assert!(read_list.is_null(3));
2395 assert!(read_list.is_valid(4));
2396
2397 let get_vals = |i: usize| -> Vec<i32> {
2398 read_list
2399 .value(i)
2400 .as_primitive::<arrow_array::types::Int32Type>()
2401 .values()
2402 .iter()
2403 .copied()
2404 .collect()
2405 };
2406 assert_eq!(get_vals(0), vec![1, 2]);
2407 assert_eq!(get_vals(2), vec![6, 7]);
2408 assert_eq!(get_vals(4), vec![10]);
2409 }
2410}