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