diff --git a/parquet/src/column/chunker/cdc.rs b/parquet/src/column/chunker/cdc.rs index f524250b1b03..95b37333dc03 100644 --- a/parquet/src/column/chunker/cdc.rs +++ b/parquet/src/column/chunker/cdc.rs @@ -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 --- @@ -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 { + let values: Vec = (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 = (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::; 4096])) as ArrayRef; diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 01c812051f24..938d4053dd5a 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -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()?;