diff --git a/crates/paimon/src/io/file_io.rs b/crates/paimon/src/io/file_io.rs index 42814fd8a..edb197e52 100644 --- a/crates/paimon/src/io/file_io.rs +++ b/crates/paimon/src/io/file_io.rs @@ -1955,11 +1955,23 @@ 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 + } + + /// Streaming writer with an explicit multipart concurrency. + /// + /// Callers that do not pass a concurrency keep the historical single + /// upload. Vector-index serialization uses eight so part uploads overlap + /// while each chunk stays 8 MiB. + pub(crate) async fn async_writer_with_concurrency( + &self, + concurrency: usize, + ) -> crate::Result> { let (op, relative_path, cache_path) = self.source.resolve(&self.path).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(), @@ -1977,6 +1989,9 @@ impl OutputFile { } } +#[cfg(test)] +pub(crate) mod multipart_test; + #[cfg(test)] mod file_action_test { use std::collections::BTreeSet; @@ -4324,7 +4339,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..7d43d477e --- /dev/null +++ b/crates/paimon/src/io/file_io/multipart_test.rs @@ -0,0 +1,637 @@ +// 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 bytes::Bytes; +use opendal::raw::{ + oio, OpCopy, OpCreateDir, OpList, OpPresign, OpRead, OpRename, OpStat, OpWrite, RpCreateDir, + RpPresign, RpRename, RpStat, Service, ServiceInfo, +}; +use opendal::{Buffer, Capability, Metadata, MetadataBuilder, OperationContext, Operator}; +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use tokio::io::AsyncWriteExt; +use tokio::sync::Notify; +use url::Url; + +const CHUNK: usize = 8 * 1024 * 1024; + +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub(crate) enum Fault { + #[default] + None, + Part, + Close, +} + +#[derive(Debug)] +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) part_attempts: Vec, + /// When false, stored part bytes are copies so payload owners can drop. + pub(crate) share_payloads: bool, + /// Delay part 0 until part 1 is accepted. Tests that retry part 0 turn this off. + pub(crate) reorder: bool, + /// Fail part 0 this many times with a temporary error, then accept it. + pub(crate) temporary_failures: usize, +} + +impl Default for UploadState { + fn default() -> Self { + Self { + fault: Fault::None, + fail_on_index: 0, + index_writes: 0, + concurrency: Vec::new(), + uploads: BTreeMap::new(), + completed_parts: Vec::new(), + part_attempts: Vec::new(), + share_payloads: true, + reorder: true, + temporary_failures: 0, + } + } +} + +/// 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; + type Composer = (); + + 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 + }; + let reorder = state.reorder && args.concurrent() > 1; + Ok(Box::new(oio::MultipartWriter::new( + ctx.executor().clone(), + Upload { + provider: self.clone(), + path: path.to_string(), + fault, + reorder, + second_part: Notify::new(), + }, + args.concurrent(), + ))) + } + + async fn create_dir( + &self, + _ctx: &OperationContext, + _path: &str, + _args: OpCreateDir, + ) -> opendal::Result { + // The in-memory map is flat. Directory creation only needs to succeed so + // index writers can prepare `index/` before the multipart upload. + Ok(RpCreateDir::default()) + } + 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, + ) -> opendal::Result { + self.memory.service().copy(ctx, from, to, args) + } + 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") +} + +fn file_metadata(size: u64) -> Metadata { + MetadataBuilder::file(size).build() +} + +impl oio::MultipartWrite for Upload { + async fn write_once(&self, size: u64, body: Buffer) -> opendal::Result { + self.provider.memory.write(&self.path, body).await?; + Ok(file_metadata(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 { + if self.reorder && part_number == 0 { + self.second_part.notified().await; + } + let mut temporary = false; + { + let mut state = self.provider.state.lock().unwrap(); + state.part_attempts.push(part_number); + state.completed_parts.push(part_number); + let stored = if state.share_payloads { + body.to_bytes() + } else { + Bytes::copy_from_slice(&body.to_bytes()) + }; + state + .uploads + .get_mut(upload_id) + .unwrap() + .insert(part_number, stored); + if part_number == 0 && state.temporary_failures > 0 { + state.temporary_failures -= 1; + temporary = true; + } + } + if part_number == 1 { + self.second_part.notify_one(); + } + if temporary { + return Err(injected_error().set_temporary()); + } + 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(); + for part in parts { + bytes.extend_from_slice(&uploaded[&part.part_number]); + } + } + let size = bytes.len() as u64; + self.provider.memory.write(&self.path, bytes).await?; + Ok(file_metadata(size)) + } + + async fn abort_part(&self, upload_id: &str) -> opendal::Result<()> { + let mut state = self.provider.state.lock().unwrap(); + 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", + directory.path().join("stream.index").display().to_string(), + ), + ] { + 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()); + } +} + +struct TrackedChunk { + data: Vec, + live: Arc, +} + +impl AsRef<[u8]> for TrackedChunk { + fn as_ref(&self) -> &[u8] { + &self.data + } +} + +impl Drop for TrackedChunk { + fn drop(&mut self) { + self.live.fetch_sub(self.data.len(), Ordering::SeqCst); + } +} + +fn rss_bytes() -> u64 { + let Ok(status) = std::fs::read_to_string("/proc/self/status") else { + return 0; + }; + for line in status.lines() { + let Some(rest) = line.strip_prefix("VmRSS:") else { + continue; + }; + let kb = rest + .split_whitespace() + .next() + .unwrap_or("0") + .parse::() + .unwrap_or(0); + return kb.saturating_mul(1024); + } + 0 +} + +async fn settle(live: &AtomicUsize) -> usize { + let mut stable = 0; + let mut last = live.load(Ordering::SeqCst); + for _ in 0..100 { + tokio::task::yield_now().await; + let now = live.load(Ordering::SeqCst); + if now == last { + stable += 1; + if stable >= 5 { + return now; + } + } else { + stable = 0; + last = now; + } + } + last +} + +/// OpenDAL 0.59 releases a successful part's input when that part task finishes, +/// instead of retaining it until `close` collects the result. Twelve 8 MiB parts +/// at concurrency 8 must therefore keep only the unsent tail, not the whole upload. +#[tokio::test] +async fn streaming_upload_releases_completed_part_buffers() { + const PARTS: usize = 12; + const CONCURRENCY: usize = 8; + let provider = MultipartProvider::new(CHUNK); + provider.state.lock().unwrap().share_payloads = false; + let (op, relative) = FileIOProvider::create(&provider, "memory:/release.index") + .await + .unwrap(); + let live = Arc::new(AtomicUsize::new(0)); + let rss_before = rss_bytes(); + let mut rss_peak = rss_before; + let mut writer = op + .writer_with(&relative) + .chunk(CHUNK) + .concurrent(CONCURRENCY) + .await + .unwrap(); + + let mut peak = 0usize; + for part in 0..PARTS { + live.fetch_add(CHUNK, Ordering::SeqCst); + let bytes = Bytes::from_owner(TrackedChunk { + data: vec![part as u8; CHUNK], + live: Arc::clone(&live), + }); + writer.write(bytes).await.unwrap(); + if part == 0 { + let held = settle(&live).await; + assert_eq!( + held, CHUNK, + "the first chunk stays buffered until the next part arrives" + ); + } + peak = peak.max(settle(&live).await); + rss_peak = rss_peak.max(rss_bytes()); + } + let before_close = settle(&live).await; + let completed = provider.state.lock().unwrap().completed_parts.len(); + writer.close().await.unwrap(); + let after_close = settle(&live).await; + rss_peak = rss_peak.max(rss_bytes()); + + let scheduled = PARTS.saturating_sub(2); + assert_eq!( + completed, scheduled, + "all but the generator tail and multipart cache should already have been uploaded" + ); + assert!( + before_close <= 4 * CHUNK, + "completed part payloads were still live before close: {before_close} bytes, peak {peak}" + ); + assert!( + peak <= (CONCURRENCY + 3) * CHUNK, + "peak live part payloads {peak} exceeded concurrency {CONCURRENCY}" + ); + assert!( + peak < PARTS * CHUNK, + "peak live part payloads {peak} retained the whole {} byte upload", + PARTS * CHUNK + ); + assert_eq!(after_close, 0, "close must drop every part payload"); + + let got = provider.memory.read(&relative).await.unwrap().to_vec(); + assert_eq!(got.len(), PARTS * CHUNK); + for part in 0..PARTS { + let start = part * CHUNK; + assert_eq!(got[start], part as u8); + assert_eq!(got[start + CHUNK - 1], part as u8); + } + eprintln!( + "part-buffer release: parts={PARTS} chunk={CHUNK} concurrency={CONCURRENCY} \ + scheduled_before_close={completed} live_before_close={before_close} \ + peak_live={peak} rss_before={rss_before} rss_peak={rss_peak}" + ); +} + +/// A temporary part failure keeps that part's bytes for the scheduler retry. +/// The caller only retries `close`; it does not submit the payload again. +#[tokio::test] +async fn streaming_upload_retries_temporary_part_failure() { + const PARTS: usize = 4; + const PART: usize = 64 * 1024; + let provider = MultipartProvider::new(PART); + { + let mut state = provider.state.lock().unwrap(); + state.reorder = false; + state.temporary_failures = 1; + state.share_payloads = false; + } + let (op, relative) = FileIOProvider::create(&provider, "memory:/retry.index") + .await + .unwrap(); + let mut writer = op + .writer_with(&relative) + .chunk(PART) + .concurrent(8) + .await + .unwrap(); + let mut data = Vec::with_capacity(PARTS * PART); + for part in 0..PARTS { + let chunk = vec![part as u8 + 1; PART]; + data.extend_from_slice(&chunk); + writer.write(chunk).await.unwrap(); + } + + let mut closes = 0; + loop { + closes += 1; + match writer.close().await { + Ok(_) => break, + Err(error) if error.is_temporary() && closes < 5 => continue, + Err(error) => panic!("multipart close failed: {error}"), + } + } + assert!( + closes >= 2, + "the first close must observe the temporary part failure, closes={closes}" + ); + let attempts = provider + .state + .lock() + .unwrap() + .part_attempts + .iter() + .filter(|part| **part == 0) + .count(); + assert_eq!( + attempts, 2, + "part 0 must be sent again from the retained input" + ); + let got = provider.memory.read(&relative).await.unwrap().to_vec(); + assert_eq!(got, data); +} + +#[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/index_build/vector.rs b/crates/paimon/src/table/index_build/vector.rs index a3640d6be..2d4d48257 100644 --- a/crates/paimon/src/table/index_build/vector.rs +++ b/crates/paimon/src/table/index_build/vector.rs @@ -276,7 +276,7 @@ pub(in crate::table) async fn serialize_writer( mut writer: VectorIndexWriter, output: OutputFile, ) -> Result<()> { - let mut output = SyncIoBridge::new(output.async_writer().await?); + let mut output = SyncIoBridge::new(output.async_writer_with_concurrency(8).await?); tokio::task::spawn_blocking(move || -> std::io::Result<()> { writer.write(&mut PosWriter::new(&mut output))?; output.shutdown() 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 137de863a..0663ff17b 100644 --- a/crates/paimon/src/table/vindex_index_build_builder/tests.rs +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -1423,6 +1423,91 @@ async fn vindex_granule_training_sees_one_oversized_row_group() { } } +#[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_global_index_build_builder() + .with_index_type(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 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.row_ids().unwrap().row_ids.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_global_index_build_builder() + .with_index_type(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!( + search.execute().await.unwrap().row_ids().unwrap().row_ids, + old_result.row_ids().unwrap().row_ids + ); + 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); + } +} + #[tokio::test] async fn vindex_build_cleans_written_shards_when_later_shard_fails() { let table_path = "memory:/test_vindex_abort_written_shard";