diff --git a/crates/paimon/src/arrow/format/mod.rs b/crates/paimon/src/arrow/format/mod.rs index c85c0578a..2b24de7e1 100644 --- a/crates/paimon/src/arrow/format/mod.rs +++ b/crates/paimon/src/arrow/format/mod.rs @@ -130,6 +130,12 @@ pub(crate) trait FormatFileWriter: Send { /// Number of bytes buffered in the current row group (not yet flushed). fn in_progress_size(&self) -> usize; + /// Estimated input bytes reserved by this file writer. This excludes + /// reservations held by other consumers of the same resource context. + fn retained_input_size(&self) -> usize { + 0 + } + /// Whether this writer still retains batch data. This can differ from /// `in_progress_size` while a format is inferring its physical schema. fn retains_batch_data(&self) -> bool { @@ -266,6 +272,10 @@ impl FormatFileWriter for ResourceFormatWriter { self.inner.in_progress_size() } + fn retained_input_size(&self) -> usize { + self.reservation.size() + } + fn retains_batch_data(&self) -> bool { self.inner.retains_batch_data() } @@ -759,6 +769,57 @@ mod tests { use crate::io::FileIOBuilder; use crate::spec::{DataType, IntType}; + struct FailingFlushWriter; + + #[async_trait] + impl FormatFileWriter for FailingFlushWriter { + async fn write(&mut self, _batch: &RecordBatch) -> crate::Result<()> { + Ok(()) + } + fn num_bytes(&self) -> usize { + 0 + } + fn in_progress_size(&self) -> usize { + 1 + } + async fn flush(&mut self) -> crate::Result<()> { + Err(Error::Unsupported { + message: "injected flush failure".into(), + }) + } + async fn close(self: Box) -> crate::Result { + Err(Error::Unsupported { + message: "injected close failure".into(), + }) + } + } + + #[tokio::test] + async fn retained_input_survives_flush_failure_until_close_or_drop() { + let batch = RecordBatch::try_from_iter(vec![( + "id", + Arc::new(arrow_array::Int32Array::from(vec![1, 2])) as arrow_array::ArrayRef, + )]) + .unwrap(); + for close in [false, true] { + let resources = crate::resource::ResourceContext::builder().build().unwrap(); + let mut writer = with_write_resources(Box::new(FailingFlushWriter), Some(&resources)); + writer.write(&batch).await.unwrap(); + assert!(writer.flush().await.is_err()); + assert_eq!(writer.retained_input_size(), batch.get_array_memory_size()); + assert_eq!( + resources.metrics().reserved_memory_bytes, + batch.get_array_memory_size() + ); + if close { + assert!(writer.close().await.is_err()); + } else { + drop(writer); + } + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + } + } + #[test] fn write_format_validation_matches_compiled_factories() { for extension in supported_write_formats() { diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index a2de1a289..3f7039155 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -4646,15 +4646,19 @@ mod tests { writer.write(&first).await.unwrap(); assert!(writer.in_progress_size() > 0); + assert_eq!(writer.retained_input_size(), first_bytes.div_ceil(5)); assert_eq!( resources.metrics().reserved_memory_bytes, first_bytes.div_ceil(5) ); writer.write(&second).await.unwrap(); + assert_eq!(writer.retained_input_size(), second_bytes.div_ceil(4)); assert_eq!( resources.metrics().reserved_memory_bytes, second_bytes.div_ceil(4) ); + writer.flush().await.unwrap(); + assert_eq!(writer.retained_input_size(), 0); writer.close().await.unwrap(); assert_eq!(resources.metrics().reserved_memory_bytes, 0); diff --git a/crates/paimon/src/table/data_file_writer.rs b/crates/paimon/src/table/data_file_writer.rs index bb7bea28e..bc191f0fc 100644 --- a/crates/paimon/src/table/data_file_writer.rs +++ b/crates/paimon/src/table/data_file_writer.rs @@ -264,7 +264,11 @@ impl DataFileWriter { } if let Some(writer) = self.current_writer.as_mut() { - if writer.in_progress_size() as i64 >= self.write_buffer_size { + // Dictionary encoding may be much smaller than the retained input. + // Only this file's reservation belongs in its flush decision. + if writer.in_progress_size().max(writer.retained_input_size()) as i64 + >= self.write_buffer_size + { writer.flush().await?; } } @@ -624,6 +628,125 @@ mod tests { use opendal::Operator; use std::sync::{Arc, Mutex}; + fn dictionary_batch(rows: usize) -> RecordBatch { + let value = "x".repeat(4096); + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "id", + ArrowDataType::Utf8, + true, + )])), + vec![Arc::new(arrow_array::StringArray::from(vec![value; rows]))], + ) + .unwrap() + } + + fn dictionary_writer(io: FileIO, path: &str, buffer_size: i64) -> DataFileWriter { + DataFileWriter::new( + io, + path.into(), + String::new(), + 0, + 0, + i64::MAX, + "zstd".into(), + 1, + buffer_size, + "parquet".into(), + vec![DataField::new( + 0, + "id".into(), + DataType::VarChar(crate::spec::VarCharType::string_type()), + )], + HashMap::new(), + None, + None, + None, + ) + .unwrap() + } + + #[tokio::test] + async fn retained_input_flushes_dictionary_batches_within_budget() { + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + // Without accounting, the same input keeps the original encoded-size + // policy. With accounting, each batch must be flushed before the next. + for account_input in [false, true] { + let io = FileIOBuilder::new("memory").build().unwrap(); + let path = "memory:/dictionary-input"; + let resources = ResourceContext::builder() + .memory_limit(6 * 1024 * 1024) + .build() + .unwrap(); + let mut writer = dictionary_writer(io.clone(), path, 1024 * 1024) + .with_resources(account_input.then(|| resources.clone())); + for _ in 0..4 { + // Each write owns distinct Arrow buffers, despite identical values. + writer.write(&dictionary_batch(1024)).await.unwrap(); + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + } + let files = writer.prepare_commit().await.unwrap(); + assert_eq!(files.len(), 1); + let bytes = io + .new_input(&format!("{path}/bucket-0/{}", files[0].file_name)) + .unwrap() + .read() + .await + .unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + assert_eq!( + reader.metadata().num_row_groups(), + if account_input { 4 } else { 1 } + ); + let mut rows = 0; + let expected_value = "x".repeat(4096); + for batch in reader.build().unwrap() { + let batch = batch.unwrap(); + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(values + .iter() + .all(|value| value == Some(expected_value.as_str()))); + rows += batch.num_rows(); + } + assert_eq!(rows, 4096); + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + assert!(resources.metrics().peak_reserved_memory_bytes <= 6 * 1024 * 1024); + } + } + + #[tokio::test] + async fn retained_input_flush_ignores_other_files_in_shared_context() { + let io = FileIOBuilder::new("memory").build().unwrap(); + let resources = ResourceContext::builder().build().unwrap(); + let mut first = dictionary_writer(io.clone(), "memory:/first", i64::MAX) + .with_resources(Some(resources.clone())); + first.write(&dictionary_batch(512)).await.unwrap(); + let first_bytes = resources.metrics().reserved_memory_bytes; + assert!(first_bytes > 1024 * 1024); + + let mut second = dictionary_writer(io, "memory:/second", 1024 * 1024) + .with_resources(Some(resources.clone())); + second.write(&dictionary_batch(1)).await.unwrap(); + assert!( + second + .current_writer + .as_ref() + .unwrap() + .retained_input_size() + > 0 + ); + assert!(resources.metrics().reserved_memory_bytes > first_bytes); + second.close_current_file().await.unwrap(); + assert_eq!(resources.metrics().reserved_memory_bytes, first_bytes); + first.abort().await; + assert_eq!(resources.metrics().reserved_memory_bytes, 0); + } + #[derive(Debug)] struct RecordingProvider { operator: Operator, diff --git a/docs/src/c-binding.md b/docs/src/c-binding.md index b32a739ea..d0cab99c3 100644 --- a/docs/src/c-binding.md +++ b/docs/src/c-binding.md @@ -802,6 +802,12 @@ batches, file indexes, caller-owned Arrow input, and allocation overhead are not fully charged. A reservation failure does not automatically flush or retry a writer; handle the error as a failed write operation. +For formats that support flushing, a data file also flushes when its own +retained-input reservation reaches `write-buffer-size`, even if encoding (for +example, a Parquet dictionary) is much smaller. Other consumers' reservations +do not trigger that file's flush. This may change row-group boundaries; it does +not make an individual batch larger than the remaining budget admissible. + The input Arrow schema must match the table schema exactly, including field count, order, names, and types. A non-nullable table field must not contain null values.