From e9fce18404a4155c7793a124fd881ee2744478e4 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 23 Sep 2026 13:49:03 +0800 Subject: [PATCH] fix: preserve ordinary export cleanup and large rows Signed-off-by: jeremyhi --- src/object-store/src/secure_fs.rs | 75 +++++++++++-------- src/operator/src/statement/copy_table_to.rs | 26 ++++--- .../export_logical_tables/writers.rs | 2 +- 3 files changed, 60 insertions(+), 43 deletions(-) diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 11bfce7b04d..1875ec806fe 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -476,14 +476,20 @@ impl oio::Write for SecureFsWriter { } return flush; } - flush?; if self.file.is_none() { - return Ok(()); + return flush; } - Err(Error::new( + let mut error = Error::new( ErrorKind::Unsupported, "filesystem writes cannot be aborted without atomic writes", - )) + ); + if let Err(flush_error) = flush { + if self.closing { + return Err(flush_error); + } + error = error.set_source(flush_error); + } + Err(error) } } @@ -875,7 +881,7 @@ mod tests { #[cfg(target_os = "linux")] #[tokio::test] - async fn test_conditional_abort_cleans_after_background_write_error() { + async fn test_abort_after_background_write_error() { use std::path::PathBuf; use opendal::options::WriteOptions; @@ -883,34 +889,41 @@ mod tests { 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 + for if_not_exists in [true, false] { + 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, + ..Default::default() + }, + ) .unwrap(); - failing_file.write_all(b"partial").await.unwrap(); - writer.file = Some(failing_file); + 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()); + let error = writer.abort().await.unwrap_err(); + if if_not_exists { + assert!(!temp_dir.path().join("partial").exists()); + } else { + assert_eq!(error.kind(), opendal::ErrorKind::Unsupported); + assert!(std::error::Error::source(&error).is_some()); + } + } } #[test] diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index d88bf4920a4..ff8af79f604 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, expand_export_batch, map_writer_error, rows_within_budget, + expand_export_batch, map_writer_error, rows_within_budget, }; // The buffer size should be greater than 5MB (minimum multipart upload size). @@ -270,10 +270,8 @@ pub(crate) async fn stream_to_managed_parquet( .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 (conversion, retained) = + ExportWriteBudget::conversion_budget(&batch, usize::MAX)?; let input = batch.clone(); // JSON escaping and number formatting can expand binary JSON. let expansion = if json { 8 } else { 1 }; @@ -415,15 +413,21 @@ mod tests { ColumnSchema::new("json", ConcreteDataType::json_datatype(), true), ])); let dictionary = DictionaryArray::::new( - Int32Array::from(vec![Some(0), None]), - Arc::new(StringArray::from(vec!["tag"])), + Int32Array::from(vec![Some(0), None, Some(1)]), + Arc::new(StringArray::from(vec!["tag", &"x".repeat(1_100_000)])), ); - let json = - datatypes::types::parse_string_to_jsonb(r#"{"value":"escaped\ntext","n":123456789}"#) - .unwrap(); + let json = datatypes::types::parse_string_to_jsonb(&format!( + r#"{{"value":"escaped\ntext{}","n":123456789}}"#, + "y".repeat(160_000) + )) + .unwrap(); let arrays = vec![ Arc::new(dictionary) as ArrayRef, - Arc::new(BinaryArray::from(vec![Some(json.as_slice()), None])), + Arc::new(BinaryArray::from(vec![ + Some(json.as_slice()), + None, + Some(json.as_slice()), + ])), ]; let batch = arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays) diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs index 75659e75811..efbac24acdd 100644 --- a/src/operator/src/statement/export_logical_tables/writers.rs +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -175,7 +175,7 @@ impl TableWriters { let result = tokio::select! { biased; _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(), - result = self.tasks.next() => result.expect("writer task queue is not empty"), + result = self.tasks.next() => result.context(error::UnexpectedSnafu { violated: "writer task queue unexpectedly empty" })?, }; result.context(error::JoinTaskSnafu)??; }