perf: share Metric export writer and payload budgets

Signed-off-by: jeremyhi <fengjiachun@gmail.com>
2026-10-13 07:22:32 +00:00 · 2026-09-22 10:57:33 +08:00
parent 683657ad32
commit 7bff9a9fef
7 changed files with 1067 additions and 224 deletions
@@ -186,9 +186,12 @@ impl ParquetFileWriter {
    }

    async fn write_bytes(&mut self, bytes: Vec<u8>) -> Result<()> {
-        if !bytes.is_empty() {
+        let bytes = Bytes::from(bytes);
+        let chunk = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize;
+        // Slices retain the complete encoded allocation until its last submission.
+        for offset in (0..bytes.len()).step_by(chunk) {
            self.sink
-                .write(bytes)
+                .write(bytes.slice(offset..(offset + chunk).min(bytes.len())))
                .await
                .context(error::WriteObjectSnafu { path: &self.path })?;
        }
@@ -328,7 +331,9 @@ mod tests {
            async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
                self.0.write(bytes).await
            }
-            async fn close(&mut self) -> object_store::Result<object_store::layers::mock::Metadata> {
+            async fn close(
+                &mut self,
+            ) -> object_store::Result<object_store::layers::mock::Metadata> {
                self.0.close().await?;
                Err(object_store::Error::new(
                    object_store::ErrorKind::Unexpected,
@@ -391,6 +396,74 @@ mod tests {
        }
    }

+    #[tokio::test]
+    async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() {
+        use std::sync::Mutex;
+
+        use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio};
+        struct ObservedWriter(oio::Writer, Arc<Mutex<Vec<usize>>>);
+        impl oio::Write for ObservedWriter {
+            async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
+                self.1.lock().unwrap().push(bytes.len());
+                self.0.write(bytes).await
+            }
+            async fn close(&mut self) -> object_store::Result<Metadata> {
+                self.0.close().await
+            }
+            async fn abort(&mut self) -> object_store::Result<()> {
+                self.0.abort().await
+            }
+        }
+        let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet");
+        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
+            .unwrap()
+            .build_operator();
+        let sizes = Arc::new(Mutex::new(Vec::new()));
+        let factory: MockWriterFactory = Arc::new({
+            let sizes = sizes.clone();
+            move |_, _, writer| Box::new(ObservedWriter(writer, sizes.clone()))
+        });
+        let store = store.layer(
+            MockLayerBuilder::default()
+                .writer_factory(factory)
+                .build()
+                .unwrap(),
+        );
+        let mut state = 17u64;
+        let values = (0..600_000)
+            .map(|_| {
+                state ^= state << 13;
+                state ^= state >> 7;
+                state ^= state << 17;
+                state as i64
+            })
+            .collect::<Vec<_>>();
+        let array = Arc::new(Int64Array::from(values)) as arrow::array::ArrayRef;
+        let batch = RecordBatch::try_from_iter([("a", array.clone()), ("b", array)]).unwrap();
+        let mut writer =
+            ParquetFileWriter::open(batch.schema(), store.clone(), "large.parquet", 1, None)
+                .await
+                .unwrap();
+        // No storage-layer chunking: observe the application's actual submissions.
+        writer.sink = store.writer("large.parquet").await.unwrap();
+        writer.write(batch.clone(), None).await.unwrap();
+        writer.finish(None).await.unwrap();
+        let sizes = sizes.lock().unwrap().clone();
+        let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize;
+        assert!(sizes.iter().sum::<usize>() > limit);
+        assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}");
+        let actual = read(&store, "large.parquet")
+            .await
+            .build()
+            .unwrap()
+            .collect::<std::result::Result<Vec<_>, _>>()
+            .unwrap();
+        assert_eq!(
+            arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(),
+            batch
+        );
+    }
+
    #[tokio::test]
    async fn row_and_byte_limits_split_batches() {
        let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
@@ -290,7 +290,7 @@ async fn test_s3_backend() -> Result<()> {
            .secret_access_key(&env::var("GT_S3_ACCESS_KEY")?)
            .region(&env::var("GT_S3_REGION")?)
            .bucket(&bucket);
-        if let Ok(endpoint) = env::var("GT_S3_ENDPOINT") {
+        if let Ok(endpoint) = env::var("GT_S3_ENDPOINT_URL") {
            builder = builder.endpoint(&endpoint);
        }

@@ -22,6 +22,7 @@ use common_datasource::file_format::csv::stream_to_csv;
 use common_datasource::file_format::json::stream_to_json;
 use common_datasource::file_format::parquet::stream_to_parquet;
 use common_datasource::object_store::build_backend_for_write_with_path;
+use common_datasource::parquet_writer::ParquetFileWriter;
 use common_query::Output;
 use common_recordbatch::adapter::DfRecordBatchStreamAdapter;
 use common_recordbatch::{
@@ -32,16 +33,22 @@ use common_telemetry::{debug, tracing};
 use datafusion::datasource::DefaultTableSource;
 use datafusion_common::TableReference as DfTableReference;
 use datafusion_expr::LogicalPlanBuilder;
+use futures::StreamExt;
 use object_store::ObjectStore;
 use session::context::QueryContextRef;
-use snafu::{OptionExt, ResultExt};
+use snafu::{OptionExt, ResultExt, ensure};
 use table::TableRef;
 use table::requests::CopyTableRequest;
 use table::table::adapter::DfTableProviderAdapter;
 use table::table_reference::TableReference;
+use tokio_util::sync::CancellationToken;

 use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result};
 use crate::statement::StatementExecutor;
+use crate::statement::export_logical_tables::writers::ExportWriteBudget;
+use crate::statement::export_logical_tables::{
+    LogicalTableExportLimits, map_writer_error, rows_within_budget,
+};

 // The buffer size should be greater than 5MB (minimum multipart upload size).
 /// Buffer size to flush data to object stores.
@@ -117,6 +124,17 @@ impl StatementExecutor {
        table: TableRef,
        req: CopyTableRequest,
        query_ctx: QueryContextRef,
+    ) -> Result<usize> {
+        self.copy_captured_table_to_managed(table, req, query_ctx, None)
+            .await
+    }
+
+    pub(crate) async fn copy_captured_table_to_managed(
+        &self,
+        table: TableRef,
+        req: CopyTableRequest,
+        query_ctx: QueryContextRef,
+        managed: Option<(&ExportWriteBudget, &CancellationToken)>,
    ) -> Result<usize> {
        let info = table.table_info();
        let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name);
@@ -165,7 +183,7 @@ impl StatementExecutor {
        } = &req;

        debug!("Copy table: {table_id} to location: {location}");
-        self.copy_to_file(&format, output, location, connection)
+        self.copy_to_file_managed(&format, output, location, connection, managed)
            .await
    }

@@ -176,9 +194,25 @@ impl StatementExecutor {
        location: &str,
        connection: &HashMap<String, String>,
    ) -> Result<usize> {
-        let output = output
-            .map_dictionary_to_values()
-            .context(error::BuildRecordBatchSnafu)?;
+        self.copy_to_file_managed(format, output, location, connection, None)
+            .await
+    }
+
+    async fn copy_to_file_managed(
+        &self,
+        format: &Format,
+        output: Output,
+        location: &str,
+        connection: &HashMap<String, String>,
+        managed: Option<(&ExportWriteBudget, &CancellationToken)>,
+    ) -> Result<usize> {
+        let output = if managed.is_none() {
+            output
+                .map_dictionary_to_values()
+                .context(error::BuildRecordBatchSnafu)?
+        } else {
+            output
+        };
        let stream = match output.data {
            OutputData::Stream(stream) => stream,
            OutputData::RecordBatches(record_batches) => record_batches.as_stream(),
@@ -192,7 +226,214 @@ impl StatementExecutor {
        let filename = backend.object_path.context(error::UnexpectedSnafu {
            violated: format!("Expected filename, path: {location}"),
        })?;
-        self.stream_to_file(stream, format, backend.object_store, &filename)
+        if let Some((budget, token)) = managed {
+            stream_to_managed_parquet(stream, backend.object_store, &filename, budget, token).await
+        } else {
+            self.stream_to_file(stream, format, backend.object_store, &filename)
+                .await
+        }
+    }
+}
+
+pub(crate) async fn stream_to_managed_parquet(
+    mut stream: SendableRecordBatchStream,
+    store: ObjectStore,
+    path: &str,
+    budget: &ExportWriteBudget,
+    token: &CancellationToken,
+) -> Result<usize> {
+    use common_recordbatch::{
+        RecordBatch, map_dictionary_to_values, map_dictionary_to_values_schema,
+    };
+    let original = stream.schema();
+    let (expanded_schema, expand) = map_dictionary_to_values_schema(original.clone());
+    let (mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone());
+    let output_schema = if json {
+        mapped_schema
+    } else {
+        expanded_schema.clone()
+    };
+    let mut writer =
+        ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None)
            .await
+            .context(error::WriteStreamToFileSnafu { path })?;
+    let mut started = false;
+    let result = async {
+        let mut rows = 0;
+        loop {
+            let batch = tokio::select! {
+                biased;
+                _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(),
+                batch = stream.next() => batch,
+            };
+            let Some(batch) = batch else { break };
+            let batch = batch
+                .context(error::BuildRecordBatchSnafu)?
+                .into_df_record_batch();
+            let mut offset = 0;
+            while offset < batch.num_rows() {
+                let (conversion, retained) = ExportWriteBudget::conversion_budget(
+                    &batch,
+                    LogicalTableExportLimits::default().conversion_bytes,
+                )?;
+                let input = batch.clone();
+                // JSON escaping and number formatting can expand binary JSON.
+                let expansion = if json { 8 } else { 1 };
+                let (len, estimated) = common_runtime::spawn_blocking_global(move || {
+                    rows_within_budget(&input, offset, input.num_rows(), conversion / expansion)
+                })
+                .await
+                .context(error::JoinTaskSnafu)??;
+                let reservation = retained.saturating_add(estimated.saturating_mul(2 * expansion));
+                let permit = budget.reserve(reservation, token).await?;
+                let batch = batch.clone();
+                let (original, expanded_schema, output_schema) = (
+                    original.clone(),
+                    expanded_schema.clone(),
+                    output_schema.clone(),
+                );
+                let (batch, len) = common_runtime::spawn_blocking_global(move || {
+                    let mut batch = RecordBatch::from_df_record_batch(
+                        original.clone(),
+                        batch.slice(offset, len),
+                    );
+                    if expand {
+                        batch = map_dictionary_to_values(batch, &original, &expanded_schema)
+                            .context(error::BuildRecordBatchSnafu)?;
+                    }
+                    if json {
+                        batch = map_json_type_to_string(batch, &expanded_schema, &output_schema)
+                            .context(error::BuildRecordBatchSnafu)?;
+                    }
+                    let batch = batch.into_df_record_batch();
+                    ensure!(
+                        batch.get_array_memory_size() <= reservation,
+                        error::LogicalTableExportResourceSnafu {
+                            reason: "converted backing buffers exceed reservation"
+                        }
+                    );
+                    Ok::<_, error::Error>((batch, len))
+                })
+                .await
+                .context(error::JoinTaskSnafu)??;
+                started = true;
+                let write = writer.write(batch, Some(token)).await;
+                drop(permit);
+                write.map_err(|error| map_writer_error(error, path))?;
+                rows += len;
+                offset += len;
+            }
+        }
+        started = true;
+        writer
+            .finish(Some(token))
+            .await
+            .map_err(|error| map_writer_error(error, path))?;
+        Ok(rows)
+    }
+    .await;
+    if result.is_err() {
+        token.cancel();
+        if started && let Err(error) = writer.abort().await {
+            common_telemetry::warn!(error; "Failed to abort ordinary export file");
+        }
+    }
+    result
+}
+
+#[cfg(test)]
+mod tests {
+    use arrow::array::{ArrayRef, BinaryArray, DictionaryArray, Int32Array, StringArray};
+    use arrow::datatypes::Int32Type;
+    use common_recordbatch::{RecordBatch, RecordBatches};
+    use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+    use datatypes::prelude::ConcreteDataType;
+    use datatypes::schema::{ColumnSchema, Schema};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn managed_copy_preserves_dictionary_json_and_empty_schema() {
+        let schema = Arc::new(Schema::new(vec![
+            ColumnSchema::new(
+                "host",
+                ConcreteDataType::dictionary_datatype(
+                    ConcreteDataType::int32_datatype(),
+                    ConcreteDataType::string_datatype(),
+                ),
+                true,
+            ),
+            ColumnSchema::new("json", ConcreteDataType::json_datatype(), true),
+        ]));
+        let dictionary = DictionaryArray::<Int32Type>::new(
+            Int32Array::from(vec![Some(0), None]),
+            Arc::new(StringArray::from(vec!["tag"])),
+        );
+        let json =
+            datatypes::types::parse_string_to_jsonb(r#"{"value":"escaped\ntext","n":123456789}"#)
+                .unwrap();
+        let arrays = vec![
+            Arc::new(dictionary) as ArrayRef,
+            Arc::new(BinaryArray::from(vec![Some(json.as_slice()), None])),
+        ];
+        let batch =
+            arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays)
+                .unwrap();
+        let batch = RecordBatch::from_df_record_batch(schema.clone(), batch);
+        for empty in [false, true] {
+            let batches = RecordBatches::try_new(
+                schema.clone(),
+                if empty { vec![] } else { vec![batch.clone()] },
+            )
+            .unwrap();
+            let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
+            let budget = ExportWriteBudget::new(1);
+            let token = CancellationToken::new();
+            let managed = stream_to_managed_parquet(
+                batches.as_stream(),
+                store.clone(),
+                "managed.parquet",
+                &budget,
+                &token,
+            )
+            .await
+            .unwrap();
+            let output = Output::new_with_record_batches(batches)
+                .map_dictionary_to_values()
+                .unwrap();
+            let OutputData::RecordBatches(batches) = output.data else {
+                panic!("expected batches");
+            };
+            let stream = SendableRecordBatchMapper::new(
+                batches.as_stream(),
+                map_json_type_to_string,
+                map_json_type_to_string_schema,
+            );
+            let ordinary = stream_to_parquet(
+                Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))),
+                store.clone(),
+                "ordinary.parquet",
+                WRITE_CONCURRENCY,
+            )
+            .await
+            .unwrap();
+            assert_eq!(managed, ordinary);
+            let mut outputs = Vec::new();
+            for name in ["managed.parquet", "ordinary.parquet"] {
+                let reader = ParquetRecordBatchReaderBuilder::try_new(
+                    store.read(name).await.unwrap().to_bytes(),
+                )
+                .unwrap();
+                let schema = reader.schema().clone();
+                let values = reader
+                    .build()
+                    .unwrap()
+                    .collect::<std::result::Result<Vec<_>, _>>()
+                    .unwrap();
+                outputs.push((schema, values));
+            }
+            assert_eq!(outputs[0], outputs[1]);
+            assert_eq!(budget.available(), (1, 64 * 1024 * 1024));
+        }
    }
 }
@@ -37,6 +37,7 @@ use crate::statement::database_copy::{
    DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory,
    validate_database_export_layout,
 };
+use crate::statement::export_logical_tables::writers::{ExportWriteBudget, retain_error};
 use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits};

 /// A validated request-scoped selection, not a metadata snapshot or an ACL token.
@@ -202,46 +203,52 @@ impl StatementExecutor {
        }
        output_files.sort();
        let req = &plan.request;
-        let rows = run_database_export_jobs(
-            plan.jobs,
-            parse_parallelism_from_option_map(&req.with),
-            cancellation,
-            |job, token| {
-                let ctx = ctx.clone();
-                async move {
-                    match job {
-                        DatabaseExportJob::Metric(unit) => self
-                            .export_logical_tables(
-                                &unit,
-                                &req.location,
-                                &req.connection,
-                                req.time_range.as_ref(),
-                                LogicalTableExportLimits::default(),
-                                &token,
-                                ctx,
-                            )
-                            .await
-                            .map(|summary| summary.rows),
-                        DatabaseExportJob::Ordinary { table, output } => {
-                            let info = table.table_info();
-                            let copy = CopyTableRequest {
-                                catalog_name: info.catalog_name.clone(),
-                                schema_name: info.schema_name.clone(),
-                                table_name: info.name.clone(),
-                                location: output.location,
-                                with: req.with.clone(),
-                                connection: req.connection.clone(),
-                                pattern: None,
-                                direction: CopyDirection::Export,
-                                timestamp_range: req.time_range,
-                                limit: None,
-                            };
-                            self.copy_captured_table_to(table, copy, ctx).await
-                        }
+        let parallelism = parse_parallelism_from_option_map(&req.with);
+        let budget = ExportWriteBudget::new(parallelism);
+        let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| {
+            let ctx = ctx.clone();
+            let budget = budget.clone();
+            async move {
+                match job {
+                    DatabaseExportJob::Metric(unit) => self
+                        .export_logical_tables_managed(
+                            &unit,
+                            &req.location,
+                            &req.connection,
+                            req.time_range.as_ref(),
+                            LogicalTableExportLimits::default(),
+                            &token,
+                            ctx,
+                            budget,
+                        )
+                        .await
+                        .map(|summary| summary.rows),
+                    DatabaseExportJob::Ordinary { table, output } => {
+                        let info = table.table_info();
+                        let copy = CopyTableRequest {
+                            catalog_name: info.catalog_name.clone(),
+                            schema_name: info.schema_name.clone(),
+                            table_name: info.name.clone(),
+                            location: output.location,
+                            with: req.with.clone(),
+                            connection: req.connection.clone(),
+                            pattern: None,
+                            direction: CopyDirection::Export,
+                            timestamp_range: req.time_range,
+                            limit: None,
+                        };
+                        let _permit = budget.writer(&token).await?;
+                        self.copy_captured_table_to_managed(
+                            table,
+                            copy,
+                            ctx,
+                            Some((&budget, &token)),
+                        )
+                        .await
                    }
                }
-            },
-        )
+            }
+        })
        .await?;
        Ok(DatabaseExportSummary { rows, output_files })
    }
@@ -261,6 +268,7 @@ async fn run_database_export_jobs<J, F: Future<Output = Result<usize>>>(
    loop {
        while first_error.is_none()
            && !cancellation.is_cancelled()
+            && !token.is_cancelled()
            && active.len() < parallelism.max(1)
        {
            let Some(job) = jobs.next() else { break };
@@ -284,11 +292,12 @@ async fn run_database_export_jobs<J, F: Future<Output = Result<usize>>>(
        };
        match result {
            Some(Ok(count)) => rows += count,
-            Some(Err(err)) if first_error.is_none() => {
-                first_error = Some(err);
+            Some(Err(err)) => {
+                if !cancellation.is_cancelled() || first_error.is_none() {
+                    retain_error(&mut first_error, err);
+                }
                token.cancel();
            }
-            Some(Err(err)) => common_telemetry::warn!(err; "Failed to drain database export job"),
            None => break,
        }
    }
@@ -19,6 +19,8 @@
 //! exclusively by the export attempt. Completion and cleanup of closed files
 //! belong to the caller's chunk protocol, not individual file completion.

+pub(crate) mod writers;
+
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::sync::Arc;

@@ -54,6 +56,7 @@ use tokio_util::sync::CancellationToken;
 use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result};
 use crate::statement::StatementExecutor;
 use crate::statement::database_copy::DatabaseExportFile;
+use crate::statement::export_logical_tables::writers::{ExportWriteBudget, Payload, TableWriters};

 /// Export preprocessing and per-file limits. Query memory and spill remain
 /// governed by the query engine.
@@ -103,7 +106,7 @@ pub struct LogicalTableExport {
    logical_tables: BTreeMap<TableId, LogicalTableProjection>,
 }

-struct LogicalTableProjection {
+pub(crate) struct LogicalTableProjection {
    output: DatabaseExportFile,
    schema: SchemaRef,
    projection: Vec<usize>,
@@ -308,6 +311,31 @@ impl StatementExecutor {
        limits: LogicalTableExportLimits,
        cancellation: &CancellationToken,
        query_ctx: QueryContextRef,
+    ) -> Result<LogicalTableExportSummary> {
+        self.export_logical_tables_managed(
+            unit,
+            directory,
+            connection,
+            time_range,
+            limits,
+            &cancellation.child_token(),
+            query_ctx,
+            ExportWriteBudget::new(1),
+        )
+        .await
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) async fn export_logical_tables_managed(
+        &self,
+        unit: &LogicalTableExport,
+        directory: &str,
+        connection: &HashMap<String, String>,
+        time_range: Option<&TimestampRange>,
+        limits: LogicalTableExportLimits,
+        cancellation: &CancellationToken,
+        query_ctx: QueryContextRef,
+        budget: Arc<ExportWriteBudget>,
    ) -> Result<LogicalTableExportSummary> {
        limits.validate()?;
        let (store, stream) = tokio::select! {
@@ -327,10 +355,11 @@ impl StatementExecutor {
                Ok((store, stream))
            } => result?,
        };
-        export_stream(unit, stream, &store, limits, cancellation).await
+        export_stream_managed(unit, stream, &store, limits, cancellation, budget).await
    }
 }

+#[cfg(test)]
 async fn export_stream(
    unit: &LogicalTableExport,
    stream: SendableRecordBatchStream,
@@ -338,19 +367,42 @@ async fn export_stream(
    limits: LogicalTableExportLimits,
    cancellation: &CancellationToken,
 ) -> Result<LogicalTableExportSummary> {
-    let mut active = None;
-    let result = write_tables(unit, stream, store, limits, cancellation, &mut active).await;
-    if result.is_err()
-        && let Some(writer) = active
-        && let Err(cleanup_error) = writer
-            .writer
-            .abort()
-            .await
-            .map_err(|error| map_writer_error(error, &writer.path))
-    {
-        common_telemetry::warn!(cleanup_error; "Failed to clean up incomplete Metric export file");
-    }
-    result
+    export_stream_managed(
+        unit,
+        stream,
+        store,
+        limits,
+        &cancellation.child_token(),
+        ExportWriteBudget::new(1),
+    )
+    .await
+}
+
+async fn export_stream_managed(
+    unit: &LogicalTableExport,
+    stream: SendableRecordBatchStream,
+    store: &ObjectStore,
+    limits: LogicalTableExportLimits,
+    cancellation: &CancellationToken,
+    budget: Arc<ExportWriteBudget>,
+) -> Result<LogicalTableExportSummary> {
+    let mut writers = TableWriters::new(budget.clone());
+    let result = write_tables(
+        unit,
+        stream,
+        store,
+        limits,
+        cancellation,
+        &mut writers,
+        &budget,
+    )
+    .await;
+    let (summary, result) = match result {
+        Ok(summary) => (summary, Ok(())),
+        Err(error) => (LogicalTableExportSummary::default(), Err(error)),
+    };
+    writers.drain(result, cancellation).await?;
+    Ok(summary)
 }

 fn check_cancelled(cancellation: &CancellationToken) -> Result<()> {
@@ -367,7 +419,8 @@ async fn write_tables(
    store: &ObjectStore,
    limits: LogicalTableExportLimits,
    cancellation: &CancellationToken,
-    active: &mut Option<ActiveWriter>,
+    writers: &mut TableWriters,
+    budget: &ExportWriteBudget,
 ) -> Result<LogicalTableExportSummary> {
    let id_index =
        stream
@@ -425,22 +478,19 @@ async fn write_tables(
            while end < batch.num_rows() && ids.value(end) == id {
                end += 1;
            }
-            if active.as_ref().is_some_and(|writer| writer.table_id != id) {
-                finish_active(active, cancellation).await?;
+            if writers.table_id().is_some_and(|last| last != id) {
+                writers.close_input();
            }
            check_cancelled(cancellation)?;
            if let Some(file) = unit.logical_tables.get(&id) {
-                if active.is_none() {
-                    *active = Some(ActiveWriter::open(id, file, store, limits).await?);
+                if writers.table_id().is_none() {
+                    writers.open(id, file, store, limits, cancellation).await?;
                    written.insert(id);
                    summary.files += 1;
                }
                let projected = batch
                    .project(&file.projection)
                    .context(error::ProjectSchemaSnafu)?;
-                let writer = active.as_mut().context(error::UnexpectedSnafu {
-                    violated: "missing logical writer",
-                })?;
                let mut offset = start;
                while offset < end {
                    let (expanded, consumed) = expand_bounded_slice(
@@ -449,15 +499,12 @@ async fn write_tables(
                        offset,
                        end,
                        limits.conversion_bytes,
+                        budget,
+                        cancellation,
                    )
                    .await?;
                    check_cancelled(cancellation)?;
-                    // Do not drop an in-flight file operation before cleanup.
-                    writer
-                        .writer
-                        .write(expanded, Some(cancellation))
-                        .await
-                        .map_err(|error| map_writer_error(error, &writer.path))?;
+                    writers.send(expanded, cancellation).await?;
                    check_cancelled(cancellation)?;
                    offset += consumed;
                    summary.rows += consumed;
@@ -468,12 +515,12 @@ async fn write_tables(
            start = end;
        }
    }
-    finish_active(active, cancellation).await?;
+    writers.close_input();
    for (&id, file) in &unit.logical_tables {
        if !written.contains(&id) {
            check_cancelled(cancellation)?;
-            *active = Some(ActiveWriter::open(id, file, store, limits).await?);
-            finish_active(active, cancellation).await?;
+            writers.open(id, file, store, limits, cancellation).await?;
+            writers.close_input();
            summary.files += 1;
        }
    }
@@ -482,20 +529,18 @@ async fn write_tables(
 }

 struct ActiveWriter {
-    table_id: u32,
    path: String,
    writer: ParquetFileWriter,
 }

 impl ActiveWriter {
    async fn open(
-        id: u32,
        table: &LogicalTableProjection,
        store: &ObjectStore,
        limits: LogicalTableExportLimits,
    ) -> Result<Self> {
        let path = table.output.path.clone();
-        let conditional = store.info().full_capability().write_with_if_not_exists;
+        let conditional = store.info().capability().write_with_if_not_exists;
        if !conditional {
            ensure!(
                !store
@@ -521,39 +566,29 @@ impl ActiveWriter {
        )
        .await
        .map_err(|error| map_writer_error(error, &path))?;
-        Ok(Self {
-            table_id: id,
-            path,
-            writer,
-        })
+        Ok(Self { path, writer })
    }
 }

-async fn finish_active(
-    active: &mut Option<ActiveWriter>,
-    cancellation: &CancellationToken,
-) -> Result<()> {
-    if let Some(writer) = active.as_mut() {
-        writer
-            .writer
-            .finish(Some(cancellation))
-            .await
-            .map_err(|error| map_writer_error(error, &writer.path))?;
-        check_cancelled(cancellation)?;
-        *active = None;
-    }
-    Ok(())
-}
-
 async fn expand_bounded_slice(
    batch: RecordBatch,
    schema: SchemaRef,
    start: usize,
    end: usize,
-    budget: usize,
-) -> Result<(RecordBatch, usize)> {
+    requested: usize,
+    budget: &ExportWriteBudget,
+    cancellation: &CancellationToken,
+) -> Result<(Payload, usize)> {
+    let (conversion, retained) = ExportWriteBudget::conversion_budget(&batch, requested)?;
+    let input = batch.clone();
+    let (len, estimated) = common_runtime::spawn_blocking_global(move || {
+        rows_within_budget(&input, start, end, conversion)
+    })
+    .await
+    .context(error::JoinTaskSnafu)??;
+    let reservation = retained.saturating_add(estimated.saturating_mul(2));
+    let permit = budget.reserve(reservation, cancellation).await?;
    common_runtime::spawn_blocking_global(move || {
-        let len = rows_within_budget(&batch, start, end, budget)?;
        let slice = batch.slice(start, len);
        let arrays = slice
            .columns()
@@ -562,13 +597,28 @@ async fn expand_bounded_slice(
            .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu))
            .collect::<Result<Vec<_>>>()?;
        let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?;
-        Ok((expanded, len))
+        ensure!(
+            expanded.get_array_memory_size() <= reservation,
+            LogicalTableExportResourceSnafu {
+                reason: "converted backing buffers exceed reservation"
+            }
+        );
+        Ok((
+            Payload {
+                batch: expanded,
+                permit,
+            },
+            len,
+        ))
    })
    .await
    .context(error::JoinTaskSnafu)?
 }

-fn map_writer_error(source: common_datasource::error::Error, path: &str) -> error::Error {
+pub(crate) fn map_writer_error(
+    source: common_datasource::error::Error,
+    path: &str,
+) -> error::Error {
    match source {
        common_datasource::error::Error::ParquetWriteCancelled {} => {
            error::LogicalTableExportCancelledSnafu.build()
@@ -639,12 +689,12 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result<usize> {
    Ok(bytes.saturating_add(16))
 }

-fn rows_within_budget(
+pub(crate) fn rows_within_budget(
    batch: &RecordBatch,
    start: usize,
    end: usize,
    budget: usize,
-) -> Result<usize> {
+) -> Result<(usize, usize)> {
    let mut bytes = 0usize;
    let mut row = start;
    while row < end {
@@ -663,7 +713,7 @@ fn rows_within_budget(
            reason: "one expanded logical row exceeds conversion byte budget"
        }
    );
-    Ok(row - start)
+    Ok((row - start, bytes))
 }

 #[cfg(test)]
@@ -128,52 +128,56 @@ async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec<RecordBatch>)

 #[tokio::test]
 async fn routes_across_batches_and_writes_empty_files() {
-    let unit = unit();
-    let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
-    let batches = vec![
-        batch(vec![Some(1025), Some(1025)], vec![Some(""), None]),
-        batch(
-            vec![Some(1025), Some(1027), Some(1028)],
-            vec![Some("a"), Some("b"), Some("unselected")],
-        ),
-    ];
-    let result = write_tables(
-        &unit,
-        stream(batches),
-        &store,
-        export_limits(),
-        &CancellationToken::new(),
-        &mut None,
-    )
-    .await
-    .unwrap();
-    assert_eq!(
-        result,
-        LogicalTableExportSummary {
-            rows: 4,
-            skipped_rows: 1,
-            files: 3
-        }
-    );
-    let (schema, cpu) = read(&store, "cpu.v1.parquet").await;
-    assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields());
-    let values: Vec<_> = cpu
-        .iter()
-        .flat_map(|batch| {
-            batch
-                .column(0)
-                .as_string::<i32>()
-                .iter()
-                .map(|s| s.map(str::to_owned))
-        })
-        .collect();
-    assert_eq!(values, vec![Some("".into()), None, Some("a".into())]);
-    let (schema, empty) = read(&store, "empty.parquet").await;
-    assert_eq!(schema.fields(), unit.logical_tables[&1026].schema.fields());
-    assert!(empty.is_empty());
-    let (schema, requests) = read(&store, "requests.parquet").await;
-    assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields());
-    assert_eq!(requests[0].column(1).null_count(), 1);
+    for parallelism in [1, 4] {
+        let unit = unit();
+        let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
+        let batches = vec![
+            batch(vec![Some(1025), Some(1025)], vec![Some(""), None]),
+            batch(
+                vec![Some(1025), Some(1027), Some(1028)],
+                vec![Some("a"), Some("b"), Some("unselected")],
+            ),
+        ];
+        let budget = ExportWriteBudget::new(parallelism);
+        let result = export_stream_managed(
+            &unit,
+            stream(batches),
+            &store,
+            export_limits(),
+            &CancellationToken::new(),
+            budget.clone(),
+        )
+        .await
+        .unwrap();
+        assert_eq!(
+            result,
+            LogicalTableExportSummary {
+                rows: 4,
+                skipped_rows: 1,
+                files: 3
+            }
+        );
+        let (schema, cpu) = read(&store, "cpu.v1.parquet").await;
+        assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields());
+        let values: Vec<_> = cpu
+            .iter()
+            .flat_map(|batch| {
+                batch
+                    .column(0)
+                    .as_string::<i32>()
+                    .iter()
+                    .map(|s| s.map(str::to_owned))
+            })
+            .collect();
+        assert_eq!(values, vec![Some("".into()), None, Some("a".into())]);
+        let (schema, empty) = read(&store, "empty.parquet").await;
+        assert_eq!(schema.fields(), unit.logical_tables[&1026].schema.fields());
+        assert!(empty.is_empty());
+        let (schema, requests) = read(&store, "requests.parquet").await;
+        assert_eq!(schema.fields(), unit.logical_tables[&1027].schema.fields());
+        assert_eq!(requests[0].column(1).null_count(), 1);
+        assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024));
+    }
 }

 #[tokio::test]
@@ -211,21 +215,16 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() {
        ),
    ] {
        let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
-        let mut active = None;
-        let err = write_tables(
+        let err = export_stream(
            &unit,
            stream(batches),
            &store,
            limits,
            &CancellationToken::new(),
-            &mut active,
        )
        .await
        .unwrap_err();
        assert!(err.to_string().contains(message), "{err}");
-        if let Some(writer) = active {
-            writer.writer.abort().await.unwrap();
-        }
    }
 }

@@ -233,17 +232,25 @@ async fn rejects_invalid_order_ids_and_resource_exhaustion() {
 async fn existing_outputs_are_not_overwritten() {
    let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
    store.write("cpu.v1.parquet", "keep").await.unwrap();
-    let err = write_tables(
+    let err = export_stream(
        &unit(),
        stream(vec![batch(vec![Some(1025)], vec![None])]),
        &store,
        LogicalTableExportLimits::default(),
        &CancellationToken::new(),
-        &mut None,
    )
    .await
    .unwrap_err();
-    assert!(err.to_string().contains("already exists"));
+    if store.info().capability().write_with_if_not_exists {
+        assert!(matches!(err, error::Error::WriteStreamToFile {
+            source: common_datasource::error::Error::WriteObject { error, .. }, ..
+        } if error.kind() == object_store::ErrorKind::ConditionNotMatch));
+    } else {
+        assert!(matches!(
+            err,
+            error::Error::InvalidLogicalTableExport { .. }
+        ));
+    }
    assert_eq!(
        store.read("cpu.v1.parquet").await.unwrap().to_bytes(),
        Bytes::from_static(b"keep")
@@ -272,10 +279,10 @@ fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() {
        ("histogram", Arc::new(histogram)),
    ])
    .unwrap();
-    assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap(), 1);
+    assert_eq!(rows_within_budget(&batch, 0, 3, 4300).unwrap().0, 1);
    // The dictionary and container overhead fit; the nested list elements do not.
    assert!(rows_within_budget(&batch, 0, 3, 4180).is_err());
-    assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap(), 3);
+    assert_eq!(rows_within_budget(&batch, 0, 3, 15000).unwrap().0, 3);
 }

 #[test]
@@ -440,60 +447,117 @@ impl object_store::layers::mock::oio::Write for PausedFileWriter {
 }

 #[tokio::test]
-async fn cancellation_waits_for_file_creation_before_cleanup() {
-    use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory};
-    let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open");
-    let access =
-        common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap();
-    let store = build_backend_for_write(
-        &format!("{}/", directory.path().display()),
-        &HashMap::new(),
-        &access,
-    )
-    .await
-    .unwrap();
-    let started = Arc::new(tokio::sync::Notify::new());
-    let release = Arc::new(tokio::sync::Notify::new());
-    let factory: MockWriterFactory = Arc::new({
-        let started = started.clone();
-        let release = release.clone();
-        move |_, _, inner| {
-            Box::new(PausedFileWriter {
-                inner: Some(inner),
-                started: started.clone(),
-                release: release.clone(),
-            })
+async fn cancellation_drains_storage_and_preserves_committed_files() {
+    for large in [false, true] {
+        use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory};
+        let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open");
+        let access =
+            common_datasource::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap();
+        let store = build_backend_for_write(
+            &format!("{}/", directory.path().display()),
+            &HashMap::new(),
+            &access,
+        )
+        .await
+        .unwrap();
+        let started = Arc::new(tokio::sync::Notify::new());
+        let release = Arc::new(tokio::sync::Notify::new());
+        let factory: MockWriterFactory = Arc::new({
+            let started = started.clone();
+            let release = release.clone();
+            move |_, _, inner| {
+                Box::new(PausedFileWriter {
+                    inner: Some(inner),
+                    started: started.clone(),
+                    release: release.clone(),
+                })
+            }
+        });
+        let store = store.layer(
+            MockLayerBuilder::default()
+                .writer_factory(factory)
+                .build()
+                .unwrap(),
+        );
+        let cancellation = CancellationToken::new();
+        let (unit, input) = if large {
+            let field = Field::new("value", DataType::Utf8, true);
+            let unit = LogicalTableExport::try_new(
+                table(
+                    1024,
+                    "phy",
+                    vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()],
+                    true,
+                ),
+                &[table(1025, "cpu.v1", vec![field], false)],
+            )
+            .unwrap();
+            let mut state = 17u64;
+            let strings = (0..512)
+                .map(|_| {
+                    (0..32768)
+                        .map(|_| {
+                            state ^= state << 13;
+                            state ^= state >> 7;
+                            state ^= state << 17;
+                            char::from(b' ' + (state % 95) as u8)
+                        })
+                        .collect::<String>()
+                })
+                .collect::<Vec<_>>();
+            let input = RecordBatch::try_from_iter([
+                (
+                    TABLE_ID,
+                    Arc::new(UInt32Array::from(vec![1025; 512])) as ArrayRef,
+                ),
+                ("value", Arc::new(StringArray::from(strings))),
+            ])
+            .unwrap();
+            (unit, input)
+        } else {
+            (unit(), batch(vec![Some(1025)], vec![Some("a")]))
+        };
+        let budget = ExportWriteBudget::new(1);
+        let mut limits = LogicalTableExportLimits::default();
+        limits.writer.flush_threshold_bytes = 1024 * 1024;
+        let export = export_stream_managed(
+            &unit,
+            stream(vec![input]),
+            &store,
+            limits,
+            &cancellation,
+            budget.clone(),
+        );
+        tokio::pin!(export);
+        tokio::select! {
+            result = &mut export => panic!("export completed before the file open: {result:?}"),
+            _ = started.notified() => {},
        }
-    });
-    let store = store.layer(
-        MockLayerBuilder::default()
-            .writer_factory(factory)
-            .build()
-            .unwrap(),
-    );
-    let cancellation = CancellationToken::new();
-    let unit = unit();
-    let export = export_stream(
-        &unit,
-        stream(vec![batch(vec![Some(1025)], vec![Some("a")])]),
-        &store,
-        LogicalTableExportLimits::default(),
-        &cancellation,
-    );
-    tokio::pin!(export);
-    tokio::select! {
-        result = &mut export => panic!("export completed before the file open: {result:?}"),
-        _ = started.notified() => {},
+        if large {
+            assert!(budget.available().1 < 64 * 1024 * 1024);
+        }
+        let held = budget.available();
+        cancellation.cancel();
+        assert!(futures::poll!(&mut export).is_pending());
+        assert_eq!(budget.available().0, held.0);
+        if large {
+            // Cancelled admission releases its pending reservation; the worker
+            // still owns payload capacity while the storage operation is paused.
+            assert!(budget.available().1 < 64 * 1024 * 1024);
+        }
+        release.notify_one();
+        let result = export.await;
+        assert!(matches!(
+            result,
+            Err(error::Error::LogicalTableExportCancelled { .. })
+        ));
+        assert_eq!(store.exists("cpu.v1.parquet").await.unwrap(), !large);
+        if !large {
+            let (_, batches) = read(&store, "cpu.v1.parquet").await;
+            assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
+        }
+        assert_eq!(budget.available(), (1, 64 * 1024 * 1024));
    }
-    cancellation.cancel();
-    assert!(futures::poll!(&mut export).is_pending());
-    release.notify_one();
-    let result = export.await;
-    assert!(matches!(
-        result,
-        Err(error::Error::LogicalTableExportCancelled { .. })
-    ));
-    assert!(!store.exists("cpu.v1.parquet").await.unwrap());
 }

 struct FailedAbortWriter(object_store::layers::mock::oio::Writer);
@@ -574,3 +638,164 @@ async fn validates_membership_by_table_route() {
        }
    }
 }
+
+#[tokio::test]
+async fn retained_backing_is_reserved_before_conversion_and_until_payload_drop() {
+    let input = batch(vec![Some(1025); 8192], vec![Some("shared"); 8192]);
+    let tiny = input.slice(0, 1);
+    let schema = Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)]));
+    let projected = tiny.project(&[2]).unwrap();
+    let full = input.project(&[2]).unwrap().get_array_memory_size();
+    assert_eq!(projected.get_array_memory_size(), full);
+    let budget = ExportWriteBudget::new(1);
+    let token = CancellationToken::new();
+    let blocker = budget.reserve(64 * 1024 * 1024, &token).await.unwrap();
+    let convert = expand_bounded_slice(projected, schema, 0, 1, 1024, &budget, &token);
+    tokio::pin!(convert);
+    assert!(futures::poll!(&mut convert).is_pending());
+    drop(blocker);
+    let (payload, rows) = convert.await.unwrap();
+    assert_eq!(rows, 1);
+    assert!(
+        64 * 1024 * 1024 - budget.available().1 >= full + payload.batch.get_array_memory_size()
+    );
+    drop(payload);
+    assert_eq!(budget.available(), (1, 64 * 1024 * 1024));
+    assert!(budget.reserve(64 * 1024 * 1024 + 1, &token).await.is_err());
+}
+
+#[tokio::test]
+async fn groups_and_ordinary_files_share_writer_admission_and_drain() {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio};
+    use tokio::sync::Semaphore;
+
+    use crate::statement::copy_table_to::stream_to_managed_parquet;
+
+    struct CountedWriter {
+        inner: oio::Writer,
+        active: Arc<AtomicUsize>,
+        started: Arc<Semaphore>,
+        release: Arc<Semaphore>,
+    }
+    impl Drop for CountedWriter {
+        fn drop(&mut self) {
+            self.active.fetch_sub(1, Ordering::SeqCst);
+        }
+    }
+    impl oio::Write for CountedWriter {
+        async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
+            self.inner.write(bytes).await
+        }
+        async fn close(&mut self) -> object_store::Result<Metadata> {
+            self.started.add_permits(1);
+            self.release.acquire().await.unwrap().forget();
+            self.inner.close().await
+        }
+        async fn abort(&mut self) -> object_store::Result<()> {
+            self.inner.abort().await
+        }
+    }
+    for parallelism in [1, 4] {
+        for cancel in [false, true] {
+            let active = Arc::new(AtomicUsize::new(0));
+            let peak = Arc::new(AtomicUsize::new(0));
+            let started = Arc::new(Semaphore::new(0));
+            let release = Arc::new(Semaphore::new(0));
+            let factory: MockWriterFactory = Arc::new({
+                let (active, peak, started, release) = (
+                    active.clone(),
+                    peak.clone(),
+                    started.clone(),
+                    release.clone(),
+                );
+                move |_, args, inner| {
+                    assert_eq!(args.concurrent(), 1);
+                    let now = active.fetch_add(1, Ordering::SeqCst) + 1;
+                    peak.fetch_max(now, Ordering::SeqCst);
+                    Box::new(CountedWriter {
+                        inner,
+                        active: active.clone(),
+                        started: started.clone(),
+                        release: release.clone(),
+                    })
+                }
+            });
+            let store = ObjectStore::new(object_store::services::Memory::default())
+                .unwrap()
+                .layer(
+                    MockLayerBuilder::default()
+                        .writer_factory(factory)
+                        .build()
+                        .unwrap(),
+                );
+            let budget = ExportWriteBudget::new(parallelism);
+            let token = CancellationToken::new();
+            let mut a = unit();
+            let mut b = unit();
+            for (group, unit) in [("a", &mut a), ("b", &mut b)] {
+                for table in unit.logical_tables.values_mut() {
+                    table.output.path = format!("{group}/{}", table.output.path);
+                }
+            }
+            let rows = || {
+                stream(vec![batch(
+                    vec![Some(1025), Some(1027)],
+                    vec![Some("a"), Some("b")],
+                )])
+            };
+            let ordinary = async {
+                let _permit = budget.writer(&token).await?;
+                stream_to_managed_parquet(
+                    rows(),
+                    store.clone(),
+                    "ordinary.parquet",
+                    &budget,
+                    &token,
+                )
+                .await
+            };
+            let work = async {
+                tokio::join!(
+                    export_stream_managed(
+                        &a,
+                        rows(),
+                        &store,
+                        export_limits(),
+                        &token,
+                        budget.clone()
+                    ),
+                    export_stream_managed(
+                        &b,
+                        rows(),
+                        &store,
+                        export_limits(),
+                        &token,
+                        budget.clone()
+                    ),
+                    ordinary,
+                )
+            };
+            tokio::pin!(work);
+            tokio::select! {
+                _ = started.acquire_many(parallelism as u32) => {},
+                result = &mut work => panic!("completed while close paused: {result:?}"),
+            }
+            assert_eq!(budget.available().0, 0);
+            assert_eq!(peak.load(Ordering::SeqCst), parallelism);
+            if cancel {
+                token.cancel();
+            }
+            assert!(futures::poll!(&mut work).is_pending());
+            release.add_permits(10);
+            let (a, b, ordinary) = work.await;
+            assert_eq!(a.is_err(), cancel);
+            assert_eq!(b.is_err(), cancel);
+            assert_eq!(ordinary.is_err(), cancel);
+            assert_eq!(active.load(Ordering::SeqCst), 0);
+            assert!(peak.load(Ordering::SeqCst) <= parallelism);
+            assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024));
+        }
+    }
+}
@@ -0,0 +1,245 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use arrow::record_batch::RecordBatch;
+use futures::stream::FuturesUnordered;
+use futures::{FutureExt, StreamExt};
+use object_store::ObjectStore;
+use snafu::{OptionExt, ResultExt, ensure};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
+use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
+
+use crate::error::{self, Result};
+use crate::statement::export_logical_tables::{
+    ActiveWriter, LogicalTableExportLimits, LogicalTableProjection, check_cancelled,
+    map_writer_error,
+};
+
+const PAYLOAD_BYTES: usize = 64 * 1024 * 1024;
+
+/// Request-wide downstream payload and writer admission, separate from query memory.
+pub(crate) struct ExportWriteBudget {
+    writers: Arc<Semaphore>,
+    bytes: Arc<Semaphore>,
+}
+
+impl ExportWriteBudget {
+    pub(crate) fn new(parallelism: usize) -> Arc<Self> {
+        Arc::new(Self {
+            writers: Arc::new(Semaphore::new(parallelism.max(1))),
+            bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)),
+        })
+    }
+
+    #[cfg(test)]
+    pub(crate) fn available(&self) -> (usize, usize) {
+        (
+            self.writers.available_permits(),
+            self.bytes.available_permits(),
+        )
+    }
+
+    pub(crate) async fn writer(&self, token: &CancellationToken) -> Result<OwnedSemaphorePermit> {
+        tokio::select! {
+            biased;
+            _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(),
+            permit = self.writers.clone().acquire_owned() => permit.map_err(|_| error::UnexpectedSnafu { violated: "writer budget closed" }.build()),
+        }
+    }
+
+    pub(crate) async fn reserve(
+        &self,
+        size: usize,
+        token: &CancellationToken,
+    ) -> Result<OwnedSemaphorePermit> {
+        ensure!(
+            size <= PAYLOAD_BYTES,
+            error::LogicalTableExportResourceSnafu {
+                reason: "batch backing buffers exceed request payload budget"
+            }
+        );
+        tokio::select! {
+            biased;
+            _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(),
+            permit = self.bytes.clone().acquire_many_owned(size as u32) => permit.map_err(|_| error::UnexpectedSnafu { violated: "payload budget closed" }.build()),
+        }
+    }
+
+    /// Includes full retained input allocations, conversion capacity growth and
+    /// array metadata. Shared backing is conservatively charged for each payload.
+    pub(crate) fn conversion_budget(
+        batch: &RecordBatch,
+        requested: usize,
+    ) -> Result<(usize, usize)> {
+        let backing = batch.get_array_memory_size();
+        let overhead = batch.num_columns().saturating_mul(1024);
+        let available = PAYLOAD_BYTES
+            .saturating_sub(backing)
+            .saturating_sub(overhead)
+            / 2;
+        let conversion = requested.min(available);
+        ensure!(
+            conversion > 0,
+            error::LogicalTableExportResourceSnafu {
+                reason: "batch backing buffers exceed request payload budget"
+            }
+        );
+        Ok((conversion, backing.saturating_add(overhead)))
+    }
+}
+
+pub(crate) struct Payload {
+    pub(crate) batch: RecordBatch,
+    pub(crate) permit: OwnedSemaphorePermit,
+}
+
+pub(crate) struct TableWriters {
+    current: Option<(u32, mpsc::Sender<Payload>)>,
+    tasks: FuturesUnordered<JoinHandle<Result<()>>>,
+    budget: Arc<ExportWriteBudget>,
+}
+
+impl TableWriters {
+    pub(crate) fn new(budget: Arc<ExportWriteBudget>) -> Self {
+        Self {
+            current: None,
+            tasks: FuturesUnordered::new(),
+            budget,
+        }
+    }
+
+    pub(crate) fn table_id(&self) -> Option<u32> {
+        self.current.as_ref().map(|(id, _)| *id)
+    }
+
+    pub(crate) fn close_input(&mut self) {
+        self.current = None;
+    }
+
+    pub(crate) async fn open(
+        &mut self,
+        id: u32,
+        table: &LogicalTableProjection,
+        store: &ObjectStore,
+        limits: LogicalTableExportLimits,
+        token: &CancellationToken,
+    ) -> Result<()> {
+        // EOF must precede acquiring the next slot, including when P is one.
+        self.close_input();
+        while let Some(Some(result)) = self.tasks.next().now_or_never() {
+            result.context(error::JoinTaskSnafu)??;
+        }
+        let permit = self.budget.writer(token).await?;
+        let writer = ActiveWriter::open(table, store, limits).await?;
+        let (sender, receiver) = mpsc::channel(2);
+        self.current = Some((id, sender));
+        let token = token.clone();
+        self.tasks.push(common_runtime::spawn_global(async move {
+            let _permit = permit;
+            let guard = token.clone().drop_guard();
+            let result = run_writer(writer, receiver, &token).await;
+            if result.is_ok() {
+                guard.disarm();
+            }
+            result
+        }));
+        Ok(())
+    }
+
+    pub(crate) async fn send(&self, payload: Payload, token: &CancellationToken) -> Result<()> {
+        let (_, sender) = self.current.as_ref().context(error::UnexpectedSnafu {
+            violated: "missing logical writer",
+        })?;
+        tokio::select! {
+            biased;
+            _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(),
+            result = sender.send(payload) => result.map_err(|_| error::LogicalTableExportCancelledSnafu.build()),
+        }
+    }
+
+    pub(crate) async fn drain(
+        &mut self,
+        result: Result<()>,
+        token: &CancellationToken,
+    ) -> Result<()> {
+        self.close_input();
+        let mut first_error = result.err();
+        if first_error.is_some() {
+            token.cancel();
+        }
+        while let Some(result) = self.tasks.next().await {
+            if let Err(err) = result.context(error::JoinTaskSnafu).and_then(|r| r) {
+                retain_error(&mut first_error, err);
+                token.cancel();
+            }
+        }
+        match first_error {
+            Some(err) => Err(err),
+            None => Ok(()),
+        }
+    }
+}
+
+pub(crate) fn retain_error(first: &mut Option<error::Error>, error: error::Error) {
+    if first.as_ref().is_none_or(|err| {
+        matches!(
+            err,
+            error::Error::LogicalTableExportCancelled { .. }
+                | error::Error::DatabaseExportCancelled { .. }
+        )
+    }) {
+        *first = Some(error);
+    }
+}
+
+async fn run_writer(
+    mut writer: ActiveWriter,
+    mut receiver: mpsc::Receiver<Payload>,
+    token: &CancellationToken,
+) -> Result<()> {
+    let result = async {
+        loop {
+            let payload = tokio::select! {
+                biased;
+                _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(),
+                payload = receiver.recv() => payload,
+            };
+            let Some(Payload { batch, permit }) = payload else {
+                break;
+            };
+            let result = writer.writer.write(batch, Some(token)).await;
+            drop(permit);
+            result.map_err(|error| map_writer_error(error, &writer.path))?;
+        }
+        writer
+            .writer
+            .finish(Some(token))
+            .await
+            .map_err(|error| map_writer_error(error, &writer.path))?;
+        check_cancelled(token)
+    }
+    .await;
+    if result.is_err() {
+        token.cancel();
+        receiver.close();
+        drop(receiver);
+        if let Err(error) = writer.writer.abort().await {
+            common_telemetry::warn!(error; "Failed to abort Metric export file");
+        }
+    }
+    result
+}