fix: preserve export output after ambiguous close

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

src/common/datasource/src/parquet_writer.rs

@@ -57,6 +57,7 @@ pub struct ParquetFileWriter {
     path: String,
     limits: Option<ParquetWriterLimits>,
     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();

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"
+                );
+            });
     }
 }

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<Metadata> {
+            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::<usize>(), 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