fix(write): flush retained input before dictionary encoding exhausts its budget

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<Self>) -> crate::Result<FormatWriteResult> {
+            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() {

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);
 

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::<arrow_array::StringArray>()
+                    .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,

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.