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