This is an automated email from the ASF dual-hosted git repository.
yihua pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/hudi-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 92e99e89 perf(core): merge copy-on-write file groups in parallel (#884)
92e99e89 is described below
commit 92e99e89894c4a465abf988ed335545b60cf03e3
Author: Y Ethan Guo <[email protected]>
AuthorDate: Fri Oct 30 09:09:38 2026 -0700
perf(core): merge copy-on-write file groups in parallel (#884)
`run_write_tasks` now spawns each per-file task onto the runtime inside the
same bounded stream, so file groups are read and merged in parallel instead of
one after another on the awaiting task. The concurrency bound and input-order
results are unchanged, and dropping the write aborts the spawned tasks. The
copy-on-write merge, now `MergePlan::merge`, runs together with the encode
inside the existing `spawn_blocking`, so its long synchronous section does not
stall the async workers. A [...]
---
crates/core/src/write/mod.rs | 122 ++++++++++++++++++++++++-
crates/core/src/write/rewrite/cow.rs | 80 ++++++++--------
crates/core/tests/table_write_runtime_tests.rs | 112 +++++++++++++++++++++++
3 files changed, 276 insertions(+), 38 deletions(-)
diff --git a/crates/core/src/write/mod.rs b/crates/core/src/write/mod.rs
index f5c69141..75a22afc 100644
--- a/crates/core/src/write/mod.rs
+++ b/crates/core/src/write/mod.rs
@@ -285,7 +285,8 @@ pub(crate) fn write_task_parallelism(table:
&crate::table::Table) -> usize {
})
}
-/// Run write tasks with bounded concurrency; results come back in input order.
+/// Run write tasks with bounded concurrency on the runtime's worker threads;
+/// results come back in input order.
pub(crate) async fn run_write_tasks<T: Send + 'static>(
tasks: Vec<futures::future::BoxFuture<'static, Result<T>>>,
parallelism: usize,
@@ -298,7 +299,12 @@ pub(crate) async fn run_write_tasks<T: Send + 'static>(
let indexed: Vec<_> = tasks
.into_iter()
.enumerate()
- .map(|(index, task)| async move { (index, task.await) })
+ .map(|(index, task)| async move {
+ let result = AbortOnDrop(tokio::spawn(task))
+ .await
+ .unwrap_or_else(|e| Err(CoreError::Write(format!("write task
failed: {e}"))));
+ (index, result)
+ })
.collect();
let mut stream =
futures::stream::iter(indexed).buffer_unordered(parallelism.max(1));
while let Some((index, result)) = stream.next().await {
@@ -310,6 +316,27 @@ pub(crate) async fn run_write_tasks<T: Send + 'static>(
.collect()
}
+/// Aborts the spawned task when dropped. The task stops at its next await;
+/// work it already handed to the blocking pool still runs to completion.
+struct AbortOnDrop<T>(tokio::task::JoinHandle<T>);
+
+impl<T> Drop for AbortOnDrop<T> {
+ fn drop(&mut self) {
+ self.0.abort();
+ }
+}
+
+impl<T> std::future::Future for AbortOnDrop<T> {
+ type Output = std::result::Result<T, tokio::task::JoinError>;
+
+ fn poll(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> std::task::Poll<Self::Output> {
+ std::pin::Pin::new(&mut self.0).poll(cx)
+ }
+}
+
/// Output of one parquet base-file or log-file write task.
pub(crate) struct FileTaskOutput {
pub size: i64,
@@ -881,4 +908,95 @@ mod tests {
let ids = batch.project(&[0]).unwrap();
assert_eq!(normalize_write_batch(&ids).unwrap(), ids);
}
+
+ /// CPU-bound tasks run on the runtime's workers, not the awaiting thread,
+ /// and results keep input order with errors in place.
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_run_write_tasks_runs_on_workers_in_input_order() {
+ let caller = std::thread::current().id();
+ let tasks: Vec<futures::future::BoxFuture<'static,
Result<std::thread::ThreadId>>> = (0
+ ..8u64)
+ .map(|i| {
+ let task: futures::future::BoxFuture<'static, _> =
Box::pin(async move {
+ // Later tasks finish first, so completion order differs
from input order.
+ let spin = std::time::Duration::from_millis(5 * (8 - i));
+ let start = std::time::Instant::now();
+ while start.elapsed() < spin {
+ std::hint::spin_loop();
+ }
+ if i == 3 {
+ Err(CoreError::Write("task 3".to_string()))
+ } else {
+ Ok(std::thread::current().id())
+ }
+ });
+ task
+ })
+ .collect();
+
+ let results = run_write_tasks(tasks, 4).await;
+
+ assert_eq!(results.len(), 8);
+ let mut threads = std::collections::HashSet::new();
+ for (i, result) in results.into_iter().enumerate() {
+ match result {
+ Err(CoreError::Write(msg)) if i == 3 => assert_eq!(msg, "task
3"),
+ Ok(thread) if i != 3 => {
+ assert_ne!(thread, caller, "task {i} ran on the awaiting
thread");
+ threads.insert(thread);
+ }
+ other => panic!("unexpected result for task {i}: {other:?}"),
+ }
+ }
+ assert!(
+ threads.len() > 1,
+ "tasks ran on {} thread(s)",
+ threads.len()
+ );
+ }
+
+ /// A panicking task surfaces as a write error; its siblings still
complete.
+ #[tokio::test]
+ async fn test_run_write_tasks_maps_panic_to_write_error() {
+ let tasks: Vec<futures::future::BoxFuture<'static, Result<usize>>> =
vec![
+ Box::pin(async { Ok(0) }),
+ Box::pin(async { panic!("boom") }),
+ Box::pin(async { Ok(2) }),
+ ];
+
+ let results = run_write_tasks(tasks, 2).await;
+
+ assert!(matches!(results[0], Ok(0)));
+ assert!(
+ matches!(&results[1], Err(CoreError::Write(msg)) if
msg.contains("panicked")),
+ "{:?}",
+ results[1]
+ );
+ assert!(matches!(results[2], Ok(2)));
+ }
+
+ /// Dropping the write future aborts the tasks it spawned.
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn test_run_write_tasks_aborts_spawned_tasks_on_drop() {
+ use futures::FutureExt;
+ let (tx, mut rx) = tokio::sync::oneshot::channel::<()>();
+ let task: futures::future::BoxFuture<'static, Result<()>> =
Box::pin(async move {
+ let _tx = tx;
+ std::future::pending::<()>().await;
+ Ok(())
+ });
+
+ assert!(run_write_tasks(vec![task], 1).now_or_never().is_none());
+
+ for _ in 0..500 {
+ if matches!(
+ rx.try_recv(),
+ Err(tokio::sync::oneshot::error::TryRecvError::Closed)
+ ) {
+ return;
+ }
+ std::thread::sleep(std::time::Duration::from_millis(10));
+ }
+ panic!("spawned write task kept running after the write was dropped");
+ }
}
diff --git a/crates/core/src/write/rewrite/cow.rs
b/crates/core/src/write/rewrite/cow.rs
index 0ef063fb..81af0db8 100644
--- a/crates/core/src/write/rewrite/cow.rs
+++ b/crates/core/src/write/rewrite/cow.rs
@@ -468,11 +468,29 @@ pub(super) fn merge_file_group_task(
let old = reader
.read_file_slice(&plan.slice, &ReadOptions::new())
.await?;
+ let storage = ctx.storage.clone();
+ let path = relative_data_path(&plan.partition, &plan.out_name);
+ // The merge and encode are long synchronous sections; on the blocking
+ // pool they cannot stall the async workers other tasks run on.
+ let (output, bytes) = tokio::task::spawn_blocking(move ||
plan.merge(ctx, old))
+ .await
+ .map_err(|e| CoreError::Write(format!("merge task panicked:
{e}")))??;
+ if let Some(bytes) = bytes {
+ storage.put_file(&path, bytes).await?;
+ }
+ Ok(output)
+ })
+}
+
+impl MergePlan {
+ /// Apply the op to the group's current rows `old`, stamp meta, and encode;
+ /// the bytes are `None` when the op matched nothing and no file is
written.
+ fn merge(self, ctx: MergeCtx, old: RecordBatch) -> Result<(MergeOutput,
Option<Vec<u8>>)> {
let key_field = MetaField::RecordKey.as_ref();
let mut rli_deletes = Vec::new();
let mut marker_deletes = 0usize;
let mut inserted = 0;
- let (merged, matched) = match plan.op {
+ let (merged, matched) = match self.op {
GroupOp::Upsert {
incoming,
partial_columns,
@@ -613,7 +631,7 @@ pub(super) fn merge_file_group_task(
}
let matched = old.num_rows() - selected.len();
if matched == 0 {
- return Ok(MergeOutput::unchanged());
+ return Ok((MergeOutput::unchanged(), None));
}
(take_batch(&old, &selected)?, matched)
}
@@ -635,7 +653,7 @@ pub(super) fn merge_file_group_task(
}
}
if matched == 0 {
- return Ok(MergeOutput::unchanged());
+ return Ok((MergeOutput::unchanged(), None));
}
(take_batch(&old, &selected)?, matched)
}
@@ -643,7 +661,7 @@ pub(super) fn merge_file_group_task(
let mask = filters_to_row_mask(std::slice::from_ref(&filter),
&old)?;
let (updated, matched) = apply_set_updates(&old, &mask, &set)?;
if matched == 0 {
- return Ok(MergeOutput::unchanged());
+ return Ok((MergeOutput::unchanged(), None));
}
let updated = restamp_matched_rows(&updated, &mask,
&ctx.instant)?;
// Setting the delete marker deletes the row, as an upsert
does.
@@ -670,39 +688,29 @@ pub(super) fn merge_file_group_task(
(updated, matched)
}
};
- let merged = restamp_file_name(&merged, &plan.out_name)?;
- let path = relative_data_path(&plan.partition, &plan.out_name);
+ let merged = restamp_file_name(&merged, &self.out_name)?;
let num_rows = merged.num_rows();
- let props = ctx.props.clone();
- let bloom = ctx.bloom;
- let collect_ranges = ctx.collect_ranges;
- let (bytes, ranges) = tokio::task::spawn_blocking(move || -> Result<_>
{
- let bytes =
- crate::write::append::write_parquet_bytes_with_props(props,
bloom, &[merged])?;
- let ranges = if collect_ranges {
- Some(
-
crate::metadata::table::column_stats::column_ranges_from_parquet_bytes(&bytes)?,
- )
- } else {
- None
- };
- Ok((bytes, ranges))
- })
- .await
- .map_err(|e| CoreError::Write(format!("merge task panicked: {e}")))??;
- let size = bytes.len() as i64;
- ctx.storage.put_file(&path, bytes).await?;
- Ok(MergeOutput {
- wrote: true,
- size,
- num_rows,
- matched,
- inserted,
- marker_deletes,
- ranges,
- rli_deletes,
- })
- })
+ let bytes =
+ crate::write::append::write_parquet_bytes_with_props(ctx.props,
ctx.bloom, &[merged])?;
+ let ranges = if ctx.collect_ranges {
+
Some(crate::metadata::table::column_stats::column_ranges_from_parquet_bytes(&bytes)?)
+ } else {
+ None
+ };
+ Ok((
+ MergeOutput {
+ wrote: true,
+ size: bytes.len() as i64,
+ num_rows,
+ matched,
+ inserted,
+ marker_deletes,
+ ranges,
+ rli_deletes,
+ },
+ Some(bytes),
+ ))
+ }
}
/// Abort a fenced-but-empty instant: delete the fencing files and marker dir
diff --git a/crates/core/tests/table_write_runtime_tests.rs
b/crates/core/tests/table_write_runtime_tests.rs
new file mode 100644
index 00000000..d229209b
--- /dev/null
+++ b/crates/core/tests/table_write_runtime_tests.rs
@@ -0,0 +1,112 @@
+/*
+ * 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.
+ */
+//! How write tasks use the tokio runtime.
+//!
+//! This file holds its own test binary because it installs a global allocator
+//! that counts allocations made on a marked thread.
+
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::cell::Cell;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use arrow_array::{Int64Array, RecordBatch, StringArray};
+use arrow_schema::{DataType, Field, Schema};
+use hudi_core::table::Table;
+use tempfile::tempdir;
+
+thread_local! {
+ static COUNT_ALLOCS: Cell<bool> = const { Cell::new(false) };
+}
+static MARKED_THREAD_ALLOCS: AtomicUsize = AtomicUsize::new(0);
+
+struct CountingAllocator;
+
+unsafe impl GlobalAlloc for CountingAllocator {
+ unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+ if COUNT_ALLOCS.try_with(Cell::get).unwrap_or(false) {
+ MARKED_THREAD_ALLOCS.fetch_add(1, Ordering::Relaxed);
+ }
+ unsafe { System.alloc(layout) }
+ }
+
+ unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+ unsafe { System.dealloc(ptr, layout) }
+ }
+}
+
+#[global_allocator]
+static ALLOCATOR: CountingAllocator = CountingAllocator;
+
+fn batch(ids: &[String], value: i64) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("value", DataType::Int64, false),
+ ]));
+ RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from_iter_values(ids)),
+ Arc::new(Int64Array::from(vec![value; ids.len()])),
+ ],
+ )
+ .unwrap()
+}
+
+/// A copy-on-write upsert runs its merge off the async workers.
+///
+/// On a one-worker runtime the worker thread is marked and its allocations are
+/// counted during the upsert. The merge builds an owned key per stored and
+/// incoming row, so a merge on the worker allocates several times per row,
+/// while the worker's own share (the async read and the put) allocates per
+/// batch. Allocation counts do not depend on scheduling, so the check holds on
+/// a loaded machine.
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_cow_upsert_merge_runs_off_the_async_worker() {
+ const GROUPS: usize = 2;
+ const ROWS: usize = 20_000;
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("trips")
+ .with_record_key_fields(["id"])
+ .with_metadata(false)
+ .with_option("hoodie.parquet.small.file.limit", "0")
+ .create()
+ .await
+ .unwrap();
+ let ids: Vec<String> = (0..GROUPS * ROWS).map(|i|
format!("key{i:08}")).collect();
+ for group in ids.chunks(ROWS) {
+ table.append([batch(group, 1)]).await.unwrap();
+ }
+
+ // Spawned tasks run only on the runtime's single worker.
+ tokio::spawn(async { COUNT_ALLOCS.with(|count| count.set(true)) })
+ .await
+ .unwrap();
+ MARKED_THREAD_ALLOCS.store(0, Ordering::SeqCst);
+ let result = table.upsert([batch(&ids, 2)]).await.unwrap();
+ let worker_allocs = MARKED_THREAD_ALLOCS.load(Ordering::SeqCst);
+
+ assert_eq!(result.num_updates, GROUPS * ROWS);
+ assert!(
+ worker_allocs < ROWS,
+ "the async worker made {worker_allocs} allocations during the upsert;
the merge \
+ of {ROWS}-row file groups ran on it"
+ );
+}
