From 8281c3490c9cc007e846d02f3bc6ef500bf7ba16 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 26 Aug 2026 22:10:11 +0000 Subject: [PATCH] feat: add destination-aware native shuffle execution --- native/shuffle/src/lib.rs | 4 +- native/shuffle/src/metrics.rs | 3 +- .../shuffle/src/partitioners/empty_schema.rs | 2 +- .../src/partitioners/multi_partition.rs | 2 +- .../src/partitioners/single_partition.rs | 24 +- native/shuffle/src/rss_execution_tests.rs | 445 ++++++++++++++++++ native/shuffle/src/shuffle_writer.rs | 217 +++++++-- native/shuffle/src/writers/mod.rs | 3 +- 8 files changed, 634 insertions(+), 66 deletions(-) create mode 100644 native/shuffle/src/rss_execution_tests.rs diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 2263ae0dac0..e9674ff8956 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -19,6 +19,8 @@ pub(crate) mod comet_partitioning; pub mod ipc; pub(crate) mod metrics; pub(crate) mod partitioners; +#[cfg(test)] +mod rss_execution_tests; mod schema_align; mod shuffle_writer; mod spark_crc32c_hasher; @@ -28,5 +30,5 @@ pub(crate) mod writers; pub use comet_partitioning::CometPartitioning; pub use ipc::read_ipc_compressed; pub use schema_align::SchemaAlignExec; -pub use shuffle_writer::ShuffleWriterExec; +pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec}; pub use writers::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/metrics.rs b/native/shuffle/src/metrics.rs index 23dbfe96a15..bcb4439b784 100644 --- a/native/shuffle/src/metrics.rs +++ b/native/shuffle/src/metrics.rs @@ -33,7 +33,8 @@ pub(crate) struct ShufflePartitionerMetrics { /// Time encoding batches to IPC format pub(crate) encode_time: Time, - /// Time spent writing to disk. Maps to "shuffleWriteTime" in Spark SQL Metrics. + /// Time spent writing encoded data to its destination. Maps to "shuffleWriteTime" in Spark + /// SQL Metrics. pub(crate) write_time: Time, /// Number of input batches diff --git a/native/shuffle/src/partitioners/empty_schema.rs b/native/shuffle/src/partitioners/empty_schema.rs index 2901b82e65a..a2d0e10d3da 100644 --- a/native/shuffle/src/partitioners/empty_schema.rs +++ b/native/shuffle/src/partitioners/empty_schema.rs @@ -26,7 +26,7 @@ use tokio::time::Instant; /// A partitioner for zero-column schemas (e.g. queries where ColumnPruning removes all columns). /// This handles shuffles for operations like COUNT(*) that produce empty-schema record batches /// but contain a valid row count. Accumulates the total row count and writes a single -/// zero-column IPC batch to partition 0. All other partitions get empty entries in the index file. +/// zero-column IPC batch to partition 0. All other partitions are finalized without data. pub(crate) struct EmptySchemaShufflePartitioner { partition_writer: T, schema: SchemaRef, diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 80bf93de3ad..37ce85d1cd2 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -554,7 +554,7 @@ impl MultiPartitionShuffleRepartitioner { pub(crate) fn spill(&mut self, unreserved_bytes: usize) -> datafusion::common::Result<()> { log::info!( - "ShuffleRepartitioner spilling shuffle data of {} to disk while inserting ({} time(s) so far)", + "ShuffleRepartitioner spilling {} bytes to its partition writer ({} previous spills)", self.used(), self.spill_count() ); diff --git a/native/shuffle/src/partitioners/single_partition.rs b/native/shuffle/src/partitioners/single_partition.rs index 532fb250085..e7fd10d99b7 100644 --- a/native/shuffle/src/partitioners/single_partition.rs +++ b/native/shuffle/src/partitioners/single_partition.rs @@ -22,19 +22,19 @@ use arrow::array::RecordBatch; use std::iter; use tokio::time::Instant; -/// A partitioner that writes all shuffle data to a single file and a single index file. +/// A partitioner that streams all shuffle data to output partition zero. /// -/// Batches are streamed straight to the long-lived `BufBatchWriter` inside the -/// [`PartitionWriter`], whose internal `BatchCoalescer` combines sub-`batch_size` batches -/// into `batch_size`-row IPC blocks across calls. A batch that is already `>= batch_size` -/// and lands on an empty coalescer buffer is passed through and written verbatim as a -/// single block, which may exceed `batch_size` (see the `BatchCoalescer` bypass in -/// `BufBatchWriter`). Block boundaries therefore depend on how the input is chunked, but -/// every row is written exactly once in order. +/// Batches are streamed straight to the underlying [`PartitionWriter`]. For local output, its +/// long-lived `BufBatchWriter` uses an internal `BatchCoalescer` to combine sub-`batch_size` +/// batches into `batch_size`-row IPC blocks across calls. A batch that is already `>= batch_size` +/// and lands on an empty coalescer buffer is passed through and written verbatim as a single +/// block, which may exceed `batch_size` (see the `BatchCoalescer` bypass in `BufBatchWriter`). +/// Remote output instead sends each encoded batch directly to its task-owned callback. In either +/// case every row is written exactly once in order. /// -/// The partitioner does no buffering or concatenation of its own. A concat layer here -/// would be redundant with the coalescer, and actively wasteful whenever its output landed -/// below `batch_size` and so missed the bypass: those rows would be copied once by the +/// The partitioner does no buffering or concatenation of its own. For local output, a concat +/// layer here would be redundant with the coalescer, and actively wasteful whenever its output +/// landed below `batch_size` and so missed the bypass: those rows would be copied once by the /// concat and again into the coalescer's builders. That happens whenever the input batch /// size does not divide `batch_size` evenly. When it does divide evenly the concat output /// hit the bypass and only one copy was made either way, so streaming costs nothing there @@ -64,7 +64,7 @@ impl ShufflePartitioner for SinglePartitionShufflePartitione self.metrics.data_size.add(batch.get_array_memory_size()); self.metrics.baseline.record_output(num_rows); - // Stream directly to the writer; its BatchCoalescer handles batching to batch_size. + // Stream directly to the writer; each destination owns its encoding and batching. self.partition_writer .write(0, &mut iter::once(Ok(batch)), &self.metrics)?; } diff --git a/native/shuffle/src/rss_execution_tests.rs b/native/shuffle/src/rss_execution_tests.rs new file mode 100644 index 00000000000..8e43323fb49 --- /dev/null +++ b/native/shuffle/src/rss_execution_tests.rs @@ -0,0 +1,445 @@ +// 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 crate::{ + read_ipc_compressed, CometPartitioning, CompressionCodec, ShuffleWriterDestination, + ShuffleWriterExec, +}; +use arrow::array::{Array, Int32Array, RecordBatch, RecordBatchOptions}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::row::{RowConverter, SortField}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::config::SessionConfig; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_plan::common::collect; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_comet_jni_bridge::ShufflePartitionPusher; +use std::error::Error; +use std::fmt::{Display, Formatter}; +use std::sync::{Arc, Mutex}; + +type RecordedFrame = (i32, Vec); + +#[derive(Default)] +struct RecordingPusher { + frames: Mutex>, +} + +impl RecordingPusher { + fn frames(&self) -> Vec { + self.frames.lock().unwrap().clone() + } +} + +impl ShufflePartitionPusher for RecordingPusher { + fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()> { + self.frames + .lock() + .unwrap() + .push((partition_id, data.to_vec())); + Ok(()) + } +} + +#[derive(Debug)] +struct CallbackSentinel; + +impl Display for CallbackSentinel { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + formatter.write_str("original typed shuffle callback failure") + } +} + +impl Error for CallbackSentinel {} + +struct ExternalFailingPusher; + +impl ShufflePartitionPusher for ExternalFailingPusher { + fn push_partition_data(&self, _partition_id: i32, _data: &[u8]) -> Result<()> { + Err(DataFusionError::External(Box::new(CallbackSentinel))) + } +} + +fn int_batch(start: i32, count: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let values = Int32Array::from_iter_values(start..start + count); + RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap() +} + +fn memory_input(batches: Vec, schema: SchemaRef) -> Arc { + let config = MemorySourceConfig::try_new(std::slice::from_ref(&batches), schema, None).unwrap(); + Arc::new(DataSourceExec::new(Arc::new(config))) +} + +fn rss_execution( + batches: Vec, + schema: SchemaRef, + partitioning: CometPartitioning, + pusher: Arc, + codec: CompressionCodec, + max_frame_size: usize, + max_buffer_bytes: Option, +) -> ShuffleWriterExec { + ShuffleWriterExec::try_new_with_destination( + memory_input(batches, schema), + partitioning, + codec, + ShuffleWriterDestination::Rss { + pusher, + max_frame_size, + }, + false, + 1024 * 1024, + max_buffer_bytes, + ) + .unwrap() +} + +fn run_execution(plan: &dyn ExecutionPlan) -> Result> { + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_limit(1024 * 1024 * 1024, 1.0) + .build() + .unwrap(), + ); + let context = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + let stream = plan.execute(0, context.task_ctx())?; + tokio::runtime::Runtime::new() + .unwrap() + .block_on(collect(stream)) +} + +fn decode_frame(frame: &[u8]) -> RecordBatch { + assert!(frame.len() >= 16, "shuffle frame must contain its header"); + let encoded_length = u64::from_le_bytes(frame[..8].try_into().unwrap()); + assert_eq!( + usize::try_from(encoded_length).unwrap() + 8, + frame.len(), + "each callback must receive one complete length-prefixed frame" + ); + read_ipc_compressed(&frame[16..]).unwrap() +} + +fn values_in_frames(frames: &[RecordedFrame]) -> Vec { + frames + .iter() + .flat_map(|(_, frame)| { + let batch = decode_frame(frame); + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() +} + +fn metric_value(execution: &ShuffleWriterExec, name: &str) -> usize { + execution + .metrics() + .unwrap() + .iter() + .find(|metric| metric.value().name() == name) + .map(|metric| metric.value().as_usize()) + .unwrap_or_default() +} + +fn assert_original_callback_failure(error: DataFusionError) { + match error { + DataFusionError::External(original) => assert!( + original.downcast_ref::().is_some(), + "the original typed callback error must remain downcastable" + ), + other => panic!("callback failure was wrapped or replaced: {other}"), + } +} + +#[test] +#[cfg_attr(miri, ignore)] +fn rss_single_partition_preserves_complete_frames_and_codecs() { + let first = int_batch(0, 16); + let second = int_batch(16, 16); + + for codec in [ + CompressionCodec::None, + CompressionCodec::Lz4Frame, + CompressionCodec::Snappy, + CompressionCodec::Zstd(1), + ] { + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + vec![first.clone(), second.clone()], + first.schema(), + CometPartitioning::SinglePartition, + pusher.clone(), + codec, + 1024 * 1024, + None, + ); + + assert!(run_execution(&execution).unwrap().is_empty()); + let frames = pusher.frames(); + assert_eq!(frames.len(), 2); + assert!(frames.iter().all(|(partition_id, _)| *partition_id == 0)); + assert_eq!(values_in_frames(&frames), (0..32).collect::>()); + assert_eq!(metric_value(&execution, "input_batches"), 2); + assert_eq!(metric_value(&execution, "output_rows"), 32); + } +} + +#[test] +#[cfg_attr(miri, ignore)] +fn rss_multi_partition_supports_hash_range_and_round_robin() { + let batch = int_batch(0, 100); + let expression = Arc::new(Column::new("value", 0)); + let ordering = + LexOrdering::new(vec![PhysicalSortExpr::new_default(expression.clone())]).unwrap(); + let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap(); + let boundary_values: Arc = Arc::new(Int32Array::from(vec![25, 50, 75])); + let boundaries = converter + .convert_columns(&[boundary_values]) + .unwrap() + .iter() + .map(|row| row.owned()) + .collect(); + + for partitioning in [ + CometPartitioning::Hash(vec![expression], 4), + CometPartitioning::RangePartitioning(ordering, 4, Arc::new(converter), boundaries), + CometPartitioning::RoundRobin(4, 0), + ] { + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + vec![batch.clone()], + batch.schema(), + partitioning, + pusher.clone(), + CompressionCodec::Lz4Frame, + 1024 * 1024, + None, + ); + + run_execution(&execution).unwrap(); + let frames = pusher.frames(); + assert!(!frames.is_empty()); + assert!(frames + .iter() + .all(|(partition_id, _)| (0..4).contains(partition_id))); + let mut values = values_in_frames(&frames); + values.sort_unstable(); + assert_eq!(values, (0..100).collect::>()); + assert_eq!(metric_value(&execution, "output_rows"), 100); + } +} + +#[test] +#[cfg_attr(miri, ignore)] +fn rss_empty_schema_preserves_row_counts_in_partition_zero() { + let schema = Arc::new(Schema::empty()); + let batch = RecordBatch::try_new_with_options( + schema.clone(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(37)), + ) + .unwrap(); + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + vec![batch.clone(), batch], + schema, + CometPartitioning::RoundRobin(4, 0), + pusher.clone(), + CompressionCodec::None, + 1024 * 1024, + None, + ); + + run_execution(&execution).unwrap(); + let frames = pusher.frames(); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].0, 0); + assert_eq!(decode_frame(&frames[0].1).num_rows(), 74); + assert_eq!(metric_value(&execution, "output_rows"), 74); +} + +#[test] +fn rss_empty_schema_without_rows_does_not_push_frames() { + let schema = Arc::new(Schema::empty()); + let batch = RecordBatch::try_new_with_options( + schema.clone(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(0)), + ) + .unwrap(); + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + vec![batch], + schema, + CometPartitioning::RoundRobin(4, 0), + pusher.clone(), + CompressionCodec::None, + 1024 * 1024, + None, + ); + + run_execution(&execution).unwrap(); + assert!(pusher.frames().is_empty()); +} + +#[test] +#[cfg_attr(miri, ignore)] +fn rss_spill_pushes_remotely_without_reporting_disk_bytes() { + let batches = (0..6) + .map(|index| int_batch(index * 512, 512)) + .collect::>(); + let schema = batches[0].schema(); + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + batches, + schema, + CometPartitioning::Hash(vec![Arc::new(Column::new("value", 0))], 4), + pusher.clone(), + CompressionCodec::Lz4Frame, + 1024 * 1024, + Some(256), + ); + + run_execution(&execution).unwrap(); + assert!(metric_value(&execution, "spill_count") > 0); + assert!(metric_value(&execution, "memory_spilled_bytes") > 0); + assert_eq!(metric_value(&execution, "spilled_bytes"), 0); + let mut values = values_in_frames(&pusher.frames()); + values.sort_unstable(); + assert_eq!(values, (0..3072).collect::>()); +} + +#[test] +fn rss_preserves_typed_callback_failures_while_inserting() { + let batch = int_batch(0, 8); + let execution = rss_execution( + vec![batch.clone()], + batch.schema(), + CometPartitioning::SinglePartition, + Arc::new(ExternalFailingPusher), + CompressionCodec::None, + 1024 * 1024, + None, + ); + + assert_original_callback_failure(run_execution(&execution).unwrap_err()); +} + +#[test] +fn rss_preserves_typed_callback_failures_while_finalizing_and_spilling() { + for max_buffer_bytes in [None, Some(1)] { + let batch = int_batch(0, 64); + let execution = rss_execution( + vec![batch.clone()], + batch.schema(), + CometPartitioning::Hash(vec![Arc::new(Column::new("value", 0))], 4), + Arc::new(ExternalFailingPusher), + CompressionCodec::None, + 1024 * 1024, + max_buffer_bytes, + ); + + assert_original_callback_failure(run_execution(&execution).unwrap_err()); + } +} + +#[test] +fn rss_rejects_invalid_and_oversized_frames_without_pushing() { + for max_frame_size in [0, 1] { + let batch = int_batch(0, 8); + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + vec![batch.clone()], + batch.schema(), + CometPartitioning::SinglePartition, + pusher.clone(), + CompressionCodec::None, + max_frame_size, + None, + ); + + let error = run_execution(&execution).unwrap_err(); + assert!(error.to_string().contains("frame")); + assert!(pusher.frames().is_empty()); + } +} + +#[test] +fn rss_callback_survives_execution_plan_child_replacement() { + let original = int_batch(0, 4); + let replacement = int_batch(10, 6); + let pusher = Arc::new(RecordingPusher::default()); + let execution = rss_execution( + vec![original.clone()], + original.schema(), + CometPartitioning::SinglePartition, + pusher.clone(), + CompressionCodec::None, + 1024 * 1024, + None, + ); + let replacement_input = memory_input(vec![replacement.clone()], replacement.schema()); + let rewritten = Arc::new(execution) + .with_new_children(vec![replacement_input]) + .unwrap(); + + run_execution(rewritten.as_ref()).unwrap(); + assert_eq!( + values_in_frames(&pusher.frames()), + (10..16).collect::>() + ); +} + +#[test] +fn explicit_local_destination_preserves_data_and_index_files() { + let batch = int_batch(0, 8); + let directory = tempfile::tempdir().unwrap(); + let data_file = directory.path().join("shuffle.data"); + let index_file = directory.path().join("shuffle.index"); + let execution = ShuffleWriterExec::try_new_with_destination( + memory_input(vec![batch.clone()], batch.schema()), + CometPartitioning::SinglePartition, + CompressionCodec::None, + ShuffleWriterDestination::Local { + output_data_file: data_file.to_str().unwrap().to_string(), + output_index_file: index_file.to_str().unwrap().to_string(), + }, + false, + 1024 * 1024, + None, + ) + .unwrap(); + + run_execution(&execution).unwrap(); + let frame = std::fs::read(data_file).unwrap(); + assert_eq!(decode_frame(&frame).num_rows(), 8); + assert_eq!(std::fs::read(index_file).unwrap().len(), 16); +} diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 46d3b592bc9..810dbd626f0 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -22,10 +22,10 @@ use crate::partitioners::{ EmptySchemaShufflePartitioner, MultiPartitionShuffleRepartitioner, ShufflePartitioner, SinglePartitionShufflePartitioner, }; -use crate::writers::LocalPartitionWriter; +use crate::writers::{LocalPartitionWriter, PartitionWriter, RssPartitionWriter}; use crate::{CometPartitioning, CompressionCodec, ShuffleBlockWriter}; use async_trait::async_trait; -use datafusion::common::exec_datafusion_err; +use datafusion::common::{exec_datafusion_err, DataFusionError}; use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::EmptyRecordBatchStream; @@ -39,6 +39,7 @@ use datafusion::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, }, }; +use datafusion_comet_jni_bridge::ShufflePartitionPusher; use futures::{StreamExt, TryStreamExt}; use std::{ fmt, @@ -46,6 +47,44 @@ use std::{ sync::Arc, }; +/// Storage destination for a native shuffle writer. +#[derive(Clone)] +pub enum ShuffleWriterDestination { + /// Writes partition data and offsets to local shuffle files. + Local { + /// Path of the local shuffle data file. + output_data_file: String, + /// Path of the local shuffle index file. + output_index_file: String, + }, + /// Pushes complete encoded partition blocks to a task-owned callback. + Rss { + /// Callback that owns the remote shuffle transport for this task. + pusher: Arc, + /// Largest complete encoded shuffle block accepted by the callback. + max_frame_size: usize, + }, +} + +impl Debug for ShuffleWriterDestination { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Local { + output_data_file, + output_index_file, + } => f + .debug_struct("Local") + .field("output_data_file", output_data_file) + .field("output_index_file", output_index_file) + .finish(), + Self::Rss { max_frame_size, .. } => f + .debug_struct("Rss") + .field("max_frame_size", max_frame_size) + .finish(), + } + } +} + /// The shuffle writer operator maps each input partition to M output partitions based on a /// partitioning scheme. No guarantees are made about the order of the resulting partitions. #[derive(Debug)] @@ -54,10 +93,8 @@ pub struct ShuffleWriterExec { input: Arc, /// Partitioning scheme to use partitioning: CometPartitioning, - /// Output data file path - output_data_file: String, - /// Output index file path - output_index_file: String, + /// Task-owned output destination for encoded partition data. + destination: ShuffleWriterDestination, /// Metrics metrics: ExecutionPlanMetricsSet, /// Cache for expensive-to-compute plan properties @@ -72,7 +109,7 @@ pub struct ShuffleWriterExec { } impl ShuffleWriterExec { - /// Create a new ShuffleWriterExec + /// Creates a shuffle writer that writes to local data and index files. #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, @@ -83,6 +120,30 @@ impl ShuffleWriterExec { tracing_enabled: bool, write_buffer_size: usize, max_buffer_bytes: Option, + ) -> Result { + Self::try_new_with_destination( + input, + partitioning, + codec, + ShuffleWriterDestination::Local { + output_data_file, + output_index_file, + }, + tracing_enabled, + write_buffer_size, + max_buffer_bytes, + ) + } + + /// Creates a shuffle writer for a local or task-owned remote destination. + pub fn try_new_with_destination( + input: Arc, + partitioning: CometPartitioning, + codec: CompressionCodec, + destination: ShuffleWriterDestination, + tracing_enabled: bool, + write_buffer_size: usize, + max_buffer_bytes: Option, ) -> Result { let cache = Arc::new(PlanProperties::new( EquivalenceProperties::new(Arc::clone(&input.schema())), @@ -95,8 +156,7 @@ impl ShuffleWriterExec { input, partitioning, metrics: ExecutionPlanMetricsSet::new(), - output_data_file, - output_index_file, + destination, cache, codec, tracing_enabled, @@ -149,12 +209,11 @@ impl ExecutionPlan for ShuffleWriterExec { children: Vec>, ) -> Result> { match children.len() { - 1 => Ok(Arc::new(ShuffleWriterExec::try_new( + 1 => Ok(Arc::new(ShuffleWriterExec::try_new_with_destination( Arc::clone(&children[0]), self.partitioning.clone(), self.codec.clone(), - self.output_data_file.clone(), - self.output_index_file.clone(), + self.destination.clone(), self.tracing_enabled, self.write_buffer_size, self.max_buffer_bytes, @@ -179,8 +238,7 @@ impl ExecutionPlan for ShuffleWriterExec { futures::stream::once(external_shuffle( input, partition, - self.output_data_file.clone(), - self.output_index_file.clone(), + self.destination.clone(), self.partitioning.clone(), metrics, context, @@ -198,8 +256,7 @@ impl ExecutionPlan for ShuffleWriterExec { async fn external_shuffle( mut input: SendableRecordBatchStream, partition: usize, - output_data_file: String, - output_index_file: String, + destination: ShuffleWriterDestination, partitioning: CometPartitioning, metrics: ShufflePartitionerMetrics, context: Arc, @@ -211,40 +268,54 @@ async fn external_shuffle( let schema = input.schema(); let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?; - let local_partition_writer = LocalPartitionWriter::try_new( - output_data_file, - output_index_file, - shuffle_block_writer, - partitioning.partition_count(), - context.session_config().batch_size(), - write_buffer_size, - context.runtime_env(), - )?; - - let mut repartitioner: Box = match &partitioning { - _ if schema.fields().is_empty() => { - log::debug!("found empty schema, overriding {partitioning:?} partitioning with EmptySchemaShufflePartitioner"); - Box::new(EmptySchemaShufflePartitioner::try_new( - local_partition_writer, + let mut repartitioner = match destination { + ShuffleWriterDestination::Local { + output_data_file, + output_index_file, + } => { + let writer = LocalPartitionWriter::try_new( + output_data_file, + output_index_file, + shuffle_block_writer, + partitioning.partition_count(), + context.session_config().batch_size(), + write_buffer_size, + context.runtime_env(), + )?; + + create_repartitioner( + writer, + partition, Arc::clone(&schema), + partitioning, + metrics, + context.as_ref(), + tracing_enabled, + max_buffer_bytes, + )? + } + ShuffleWriterDestination::Rss { + pusher, + max_frame_size, + } => { + let writer = RssPartitionWriter::try_new( + shuffle_block_writer, + pusher, partitioning.partition_count(), + max_frame_size, + )?; + + create_repartitioner( + writer, + partition, + Arc::clone(&schema), + partitioning, metrics, - )?) + context.as_ref(), + tracing_enabled, + max_buffer_bytes, + )? } - any if any.partition_count() == 1 => Box::new(SinglePartitionShufflePartitioner::new( - local_partition_writer, - metrics, - )), - _ => Box::new(MultiPartitionShuffleRepartitioner::try_new( - partition, - local_partition_writer, - partitioning, - metrics, - context.runtime_env(), - context.session_config().batch_size(), - tracing_enabled, - max_buffer_bytes, - )?), }; while let Some(batch) = input.next().await { @@ -255,17 +326,67 @@ async fn external_shuffle( repartitioner .insert_batch(batch?) .await - .map_err(|err| exec_datafusion_err!("Error inserting batch: {err}"))?; + .map_err(|error| contextualize_shuffle_error(error, "inserting batch"))?; } repartitioner .shuffle_write() - .map_err(|err| exec_datafusion_err!("Error in shuffle write: {err}"))?; + .map_err(|error| contextualize_shuffle_error(error, "in shuffle write"))?; // shuffle writer always has empty output Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) as SendableRecordBatchStream) } +/// Constructs the existing schema-appropriate partitioner for either writer backend. +#[allow(clippy::too_many_arguments)] +fn create_repartitioner( + writer: T, + partition: usize, + schema: SchemaRef, + partitioning: CometPartitioning, + metrics: ShufflePartitionerMetrics, + context: &TaskContext, + tracing_enabled: bool, + max_buffer_bytes: Option, +) -> Result> { + let partition_count = partitioning.partition_count(); + + if schema.fields().is_empty() { + log::debug!( + "found empty schema, overriding {partitioning:?} partitioning with EmptySchemaShufflePartitioner" + ); + Ok(Box::new(EmptySchemaShufflePartitioner::try_new( + writer, + schema, + partition_count, + metrics, + )?)) + } else if partition_count == 1 { + Ok(Box::new(SinglePartitionShufflePartitioner::new( + writer, metrics, + ))) + } else { + Ok(Box::new(MultiPartitionShuffleRepartitioner::try_new( + partition, + writer, + partitioning, + metrics, + context.runtime_env(), + context.session_config().batch_size(), + tracing_enabled, + max_buffer_bytes, + )?)) + } +} + +/// Preserves Java throwables while retaining existing context for ordinary shuffle failures. +fn contextualize_shuffle_error(error: DataFusionError, phase: &str) -> DataFusionError { + match error { + external @ DataFusionError::External(_) => external, + other => exec_datafusion_err!("Error {phase}: {other}"), + } +} + #[cfg(test)] mod test { use super::*; diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index cc2722ca279..fb3af2c991a 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -19,8 +19,6 @@ mod buf_batch_writer; mod checksum; mod local; mod partition_writer; -// Remote shuffle execution will use this writer in a subsequent change. -#[allow(dead_code)] mod rss; mod shuffle_block_writer; @@ -28,4 +26,5 @@ pub(crate) use buf_batch_writer::BufBatchWriter; pub(crate) use checksum::Checksum; pub(crate) use local::local_partition_writer::LocalPartitionWriter; pub(crate) use partition_writer::PartitionWriter; +pub(crate) use rss::rss_partition_writer::RssPartitionWriter; pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter};