From 683657ad32467145c336ae19ca9b12ed491b2e6d Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 10:43:16 +0800 Subject: [PATCH 1/8] fix: preserve objects after conditional Metric export collisions Signed-off-by: jeremyhi --- src/common/datasource/src/parquet_writer.rs | 113 +++++++++++++++++- src/object-store/src/secure_fs.rs | 40 +++++++ src/object-store/tests/object_store_test.rs | 74 +++++++++++- .../src/statement/export_logical_tables.rs | 32 +++-- 4 files changed, 242 insertions(+), 17 deletions(-) diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index ae1fcec04a9..955df4032cb 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -29,6 +29,13 @@ use tokio_util::sync::CancellationToken; use crate::DEFAULT_WRITE_BUFFER_SIZE; use crate::error::{self, Result}; +/// Destination creation policy; conditional failures never authorize path deletion. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ParquetCreationPolicy { + Overwrite, + IfNotExists, +} + /// Limits for one Parquet file. Flush thresholds are not hard memory caps. #[derive(Clone, Copy, Debug)] pub struct ParquetWriterLimits { @@ -49,6 +56,7 @@ pub struct ParquetFileWriter { store: ObjectStore, path: String, limits: Option, + creation: ParquetCreationPolicy, } impl ParquetFileWriter { @@ -60,6 +68,26 @@ impl ParquetFileWriter { path: &str, concurrency: usize, limits: Option, + ) -> Result { + Self::open_with_creation( + schema, + store, + path, + concurrency, + limits, + ParquetCreationPolicy::Overwrite, + ) + .await + } + + /// Open with a conditional policy only when the caller verified backend support. + pub async fn open_with_creation( + schema: SchemaRef, + store: ObjectStore, + path: &str, + concurrency: usize, + limits: Option, + creation: ParquetCreationPolicy, ) -> Result { let mut props = WriterProperties::builder() .set_compression(Compression::ZSTD(ZstdLevel::default())) @@ -90,6 +118,7 @@ impl ParquetFileWriter { .writer_with(path) .concurrent(concurrency) .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .if_not_exists(creation == ParquetCreationPolicy::IfNotExists) .await .context(error::WriteObjectSnafu { path })?; Ok(Self { @@ -98,6 +127,7 @@ impl ParquetFileWriter { store, path: path.to_owned(), limits, + creation, }) } @@ -191,13 +221,14 @@ impl ParquetFileWriter { Ok(()) } - /// Abort an exclusively owned file after all in-flight operations have completed. - /// If the backend cannot abort, deletion assumes this attempt owns the path. + /// Abort after all in-flight operations complete. Overwrite callers own the + /// path; conditional callers delegate cleanup exclusively to the backend. pub async fn abort(mut self) -> Result<()> { let result = self.sink.abort().await; - if result - .as_ref() - .is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported) + if self.creation == ParquetCreationPolicy::Overwrite + && result + .as_ref() + .is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported) { let store = self.store.clone(); let path = self.path.clone(); @@ -288,6 +319,78 @@ mod tests { .unwrap() } + #[tokio::test] + async fn conditional_abort_preserves_collisions_and_ambiguous_commits() { + use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio}; + + struct AmbiguousCommit(oio::Writer); + impl oio::Write for AmbiguousCommit { + 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 { + self.0.close().await?; + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "lost close reply", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "cannot abort", + )) + } + } + + for existing in [false, true] { + let directory = common_test_util::temp_dir::create_temp_dir("conditional_parquet"); + let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) + .unwrap() + .build_operator(); + let path = "conditional.parquet"; + if existing { + store.write(path, "original").await.unwrap(); + } + let factory: MockWriterFactory = + Arc::new(|_, _, writer| Box::new(AmbiguousCommit(writer))); + let store = store.layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + let mut writer = ParquetFileWriter::open_with_creation( + batch().schema(), + store.clone(), + path, + 1, + None, + ParquetCreationPolicy::IfNotExists, + ) + .await + .unwrap(); + writer.write(batch(), None).await.unwrap(); + assert!(writer.finish(None).await.is_err()); + assert!(writer.abort().await.is_err()); + if existing { + assert_eq!( + store.read(path).await.unwrap().to_bytes(), + Bytes::from_static(b"original") + ); + } else { + assert_eq!( + read(&store, path) + .await + .metadata() + .file_metadata() + .num_rows(), + 4 + ); + } + } + } + #[tokio::test] async fn row_and_byte_limits_split_batches() { let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 18afa3ac820..5b81d286e50 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -274,6 +274,7 @@ impl Service for SecureFsBackend { path, args, file: None, + closing: false, }) } @@ -387,6 +388,7 @@ struct SecureFsWriter { path: PathBuf, args: OpWrite, file: Option, + closing: bool, } impl SecureFsWriter { @@ -438,6 +440,7 @@ impl oio::Write for SecureFsWriter { } async fn close(&mut self) -> Result { + self.closing = true; let file = self.ensure_file().await?; file.flush().await.map_err(new_std_io_error)?; file.sync_all().await.map_err(new_std_io_error)?; @@ -450,6 +453,23 @@ impl oio::Write for SecureFsWriter { } async fn abort(&mut self) -> Result<()> { + if self.args.if_not_exists() { + // A failed exclusive create owns no file. Once close starts, preserve + // potentially committed output for the caller's deliberate retry. + if let Some(mut file) = self.file.take() { + file.flush().await.map_err(new_std_io_error)?; + drop(file); + if !self.closing { + let root = self.root.clone(); + let path = self.path.clone(); + common_runtime::spawn_blocking_global(move || root.dir.remove_file(path)) + .await + .map_err(new_task_join_error)? + .map_err(new_std_io_error)?; + } + } + return Ok(()); + } Err(Error::new( ErrorKind::Unsupported, "filesystem writes cannot be aborted without atomic writes", @@ -799,6 +819,26 @@ mod tests { assert!(!temp_dir.path().join("nested").exists()); } + #[tokio::test] + async fn test_conditional_abort_only_removes_owned_partial_file() { + let temp_dir = create_temp_dir("secure_fs_conditional_abort"); + let operator = SecureFsRoot::open(temp_dir.path()) + .unwrap() + .build_operator(); + for started in [false, true] { + let mut writer = operator + .writer_with("partial") + .if_not_exists(true) + .await + .unwrap(); + if started { + writer.write(Bytes::from_static(b"partial")).await.unwrap(); + } + writer.abort().await.unwrap(); + assert!(!operator.exists("partial").await.unwrap()); + } + } + #[tokio::test] async fn test_writer_abort_is_unsupported_without_atomic_write() { let temp_dir = create_temp_dir("secure_fs_writer_abort"); diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index 9cf03fe42c5..c1537b5b230 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -284,12 +284,15 @@ async fn test_s3_backend() -> Result<()> { let root = uuid::Uuid::new_v4().to_string(); - let builder = S3::default() + let mut builder = S3::default() .root(&root) .access_key_id(&env::var("GT_S3_ACCESS_KEY_ID")?) .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") { + builder = builder.endpoint(&endpoint); + } let store = ObjectStore::new(builder).unwrap(); let store = object_store::util::with_instrument_layers(store, false); @@ -298,6 +301,7 @@ async fn test_s3_backend() -> Result<()> { test_object_crud(&store).await?; test_object_list(&store).await?; test_object_list_start_after(&store).await?; + test_conditional_creation(&store).await?; assert_opendal_metrics(); guard.remove_all().await?; } @@ -305,6 +309,74 @@ async fn test_s3_backend() -> Result<()> { Ok(()) } +async fn test_conditional_creation(store: &ObjectStore) -> Result<()> { + const PART: usize = 8 * 1024 * 1024; + for size in [32, 2 * PART + 17] { + let payload = Bytes::from(vec![42; size]); + for preexisting in [false, true] { + let path = format!("conditional-{size}-{preexisting}"); + if preexisting { + store.write(&path, "original").await?; + } + let attempt = |value: Bytes| { + let path = &path; + async move { + let mut writer = store + .writer_with(path) + .if_not_exists(true) + .concurrent(1) + .chunk(PART) + .await + .unwrap(); + let result = async { + for offset in (0..value.len()).step_by(PART) { + writer + .write(value.slice(offset..(offset + PART).min(value.len()))) + .await?; + } + writer.close().await + } + .await; + if result.is_err() { + writer.abort().await.unwrap(); + } + result + } + }; + let (a, b) = tokio::join!(attempt(payload.clone()), attempt(payload.clone())); + assert_eq!( + usize::from(a.is_ok()) + usize::from(b.is_ok()), + usize::from(!preexisting) + ); + for error in [a.err(), b.err()].into_iter().flatten() { + assert!( + matches!( + error.kind(), + object_store::ErrorKind::ConditionNotMatch + | object_store::ErrorKind::AlreadyExists + ), + "{error:?}" + ); + } + let expected = if preexisting { + Bytes::from_static(b"original") + } else { + payload.clone() + }; + assert_eq!(store.read(&path).await?.to_bytes(), expected); + store.delete(&path).await?; + } + } + Ok(()) +} + +#[tokio::test] +async fn test_secure_fs_conditional_creation() -> Result<()> { + let dir = TempDir::new()?; + let store = object_store::secure_fs::SecureFsRoot::open(dir.path())?.build_operator(); + test_conditional_creation(&store).await +} + #[tokio::test] async fn test_oss_backend() -> Result<()> { common_telemetry::init_default_ut_logging(); diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 7547a3f54de..f52f4e96d90 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -28,7 +28,9 @@ use arrow::datatypes::{DataType, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; use common_datasource::object_store::build_backend_for_write; -use common_datasource::parquet_writer::{ParquetFileWriter, ParquetWriterLimits}; +use common_datasource::parquet_writer::{ + ParquetCreationPolicy, ParquetFileWriter, ParquetWriterLimits, +}; use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; use common_query::OutputData; use common_recordbatch::SendableRecordBatchStream; @@ -493,21 +495,29 @@ impl ActiveWriter { limits: LogicalTableExportLimits, ) -> Result { let path = table.output.path.clone(); - ensure!( - !store - .exists(&path) - .await - .context(error::ReadObjectSnafu { path: &path })?, - InvalidLogicalTableExportSnafu { - reason: format!("output already exists: {path}") - } - ); - let writer = ParquetFileWriter::open( + let conditional = store.info().full_capability().write_with_if_not_exists; + if !conditional { + ensure!( + !store + .exists(&path) + .await + .context(error::ReadObjectSnafu { path: &path })?, + InvalidLogicalTableExportSnafu { + reason: format!("output already exists: {path}") + } + ); + } + let writer = ParquetFileWriter::open_with_creation( table.schema.clone(), store.clone(), &path, 1, Some(limits.writer), + if conditional { + ParquetCreationPolicy::IfNotExists + } else { + ParquetCreationPolicy::Overwrite + }, ) .await .map_err(|error| map_writer_error(error, &path))?; From 7bff9a9fef0929d033cef633c9caa79370fe728f Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 10:57:33 +0800 Subject: [PATCH 2/8] perf: share Metric export writer and payload budgets Signed-off-by: jeremyhi --- src/common/datasource/src/parquet_writer.rs | 79 +++- src/object-store/tests/object_store_test.rs | 2 +- src/operator/src/statement/copy_table_to.rs | 253 +++++++++- src/operator/src/statement/export_database.rs | 91 ++-- .../src/statement/export_logical_tables.rs | 178 ++++--- .../statement/export_logical_tables/tests.rs | 443 +++++++++++++----- .../export_logical_tables/writers.rs | 245 ++++++++++ 7 files changed, 1067 insertions(+), 224 deletions(-) create mode 100644 src/operator/src/statement/export_logical_tables/writers.rs diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index 955df4032cb..a760de187d2 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -186,9 +186,12 @@ impl ParquetFileWriter { } async fn write_bytes(&mut self, bytes: Vec) -> 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 { + async fn close( + &mut self, + ) -> object_store::Result { 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>>); + 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 { + 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::>(); + 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::() > limit); + assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}"); + let actual = read(&store, "large.parquet") + .await + .build() + .unwrap() + .collect::, _>>() + .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(); diff --git a/src/object-store/tests/object_store_test.rs b/src/object-store/tests/object_store_test.rs index c1537b5b230..07b185de6cd 100644 --- a/src/object-store/tests/object_store_test.rs +++ b/src/object-store/tests/object_store_test.rs @@ -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); } diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 21b82f51a3b..a415a7c3904 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -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 { + 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 { 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, ) -> Result { - 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, + managed: Option<(&ExportWriteBudget, &CancellationToken)>, + ) -> Result { + 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 { + 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::::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::, _>>() + .unwrap(); + outputs.push((schema, values)); + } + assert_eq!(outputs[0], outputs[1]); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + } } } diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index 7c5de554487..ec0d8f08b48 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -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>>( 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>>( }; 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, } } diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index f52f4e96d90..a3f025d3627 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -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, } -struct LogicalTableProjection { +pub(crate) struct LogicalTableProjection { output: DatabaseExportFile, schema: SchemaRef, projection: Vec, @@ -308,6 +311,31 @@ impl StatementExecutor { limits: LogicalTableExportLimits, cancellation: &CancellationToken, query_ctx: QueryContextRef, + ) -> Result { + 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, + time_range: Option<&TimestampRange>, + limits: LogicalTableExportLimits, + cancellation: &CancellationToken, + query_ctx: QueryContextRef, + budget: Arc, ) -> Result { 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 { - 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, +) -> Result { + 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, + writers: &mut TableWriters, + budget: &ExportWriteBudget, ) -> Result { 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 { 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, - 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::>>()?; 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 { Ok(bytes.saturating_add(16)) } -fn rows_within_budget( +pub(crate) fn rows_within_budget( batch: &RecordBatch, start: usize, end: usize, budget: usize, -) -> Result { +) -> 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)] diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 4576c1e395d..74311963828 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -128,52 +128,56 @@ async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec) #[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::() - .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::() + .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::() + }) + .collect::>(); + 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::(), 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, + started: Arc, + release: Arc, + } + 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 { + 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)); + } + } +} diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs new file mode 100644 index 00000000000..c19ee4824fd --- /dev/null +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -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, + bytes: Arc, +} + +impl ExportWriteBudget { + pub(crate) fn new(parallelism: usize) -> Arc { + 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 { + 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 { + 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)>, + tasks: FuturesUnordered>>, + budget: Arc, +} + +impl TableWriters { + pub(crate) fn new(budget: Arc) -> Self { + Self { + current: None, + tasks: FuturesUnordered::new(), + budget, + } + } + + pub(crate) fn table_id(&self) -> Option { + 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) { + 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, + 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 +} From c8244ab9c7eb81973dcf70d9790fd837800ce17b Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 11:04:29 +0800 Subject: [PATCH 3/8] test: cover bounded Metric export conversion and mixed roundtrips Signed-off-by: jeremyhi --- src/operator/src/statement/copy_table_to.rs | 17 ++-- .../src/statement/export_logical_tables.rs | 81 ++++++++++++++++--- .../statement/export_logical_tables/tests.rs | 79 ++++++++++++++++++ .../export_logical_tables/writers.rs | 7 +- .../tests/export_logical_tables.rs | 41 ++++++---- 5 files changed, 189 insertions(+), 36 deletions(-) diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index a415a7c3904..5f5b73003db 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -47,7 +47,7 @@ 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, + LogicalTableExportLimits, expand_export_batch, map_writer_error, rows_within_budget, }; // The buffer size should be greater than 5MB (minimum multipart upload size). @@ -242,9 +242,7 @@ pub(crate) async fn stream_to_managed_parquet( budget: &ExportWriteBudget, token: &CancellationToken, ) -> Result { - use common_recordbatch::{ - RecordBatch, map_dictionary_to_values, map_dictionary_to_values_schema, - }; + use common_recordbatch::{RecordBatch, 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()); @@ -284,7 +282,7 @@ pub(crate) async fn stream_to_managed_parquet( }) .await .context(error::JoinTaskSnafu)??; - let reservation = retained.saturating_add(estimated.saturating_mul(2 * expansion)); + let reservation = retained.saturating_add(estimated.saturating_mul(4 * expansion)); let permit = budget.reserve(reservation, token).await?; let batch = batch.clone(); let (original, expanded_schema, output_schema) = ( @@ -298,8 +296,13 @@ pub(crate) async fn stream_to_managed_parquet( batch.slice(offset, len), ); if expand { - batch = map_dictionary_to_values(batch, &original, &expanded_schema) - .context(error::BuildRecordBatchSnafu)?; + batch = RecordBatch::from_df_record_batch( + expanded_schema.clone(), + expand_export_batch( + &batch.into_df_record_batch(), + expanded_schema.arrow_schema().clone(), + )?, + ); } if json { batch = map_json_type_to_string(batch, &expanded_schema, &output_schema) diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index a3f025d3627..1b8d012fc65 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -24,8 +24,9 @@ pub(crate) mod writers; use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; -use arrow::array::{Array, AsArray, UInt32Array}; -use arrow::compute::cast; +use arrow::array::{Array, ArrayRef, AsArray, ListArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{cast, take}; use arrow::datatypes::{DataType, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; @@ -586,17 +587,11 @@ async fn expand_bounded_slice( }) .await .context(error::JoinTaskSnafu)??; - let reservation = retained.saturating_add(estimated.saturating_mul(2)); + let reservation = retained.saturating_add(estimated.saturating_mul(4)); let permit = budget.reserve(reservation, cancellation).await?; common_runtime::spawn_blocking_global(move || { let slice = batch.slice(start, len); - let arrays = slice - .columns() - .iter() - .zip(schema.fields()) - .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu)) - .collect::>>()?; - let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?; + let expanded = expand_export_batch(&slice, schema)?; ensure!( expanded.get_array_memory_size() <= reservation, LogicalTableExportResourceSnafu { @@ -615,6 +610,70 @@ async fn expand_bounded_slice( .context(error::JoinTaskSnafu)? } +/// Limits nested casts to selected children: Arrow's List cast otherwise expands +/// the entire values array, even when the parent has been sliced. +pub(crate) fn expand_export_batch(batch: &RecordBatch, schema: SchemaRef) -> Result { + let arrays = batch + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu) +} + +fn expand_export_array(array: &ArrayRef, target: &DataType) -> Result { + if array.data_type() == target { + return Ok(array.clone()); + } + match (array.data_type(), target) { + (DataType::Dictionary(_, _), _) => { + downcast_dictionary_array! { + array => { + // Select dictionary entries before recursively converting nested values. + let selected = take(array.values().as_ref(), array.keys(), None) + .context(error::ComputeArrowSnafu)?; + expand_export_array(&selected, target) + }, + _ => error::UnexpectedSnafu { violated: "invalid dictionary array" }.fail(), + } + } + (DataType::List(_), DataType::List(field)) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let start = offsets[0]; + let end = offsets[list.len()]; + let values = list.values().slice(start as usize, (end - start) as usize); + let values = expand_export_array(&values, field.data_type())?; + let offsets = OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - start) + .collect::>() + .into(), + ); + Ok(Arc::new( + ListArray::try_new(field.clone(), offsets, values, list.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + (DataType::Struct(_), DataType::Struct(fields)) => { + let array = array.as_struct(); + let columns = array + .columns() + .iter() + .zip(fields) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + Ok(Arc::new( + StructArray::try_new(fields.clone(), columns, array.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + _ => cast(array, target).context(error::ComputeArrowSnafu), + } +} + pub(crate) fn map_writer_error( source: common_datasource::error::Error, path: &str, @@ -639,7 +698,7 @@ pub(crate) fn map_writer_error( // Count only selected logical values before dictionary expansion, including nested // histogram lists/structs. Offset and validity overhead is charged per value. fn estimate_value_size(array: &dyn Array, row: usize) -> Result { - if array.is_null(row) { + if array.is_null(row) && !matches!(array.data_type(), DataType::Struct(_) | DataType::List(_)) { return Ok(32); } let bytes = match array.data_type() { diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 74311963828..dc0e71c9cb8 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -799,3 +799,82 @@ async fn groups_and_ordinary_files_share_writer_admission_and_drain() { } } } + +#[tokio::test] +async fn nested_dictionary_conversion_limits_child_ranges_and_charges_null_parents() { + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0; 1024]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let list = ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0i32, 1023, 1024].into()), + values, + None, + ); + let input = RecordBatch::try_from_iter([("list", Arc::new(list) as ArrayRef)]).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "list", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + )])); + let budget = ExportWriteBudget::new(1); + let (payload, rows) = expand_bounded_slice( + input, + schema, + 1, + 2, + 8192, + &budget, + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(rows, 1); + assert!(payload.batch.get_array_memory_size() < 16384); + let list = payload.batch.column(0).as_list::(); + assert_eq!(list.value_offsets(), &[0, 1]); + assert_eq!(list.values().as_string::().value(0), "x".repeat(4096)); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let structure = StructArray::new( + vec![Arc::new(Field::new( + "child", + values.data_type().clone(), + true, + ))] + .into(), + vec![values], + Some(arrow::buffer::NullBuffer::from(vec![false])), + ); + let input = RecordBatch::try_from_iter([("struct", Arc::new(structure) as ArrayRef)]).unwrap(); + assert!(rows_within_budget(&input, 0, 1, 128).is_err()); +} + +#[tokio::test] +async fn completed_table_workers_are_reaped_during_admission() { + for parallelism in [1, 4] { + let mut unit = unit(); + let file = unit.logical_tables.get_mut(&1025).unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + for id in 0..64 { + file.output.path = format!("empty-{id}.parquet"); + writers + .open(id, file, &store, export_limits(), &token) + .await + .unwrap(); + assert!(writers.pending_tasks() <= parallelism + 1); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(writers.pending_tasks(), 0); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } +} diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs index c19ee4824fd..0c818b9119c 100644 --- a/src/operator/src/statement/export_logical_tables/writers.rs +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -90,7 +90,7 @@ impl ExportWriteBudget { let available = PAYLOAD_BYTES .saturating_sub(backing) .saturating_sub(overhead) - / 2; + / 4; let conversion = requested.min(available); ensure!( conversion > 0, @@ -122,6 +122,11 @@ impl TableWriters { } } + #[cfg(test)] + pub(crate) fn pending_tasks(&self) -> usize { + self.tasks.len() + } + pub(crate) fn table_id(&self) -> Option { self.current.as_ref().map(|(id, _)| *id) } diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 92888cced3c..7930abb7c39 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -228,7 +228,7 @@ fn database_export_request(directory: &std::path::Path) -> table::requests::Copy } } -async fn database_export_roundtrip(instance: &Arc) { +async fn database_export_roundtrip(instance: &Arc, parallelism: usize) { let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); let (first_logical_table_names, _, renamed_physical_table) = create_metric_export_source_tables(instance, "db_a", "dense").await; @@ -253,7 +253,9 @@ async fn database_export_roundtrip(instance: &Arc) { ]; let mut names = selected.clone(); names.extend([renamed_physical_table, "dashboard".into()]); - let req = database_export_request(&destination.path().join("data")); + let mut req = database_export_request(&destination.path().join("data")); + req.with + .insert("parallelism".into(), parallelism.to_string()); let executor = instance.statement_executor(); let captured = executor .capture_database_export_tables(&req, None, &QueryContext::arc()) @@ -445,26 +447,31 @@ async fn database_export_roundtrip(instance: &Arc) { #[tokio::test(flavor = "multi_thread")] async fn database_export_standalone_roundtrip() { - let standalone = GreptimeDbStandaloneBuilder::new("database_export") - .build() - .await; - database_export_roundtrip(standalone.fe_instance()).await; + for parallelism in [1, 4] { + let standalone = + GreptimeDbStandaloneBuilder::new(&format!("database_export_{parallelism}")) + .build() + .await; + database_export_roundtrip(standalone.fe_instance(), parallelism).await; + } } #[tokio::test(flavor = "multi_thread")] async fn database_export_distributed_roundtrip() { - let cluster = GreptimeDbClusterBuilder::new("database_export") - .await - .with_datanodes(2) - .with_local_file_access( - common_datasource::object_store::LocalFileAccess::sandboxed( - common_test_util::find_workspace_path("."), + for parallelism in [1, 4] { + let cluster = GreptimeDbClusterBuilder::new(&format!("database_export_{parallelism}")) + .await + .with_datanodes(2) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed( + common_test_util::find_workspace_path("."), + ) + .unwrap(), ) - .unwrap(), - ) - .build(false) - .await; - database_export_roundtrip(cluster.fe_instance()).await; + .build(false) + .await; + database_export_roundtrip(cluster.fe_instance(), parallelism).await; + } } #[tokio::test(flavor = "multi_thread")] From fae799a98c22d9aebe0bb53b34a1c85e27698c6c Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 11:30:55 +0800 Subject: [PATCH 4/8] fix: preserve export output after ambiguous close Signed-off-by: jeremyhi --- src/common/datasource/src/parquet_writer.rs | 18 +++-- src/object-store/src/secure_fs.rs | 62 ++++++++++++---- .../statement/export_logical_tables/tests.rs | 74 +++++++++++++++++++ 3 files changed, 136 insertions(+), 18 deletions(-) diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index a760de187d2..a3b5680b8d3 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -57,6 +57,7 @@ pub struct ParquetFileWriter { path: String, limits: Option, creation: ParquetCreationPolicy, + close_started: bool, } impl ParquetFileWriter { @@ -128,6 +129,7 @@ impl ParquetFileWriter { path: path.to_owned(), limits, creation, + close_started: false, }) } @@ -216,6 +218,7 @@ impl ParquetFileWriter { .context(error::JoinHandleSnafu)??; self.write_bytes(bytes).await?; check_cancelled(cancellation)?; + self.close_started = true; self.sink .close() .await @@ -224,11 +227,12 @@ impl ParquetFileWriter { Ok(()) } - /// Abort after all in-flight operations complete. Overwrite callers own the - /// path; conditional callers delegate cleanup exclusively to the backend. + /// Abort after all in-flight operations complete. Preserve ambiguous commits; + /// conditional callers delegate cleanup exclusively to the backend. pub async fn abort(mut self) -> Result<()> { let result = self.sink.abort().await; if self.creation == ParquetCreationPolicy::Overwrite + && !self.close_started && result .as_ref() .is_err_and(|e| e.kind() == object_store::ErrorKind::Unsupported) @@ -323,7 +327,7 @@ mod tests { } #[tokio::test] - async fn conditional_abort_preserves_collisions_and_ambiguous_commits() { + async fn abort_preserves_collisions_and_ambiguous_commits() { use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio}; struct AmbiguousCommit(oio::Writer); @@ -348,7 +352,11 @@ mod tests { } } - for existing in [false, true] { + for (creation, existing) in [ + (ParquetCreationPolicy::Overwrite, false), + (ParquetCreationPolicy::IfNotExists, false), + (ParquetCreationPolicy::IfNotExists, true), + ] { let directory = common_test_util::temp_dir::create_temp_dir("conditional_parquet"); let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() @@ -371,7 +379,7 @@ mod tests { path, 1, None, - ParquetCreationPolicy::IfNotExists, + creation, ) .await .unwrap(); diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 5b81d286e50..d4f6b156046 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -453,11 +453,14 @@ impl oio::Write for SecureFsWriter { } async fn abort(&mut self) -> Result<()> { + // Tokio writes may finish in the blocking pool after write_all returns. + if let Some(file) = self.file.as_mut() { + file.flush().await.map_err(new_std_io_error)?; + } if self.args.if_not_exists() { // A failed exclusive create owns no file. Once close starts, preserve // potentially committed output for the caller's deliberate retry. - if let Some(mut file) = self.file.take() { - file.flush().await.map_err(new_std_io_error)?; + if let Some(file) = self.file.take() { drop(file); if !self.closing { let root = self.root.clone(); @@ -839,17 +842,50 @@ mod tests { } } - #[tokio::test] - async fn test_writer_abort_is_unsupported_without_atomic_write() { - let temp_dir = create_temp_dir("secure_fs_writer_abort"); - let operator = SecureFsRoot::open(temp_dir.path()) + #[test] + fn test_writer_abort_drains_before_reporting_unsupported() { + use std::path::PathBuf; + + use opendal::Buffer; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + + use super::SecureFsWriter; + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() .unwrap() - .build_operator(); - let mut writer = operator.writer("partial").await.unwrap(); - writer.write(Bytes::from_static(b"partial")).await.unwrap(); - - let error = writer.abort().await.unwrap_err(); - - assert_eq!(ErrorKind::Unsupported, error.kind()); + .block_on(async { + let temp_dir = create_temp_dir("secure_fs_writer_abort"); + let mut writer = SecureFsWriter { + root: SecureFsRoot::open(temp_dir.path()).unwrap(), + path: PathBuf::from("partial"), + args: OpWrite::default(), + file: None, + closing: false, + }; + writer.ensure_file().await.unwrap(); + let (started, ready) = tokio::sync::oneshot::channel(); + let (release, blocked) = std::sync::mpsc::channel(); + let blocking = tokio::task::spawn_blocking(move || { + started.send(()).unwrap(); + blocked.recv().unwrap(); + }); + ready.await.unwrap(); + writer.write(Buffer::from("partial")).await.unwrap(); + let abort = writer.abort(); + tokio::pin!(abort); + let pending = futures::poll!(&mut abort).is_pending(); + release.send(()).unwrap(); + assert!(pending); + assert_eq!(ErrorKind::Unsupported, abort.await.unwrap_err().kind()); + blocking.await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join("partial")).unwrap(), + b"partial" + ); + }); } } diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index dc0e71c9cb8..964489e7ea6 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -562,6 +562,80 @@ async fn cancellation_drains_storage_and_preserves_committed_files() { struct FailedAbortWriter(object_store::layers::mock::oio::Writer); +#[tokio::test] +async fn metric_fallback_and_ordinary_preserve_ambiguous_commits() { + use object_store::layers::CapabilityOverrideLayer; + use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + + use crate::statement::copy_table_to::stream_to_managed_parquet; + + struct AmbiguousCommit(oio::Writer); + impl oio::Write for AmbiguousCommit { + 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 { + self.0.close().await?; + Err(object_store::Error::new( + object_store::ErrorKind::Unexpected, + "lost close reply", + )) + } + async fn abort(&mut self) -> object_store::Result<()> { + Err(object_store::Error::new( + object_store::ErrorKind::Unsupported, + "cannot abort", + )) + } + } + let factory: MockWriterFactory = Arc::new(|_, args, inner| { + assert!(!args.if_not_exists()); + Box::new(AmbiguousCommit(inner)) + }); + let store = ObjectStore::new(object_store::services::Memory::default()) + .unwrap() + .layer(CapabilityOverrideLayer::new(|mut capability| { + capability.write_with_if_not_exists = false; + capability + })) + .layer( + MockLayerBuilder::default() + .writer_factory(factory) + .build() + .unwrap(), + ); + assert!(!store.info().capability().write_with_if_not_exists); + let rows = || stream(vec![batch(vec![Some(1025)], vec![Some("a")])]); + assert!( + export_stream( + &unit(), + rows(), + &store, + export_limits(), + &CancellationToken::new() + ) + .await + .is_err() + ); + let budget = ExportWriteBudget::new(1); + assert!( + stream_to_managed_parquet( + rows(), + store.clone(), + "ordinary.parquet", + &budget, + &CancellationToken::new() + ) + .await + .is_err() + ); + for path in ["cpu.v1.parquet", "ordinary.parquet"] { + let (_, batches) = read(&store, path).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); +} + impl object_store::layers::mock::oio::Write for FailedAbortWriter { async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { self.0.write(bytes).await From 0690ff2c55439ce1edbd081fece8976336675d0c Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 12:56:28 +0800 Subject: [PATCH 5/8] test: align HTTP export timeout with managed cancellation Signed-off-by: jeremyhi --- src/frontend/src/instance.rs | 33 +++++++------------ src/operator/src/statement/export_database.rs | 2 +- 2 files changed, 12 insertions(+), 23 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 22d2afaa7bc..73d5c370ece 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -2760,12 +2760,12 @@ mod tests { ) } - struct ReleasedExportSource { + struct CancellableExportSource { schema: GtSchemaRef, channels: std::sync::Mutex, oneshot::Receiver<()>)>>, } - impl DataSource for ReleasedExportSource { + impl DataSource for CancellableExportSource { fn get_stream( &self, _request: ScanRequest, @@ -2785,15 +2785,15 @@ mod tests { } #[tokio::test] - async fn test_metric_export_http_timeout_drains_ordinary_writer() { + async fn test_metric_export_http_timeout_cancels_ordinary_source() { let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout"); let (started_tx, started_rx) = oneshot::channel(); - let (release_tx, release_rx) = oneshot::channel(); + let (mut release_tx, release_rx) = oneshot::channel(); let info = test_table_info(1024, "source").unwrap(); let source = Arc::new(Table::new( Arc::new(info.clone()), FilterPushDownType::Unsupported, - Arc::new(ReleasedExportSource { + Arc::new(CancellableExportSource { schema: info.meta.schema.clone(), channels: std::sync::Mutex::new(Some((started_tx, release_rx))), }), @@ -2831,9 +2831,9 @@ mod tests { let server_task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); + let uri = reqwest::Url::from_directory_path(destination.path()).unwrap(); let sql = format!( - "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')", - destination.path().display() + "COPY DATABASE greptime.public TO '{uri}' WITH (experimental_metric_export='true')" ); let request = tokio::spawn(async move { reqwest::Client::builder() @@ -2852,21 +2852,10 @@ mod tests { .unwrap(); let response = request.await.unwrap(); assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT); - // A dropped ordinary stream would close this receiver before release. - release_tx.send(()).unwrap(); - let file = destination.path().join("source.parquet"); - tokio::time::timeout(Duration::from_secs(5), async { - loop { - if std::fs::read(&file) - .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1")) - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .unwrap(); + tokio::time::timeout(Duration::from_secs(5), release_tx.closed()) + .await + .unwrap(); + assert!(!destination.path().join("source.parquet").exists()); server_task.abort(); } diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs index ec0d8f08b48..dbe18ba763a 100644 --- a/src/operator/src/statement/export_database.rs +++ b/src/operator/src/statement/export_database.rs @@ -348,7 +348,7 @@ mod tests { } .fail(); } - // Ordinary COPY continues its I/O even when Metric jobs cancel. + // Already-started ordinary I/O drains after cancellation. token.cancelled().await; started.send(10).unwrap(); finish_io.acquire().await.unwrap().forget(); From 0c63fb8677ac82dd078124251cb1c20b7544ba04 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 16:51:30 +0800 Subject: [PATCH 6/8] test: trim redundant Metric export roundtrips Signed-off-by: jeremyhi --- .../statement/export_logical_tables/tests.rs | 98 +++++++++---------- .../tests/export_logical_tables.rs | 35 +++---- 2 files changed, 63 insertions(+), 70 deletions(-) diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 964489e7ea6..e883bbb89f6 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -128,56 +128,54 @@ async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec) #[tokio::test] async fn routes_across_batches_and_writes_empty_files() { - 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::() - .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)); - } + 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(4); + 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::() + .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(), (4, 64 * 1024 * 1024)); } #[tokio::test] diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 7930abb7c39..cc6354808a8 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -447,31 +447,26 @@ async fn database_export_roundtrip(instance: &Arc, parallelism: usize) #[tokio::test(flavor = "multi_thread")] async fn database_export_standalone_roundtrip() { - for parallelism in [1, 4] { - let standalone = - GreptimeDbStandaloneBuilder::new(&format!("database_export_{parallelism}")) - .build() - .await; - database_export_roundtrip(standalone.fe_instance(), parallelism).await; - } + let standalone = GreptimeDbStandaloneBuilder::new("database_export") + .build() + .await; + database_export_roundtrip(standalone.fe_instance(), 1).await; } #[tokio::test(flavor = "multi_thread")] async fn database_export_distributed_roundtrip() { - for parallelism in [1, 4] { - let cluster = GreptimeDbClusterBuilder::new(&format!("database_export_{parallelism}")) - .await - .with_datanodes(2) - .with_local_file_access( - common_datasource::object_store::LocalFileAccess::sandboxed( - common_test_util::find_workspace_path("."), - ) - .unwrap(), + let cluster = GreptimeDbClusterBuilder::new("database_export") + .await + .with_datanodes(2) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed( + common_test_util::find_workspace_path("."), ) - .build(false) - .await; - database_export_roundtrip(cluster.fe_instance(), parallelism).await; - } + .unwrap(), + ) + .build(false) + .await; + database_export_roundtrip(cluster.fe_instance(), 4).await; } #[tokio::test(flavor = "multi_thread")] From fac310a6c36b9744783d16e257deeee62864ebf1 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 23 Sep 2026 10:06:36 +0800 Subject: [PATCH 7/8] fix: bound pending Metric export writer handles Signed-off-by: jeremyhi --- .../export_logical_tables/writers.rs | 68 +++++++++++++++++-- 1 file changed, 64 insertions(+), 4 deletions(-) diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs index 0c818b9119c..75659e75811 100644 --- a/src/operator/src/statement/export_logical_tables/writers.rs +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -35,13 +35,16 @@ const PAYLOAD_BYTES: usize = 64 * 1024 * 1024; pub(crate) struct ExportWriteBudget { writers: Arc, bytes: Arc, + max_writers: usize, } impl ExportWriteBudget { pub(crate) fn new(parallelism: usize) -> Arc { + let max_writers = parallelism.max(1); Arc::new(Self { - writers: Arc::new(Semaphore::new(parallelism.max(1))), + writers: Arc::new(Semaphore::new(max_writers)), bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)), + max_writers, }) } @@ -145,10 +148,8 @@ impl TableWriters { ) -> 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?; + self.reap_for_admission(token).await?; let writer = ActiveWriter::open(table, store, limits).await?; let (sender, receiver) = mpsc::channel(2); self.current = Some((id, sender)); @@ -165,6 +166,22 @@ impl TableWriters { Ok(()) } + async fn reap_for_admission(&mut self, token: &CancellationToken) -> Result<()> { + while let Some(Some(result)) = self.tasks.next().now_or_never() { + result.context(error::JoinTaskSnafu)??; + } + // A worker can release its permit before its JoinHandle becomes ready. + while self.tasks.len() > self.budget.max_writers { + let result = tokio::select! { + biased; + _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), + result = self.tasks.next() => result.expect("writer task queue is not empty"), + }; + result.context(error::JoinTaskSnafu)??; + } + 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", @@ -248,3 +265,46 @@ async fn run_writer( } result } + +#[cfg(test)] +mod tests { + use tokio::sync::oneshot; + + use super::*; + + #[tokio::test] + async fn released_permit_does_not_hide_pending_join_handles() { + for parallelism in [1, 4] { + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + let mut resume = Vec::new(); + for _ in 0..=parallelism { + let permit = budget.writer(&token).await.unwrap(); + let (released_tx, released_rx) = oneshot::channel(); + let (resume_tx, resume_rx) = oneshot::channel(); + writers.tasks.push(common_runtime::spawn_global(async move { + drop(permit); + released_tx.send(()).unwrap(); + resume_rx.await.unwrap(); + Ok(()) + })); + released_rx.await.unwrap(); + resume.push(resume_tx); + } + assert_eq!(budget.available().0, parallelism); + + let mut reap = Box::pin(writers.reap_for_admission(&token)); + assert!(reap.as_mut().now_or_never().is_none()); + resume.pop().unwrap().send(()).unwrap(); + reap.await.unwrap(); + assert_eq!(writers.pending_tasks(), parallelism); + + for tx in resume { + tx.send(()).unwrap(); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(budget.available().0, parallelism); + } + } +} From bddd7cfff6126b3df62469197676d9103756f585 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 23 Sep 2026 12:49:05 +0800 Subject: [PATCH 8/8] fix: protect export files on failed writes Signed-off-by: jeremyhi --- src/object-store/src/secure_fs.rs | 85 +++++++++++++++++++-- src/operator/src/statement/copy_table_to.rs | 46 +++++++++++ src/operator/src/statement/database_copy.rs | 9 ++- 3 files changed, 132 insertions(+), 8 deletions(-) diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index d4f6b156046..11bfce7b04d 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -454,9 +454,10 @@ impl oio::Write for SecureFsWriter { async fn abort(&mut self) -> Result<()> { // Tokio writes may finish in the blocking pool after write_all returns. - if let Some(file) = self.file.as_mut() { - file.flush().await.map_err(new_std_io_error)?; - } + let flush = match self.file.as_mut() { + Some(file) => file.flush().await.map_err(new_std_io_error), + None => Ok(()), + }; if self.args.if_not_exists() { // A failed exclusive create owns no file. Once close starts, preserve // potentially committed output for the caller's deliberate retry. @@ -465,12 +466,18 @@ impl oio::Write for SecureFsWriter { if !self.closing { let root = self.root.clone(); let path = self.path.clone(); - common_runtime::spawn_blocking_global(move || root.dir.remove_file(path)) - .await - .map_err(new_task_join_error)? - .map_err(new_std_io_error)?; + let cleanup = + common_runtime::spawn_blocking_global(move || root.dir.remove_file(path)) + .await + .map_err(new_task_join_error) + .and_then(|result| result.map_err(new_std_io_error)); + return flush.and(cleanup); } } + return flush; + } + flush?; + if self.file.is_none() { return Ok(()); } Err(Error::new( @@ -842,6 +849,70 @@ mod tests { } } + #[tokio::test] + async fn test_overwrite_abort_preserves_unopened_destination() { + use std::path::PathBuf; + + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + + let temp_dir = create_temp_dir("secure_fs_unopened_abort"); + std::fs::write(temp_dir.path().join("existing"), b"original").unwrap(); + let mut writer = super::SecureFsWriter { + root: SecureFsRoot::open(temp_dir.path()).unwrap(), + path: PathBuf::from("existing"), + args: OpWrite::default(), + file: None, + closing: false, + }; + + writer.abort().await.unwrap(); + assert_eq!( + std::fs::read(temp_dir.path().join("existing")).unwrap(), + b"original" + ); + } + + #[cfg(target_os = "linux")] + #[tokio::test] + async fn test_conditional_abort_cleans_after_background_write_error() { + use std::path::PathBuf; + + use opendal::options::WriteOptions; + use opendal::raw::OpWrite; + use opendal::raw::oio::Write; + use tokio::io::AsyncWriteExt; + + let temp_dir = create_temp_dir("secure_fs_flush_error_abort"); + let root = SecureFsRoot::open(temp_dir.path()).unwrap(); + let (args, _) = OpWrite::from_options( + &root.build_operator().info().capability(), + WriteOptions { + if_not_exists: true, + ..Default::default() + }, + ) + .unwrap(); + let mut writer = super::SecureFsWriter { + root, + path: PathBuf::from("partial"), + args, + file: None, + closing: false, + }; + writer.ensure_file().await.unwrap(); + let mut failing_file = tokio::fs::OpenOptions::new() + .write(true) + .open("/dev/full") + .await + .unwrap(); + failing_file.write_all(b"partial").await.unwrap(); + writer.file = Some(failing_file); + + assert!(writer.abort().await.is_err()); + assert!(!temp_dir.path().join("partial").exists()); + } + #[test] fn test_writer_abort_drains_before_reporting_unsupported() { use std::path::PathBuf; diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index 5f5b73003db..d88bf4920a4 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -355,6 +355,52 @@ mod tests { use super::*; + #[tokio::test] + async fn managed_copy_preserves_existing_file_before_sink_open() { + let temp_dir = common_test_util::temp_dir::create_temp_dir("managed_copy_existing_file"); + let store = object_store::secure_fs::SecureFsRoot::open(temp_dir.path()) + .unwrap() + .build_operator(); + let path = "existing.parquet"; + store.write(path, "original").await.unwrap(); + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = arrow::record_batch::RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let batches = async_stream::stream! { + yield Ok(batch); + yield Err(datafusion::error::DataFusionError::Execution("source failed".into())); + }; + let stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + schema.arrow_schema().clone(), + batches, + ); + let stream = + common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(stream)) + .unwrap(); + + let result = stream_to_managed_parquet( + Box::pin(stream), + store.clone(), + path, + &ExportWriteBudget::new(1), + &CancellationToken::new(), + ) + .await; + + assert!(result.is_err()); + assert_eq!( + store.read(path).await.unwrap().to_bytes().as_ref(), + b"original" + ); + } + #[tokio::test] async fn managed_copy_preserves_dictionary_json_and_empty_schema() { let schema = Arc::new(Schema::new(vec![ diff --git a/src/operator/src/statement/database_copy.rs b/src/operator/src/statement/database_copy.rs index 117104b9d8f..9878560afaf 100644 --- a/src/operator/src/statement/database_copy.rs +++ b/src/operator/src/statement/database_copy.rs @@ -27,6 +27,7 @@ use table::TableRef; use table::metadata::TableType; use table::requests::CopyDatabaseRequest; use table::table_reference::TableReference; +use tokio::sync::Semaphore; use url::Url; use crate::error::{self, Result}; @@ -56,7 +57,7 @@ pub(crate) fn parse_parallelism_from_option_map(options: &HashMap().ok()) .unwrap_or_else(get_total_cpu_cores) - .max(1) + .clamp(1, Semaphore::MAX_PERMITS) } /// Rejects import-only layouts before either database export path creates output. @@ -321,5 +322,11 @@ mod tests { let options = HashMap::from([("parallelism".to_string(), "0".to_string())]); assert_eq!(parse_parallelism_from_option_map(&options), 1); + + let options = HashMap::from([("parallelism".to_string(), usize::MAX.to_string())]); + assert_eq!( + parse_parallelism_from_option_map(&options), + Semaphore::MAX_PERMITS + ); } }