From a7a18246bbd138076ec40e96b41a221a26231d51 Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 9 Sep 2026 10:35:21 +0800 Subject: [PATCH 1/2] perf(vindex): upload index files with concurrency of 8 Keep 8 MiB chunks and default FileIO upload concurrency at 1. Add multipart ordering, round-trip, and index build failure coverage. Validation: formatting and diff checks passed. Runtime tests and benchmarks remain blocked by a Cargo registry connection timeout. --- crates/paimon/src/io/file_io.rs | 14 +- .../paimon/src/io/file_io/multipart_test.rs | 406 ++++++++++++++++++ .../table/vindex_index_build_builder/tests.rs | 99 +++++ .../vindex_index_build_builder/writer.rs | 2 +- 4 files changed, 518 insertions(+), 3 deletions(-) create mode 100644 crates/paimon/src/io/file_io/multipart_test.rs diff --git a/crates/paimon/src/io/file_io.rs b/crates/paimon/src/io/file_io.rs index 1257f6328..0d82c8859 100644 --- a/crates/paimon/src/io/file_io.rs +++ b/crates/paimon/src/io/file_io.rs @@ -838,11 +838,18 @@ impl OutputFile { /// Get an async streaming writer for format-level writes (e.g. parquet). pub(crate) async fn async_writer(&self) -> crate::Result> { + self.async_writer_with_concurrency(1).await + } + + pub(crate) async fn async_writer_with_concurrency( + &self, + concurrency: usize, + ) -> crate::Result> { let (op, relative_path, cache_path) = self.resolve().await?; let writer: Box = Box::new( op.writer_with(&relative_path) .chunk(8 * 1024 * 1024) - .concurrent(1) + .concurrent(concurrency) .await? .into_futures_async_write() .compat_write(), @@ -860,6 +867,9 @@ impl OutputFile { } } +#[cfg(test)] +pub(crate) mod multipart_test; + #[cfg(test)] mod file_action_test { use std::collections::BTreeSet; @@ -1937,7 +1947,7 @@ mod input_output_test { let mut writer = file_io .new_output(&location) .unwrap() - .async_writer() + .async_writer_with_concurrency(8) .await .unwrap(); writer.write_all(b"new metadata").await.unwrap(); diff --git a/crates/paimon/src/io/file_io/multipart_test.rs b/crates/paimon/src/io/file_io/multipart_test.rs new file mode 100644 index 000000000..bef45c611 --- /dev/null +++ b/crates/paimon/src/io/file_io/multipart_test.rs @@ -0,0 +1,406 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 super::*; +use crate::io::FileIOBuilder; +use opendal::raw::{ + oio, OpCopier, OpCopy, OpCreateDir, OpList, OpPresign, OpRead, OpRename, OpStat, OpWrite, + RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, +}; +use opendal::{Buffer, Capability, Metadata, OperationContext}; +use std::collections::BTreeMap; +use std::sync::Mutex; +use tokio::io::AsyncWriteExt; +use tokio::sync::Notify; + +const CHUNK: usize = 8 * 1024 * 1024; + +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub(crate) enum Fault { + #[default] + None, + Part, + Close, +} + +#[derive(Debug, Default)] +pub(crate) struct UploadState { + pub(crate) fault: Fault, + pub(crate) fail_on_index: usize, + pub(crate) index_writes: usize, + pub(crate) concurrency: Vec, + pub(crate) uploads: BTreeMap>, + pub(crate) completed_parts: Vec, + pub(crate) aborts: usize, +} + +/// Exercise OpenDAL's real multipart scheduler, storing completed objects in memory. +#[derive(Clone, Debug)] +pub(crate) struct MultipartProvider { + memory: Operator, + part_size: usize, + pub(crate) state: Arc>, +} + +impl MultipartProvider { + pub(crate) fn new(part_size: usize) -> Self { + Self { + memory: Operator::via_iter(opendal::services::MEMORY_SCHEME, []).unwrap(), + part_size, + state: Arc::default(), + } + } + + pub(crate) fn file_io(&self) -> FileIO { + FileIOBuilder::new("memory") + .build() + .unwrap() + .with_provider(Arc::new(self.clone())) + } +} + +#[async_trait::async_trait] +impl FileIOProvider for MultipartProvider { + async fn create(&self, path: &str) -> crate::Result<(Operator, String)> { + Ok(( + Operator::from_parts(OperationContext::default(), Arc::new(self.clone())), + Url::parse(path) + .unwrap() + .path() + .trim_start_matches('/') + .to_string(), + )) + } +} + +impl Service for MultipartProvider { + type Reader = oio::Reader; + type Writer = oio::Writer; + type Lister = oio::Lister; + type Deleter = oio::Deleter; + type Copier = oio::Copier; + + fn info(&self) -> ServiceInfo { + self.memory.service().info() + } + + fn capability(&self) -> Capability { + Capability { + write_can_multi: true, + write_multi_max_size: Some(self.part_size), + ..self.memory.service().capability() + } + } + + fn write( + &self, + ctx: &OperationContext, + path: &str, + args: OpWrite, + ) -> opendal::Result { + let mut state = self.state.lock().unwrap(); + state.concurrency.push(args.concurrent()); + if !path.ends_with(".index") { + return self.memory.service().write(ctx, path, args); + } + state.index_writes += 1; + let fault = if state.index_writes == state.fail_on_index { + state.fault + } else { + Fault::None + }; + Ok(Box::new(oio::MultipartWriter::new( + ctx.executor().clone(), + Upload { + provider: self.clone(), + path: path.to_string(), + fault, + reorder: args.concurrent() > 1, + second_part: Notify::new(), + }, + args.concurrent(), + ))) + } + + async fn create_dir( + &self, + ctx: &OperationContext, + path: &str, + args: OpCreateDir, + ) -> opendal::Result { + self.memory.service().create_dir(ctx, path, args).await + } + async fn stat( + &self, + ctx: &OperationContext, + path: &str, + args: OpStat, + ) -> opendal::Result { + self.memory.service().stat(ctx, path, args).await + } + fn read( + &self, + ctx: &OperationContext, + path: &str, + args: OpRead, + ) -> opendal::Result { + self.memory.service().read(ctx, path, args) + } + fn delete(&self, ctx: &OperationContext) -> opendal::Result { + self.memory.service().delete(ctx) + } + fn list( + &self, + ctx: &OperationContext, + path: &str, + args: OpList, + ) -> opendal::Result { + self.memory.service().list(ctx, path, args) + } + fn copy( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpCopy, + opts: OpCopier, + ) -> opendal::Result { + self.memory.service().copy(ctx, from, to, args, opts) + } + async fn rename( + &self, + ctx: &OperationContext, + from: &str, + to: &str, + args: OpRename, + ) -> opendal::Result { + self.memory.service().rename(ctx, from, to, args).await + } + async fn presign( + &self, + ctx: &OperationContext, + path: &str, + args: OpPresign, + ) -> opendal::Result { + self.memory.service().presign(ctx, path, args).await + } +} + +struct Upload { + provider: MultipartProvider, + path: String, + fault: Fault, + reorder: bool, + second_part: Notify, +} + +fn injected_error() -> opendal::Error { + opendal::Error::new(opendal::ErrorKind::Unexpected, "injected multipart failure") +} + +impl oio::MultipartWrite for Upload { + async fn write_once(&self, size: u64, body: Buffer) -> opendal::Result { + assert_eq!( + self.fault, + Fault::None, + "failure test must exercise multipart" + ); + self.provider.memory.write(&self.path, body).await?; + Ok(Metadata::default().with_content_length(size)) + } + + async fn initiate_part(&self) -> opendal::Result { + self.provider + .state + .lock() + .unwrap() + .uploads + .insert(self.path.clone(), BTreeMap::new()); + Ok(self.path.clone()) + } + + async fn write_part( + &self, + upload_id: &str, + part_number: usize, + size: u64, + body: Buffer, + ) -> opendal::Result { + assert_eq!(size as usize, body.len()); + if self.reorder && part_number == 0 { + self.second_part.notified().await; + } + { + let mut state = self.provider.state.lock().unwrap(); + state.completed_parts.push(part_number); + state + .uploads + .get_mut(upload_id) + .unwrap() + .insert(part_number, body.to_bytes()); + } + if part_number == 1 { + self.second_part.notify_one(); + } + if self.fault == Fault::Part && part_number == 0 { + return Err(injected_error()); + } + Ok(oio::MultipartPart { + part_number, + etag: part_number.to_string(), + checksum: None, + size: Some(size), + }) + } + + async fn complete_part( + &self, + upload_id: &str, + parts: &[oio::MultipartPart], + ) -> opendal::Result { + if self.fault == Fault::Close { + return Err(injected_error()); + } + let mut bytes = Vec::new(); + { + let mut state = self.provider.state.lock().unwrap(); + let uploaded = state.uploads.remove(upload_id).unwrap(); + assert_eq!(parts.len(), uploaded.len()); + for (expected, part) in parts.iter().enumerate() { + assert_eq!(part.part_number, expected); + bytes.extend_from_slice(&uploaded[&part.part_number]); + } + } + let size = bytes.len() as u64; + self.provider.memory.write(&self.path, bytes).await?; + Ok(Metadata::default().with_content_length(size)) + } + + async fn abort_part(&self, upload_id: &str) -> opendal::Result<()> { + let mut state = self.provider.state.lock().unwrap(); + state.aborts += 1; + state.uploads.remove(upload_id); + Ok(()) + } +} + +async fn round_trip(file_io: &FileIO, path: &str, concurrency: usize) { + // Different bytes within and between parts expose truncation and reordering. + let data: Vec = (0..9 * CHUNK + 12345) + .map(|i| ((i / CHUNK * 37 + i % 251) % 256) as u8) + .collect(); + let output = file_io.new_output(path).unwrap(); + let mut writer = if concurrency == 1 { + output.async_writer().await.unwrap() + } else { + output + .async_writer_with_concurrency(concurrency) + .await + .unwrap() + }; + writer.write_all(&data).await.unwrap(); + writer.shutdown().await.unwrap(); + assert_eq!( + file_io.get_status(path).await.unwrap().size, + data.len() as u64 + ); + assert_eq!( + file_io + .new_input(path) + .unwrap() + .read() + .await + .unwrap() + .as_ref(), + data + ); +} + +#[tokio::test] +async fn streaming_upload_memory_and_fs() { + let directory = tempfile::tempdir().unwrap(); + for (scheme, path) in [ + ("memory", "memory:/stream.index".to_string()), + ( + "file", + format!("file:{}/stream.index", directory.path().display()), + ), + ] { + let file_io = FileIOBuilder::new(scheme).build().unwrap(); + for concurrency in [1, 8] { + round_trip(&file_io, &path, concurrency).await; + } + } +} + +#[tokio::test] +async fn streaming_upload_preserves_out_of_order_parts() { + for concurrency in [1, 8] { + let provider = MultipartProvider::new(CHUNK); + tokio::time::timeout( + std::time::Duration::from_secs(30), + round_trip(&provider.file_io(), "memory:/ordered.index", concurrency), + ) + .await + .unwrap(); + let state = provider.state.lock().unwrap(); + assert_eq!(state.concurrency, [concurrency]); + assert_eq!(state.completed_parts.len(), 10); + if concurrency == 8 { + assert!( + state.completed_parts.iter().position(|&p| p == 1).unwrap() + < state.completed_parts.iter().position(|&p| p == 0).unwrap() + ); + } + assert!(state.uploads.is_empty()); + } +} + +#[tokio::test] +#[ignore = "requires PAIMON_CATALOG_OPTIONS and PAIMON_UPLOAD_TEST_DATABASE/TABLE for OSS"] +async fn streaming_upload_oss() { + use crate::catalog::Identifier; + use crate::common::Options; + use crate::CatalogFactory; + use futures::FutureExt; + + let options = serde_json::from_str(&std::env::var("PAIMON_CATALOG_OPTIONS").unwrap()).unwrap(); + let catalog = CatalogFactory::create(Options::from_map(options)) + .await + .unwrap(); + let table = catalog + .get_table(&Identifier::new( + std::env::var("PAIMON_UPLOAD_TEST_DATABASE").unwrap(), + std::env::var("PAIMON_UPLOAD_TEST_TABLE").unwrap(), + )) + .await + .unwrap(); + assert!(table.location().starts_with("oss://")); + let path = format!( + "{}/index/upload-check-{}.index", + table.location(), + uuid::Uuid::new_v4() + ); + eprintln!("OSS multipart byte-for-byte check: {path}"); + let result = std::panic::AssertUnwindSafe(round_trip(table.file_io(), &path, 8)) + .catch_unwind() + .await; + table.file_io().delete_file(&path).await.unwrap(); + assert!(!table.file_io().exists(&path).await.unwrap()); + eprintln!("OSS scratch object removed: {path}"); + result.unwrap(); +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs b/crates/paimon/src/table/vindex_index_build_builder/tests.rs index a4636b578..5eaff7fe1 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -643,6 +643,105 @@ async fn vindex_incremental_build_indexes_only_new_rows() { } } +#[tokio::test] +async fn vindex_upload_failure_preserves_committed_index() { + use crate::io::multipart_test::{Fault, MultipartProvider}; + + for fault in [Fault::Part, Fault::Close] { + // A small backend part limit makes the tiny test index use multipart. + let provider = MultipartProvider::new(128); + let table_path = "memory:/test_vindex_upload_failure"; + let table = test_table_with_io( + provider.file_io(), + table_path, + vindex_schema_builder(vindex_e2e_options("3")) + .build() + .unwrap(), + ); + setup_dirs(table.file_io(), table_path).await; + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + ) + .await; + table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .unwrap(); + let existing = latest_vindex_index_files(&table).await; + let old_path = format!("{table_path}/{INDEX_DIR}/{}", existing[0].file_name); + let old_bytes = table + .file_io() + .new_input(&old_path) + .unwrap() + .read() + .await + .unwrap(); + let mut search = table.new_vector_search_builder(); + search + .with_vector_column("embedding") + .with_query_vector(vec![1.0, 0.0]) + .with_limit(1); + let old_result = search.execute().await.unwrap(); + assert!(!old_result.is_empty()); + + write_vectors(&table, vec![4, 5, 6, 7, 8, 9], vec![vec![-1.0, 0.0]; 6]).await; + let snapshots = SnapshotManager::new(table.file_io().clone(), table_path.to_string()); + let before = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + { + let mut state = provider.state.lock().unwrap(); + state.fault = fault; + state.fail_on_index = state.index_writes + 2; + state.concurrency.clear(); + } + let error = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .expect_err("injected upload must fail"); + assert!( + error.to_string().contains("injected multipart failure"), + "{error}" + ); + let after = snapshots.get_latest_snapshot().await.unwrap().unwrap(); + assert_eq!(before.id(), after.id()); + assert_eq!(before.index_manifest(), after.index_manifest()); + assert_eq!(latest_vindex_index_files(&table).await, existing); + assert_eq!( + table + .file_io() + .new_input(&old_path) + .unwrap() + .read() + .await + .unwrap(), + old_bytes + ); + assert_eq!(search.execute().await.unwrap(), old_result); + let files = table + .file_io() + .list_status(&format!("{table_path}/{INDEX_DIR}/")) + .await + .unwrap(); + assert_eq!( + files.len(), + 1, + "failed and previously written new shards must be removed" + ); + assert!(table.file_io().exists(&old_path).await.unwrap()); + let state = provider.state.lock().unwrap(); + assert_eq!(state.concurrency, [8, 8]); + // The AsyncWrite adapter cannot abort: object deletion leaves the upload + // for storage lifecycle cleanup. Keep this distinct from committed files. + assert_eq!(state.uploads.len(), 1); + assert_eq!(state.aborts, 0); + } +} + #[tokio::test] async fn vindex_build_cleans_written_shards_when_later_shard_fails() { let table_path = "memory:/test_vindex_abort_written_shard"; diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs b/crates/paimon/src/table/vindex_index_build_builder/writer.rs index 803c08887..f3f2f9306 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/writer.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -350,7 +350,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { .table .file_io() .new_output(&index_path)? - .async_writer() + .async_writer_with_concurrency(8) .await?; let mut output = SyncIoBridge::new(async_writer); tokio::task::spawn_blocking(move || -> std::io::Result<()> { From c5e0c6502810dee07834863dca2a40fa0378ca79 Mon Sep 17 00:00:00 2001 From: yantian Date: Wed, 9 Sep 2026 13:52:04 +0800 Subject: [PATCH 2/2] test(io): fix multipart upload paths on Windows Use the native temporary file path so FileIO preserves the Windows drive prefix. Simplify redundant multipart and index upload failure assertions. --- .../paimon/src/io/file_io/multipart_test.rs | 14 ++------------ .../table/vindex_index_build_builder/tests.rs | 19 ------------------- 2 files changed, 2 insertions(+), 31 deletions(-) diff --git a/crates/paimon/src/io/file_io/multipart_test.rs b/crates/paimon/src/io/file_io/multipart_test.rs index bef45c611..3d49df2d3 100644 --- a/crates/paimon/src/io/file_io/multipart_test.rs +++ b/crates/paimon/src/io/file_io/multipart_test.rs @@ -45,7 +45,6 @@ pub(crate) struct UploadState { pub(crate) concurrency: Vec, pub(crate) uploads: BTreeMap>, pub(crate) completed_parts: Vec, - pub(crate) aborts: usize, } /// Exercise OpenDAL's real multipart scheduler, storing completed objects in memory. @@ -214,11 +213,6 @@ fn injected_error() -> opendal::Error { impl oio::MultipartWrite for Upload { async fn write_once(&self, size: u64, body: Buffer) -> opendal::Result { - assert_eq!( - self.fault, - Fault::None, - "failure test must exercise multipart" - ); self.provider.memory.write(&self.path, body).await?; Ok(Metadata::default().with_content_length(size)) } @@ -240,7 +234,6 @@ impl oio::MultipartWrite for Upload { size: u64, body: Buffer, ) -> opendal::Result { - assert_eq!(size as usize, body.len()); if self.reorder && part_number == 0 { self.second_part.notified().await; } @@ -279,9 +272,7 @@ impl oio::MultipartWrite for Upload { { let mut state = self.provider.state.lock().unwrap(); let uploaded = state.uploads.remove(upload_id).unwrap(); - assert_eq!(parts.len(), uploaded.len()); - for (expected, part) in parts.iter().enumerate() { - assert_eq!(part.part_number, expected); + for part in parts { bytes.extend_from_slice(&uploaded[&part.part_number]); } } @@ -292,7 +283,6 @@ impl oio::MultipartWrite for Upload { async fn abort_part(&self, upload_id: &str) -> opendal::Result<()> { let mut state = self.provider.state.lock().unwrap(); - state.aborts += 1; state.uploads.remove(upload_id); Ok(()) } @@ -337,7 +327,7 @@ async fn streaming_upload_memory_and_fs() { ("memory", "memory:/stream.index".to_string()), ( "file", - format!("file:{}/stream.index", directory.path().display()), + directory.path().join("stream.index").display().to_string(), ), ] { let file_io = FileIOBuilder::new(scheme).build().unwrap(); diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs b/crates/paimon/src/table/vindex_index_build_builder/tests.rs index 5eaff7fe1..36318a487 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -673,13 +673,6 @@ async fn vindex_upload_failure_preserves_committed_index() { .unwrap(); let existing = latest_vindex_index_files(&table).await; let old_path = format!("{table_path}/{INDEX_DIR}/{}", existing[0].file_name); - let old_bytes = table - .file_io() - .new_input(&old_path) - .unwrap() - .read() - .await - .unwrap(); let mut search = table.new_vector_search_builder(); search .with_vector_column("embedding") @@ -710,17 +703,6 @@ async fn vindex_upload_failure_preserves_committed_index() { let after = snapshots.get_latest_snapshot().await.unwrap().unwrap(); assert_eq!(before.id(), after.id()); assert_eq!(before.index_manifest(), after.index_manifest()); - assert_eq!(latest_vindex_index_files(&table).await, existing); - assert_eq!( - table - .file_io() - .new_input(&old_path) - .unwrap() - .read() - .await - .unwrap(), - old_bytes - ); assert_eq!(search.execute().await.unwrap(), old_result); let files = table .file_io() @@ -738,7 +720,6 @@ async fn vindex_upload_failure_preserves_committed_index() { // The AsyncWrite adapter cannot abort: object deletion leaves the upload // for storage lifecycle cleanup. Keep this distinct from committed files. assert_eq!(state.uploads.len(), 1); - assert_eq!(state.aborts, 0); } }