Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 83 additions & 1 deletion parquet/src/column/chunker/cdc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -754,7 +754,7 @@ mod arrow_tests {
use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use crate::arrow::arrow_writer::ArrowWriter;
use crate::column::writer::LevelDataRef;
use crate::file::properties::{CdcOptions, WriterProperties};
use crate::file::properties::{CdcOptions, WriterProperties, WriterVersion};
use crate::file::reader::{FileReader, SerializedFileReader};

// --- Constants matching C++ TestCDCSingleRowGroup ---
Expand Down Expand Up @@ -1134,6 +1134,88 @@ mod arrow_tests {
buf
}

/// Write a single-column batch with CDC plus the page-size / row-count
/// limits that make a forced chunk boundary flush an already-empty page.
fn write_boolean_cdc_with_small_pages(num_rows: usize) -> Vec<u8> {
let values: Vec<bool> = (0..num_rows).map(|i| i % 7 == 0).collect();
let col = Arc::new(BooleanArray::from(values)) as ArrayRef;
let batch = RecordBatch::try_from_iter([("flag", col)]).unwrap();

let props = WriterProperties::builder()
.set_writer_version(WriterVersion::PARQUET_2_0)
.set_data_page_size_limit(1024)
.set_content_defined_chunking(Some(CdcOptions {
min_chunk_size: 8 * 1024,
max_chunk_size: 16 * 1024,
norm_level: 0,
}))
.build();

let mut buf = Vec::new();
let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
buf
}

fn count_data_pages(data: &[u8]) -> (usize, usize) {
let info = get_column_info(data, 0);
let total = info.iter().map(|c| c.page_lengths.len()).sum();
let empty = info
.iter()
.flat_map(|c| c.page_lengths.iter())
.filter(|&&n| n == 0)
.count();
(total, empty)
}

#[test]
#[cfg_attr(miri, ignore)] // large enough to hit CDC chunk boundaries
fn cdc_boolean_small_pages_does_not_panic() {
let data = write_boolean_cdc_with_small_pages(500_000);
let (total, empty) = count_data_pages(&data);
assert!(
total > 1,
"CDC should produce multiple data pages, got {total}"
);
assert_eq!(empty, 0, "forced CDC page breaks must not emit empty pages");

let readback = read_batches(&data);
let rows: usize = readback.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, 500_000);
}

#[test]
#[cfg_attr(miri, ignore)]
fn cdc_int32_row_count_limit_does_not_emit_empty_pages() {
let values: Vec<i32> = (0..500_000).map(|i| i % 97).collect();
let col = Arc::new(Int32Array::from(values)) as ArrayRef;
let batch = RecordBatch::try_from_iter([("a", col)]).unwrap();

let props = WriterProperties::builder()
.set_writer_version(WriterVersion::PARQUET_2_0)
.set_dictionary_enabled(false)
.set_data_page_row_count_limit(128)
.set_content_defined_chunking(Some(CdcOptions {
min_chunk_size: 8 * 1024,
max_chunk_size: 16 * 1024,
norm_level: 0,
}))
.build();

let mut buf = Vec::new();
let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();

let (total, empty) = count_data_pages(&buf);
assert!(
total > 1,
"CDC should produce multiple data pages, got {total}"
);
assert_eq!(empty, 0, "forced CDC page breaks must not emit empty pages");
}

#[test]
fn cdc_all_null_arrow_column_writes_data_pages() {
let array = Arc::new(Int32Array::from(vec![None::<i32>; 4096])) as ArrayRef;
Expand Down
10 changes: 10 additions & 0 deletions parquet/src/column/writer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1466,7 +1466,17 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {

/// Adds data page.
/// Data page is either buffered in case of dictionary encoding or written directly.
///
/// A call with nothing buffered is a no-op. Content-defined chunking forces a page
/// break at chunk boundaries, and writing the chunk can already have flushed the
/// page (the chunk hits `data_page_size_limit` or `data_page_row_count_limit`).
/// Flushing an empty BOOLEAN RLE encoder panics; other encodings would emit a
/// zero-value data page.
pub(crate) fn add_data_page(&mut self) -> Result<()> {
if self.page_metrics.num_buffered_values == 0 {
return Ok(());
}

// Extract encoded values
let values_data = self.encoder.flush_data_page()?;

Expand Down
Loading