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::{
758 CdcOptions, WriterProperties, WriterPropertiesBuilder, WriterVersion,
759 };
760 use crate::file::reader::{FileReader, SerializedFileReader};
761
762 const CDC_MIN_CHUNK_SIZE: usize = 4 * 1024;
765 const CDC_MAX_CHUNK_SIZE: usize = 16 * 1024;
766 const CDC_PART_SIZE: usize = 128 * 1024;
767 const CDC_EDIT_SIZE: usize = 128;
768 const CDC_ROW_GROUP_LENGTH: usize = 1024 * 1024;
769
770 fn test_hash(seed: u64, index: u64) -> u64 {
774 let mut h = (index.wrapping_add(seed)).wrapping_mul(0xc4ceb9fe1a85ec53u64);
775 h ^= h >> 33;
776 h = h.wrapping_mul(0xff51afd7ed558ccdu64);
777 h ^= h >> 33;
778 h = h.wrapping_mul(0xc4ceb9fe1a85ec53u64);
779 h ^= h >> 33;
780 h
781 }
782
783 fn generate_array(dtype: &DataType, nullable: bool, length: usize, seed: u64) -> ArrayRef {
785 macro_rules! gen_primitive {
786 ($array_type:ty, $cast:expr) => {{
787 if nullable {
788 let arr: $array_type = (0..length)
789 .map(|i| {
790 let val = test_hash(seed, i as u64);
791 if val % 10 == 0 {
792 None
793 } else {
794 Some($cast(val))
795 }
796 })
797 .collect();
798 Arc::new(arr) as ArrayRef
799 } else {
800 let arr: $array_type = (0..length)
801 .map(|i| Some($cast(test_hash(seed, i as u64))))
802 .collect();
803 Arc::new(arr) as ArrayRef
804 }
805 }};
806 }
807
808 match dtype {
809 DataType::Boolean => {
810 if nullable {
811 let arr: BooleanArray = (0..length)
812 .map(|i| {
813 let val = test_hash(seed, i as u64);
814 if val.is_multiple_of(10) {
815 None
816 } else {
817 Some(val.is_multiple_of(2))
818 }
819 })
820 .collect();
821 Arc::new(arr)
822 } else {
823 let arr: BooleanArray = (0..length)
824 .map(|i| Some(test_hash(seed, i as u64).is_multiple_of(2)))
825 .collect();
826 Arc::new(arr)
827 }
828 }
829 DataType::Int32 => gen_primitive!(Int32Array, |v: u64| v as i32),
830 DataType::Int64 => {
831 gen_primitive!(arrow_array::Int64Array, |v: u64| v as i64)
832 }
833 DataType::Float64 => {
834 gen_primitive!(arrow_array::Float64Array, |v: u64| (v % 100000) as f64
835 / 1000.0)
836 }
837 DataType::Utf8 => {
838 let arr: arrow_array::StringArray = if nullable {
839 (0..length)
840 .map(|i| {
841 let val = test_hash(seed, i as u64);
842 if val.is_multiple_of(10) {
843 None
844 } else {
845 Some(format!("str_{val}"))
846 }
847 })
848 .collect()
849 } else {
850 (0..length)
851 .map(|i| Some(format!("str_{}", test_hash(seed, i as u64))))
852 .collect()
853 };
854 Arc::new(arr)
855 }
856 DataType::Binary => {
857 let arr: arrow_array::BinaryArray = if nullable {
858 (0..length)
859 .map(|i| {
860 let val = test_hash(seed, i as u64);
861 if val.is_multiple_of(10) {
862 None
863 } else {
864 Some(format!("bin_{val}").into_bytes())
865 }
866 })
867 .collect()
868 } else {
869 (0..length)
870 .map(|i| Some(format!("bin_{}", test_hash(seed, i as u64)).into_bytes()))
871 .collect()
872 };
873 Arc::new(arr)
874 }
875 DataType::FixedSizeBinary(size) => {
876 let size = *size;
877 let mut builder = arrow_array::builder::FixedSizeBinaryBuilder::new(size);
878 for i in 0..length {
879 let val = test_hash(seed, i as u64);
880 if nullable && val.is_multiple_of(10) {
881 builder.append_null();
882 } else {
883 let s = format!("bin_{val}");
884 let bytes = s.as_bytes();
885 let mut buf = vec![0u8; size as usize];
886 let copy_len = bytes.len().min(size as usize);
887 buf[..copy_len].copy_from_slice(&bytes[..copy_len]);
888 builder.append_value(&buf).unwrap();
889 }
890 }
891 Arc::new(builder.finish())
892 }
893 DataType::Date32 => {
894 gen_primitive!(arrow_array::Date32Array, |v: u64| v as i32)
895 }
896 DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => {
897 gen_primitive!(arrow_array::TimestampNanosecondArray, |v: u64| v as i64)
898 }
899 _ => panic!("Unsupported test data type: {dtype:?}"),
900 }
901 }
902
903 fn generate_table(schema: &Arc<Schema>, length: usize, seed: u64) -> RecordBatch {
905 let arrays: Vec<ArrayRef> = schema
906 .fields()
907 .iter()
908 .enumerate()
909 .map(|(i, field)| {
910 generate_array(
911 field.data_type(),
912 field.is_nullable(),
913 length,
914 seed + i as u64 * 10,
915 )
916 })
917 .collect();
918 RecordBatch::try_new(schema.clone(), arrays).unwrap()
919 }
920
921 fn cdc_byte_width(dtype: &DataType) -> usize {
924 match dtype {
925 DataType::Boolean => 1,
926 DataType::Int8 | DataType::UInt8 => 1,
927 DataType::Int16 | DataType::UInt16 | DataType::Float16 => 2,
928 DataType::Int32
929 | DataType::UInt32
930 | DataType::Float32
931 | DataType::Date32
932 | DataType::Time32(_) => 4,
933 DataType::Int64
934 | DataType::UInt64
935 | DataType::Float64
936 | DataType::Date64
937 | DataType::Time64(_)
938 | DataType::Timestamp(_, _)
939 | DataType::Duration(_) => 8,
940 DataType::Decimal128(_, _) => 16,
941 DataType::Decimal256(_, _) => 32,
942 DataType::FixedSizeBinary(n) => *n as usize,
943 _ => 0, }
945 }
946
947 fn bytes_per_record(dtype: &DataType, nullable: bool) -> usize {
949 let bw = cdc_byte_width(dtype);
950 if bw > 0 {
951 if nullable { bw + 2 } else { bw }
952 } else {
953 16 }
955 }
956
957 fn calculate_cdc_size(array: &dyn Array, nullable: bool) -> i64 {
959 let dtype = array.data_type();
960 let bw = cdc_byte_width(dtype);
961 let result = if bw > 0 {
962 let valid_count = array.len() - array.null_count();
964 (valid_count * bw) as i64
965 } else {
966 match dtype {
968 DataType::Utf8 => {
969 let a = array.as_string::<i32>();
970 (0..a.len())
971 .filter(|&i| a.is_valid(i))
972 .map(|i| a.value(i).len() as i64)
973 .sum()
974 }
975 DataType::Binary => {
976 let a = array.as_binary::<i32>();
977 (0..a.len())
978 .filter(|&i| a.is_valid(i))
979 .map(|i| a.value(i).len() as i64)
980 .sum()
981 }
982 DataType::LargeBinary => {
983 let a = array.as_binary::<i64>();
984 (0..a.len())
985 .filter(|&i| a.is_valid(i))
986 .map(|i| a.value(i).len() as i64)
987 .sum()
988 }
989 _ => panic!("CDC size calculation not implemented for {dtype:?}"),
990 }
991 };
992
993 if nullable {
994 result + array.len() as i64 * 2
996 } else {
997 result
998 }
999 }
1000
1001 struct ColumnInfo {
1003 page_lengths: Vec<i64>,
1004 has_dictionary_page: bool,
1005 }
1006
1007 fn get_column_info(data: &[u8], column_index: usize) -> Vec<ColumnInfo> {
1009 let reader = SerializedFileReader::new(bytes::Bytes::from(data.to_vec())).unwrap();
1010 let metadata = reader.metadata();
1011 let mut result = Vec::new();
1012 for rg in 0..metadata.num_row_groups() {
1013 let rg_reader = reader.get_row_group(rg).unwrap();
1014 let col_reader = rg_reader.get_column_page_reader(column_index).unwrap();
1015 let mut info = ColumnInfo {
1016 page_lengths: Vec::new(),
1017 has_dictionary_page: false,
1018 };
1019 for page in col_reader {
1020 let page = page.unwrap();
1021 match page.page_type() {
1022 crate::basic::PageType::DATA_PAGE | crate::basic::PageType::DATA_PAGE_V2 => {
1023 info.page_lengths.push(page.num_values() as i64);
1024 }
1025 crate::basic::PageType::DICTIONARY_PAGE => {
1026 info.has_dictionary_page = true;
1027 }
1028 crate::basic::PageType::INDEX_PAGE => {}
1029 }
1030 }
1031 result.push(info);
1032 }
1033 result
1034 }
1035
1036 fn assert_cdc_chunk_sizes(
1039 array: &ArrayRef,
1040 info: &ColumnInfo,
1041 nullable: bool,
1042 min_chunk_size: usize,
1043 max_chunk_size: usize,
1044 expect_dictionary_page: bool,
1045 ) {
1046 let expect_dict = match array.data_type() {
1048 DataType::Boolean => false,
1049 _ => expect_dictionary_page,
1050 };
1051 assert_eq!(
1052 info.has_dictionary_page,
1053 expect_dict,
1054 "dictionary page mismatch for {:?}",
1055 array.data_type()
1056 );
1057
1058 let page_lengths = &info.page_lengths;
1059 assert!(
1060 page_lengths.len() > 1,
1061 "CDC should produce multiple pages, got {page_lengths:?}"
1062 );
1063
1064 let bw = cdc_byte_width(array.data_type());
1065 if bw > 0
1067 || matches!(
1068 array.data_type(),
1069 DataType::Utf8 | DataType::Binary | DataType::LargeBinary
1070 )
1071 {
1072 let mut offset = 0i64;
1073 for (i, &page_len) in page_lengths.iter().enumerate() {
1074 let slice = array.slice(offset as usize, page_len as usize);
1075 let cdc_size = calculate_cdc_size(slice.as_ref(), nullable);
1076 if i < page_lengths.len() - 1 {
1077 assert!(
1078 cdc_size >= min_chunk_size as i64,
1079 "Page {i}: CDC size {cdc_size} < min {min_chunk_size}, pages={page_lengths:?}"
1080 );
1081 }
1082 assert!(
1083 cdc_size <= max_chunk_size as i64,
1084 "Page {i}: CDC size {cdc_size} > max {max_chunk_size}, pages={page_lengths:?}"
1085 );
1086 offset += page_len;
1087 }
1088 assert_eq!(
1089 offset,
1090 array.len() as i64,
1091 "page lengths must sum to array length"
1092 );
1093 }
1094 }
1095
1096 fn write_with_cdc_options(
1099 batches: &[&RecordBatch],
1100 min_chunk_size: usize,
1101 max_chunk_size: usize,
1102 max_row_group_rows: Option<usize>,
1103 enable_dictionary: bool,
1104 ) -> Vec<u8> {
1105 assert!(!batches.is_empty());
1106 let schema = batches[0].schema();
1107 let mut builder = WriterProperties::builder()
1108 .set_dictionary_enabled(enable_dictionary)
1109 .set_content_defined_chunking(Some(CdcOptions {
1110 min_chunk_size,
1111 max_chunk_size,
1112 norm_level: 0,
1113 }));
1114 if let Some(max_rows) = max_row_group_rows {
1115 builder = builder.set_max_row_group_row_count(Some(max_rows));
1116 }
1117 let props = builder.build();
1118 let mut buf = Vec::new();
1119 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap();
1120 for batch in batches {
1121 writer.write(batch).unwrap();
1122 }
1123 writer.close().unwrap();
1124
1125 let readback = read_batches(&buf);
1127 let original_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
1128 let readback_rows: usize = readback.iter().map(|b| b.num_rows()).sum();
1129 assert_eq!(original_rows, readback_rows, "Roundtrip row count mismatch");
1130 if original_rows > 0 {
1131 let original = concat_batches(batches.iter().copied());
1132 let roundtrip = concat_batches(&readback);
1133 assert_eq!(original, roundtrip, "Roundtrip validation failed");
1134 }
1135
1136 buf
1137 }
1138
1139 fn write_cdc_batch(batch: &RecordBatch, builder: WriterPropertiesBuilder) -> Vec<u8> {
1141 let props = builder
1142 .set_writer_version(WriterVersion::PARQUET_2_0)
1143 .set_content_defined_chunking(Some(CdcOptions {
1144 min_chunk_size: 8 * 1024,
1145 max_chunk_size: 16 * 1024,
1146 norm_level: 0,
1147 }))
1148 .build();
1149
1150 let mut buf = Vec::new();
1151 let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
1152 writer.write(batch).unwrap();
1153 writer.close().unwrap();
1154 buf
1155 }
1156
1157 fn assert_multiple_nonempty_data_pages(data: &[u8]) {
1158 let info = get_column_info(data, 0);
1159 let total: usize = info.iter().map(|c| c.page_lengths.len()).sum();
1160 let empty = info
1161 .iter()
1162 .flat_map(|c| c.page_lengths.iter())
1163 .filter(|&&n| n == 0)
1164 .count();
1165
1166 assert!(
1167 total > 1,
1168 "CDC should produce multiple data pages, got {total}"
1169 );
1170 assert_eq!(empty, 0, "forced CDC page breaks must not emit empty pages");
1171 }
1172
1173 #[test]
1174 #[cfg_attr(miri, ignore)] fn cdc_boolean_small_pages_does_not_panic() {
1176 let num_rows = 32_768;
1177 let values: Vec<bool> = (0..num_rows).map(|i| i % 7 == 0).collect();
1178 let col = Arc::new(BooleanArray::from(values)) as ArrayRef;
1179 let batch = RecordBatch::try_from_iter([("flag", col)]).unwrap();
1180
1181 let props = WriterProperties::builder().set_data_page_size_limit(1024);
1182 let data = write_cdc_batch(&batch, props);
1183 assert_multiple_nonempty_data_pages(&data);
1184
1185 let readback = read_batches(&data);
1186 let rows: usize = readback.iter().map(|b| b.num_rows()).sum();
1187 assert_eq!(rows, num_rows);
1188 }
1189
1190 #[test]
1191 #[cfg_attr(miri, ignore)] fn cdc_int32_row_count_limit_does_not_emit_empty_pages() {
1193 let values: Vec<i32> = (0..8_192).map(|i| i % 97).collect();
1194 let col = Arc::new(Int32Array::from(values)) as ArrayRef;
1195 let batch = RecordBatch::try_from_iter([("a", col)]).unwrap();
1196
1197 let props = WriterProperties::builder()
1198 .set_dictionary_enabled(false)
1199 .set_data_page_row_count_limit(128);
1200 let data = write_cdc_batch(&batch, props);
1201 assert_multiple_nonempty_data_pages(&data);
1202 }
1203
1204 #[test]
1205 fn cdc_all_null_arrow_column_writes_data_pages() {
1206 let array = Arc::new(Int32Array::from(vec![None::<i32>; 4096])) as ArrayRef;
1207 let schema = Arc::new(Schema::new(vec![Field::new("f0", DataType::Int32, true)]));
1208 let batch = RecordBatch::try_new(schema, vec![array.clone()]).unwrap();
1209
1210 let data = write_with_cdc_options(&[&batch], 64, 256, Some(4096), false);
1211 let info = get_column_info(&data, 0);
1212
1213 assert_eq!(info.len(), 1);
1214 assert!(
1215 !info[0].page_lengths.is_empty(),
1216 "all-null CDC write should still emit data pages"
1217 );
1218 assert_eq!(
1219 info[0].page_lengths.iter().sum::<i64>(),
1220 array.len() as i64,
1221 "all-null CDC pages should account for every input row"
1222 );
1223 }
1224
1225 fn read_batches(data: &[u8]) -> Vec<RecordBatch> {
1226 let reader = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(data.to_vec()))
1227 .unwrap()
1228 .build()
1229 .unwrap();
1230 reader.collect::<std::result::Result<Vec<_>, _>>().unwrap()
1231 }
1232
1233 fn concat_batches(batches: impl IntoIterator<Item = impl Borrow<RecordBatch>>) -> RecordBatch {
1234 let batches: Vec<_> = batches.into_iter().collect();
1235 let schema = batches[0].borrow().schema();
1236 let batches = batches.iter().map(|b| b.borrow());
1237 arrow_select::concat::concat_batches(&schema, batches).unwrap()
1238 }
1239
1240 fn find_differences(first: &[i64], second: &[i64]) -> Vec<(Vec<i64>, Vec<i64>)> {
1243 let n = first.len();
1244 let m = second.len();
1245 let mut dp = vec![vec![0usize; m + 1]; n + 1];
1246 for i in 0..n {
1247 for j in 0..m {
1248 if first[i] == second[j] {
1249 dp[i + 1][j + 1] = dp[i][j] + 1;
1250 } else {
1251 dp[i + 1][j + 1] = dp[i + 1][j].max(dp[i][j + 1]);
1252 }
1253 }
1254 }
1255 let mut common = Vec::new();
1256 let (mut i, mut j) = (n, m);
1257 while i > 0 && j > 0 {
1258 if first[i - 1] == second[j - 1] {
1259 common.push((i - 1, j - 1));
1260 i -= 1;
1261 j -= 1;
1262 } else if dp[i - 1][j] >= dp[i][j - 1] {
1263 i -= 1;
1264 } else {
1265 j -= 1;
1266 }
1267 }
1268 common.reverse();
1269
1270 let mut result = Vec::new();
1271 let (mut last_i, mut last_j) = (0usize, 0usize);
1272 for (ci, cj) in &common {
1273 if *ci > last_i || *cj > last_j {
1274 result.push((first[last_i..*ci].to_vec(), second[last_j..*cj].to_vec()));
1275 }
1276 last_i = ci + 1;
1277 last_j = cj + 1;
1278 }
1279 if last_i < n || last_j < m {
1280 result.push((first[last_i..].to_vec(), second[last_j..].to_vec()));
1281 }
1282
1283 let mut merged: Vec<(Vec<i64>, Vec<i64>)> = Vec::new();
1285 for diff in result {
1286 if let Some(prev) = merged.last_mut() {
1287 if prev.0.is_empty() && diff.1.is_empty() {
1288 prev.0 = diff.0;
1289 continue;
1290 } else if prev.1.is_empty() && diff.0.is_empty() {
1291 prev.1 = diff.1;
1292 continue;
1293 }
1294 }
1295 merged.push(diff);
1296 }
1297 merged
1298 }
1299
1300 fn assert_page_length_differences(
1303 original: &ColumnInfo,
1304 modified: &ColumnInfo,
1305 exact_equal_diffs: usize,
1306 exact_larger_diffs: usize,
1307 exact_smaller_diffs: usize,
1308 edit_length: i64,
1309 ) {
1310 let diffs = find_differences(&original.page_lengths, &modified.page_lengths);
1311 let expected = exact_equal_diffs + exact_larger_diffs + exact_smaller_diffs;
1312
1313 if diffs.len() != expected {
1314 eprintln!("Original: {:?}", original.page_lengths);
1315 eprintln!("Modified: {:?}", modified.page_lengths);
1316 for d in &diffs {
1317 eprintln!(" Diff: {:?} vs {:?}", d.0, d.1);
1318 }
1319 }
1320 assert_eq!(
1321 diffs.len(),
1322 expected,
1323 "Expected {expected} diffs, got {}",
1324 diffs.len()
1325 );
1326
1327 let (mut eq, mut larger, mut smaller) = (0usize, 0usize, 0usize);
1328 for (left, right) in &diffs {
1329 let left_sum: i64 = left.iter().sum();
1330 let right_sum: i64 = right.iter().sum();
1331 match left_sum.cmp(&right_sum) {
1332 Ordering::Equal => eq += 1,
1333 Ordering::Less => {
1334 larger += 1;
1335 assert_eq!(
1336 left_sum + edit_length,
1337 right_sum,
1338 "Larger diff mismatch: {left_sum} + {edit_length} != {right_sum}"
1339 );
1340 }
1341 Ordering::Greater => {
1342 smaller += 1;
1343 assert_eq!(
1344 left_sum,
1345 right_sum + edit_length,
1346 "Smaller diff mismatch: {left_sum} != {right_sum} + {edit_length}"
1347 );
1348 }
1349 }
1350 }
1351
1352 assert_eq!(eq, exact_equal_diffs, "equal diffs count");
1353 assert_eq!(larger, exact_larger_diffs, "larger diffs count");
1354 assert_eq!(smaller, exact_smaller_diffs, "smaller diffs count");
1355 }
1356
1357 fn assert_page_length_differences_update(
1360 original: &ColumnInfo,
1361 modified: &ColumnInfo,
1362 max_equal_diffs: usize,
1363 ) {
1364 let diffs = find_differences(&original.page_lengths, &modified.page_lengths);
1365 assert!(
1366 diffs.len() <= max_equal_diffs,
1367 "Expected at most {max_equal_diffs} diffs, got {}",
1368 diffs.len()
1369 );
1370 for (left, right) in &diffs {
1371 let left_sum: i64 = left.iter().sum();
1372 let right_sum: i64 = right.iter().sum();
1373 assert_eq!(
1374 left_sum, right_sum,
1375 "Update diff should not change total row count"
1376 );
1377 }
1378 }
1379
1380 #[test]
1383 fn test_find_differences_basic() {
1384 let diffs = find_differences(&[1, 2, 3, 4, 5], &[1, 7, 8, 4, 5]);
1385 assert_eq!(diffs.len(), 1);
1386 assert_eq!(diffs[0].0, vec![2, 3]);
1387 assert_eq!(diffs[0].1, vec![7, 8]);
1388 }
1389
1390 #[test]
1391 fn test_find_differences_multiple() {
1392 let diffs = find_differences(&[1, 2, 3, 4, 5, 6, 7], &[1, 8, 9, 4, 10, 6, 11]);
1393 assert_eq!(diffs.len(), 3);
1394 assert_eq!(diffs[0].0, vec![2, 3]);
1395 assert_eq!(diffs[0].1, vec![8, 9]);
1396 assert_eq!(diffs[1].0, vec![5]);
1397 assert_eq!(diffs[1].1, vec![10]);
1398 assert_eq!(diffs[2].0, vec![7]);
1399 assert_eq!(diffs[2].1, vec![11]);
1400 }
1401
1402 #[test]
1403 fn test_find_differences_different_lengths() {
1404 let diffs = find_differences(&[1, 2, 3], &[1, 2, 3, 4, 5]);
1405 assert_eq!(diffs.len(), 1);
1406 assert!(diffs[0].0.is_empty());
1407 assert_eq!(diffs[0].1, vec![4, 5]);
1408 }
1409
1410 #[test]
1411 fn test_find_differences_empty() {
1412 let diffs = find_differences(&[], &[]);
1413 assert!(diffs.is_empty());
1414 }
1415
1416 #[test]
1417 fn test_find_differences_changes_at_both_ends() {
1418 let diffs = find_differences(&[1, 2, 3, 4, 5, 6, 7, 8, 9], &[0, 0, 2, 3, 4, 5, 7, 7, 8]);
1419 assert_eq!(diffs.len(), 3);
1420 assert_eq!(diffs[0].0, vec![1]);
1421 assert_eq!(diffs[0].1, vec![0, 0]);
1422 assert_eq!(diffs[1].0, vec![6]);
1423 assert_eq!(diffs[1].1, vec![7]);
1424 assert_eq!(diffs[2].0, vec![9]);
1425 assert!(diffs[2].1.is_empty());
1426 }
1427
1428 #[test]
1429 fn test_find_differences_additional() {
1430 let diffs = find_differences(
1431 &[445, 312, 393, 401, 410, 138, 558, 457],
1432 &[445, 312, 393, 393, 410, 138, 558, 457],
1433 );
1434 assert_eq!(diffs.len(), 1);
1435 assert_eq!(diffs[0].0, vec![401]);
1436 assert_eq!(diffs[0].1, vec![393]);
1437 }
1438
1439 macro_rules! cdc_single_rg_tests {
1442 ($mod_name:ident, $dtype:expr, $nullable:expr) => {
1443 mod $mod_name {
1444 use super::*;
1445
1446 fn config() -> (DataType, bool, usize, usize) {
1447 let dtype: DataType = $dtype;
1448 let nullable: bool = $nullable;
1449 let bpr = bytes_per_record(&dtype, nullable);
1450 let part_length = CDC_PART_SIZE / bpr;
1451 let edit_length = CDC_EDIT_SIZE / bpr;
1452 (dtype, nullable, part_length, edit_length)
1453 }
1454
1455 fn make_schema(dtype: &DataType, nullable: bool) -> Arc<Schema> {
1456 Arc::new(Schema::new(vec![Field::new("f0", dtype.clone(), nullable)]))
1457 }
1458
1459 #[test]
1460 #[cfg_attr(miri, ignore)] fn delete_once() {
1462 let (dtype, nullable, part_length, edit_length) = config();
1463 let schema = make_schema(&dtype, nullable);
1464
1465 let part1 = generate_table(&schema, part_length, 0);
1466 let part2 = generate_table(&schema, edit_length, 1);
1467 let part3 = generate_table(&schema, part_length, part_length as u64);
1468
1469 let base = concat_batches([&part1, &part2, &part3]);
1470 let modified = concat_batches([&part1, &part3]);
1471
1472 for enable_dictionary in [false, true] {
1473 let base_data = write_with_cdc_options(
1474 &[&base],
1475 CDC_MIN_CHUNK_SIZE,
1476 CDC_MAX_CHUNK_SIZE,
1477 Some(CDC_ROW_GROUP_LENGTH),
1478 enable_dictionary,
1479 );
1480 let mod_data = write_with_cdc_options(
1481 &[&modified],
1482 CDC_MIN_CHUNK_SIZE,
1483 CDC_MAX_CHUNK_SIZE,
1484 Some(CDC_ROW_GROUP_LENGTH),
1485 enable_dictionary,
1486 );
1487
1488 let base_info = get_column_info(&base_data, 0);
1489 let mod_info = get_column_info(&mod_data, 0);
1490 assert_eq!(base_info.len(), 1);
1491 assert_eq!(mod_info.len(), 1);
1492
1493 assert_cdc_chunk_sizes(
1494 &base.column(0).clone(),
1495 &base_info[0],
1496 nullable,
1497 CDC_MIN_CHUNK_SIZE,
1498 CDC_MAX_CHUNK_SIZE,
1499 enable_dictionary,
1500 );
1501 assert_cdc_chunk_sizes(
1502 &modified.column(0).clone(),
1503 &mod_info[0],
1504 nullable,
1505 CDC_MIN_CHUNK_SIZE,
1506 CDC_MAX_CHUNK_SIZE,
1507 enable_dictionary,
1508 );
1509
1510 assert_page_length_differences(
1511 &base_info[0],
1512 &mod_info[0],
1513 0,
1514 0,
1515 1,
1516 edit_length as i64,
1517 );
1518 }
1519 }
1520
1521 #[test]
1522 #[cfg_attr(miri, ignore)] fn delete_twice() {
1524 let (dtype, nullable, part_length, edit_length) = config();
1525 let schema = make_schema(&dtype, nullable);
1526
1527 let part1 = generate_table(&schema, part_length, 0);
1528 let part2 = generate_table(&schema, edit_length, 1);
1529 let part3 = generate_table(&schema, part_length, part_length as u64);
1530 let part4 = generate_table(&schema, edit_length, 2);
1531 let part5 = generate_table(&schema, part_length, 2 * part_length as u64);
1532
1533 let base = concat_batches([&part1, &part2, &part3, &part4, &part5]);
1534 let modified = concat_batches([&part1, &part3, &part5]);
1535
1536 for enable_dictionary in [false, true] {
1537 let base_data = write_with_cdc_options(
1538 &[&base],
1539 CDC_MIN_CHUNK_SIZE,
1540 CDC_MAX_CHUNK_SIZE,
1541 Some(CDC_ROW_GROUP_LENGTH),
1542 enable_dictionary,
1543 );
1544 let mod_data = write_with_cdc_options(
1545 &[&modified],
1546 CDC_MIN_CHUNK_SIZE,
1547 CDC_MAX_CHUNK_SIZE,
1548 Some(CDC_ROW_GROUP_LENGTH),
1549 enable_dictionary,
1550 );
1551
1552 let base_info = get_column_info(&base_data, 0);
1553 let mod_info = get_column_info(&mod_data, 0);
1554 assert_eq!(base_info.len(), 1);
1555 assert_eq!(mod_info.len(), 1);
1556
1557 assert_cdc_chunk_sizes(
1558 &base.column(0).clone(),
1559 &base_info[0],
1560 nullable,
1561 CDC_MIN_CHUNK_SIZE,
1562 CDC_MAX_CHUNK_SIZE,
1563 enable_dictionary,
1564 );
1565 assert_cdc_chunk_sizes(
1566 &modified.column(0).clone(),
1567 &mod_info[0],
1568 nullable,
1569 CDC_MIN_CHUNK_SIZE,
1570 CDC_MAX_CHUNK_SIZE,
1571 enable_dictionary,
1572 );
1573
1574 assert_page_length_differences(
1575 &base_info[0],
1576 &mod_info[0],
1577 0,
1578 0,
1579 2,
1580 edit_length as i64,
1581 );
1582 }
1583 }
1584
1585 #[test]
1586 #[cfg_attr(miri, ignore)] fn insert_once() {
1588 let (dtype, nullable, part_length, edit_length) = config();
1589 let schema = make_schema(&dtype, nullable);
1590
1591 let part1 = generate_table(&schema, part_length, 0);
1592 let part2 = generate_table(&schema, edit_length, 1);
1593 let part3 = generate_table(&schema, part_length, part_length as u64);
1594
1595 let base = concat_batches([&part1, &part3]);
1596 let modified = concat_batches([&part1, &part2, &part3]);
1597
1598 for enable_dictionary in [false, true] {
1599 let base_data = write_with_cdc_options(
1600 &[&base],
1601 CDC_MIN_CHUNK_SIZE,
1602 CDC_MAX_CHUNK_SIZE,
1603 Some(CDC_ROW_GROUP_LENGTH),
1604 enable_dictionary,
1605 );
1606 let mod_data = write_with_cdc_options(
1607 &[&modified],
1608 CDC_MIN_CHUNK_SIZE,
1609 CDC_MAX_CHUNK_SIZE,
1610 Some(CDC_ROW_GROUP_LENGTH),
1611 enable_dictionary,
1612 );
1613
1614 let base_info = get_column_info(&base_data, 0);
1615 let mod_info = get_column_info(&mod_data, 0);
1616 assert_eq!(base_info.len(), 1);
1617 assert_eq!(mod_info.len(), 1);
1618
1619 assert_cdc_chunk_sizes(
1620 &base.column(0).clone(),
1621 &base_info[0],
1622 nullable,
1623 CDC_MIN_CHUNK_SIZE,
1624 CDC_MAX_CHUNK_SIZE,
1625 enable_dictionary,
1626 );
1627 assert_cdc_chunk_sizes(
1628 &modified.column(0).clone(),
1629 &mod_info[0],
1630 nullable,
1631 CDC_MIN_CHUNK_SIZE,
1632 CDC_MAX_CHUNK_SIZE,
1633 enable_dictionary,
1634 );
1635
1636 assert_page_length_differences(
1637 &base_info[0],
1638 &mod_info[0],
1639 0,
1640 1,
1641 0,
1642 edit_length as i64,
1643 );
1644 }
1645 }
1646
1647 #[test]
1648 #[cfg_attr(miri, ignore)] fn insert_twice() {
1650 let (dtype, nullable, part_length, edit_length) = config();
1651 let schema = make_schema(&dtype, nullable);
1652
1653 let part1 = generate_table(&schema, part_length, 0);
1654 let part2 = generate_table(&schema, edit_length, 1);
1655 let part3 = generate_table(&schema, part_length, part_length as u64);
1656 let part4 = generate_table(&schema, edit_length, 2);
1657 let part5 = generate_table(&schema, part_length, 2 * part_length as u64);
1658
1659 let base = concat_batches([&part1, &part3, &part5]);
1660 let modified = concat_batches([&part1, &part2, &part3, &part4, &part5]);
1661
1662 for enable_dictionary in [false, true] {
1663 let base_data = write_with_cdc_options(
1664 &[&base],
1665 CDC_MIN_CHUNK_SIZE,
1666 CDC_MAX_CHUNK_SIZE,
1667 Some(CDC_ROW_GROUP_LENGTH),
1668 enable_dictionary,
1669 );
1670 let mod_data = write_with_cdc_options(
1671 &[&modified],
1672 CDC_MIN_CHUNK_SIZE,
1673 CDC_MAX_CHUNK_SIZE,
1674 Some(CDC_ROW_GROUP_LENGTH),
1675 enable_dictionary,
1676 );
1677
1678 let base_info = get_column_info(&base_data, 0);
1679 let mod_info = get_column_info(&mod_data, 0);
1680 assert_eq!(base_info.len(), 1);
1681 assert_eq!(mod_info.len(), 1);
1682
1683 assert_cdc_chunk_sizes(
1684 &base.column(0).clone(),
1685 &base_info[0],
1686 nullable,
1687 CDC_MIN_CHUNK_SIZE,
1688 CDC_MAX_CHUNK_SIZE,
1689 enable_dictionary,
1690 );
1691 assert_cdc_chunk_sizes(
1692 &modified.column(0).clone(),
1693 &mod_info[0],
1694 nullable,
1695 CDC_MIN_CHUNK_SIZE,
1696 CDC_MAX_CHUNK_SIZE,
1697 enable_dictionary,
1698 );
1699
1700 assert_page_length_differences(
1701 &base_info[0],
1702 &mod_info[0],
1703 0,
1704 2,
1705 0,
1706 edit_length as i64,
1707 );
1708 }
1709 }
1710
1711 #[test]
1712 #[cfg_attr(miri, ignore)] fn update_once() {
1714 let (dtype, nullable, part_length, edit_length) = config();
1715 let schema = make_schema(&dtype, nullable);
1716
1717 let part1 = generate_table(&schema, part_length, 0);
1718 let part2 = generate_table(&schema, edit_length, 1);
1719 let part3 = generate_table(&schema, part_length, part_length as u64);
1720 let part4 = generate_table(&schema, edit_length, 2);
1721
1722 let base = concat_batches([&part1, &part2, &part3]);
1723 let modified = concat_batches([&part1, &part4, &part3]);
1724
1725 for enable_dictionary in [false, true] {
1726 let base_data = write_with_cdc_options(
1727 &[&base],
1728 CDC_MIN_CHUNK_SIZE,
1729 CDC_MAX_CHUNK_SIZE,
1730 Some(CDC_ROW_GROUP_LENGTH),
1731 enable_dictionary,
1732 );
1733 let mod_data = write_with_cdc_options(
1734 &[&modified],
1735 CDC_MIN_CHUNK_SIZE,
1736 CDC_MAX_CHUNK_SIZE,
1737 Some(CDC_ROW_GROUP_LENGTH),
1738 enable_dictionary,
1739 );
1740
1741 let base_info = get_column_info(&base_data, 0);
1742 let mod_info = get_column_info(&mod_data, 0);
1743 assert_eq!(base_info.len(), 1);
1744 assert_eq!(mod_info.len(), 1);
1745
1746 assert_cdc_chunk_sizes(
1747 &base.column(0).clone(),
1748 &base_info[0],
1749 nullable,
1750 CDC_MIN_CHUNK_SIZE,
1751 CDC_MAX_CHUNK_SIZE,
1752 enable_dictionary,
1753 );
1754 assert_cdc_chunk_sizes(
1755 &modified.column(0).clone(),
1756 &mod_info[0],
1757 nullable,
1758 CDC_MIN_CHUNK_SIZE,
1759 CDC_MAX_CHUNK_SIZE,
1760 enable_dictionary,
1761 );
1762
1763 assert_page_length_differences_update(&base_info[0], &mod_info[0], 1);
1764 }
1765 }
1766
1767 #[test]
1768 #[cfg_attr(miri, ignore)] fn update_twice() {
1770 let (dtype, nullable, part_length, edit_length) = config();
1771 let schema = make_schema(&dtype, nullable);
1772
1773 let part1 = generate_table(&schema, part_length, 0);
1774 let part2 = generate_table(&schema, edit_length, 1);
1775 let part3 = generate_table(&schema, part_length, part_length as u64);
1776 let part4 = generate_table(&schema, edit_length, 2);
1777 let part5 = generate_table(&schema, part_length, 2 * part_length as u64);
1778 let part6 = generate_table(&schema, edit_length, 3);
1779 let part7 = generate_table(&schema, edit_length, 4);
1780
1781 let base = concat_batches([&part1, &part2, &part3, &part4, &part5]);
1782 let modified = concat_batches([&part1, &part6, &part3, &part7, &part5]);
1783
1784 for enable_dictionary in [false, true] {
1785 let base_data = write_with_cdc_options(
1786 &[&base],
1787 CDC_MIN_CHUNK_SIZE,
1788 CDC_MAX_CHUNK_SIZE,
1789 Some(CDC_ROW_GROUP_LENGTH),
1790 enable_dictionary,
1791 );
1792 let mod_data = write_with_cdc_options(
1793 &[&modified],
1794 CDC_MIN_CHUNK_SIZE,
1795 CDC_MAX_CHUNK_SIZE,
1796 Some(CDC_ROW_GROUP_LENGTH),
1797 enable_dictionary,
1798 );
1799
1800 let base_info = get_column_info(&base_data, 0);
1801 let mod_info = get_column_info(&mod_data, 0);
1802 assert_eq!(base_info.len(), 1);
1803 assert_eq!(mod_info.len(), 1);
1804
1805 assert_cdc_chunk_sizes(
1806 &base.column(0).clone(),
1807 &base_info[0],
1808 nullable,
1809 CDC_MIN_CHUNK_SIZE,
1810 CDC_MAX_CHUNK_SIZE,
1811 enable_dictionary,
1812 );
1813 assert_cdc_chunk_sizes(
1814 &modified.column(0).clone(),
1815 &mod_info[0],
1816 nullable,
1817 CDC_MIN_CHUNK_SIZE,
1818 CDC_MAX_CHUNK_SIZE,
1819 enable_dictionary,
1820 );
1821
1822 assert_page_length_differences_update(&base_info[0], &mod_info[0], 2);
1823 }
1824 }
1825
1826 #[test]
1827 #[cfg_attr(miri, ignore)] fn prepend() {
1829 let (dtype, nullable, part_length, edit_length) = config();
1830 let schema = make_schema(&dtype, nullable);
1831
1832 let part1 = generate_table(&schema, part_length, 0);
1833 let part2 = generate_table(&schema, edit_length, 1);
1834 let part3 = generate_table(&schema, part_length, part_length as u64);
1835 let part4 = generate_table(&schema, edit_length, 2);
1836
1837 let base = concat_batches([&part1, &part2, &part3]);
1838 let modified = concat_batches([&part4, &part1, &part2, &part3]);
1839
1840 for enable_dictionary in [false, true] {
1841 let base_data = write_with_cdc_options(
1842 &[&base],
1843 CDC_MIN_CHUNK_SIZE,
1844 CDC_MAX_CHUNK_SIZE,
1845 Some(CDC_ROW_GROUP_LENGTH),
1846 enable_dictionary,
1847 );
1848 let mod_data = write_with_cdc_options(
1849 &[&modified],
1850 CDC_MIN_CHUNK_SIZE,
1851 CDC_MAX_CHUNK_SIZE,
1852 Some(CDC_ROW_GROUP_LENGTH),
1853 enable_dictionary,
1854 );
1855
1856 let base_info = get_column_info(&base_data, 0);
1857 let mod_info = get_column_info(&mod_data, 0);
1858 assert_eq!(base_info.len(), 1);
1859 assert_eq!(mod_info.len(), 1);
1860
1861 assert_cdc_chunk_sizes(
1862 &base.column(0).clone(),
1863 &base_info[0],
1864 nullable,
1865 CDC_MIN_CHUNK_SIZE,
1866 CDC_MAX_CHUNK_SIZE,
1867 enable_dictionary,
1868 );
1869 assert_cdc_chunk_sizes(
1870 &modified.column(0).clone(),
1871 &mod_info[0],
1872 nullable,
1873 CDC_MIN_CHUNK_SIZE,
1874 CDC_MAX_CHUNK_SIZE,
1875 enable_dictionary,
1876 );
1877
1878 assert!(
1879 mod_info[0].page_lengths.len() >= base_info[0].page_lengths.len(),
1880 "Modified should have same or more pages"
1881 );
1882
1883 assert_page_length_differences(
1884 &base_info[0],
1885 &mod_info[0],
1886 0,
1887 1,
1888 0,
1889 edit_length as i64,
1890 );
1891 }
1892 }
1893
1894 #[test]
1895 #[cfg_attr(miri, ignore)] fn append() {
1897 let (dtype, nullable, part_length, edit_length) = config();
1898 let schema = make_schema(&dtype, nullable);
1899
1900 let part1 = generate_table(&schema, part_length, 0);
1901 let part2 = generate_table(&schema, edit_length, 1);
1902 let part3 = generate_table(&schema, part_length, part_length as u64);
1903 let part4 = generate_table(&schema, edit_length, 2);
1904
1905 let base = concat_batches([&part1, &part2, &part3]);
1906 let modified = concat_batches([&part1, &part2, &part3, &part4]);
1907
1908 for enable_dictionary in [false, true] {
1909 let base_data = write_with_cdc_options(
1910 &[&base],
1911 CDC_MIN_CHUNK_SIZE,
1912 CDC_MAX_CHUNK_SIZE,
1913 Some(CDC_ROW_GROUP_LENGTH),
1914 enable_dictionary,
1915 );
1916 let mod_data = write_with_cdc_options(
1917 &[&modified],
1918 CDC_MIN_CHUNK_SIZE,
1919 CDC_MAX_CHUNK_SIZE,
1920 Some(CDC_ROW_GROUP_LENGTH),
1921 enable_dictionary,
1922 );
1923
1924 let base_info = get_column_info(&base_data, 0);
1925 let mod_info = get_column_info(&mod_data, 0);
1926 assert_eq!(base_info.len(), 1);
1927 assert_eq!(mod_info.len(), 1);
1928
1929 assert_cdc_chunk_sizes(
1930 &base.column(0).clone(),
1931 &base_info[0],
1932 nullable,
1933 CDC_MIN_CHUNK_SIZE,
1934 CDC_MAX_CHUNK_SIZE,
1935 enable_dictionary,
1936 );
1937 assert_cdc_chunk_sizes(
1938 &modified.column(0).clone(),
1939 &mod_info[0],
1940 nullable,
1941 CDC_MIN_CHUNK_SIZE,
1942 CDC_MAX_CHUNK_SIZE,
1943 enable_dictionary,
1944 );
1945
1946 let bp = &base_info[0].page_lengths;
1947 let mp = &mod_info[0].page_lengths;
1948 assert!(mp.len() >= bp.len());
1949 for i in 0..bp.len() - 1 {
1950 assert_eq!(bp[i], mp[i], "Page {i} should be identical");
1951 }
1952 assert!(mp[bp.len() - 1] >= bp[bp.len() - 1]);
1953 }
1954 }
1955
1956 #[test]
1957 fn empty_table() {
1958 let (dtype, nullable, _, _) = config();
1959 let schema = make_schema(&dtype, nullable);
1960
1961 let empty = RecordBatch::new_empty(schema);
1962 for enable_dictionary in [false, true] {
1963 let data = write_with_cdc_options(
1964 &[&empty],
1965 CDC_MIN_CHUNK_SIZE,
1966 CDC_MAX_CHUNK_SIZE,
1967 Some(CDC_ROW_GROUP_LENGTH),
1968 enable_dictionary,
1969 );
1970 let info = get_column_info(&data, 0);
1971 if !info.is_empty() {
1973 assert!(info[0].page_lengths.is_empty());
1974 }
1975 }
1976 }
1977
1978 #[test]
1979 #[cfg_attr(miri, ignore)] fn array_offsets() {
1981 let (dtype, nullable, part_length, edit_length) = config();
1982 let schema = make_schema(&dtype, nullable);
1983
1984 let table = concat_batches([
1985 &generate_table(&schema, part_length, 0),
1986 &generate_table(&schema, edit_length, 1),
1987 &generate_table(&schema, part_length, part_length as u64),
1988 ]);
1989
1990 for offset in [0usize, 512, 1024] {
1991 if offset >= table.num_rows() {
1992 continue;
1993 }
1994 let sliced = table.slice(offset, table.num_rows() - offset);
1995 let data = write_with_cdc_options(
1996 &[&sliced],
1997 CDC_MIN_CHUNK_SIZE,
1998 CDC_MAX_CHUNK_SIZE,
1999 Some(CDC_ROW_GROUP_LENGTH),
2000 true,
2001 );
2002 let info = get_column_info(&data, 0);
2003 assert_eq!(info.len(), 1);
2004
2005 assert_cdc_chunk_sizes(
2007 &sliced.column(0).clone(),
2008 &info[0],
2009 nullable,
2010 CDC_MIN_CHUNK_SIZE,
2011 CDC_MAX_CHUNK_SIZE,
2012 true,
2013 );
2014 }
2015 }
2016 }
2017 };
2018 }
2019
2020 cdc_single_rg_tests!(cdc_bool_non_null, DataType::Boolean, false);
2022 cdc_single_rg_tests!(cdc_i32_non_null, DataType::Int32, false);
2023 cdc_single_rg_tests!(cdc_i64_nullable, DataType::Int64, true);
2024 cdc_single_rg_tests!(cdc_f64_nullable, DataType::Float64, true);
2025 cdc_single_rg_tests!(cdc_utf8_non_null, DataType::Utf8, false);
2026 cdc_single_rg_tests!(cdc_binary_nullable, DataType::Binary, true);
2027 cdc_single_rg_tests!(cdc_fsb16_nullable, DataType::FixedSizeBinary(16), true);
2028 cdc_single_rg_tests!(cdc_date32_non_null, DataType::Date32, false);
2029 cdc_single_rg_tests!(
2030 cdc_timestamp_nullable,
2031 DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, None),
2032 true
2033 );
2034
2035 mod cdc_multiple_row_groups {
2038 use super::*;
2039
2040 const PART_LENGTH: usize = 128 * 1024;
2041 const EDIT_LENGTH: usize = 128;
2042 const ROW_GROUP_LENGTH: usize = 64 * 1024;
2043
2044 fn schema() -> Arc<Schema> {
2045 Arc::new(Schema::new(vec![
2046 Field::new("int32", DataType::Int32, true),
2047 Field::new("float64", DataType::Float64, true),
2048 Field::new("bool", DataType::Boolean, false),
2049 ]))
2050 }
2051
2052 #[test]
2053 #[cfg_attr(miri, ignore)] fn insert_once() {
2055 let s = schema();
2056 let part1 = generate_table(&s, PART_LENGTH, 0);
2057 let part2 = generate_table(&s, PART_LENGTH, 2);
2058 let part3 = generate_table(&s, PART_LENGTH, 4);
2059 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2060 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2061
2062 let base = concat_batches([&part1, &edit1, &part2, &part3]);
2063 let modified = concat_batches([&part1, &edit1, &edit2, &part2, &part3]);
2064 assert_eq!(modified.num_rows(), base.num_rows() + EDIT_LENGTH);
2065
2066 let base_data = write_with_cdc_options(
2067 &[&base],
2068 CDC_MIN_CHUNK_SIZE,
2069 CDC_MAX_CHUNK_SIZE,
2070 Some(ROW_GROUP_LENGTH),
2071 false,
2072 );
2073 let mod_data = write_with_cdc_options(
2074 &[&modified],
2075 CDC_MIN_CHUNK_SIZE,
2076 CDC_MAX_CHUNK_SIZE,
2077 Some(ROW_GROUP_LENGTH),
2078 false,
2079 );
2080
2081 for col in 0..s.fields().len() {
2082 let base_info = get_column_info(&base_data, col);
2083 let mod_info = get_column_info(&mod_data, col);
2084
2085 assert_eq!(base_info.len(), 7, "expected 7 row groups for col {col}");
2086 assert_eq!(mod_info.len(), 7);
2087
2088 assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths);
2090 assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths);
2091
2092 for i in 2..mod_info.len() - 1 {
2094 assert_page_length_differences(
2095 &base_info[i],
2096 &mod_info[i],
2097 0,
2098 1,
2099 1,
2100 EDIT_LENGTH as i64,
2101 );
2102 }
2103 assert_page_length_differences(
2105 base_info.last().unwrap(),
2106 mod_info.last().unwrap(),
2107 0,
2108 1,
2109 0,
2110 EDIT_LENGTH as i64,
2111 );
2112 }
2113 }
2114
2115 #[test]
2116 #[cfg_attr(miri, ignore)] fn delete_once() {
2118 let s = schema();
2119 let part1 = generate_table(&s, PART_LENGTH, 0);
2120 let part2 = generate_table(&s, PART_LENGTH, 2);
2121 let part3 = generate_table(&s, PART_LENGTH, 4);
2122 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2123 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2124
2125 let base = concat_batches([&part1, &edit1, &part2, &part3, &edit2]);
2126 let modified = concat_batches([&part1, &part2, &part3, &edit2]);
2127
2128 let base_data = write_with_cdc_options(
2129 &[&base],
2130 CDC_MIN_CHUNK_SIZE,
2131 CDC_MAX_CHUNK_SIZE,
2132 Some(ROW_GROUP_LENGTH),
2133 false,
2134 );
2135 let mod_data = write_with_cdc_options(
2136 &[&modified],
2137 CDC_MIN_CHUNK_SIZE,
2138 CDC_MAX_CHUNK_SIZE,
2139 Some(ROW_GROUP_LENGTH),
2140 false,
2141 );
2142
2143 for col in 0..s.fields().len() {
2144 let base_info = get_column_info(&base_data, col);
2145 let mod_info = get_column_info(&mod_data, col);
2146
2147 assert_eq!(base_info.len(), 7);
2148 assert_eq!(mod_info.len(), 7);
2149
2150 assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths);
2151 assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths);
2152
2153 for i in 2..mod_info.len() - 1 {
2154 assert_page_length_differences(
2155 &base_info[i],
2156 &mod_info[i],
2157 0,
2158 1,
2159 1,
2160 EDIT_LENGTH as i64,
2161 );
2162 }
2163 assert_page_length_differences(
2164 base_info.last().unwrap(),
2165 mod_info.last().unwrap(),
2166 0,
2167 0,
2168 1,
2169 EDIT_LENGTH as i64,
2170 );
2171 }
2172 }
2173
2174 #[test]
2175 #[cfg_attr(miri, ignore)] fn update_once() {
2177 let s = schema();
2178 let part1 = generate_table(&s, PART_LENGTH, 0);
2179 let part2 = generate_table(&s, PART_LENGTH, 2);
2180 let part3 = generate_table(&s, PART_LENGTH, 4);
2181 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2182 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2183 let edit3 = generate_table(&s, EDIT_LENGTH, 5);
2184
2185 let base = concat_batches([&part1, &edit1, &part2, &part3, &edit2]);
2186 let modified = concat_batches([&part1, &edit3, &part2, &part3, &edit2]);
2187
2188 let base_data = write_with_cdc_options(
2189 &[&base],
2190 CDC_MIN_CHUNK_SIZE,
2191 CDC_MAX_CHUNK_SIZE,
2192 Some(ROW_GROUP_LENGTH),
2193 false,
2194 );
2195 let mod_data = write_with_cdc_options(
2196 &[&modified],
2197 CDC_MIN_CHUNK_SIZE,
2198 CDC_MAX_CHUNK_SIZE,
2199 Some(ROW_GROUP_LENGTH),
2200 false,
2201 );
2202
2203 for col in 0..s.fields().len() {
2204 let nullable = s.field(col).is_nullable();
2205 let base_info = get_column_info(&base_data, col);
2206 let mod_info = get_column_info(&mod_data, col);
2207
2208 assert_eq!(base_info.len(), 7);
2209 assert_eq!(mod_info.len(), 7);
2210
2211 assert_cdc_chunk_sizes(
2213 &base.column(col).slice(0, ROW_GROUP_LENGTH),
2214 &base_info[0],
2215 nullable,
2216 CDC_MIN_CHUNK_SIZE,
2217 CDC_MAX_CHUNK_SIZE,
2218 false,
2219 );
2220
2221 assert_eq!(base_info[0].page_lengths, mod_info[0].page_lengths);
2222 assert_eq!(base_info[1].page_lengths, mod_info[1].page_lengths);
2223
2224 assert_page_length_differences_update(&base_info[2], &mod_info[2], 1);
2226
2227 for i in 3..mod_info.len() {
2229 assert_eq!(base_info[i].page_lengths, mod_info[i].page_lengths);
2230 }
2231 }
2232 }
2233
2234 #[test]
2235 #[cfg_attr(miri, ignore)] fn append() {
2237 let s = schema();
2238 let part1 = generate_table(&s, PART_LENGTH, 0);
2239 let part2 = generate_table(&s, PART_LENGTH, 2);
2240 let part3 = generate_table(&s, PART_LENGTH, 4);
2241 let edit1 = generate_table(&s, EDIT_LENGTH, 1);
2242 let edit2 = generate_table(&s, EDIT_LENGTH, 3);
2243
2244 let base = concat_batches([&part1, &edit1, &part2, &part3]);
2245 let modified = concat_batches([&part1, &edit1, &part2, &part3, &edit2]);
2246
2247 let base_data = write_with_cdc_options(
2248 &[&base],
2249 CDC_MIN_CHUNK_SIZE,
2250 CDC_MAX_CHUNK_SIZE,
2251 Some(ROW_GROUP_LENGTH),
2252 false,
2253 );
2254 let mod_data = write_with_cdc_options(
2255 &[&modified],
2256 CDC_MIN_CHUNK_SIZE,
2257 CDC_MAX_CHUNK_SIZE,
2258 Some(ROW_GROUP_LENGTH),
2259 false,
2260 );
2261
2262 for col in 0..s.fields().len() {
2263 let nullable = s.field(col).is_nullable();
2264 let base_info = get_column_info(&base_data, col);
2265 let mod_info = get_column_info(&mod_data, col);
2266
2267 assert_eq!(base_info.len(), 7);
2268 assert_eq!(mod_info.len(), 7);
2269
2270 assert_cdc_chunk_sizes(
2272 &base.column(col).slice(0, ROW_GROUP_LENGTH),
2273 &base_info[0],
2274 nullable,
2275 CDC_MIN_CHUNK_SIZE,
2276 CDC_MAX_CHUNK_SIZE,
2277 false,
2278 );
2279
2280 for i in 0..base_info.len() - 1 {
2282 assert_eq!(base_info[i].page_lengths, mod_info[i].page_lengths);
2283 }
2284
2285 let bp = &base_info.last().unwrap().page_lengths;
2287 let mp = &mod_info.last().unwrap().page_lengths;
2288 assert!(mp.len() >= bp.len());
2289 for i in 0..bp.len() - 1 {
2290 assert_eq!(bp[i], mp[i]);
2291 }
2292 }
2293 }
2294 }
2295
2296 #[test]
2299 #[cfg_attr(miri, ignore)] fn test_cdc_array_offsets_direct() {
2301 use crate::basic::Type as PhysicalType;
2302 use crate::schema::types::{ColumnDescriptor, ColumnPath, Type};
2303
2304 let options = CdcOptions {
2305 min_chunk_size: CDC_MIN_CHUNK_SIZE,
2306 max_chunk_size: CDC_MAX_CHUNK_SIZE,
2307 norm_level: 0,
2308 };
2309 let desc = {
2310 let tp = Type::primitive_type_builder("col", PhysicalType::INT32)
2311 .build()
2312 .unwrap();
2313 ColumnDescriptor::new(Arc::new(tp), 0, 0, ColumnPath::new(vec![]))
2314 };
2315
2316 let bpr = bytes_per_record(&DataType::Int32, false);
2317 let n = CDC_PART_SIZE / bpr;
2318 let offset = 10usize;
2319
2320 let array: Int32Array = (0..n).map(|i| test_hash(0, i as u64) as i32).collect();
2321 let mut chunker = super::ContentDefinedChunker::new(&desc, &options).unwrap();
2322 let chunks = chunker
2323 .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &array)
2324 .unwrap();
2325
2326 let sliced = array.slice(offset, n - offset);
2327 let mut chunker2 = super::ContentDefinedChunker::new(&desc, &options).unwrap();
2328 let chunks2 = chunker2
2329 .get_arrow_chunks(LevelDataRef::Absent, LevelDataRef::Absent, &sliced)
2330 .unwrap();
2331
2332 let values: Vec<usize> = chunks.iter().map(|c| c.num_values).collect();
2333 let values2: Vec<usize> = chunks2.iter().map(|c| c.num_values).collect();
2334
2335 assert!(values.len() > 1, "expected multiple chunks, got {values:?}");
2336 assert_eq!(values.len(), values2.len(), "chunk count must match");
2337
2338 assert_eq!(
2339 values[0] - values2[0],
2340 offset,
2341 "offsetted first chunk should be {offset} values shorter"
2342 );
2343 assert_eq!(
2344 &values[1..],
2345 &values2[1..],
2346 "all chunks after the first must be identical"
2347 );
2348 }
2349
2350 #[test]
2355 #[cfg_attr(miri, ignore)] fn test_cdc_list_roundtrip() {
2357 let schema = Arc::new(Schema::new(vec![
2358 Field::new(
2359 "_1",
2360 DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))),
2361 true,
2362 ),
2363 Field::new(
2364 "_2",
2365 DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))),
2366 true,
2367 ),
2368 Field::new(
2369 "_3",
2370 DataType::LargeList(Arc::new(Field::new_list_field(DataType::Utf8, true))),
2371 true,
2372 ),
2373 ]));
2374 let batch = create_random_batch(schema, 10_000, 0.25, 0.75).unwrap();
2375 write_with_cdc_options(
2376 &[&batch],
2377 CDC_MIN_CHUNK_SIZE,
2378 CDC_MAX_CHUNK_SIZE,
2379 None,
2380 true,
2381 );
2382 }
2383
2384 #[test]
2386 #[cfg_attr(miri, ignore)] fn test_cdc_deeply_nested_roundtrip() {
2388 let inner_field = Field::new_list_field(DataType::Int32, true);
2389 let inner_type = DataType::List(Arc::new(inner_field));
2390 let outer_field = Field::new_list_field(inner_type.clone(), true);
2391 let list_list_type = DataType::List(Arc::new(outer_field));
2392
2393 let struct_inner_field = Field::new_list_field(DataType::Int32, true);
2394 let struct_inner_type = DataType::List(Arc::new(struct_inner_field));
2395 let struct_fields = Fields::from(vec![Field::new("a", struct_inner_type, true)]);
2396 let struct_type = DataType::Struct(struct_fields);
2397 let struct_list_field = Field::new_list_field(struct_type, true);
2398 let list_struct_type = DataType::List(Arc::new(struct_list_field));
2399
2400 let schema = Arc::new(Schema::new(vec![
2401 Field::new("list_list", list_list_type, true),
2402 Field::new("list_struct_list", list_struct_type, true),
2403 ]));
2404 let batch = create_random_batch(schema, 10_000, 0.25, 0.75).unwrap();
2405 write_with_cdc_options(
2406 &[&batch],
2407 CDC_MIN_CHUNK_SIZE,
2408 CDC_MAX_CHUNK_SIZE,
2409 None,
2410 true,
2411 );
2412 }
2413
2414 #[test]
2420 fn test_cdc_list_non_empty_null_segments() {
2421 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2428 let offsets = Buffer::from_iter([0_i32, 2, 5, 7, 9, 10]);
2429 let null_bitmap = Buffer::from([0b00010101]); let list_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
2432 let list_data = unsafe {
2433 ArrayData::new_unchecked(
2434 list_type.clone(),
2435 5,
2436 None,
2437 Some(null_bitmap),
2438 0,
2439 vec![offsets],
2440 vec![values.to_data()],
2441 )
2442 };
2443 let list_array = arrow_array::make_array(list_data);
2444
2445 let schema = Arc::new(Schema::new(vec![Field::new("col", list_type, true)]));
2446 let batch = RecordBatch::try_new(schema, vec![list_array]).unwrap();
2447
2448 let buf = write_with_cdc_options(
2449 &[&batch],
2450 CDC_MIN_CHUNK_SIZE,
2451 CDC_MAX_CHUNK_SIZE,
2452 None,
2453 true,
2454 );
2455 let read = concat_batches(read_batches(&buf));
2456 let read_list = read.column(0).as_list::<i32>();
2457 assert_eq!(read_list.len(), 5);
2458 assert!(read_list.is_valid(0));
2459 assert!(read_list.is_null(1));
2460 assert!(read_list.is_valid(2));
2461 assert!(read_list.is_null(3));
2462 assert!(read_list.is_valid(4));
2463
2464 let get_vals = |i: usize| -> Vec<i32> {
2465 read_list
2466 .value(i)
2467 .as_primitive::<arrow_array::types::Int32Type>()
2468 .values()
2469 .iter()
2470 .copied()
2471 .collect()
2472 };
2473 assert_eq!(get_vals(0), vec![1, 2]);
2474 assert_eq!(get_vals(2), vec![6, 7]);
2475 assert_eq!(get_vals(4), vec![10]);
2476 }
2477}