diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 7142d12461..3fdfbdccbe 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -60,6 +60,7 @@ message BallistaPhysicalPlanNode { RangeShuffleReaderExecNode range_shuffle_reader = 12; RangeFilterExecNode range_filter = 13; PrefixMergeExecNode prefix_merge = 14; + RangeShuffleWriterExecNode range_shuffle_writer = 15; } } @@ -320,6 +321,15 @@ message SortShuffleWriterExecNode { optional uint64 memory_limit_per_task_bytes = 9; } +// Passthrough shuffle writer that emits the seekable Arrow IPC file format. +// Carries no output partitioning: like `ShuffleWriterExecNode`'s passthrough +// case it writes its input's partitioning through, one file per partition. +message RangeShuffleWriterExecNode { + string job_id = 1; + uint32 stage_id = 2; + datafusion.PhysicalPlanNode input = 3; +} + message UnresolvedShuffleExecNode { uint32 stage_id = 1; datafusion_common.Schema schema = 2; @@ -359,6 +369,11 @@ message RangeShuffleReaderExecNode { repeated datafusion.PhysicalSortExprNode merge_ordering = 4; // Row limit pushed down by a consuming merge. Absent means read everything. optional uint64 fetch = 5; + // Half-open value range each output partition covers, halo already applied. + // One per output partition, or empty when the scheduler had no cuts. Lets + // the reader skip the parts of an indexed source that cannot hold a row the + // partition wants; the consuming RangeFilterExec still does the exact trim. + repeated RangeBound bounds = 6; } // CoalescePartitionsRule output: groups upstream partitions into coalesced output partitions. @@ -507,6 +522,32 @@ message ExecutePartition { datafusion.PhysicalHashRepartition output_partitioning = 6; } +// How a shuffle writer laid its output out on disk. Selects the path shape, +// and with it how a client parses the index beside the data. It says nothing +// about the framing of the bytes inside the data file. +enum ShuffleLayout { + // {stage_id}/{partition_id}/data-{file_id}.arrow — one file per output + // partition, written by the passthrough and range writers. + SHUFFLE_LAYOUT_PASSTHROUGH = 0; + // {stage_id}/{file_id}/data.arrow — one file per task holding every + // partition, written by the sort-based writer. + SHUFFLE_LAYOUT_SORT = 1; +} + +// Which of the files making up a shuffle output is being asked for. Which +// index sits beside the data, and how to read it, follows from the layout, so +// it is not spelled out here. +enum ShuffleFileKind { + SHUFFLE_FILE_KIND_DATA = 0; + SHUFFLE_FILE_KIND_INDEX = 1; +} + +// A half-open byte range of a file: `[offset, offset + length)`. +message ByteRange { + uint64 offset = 1; + uint64 length = 2; +} + message FetchPartition { string job_id = 1; uint32 stage_id = 2; @@ -514,11 +555,32 @@ message FetchPartition { string host = 5; uint32 port = 6; optional uint64 file_id = 7; - bool is_sort_shuffle = 8; - + // How the producing writer laid its output out, which is what turns the + // identifiers above into a path. + ShuffleLayout layout = 9; + // Which file beside that partition to serve. + ShuffleFileKind file_kind = 10; + // Ranges to return, concatenated in request order, as absolute offsets into + // the file. + // + // Empty asks for whatever the identifiers above address. Under the sort + // layout that is the named partition's slice, which the executor resolves + // through its own index — one round trip, as it has always been. Under the + // passthrough layout it is the whole file. + // + // With ranges given the executor serves bytes and nothing else: no index + // read, no comparison, no notion of what a range means. That is what a + // consumer needs against object storage, where there is no executor to ask, + // and it is why the range shuffle's index carries byte offsets at all. + repeated ByteRange byte_ranges = 11; + // reserved after removing deprecated `path`. reserved 4; reserved "path"; + // reserved after `is_sort_shuffle` became `layout`: the bool selected a path + // shape, which reads as data-vs-index once a file_kind exists beside it. + reserved 8; + reserved "is_sort_shuffle"; } message PartitionLocation { diff --git a/ballista/core/src/client.rs b/ballista/core/src/client.rs index 272098a917..3251ed19e1 100644 --- a/ballista/core/src/client.rs +++ b/ballista/core/src/client.rs @@ -20,7 +20,9 @@ use crate::error::{BallistaError, Result as BResult}; use crate::extension::BallistaConfigGrpcEndpoint; use crate::serde::protobuf; -use crate::serde::scheduler::{Action, PartitionId}; +use crate::serde::scheduler::{ + Action, ByteRange, PartitionId, ShuffleFileKind, ShuffleLayout, +}; use crate::utils::create_grpc_client_endpoint; use arrow_flight; use arrow_flight::Ticket; @@ -157,7 +159,7 @@ impl BallistaClient { executor_id: &str, partition_id: &PartitionId, file_id: Option, - is_sort_shuffle: bool, + layout: ShuffleLayout, flight_transport: bool, ) -> BResult { let host = self.host.to_owned(); @@ -166,7 +168,7 @@ impl BallistaClient { executor_id, partition_id, file_id, - is_sort_shuffle, + layout, &host, port, flight_transport, @@ -174,6 +176,42 @@ impl BallistaClient { .await } + /// Fetch byte ranges of a shuffle file, as absolute offsets into it. + /// + /// The executor returns exactly those bytes concatenated in request order, + /// having resolved nothing: a caller that knows which bytes it wants — from + /// an index it fetched itself — gets them without the executor reading an + /// index or comparing a value. Block transport only, since decoding is the + /// caller's business. + #[allow(clippy::too_many_arguments)] + pub async fn fetch_byte_ranges( + &mut self, + executor_id: &str, + partition_id: &PartitionId, + file_id: Option, + layout: ShuffleLayout, + file_kind: ShuffleFileKind, + byte_ranges: Vec, + header: Option>, + ) -> BResult { + let host = self.host.to_owned(); + let port = self.port; + let action = Action::FetchPartition { + job_id: partition_id.job_id.clone(), + stage_id: partition_id.stage_id, + partition_id: partition_id.partition_id, + host, + port, + file_id, + layout, + file_kind, + byte_ranges, + }; + self.execute_do_action(&action, header) + .await + .map_err(|error| Self::as_fetch_failed(executor_id, partition_id, error)) + } + /// Retrieves a partition from an executor. /// /// Depending on the value of the `flight_transport` parameter, this method will utilize either @@ -187,11 +225,14 @@ impl BallistaClient { executor_id: &str, partition_id: &PartitionId, file_id: Option, - is_sort_shuffle: bool, + layout: ShuffleLayout, host: &str, port: u16, flight_transport: bool, ) -> BResult { + // No ranges: asks for whatever these identifiers address, which under + // the sort layout is the named partition's slice and under passthrough + // is the whole file. let action = Action::FetchPartition { job_id: partition_id.job_id.clone(), stage_id: partition_id.stage_id, @@ -199,36 +240,47 @@ impl BallistaClient { host: host.to_owned(), port, file_id, - is_sort_shuffle, + layout, + file_kind: ShuffleFileKind::Data, + byte_ranges: vec![], }; let result = if flight_transport { self.execute_do_get(&action).await } else { - self.execute_do_action(&action).await + self.execute_do_action(&action, None).await }; - result - .map_err(|error| match error { - // map grpc connection error to partition fetch error. - BallistaError::GrpcActionError(msg) => { - log::warn!( - "grpc client failed to fetch partition: {partition_id:?} , message: {msg:?}" - ); - BallistaError::FetchFailed( - executor_id.to_owned(), - partition_id.stage_id, - partition_id.partition_id, - msg, - ) - } - error => { - log::warn!( - "grpc client failed to fetch partition: {partition_id:?} , error: {error:?}" - ); - error - } - }) + result.map_err(|error| Self::as_fetch_failed(executor_id, partition_id, error)) + } + + /// Report a transport failure as a partition fetch failure, which is what + /// lets the scheduler retry the task rather than fail the query. + fn as_fetch_failed( + executor_id: &str, + partition_id: &PartitionId, + error: BallistaError, + ) -> BallistaError { + match error { + // map grpc connection error to partition fetch error. + BallistaError::GrpcActionError(msg) => { + log::warn!( + "grpc client failed to fetch partition: {partition_id:?} , message: {msg:?}" + ); + BallistaError::FetchFailed( + executor_id.to_owned(), + partition_id.stage_id, + partition_id.partition_id, + msg, + ) + } + error => { + log::warn!( + "grpc client failed to fetch partition: {partition_id:?} , error: {error:?}" + ); + error + } + } } #[allow(rustdoc::private_intra_doc_links)] @@ -321,6 +373,7 @@ impl BallistaClient { pub async fn execute_do_action( &mut self, action: &Action, + header: Option>, ) -> BResult { let serialized_action: protobuf::Action = action.to_owned().try_into()?; @@ -375,6 +428,19 @@ impl BallistaClient { }) }); + // A caller fetching byte ranges gets batch messages with no schema + // ahead of them, because the schema is not in the bytes it asked + // for. It supplies the header it already knows. + let stream = + match header.clone() { + Some(header) => futures::stream::once(async move { + Ok(prost::bytes::Bytes::from(header)) + }) + .chain(stream) + .boxed(), + None => stream.boxed(), + }; + return Ok(Box::pin(BlockDataStream::try_new(stream).await?)); } unreachable!("Did not receive schema batch from flight server"); diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index d7fd6a4d18..98bcd4f6f5 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -27,6 +27,7 @@ use crate::serde::protobuf::{ scheduler_grpc_client::SchedulerGrpcClient, }; use crate::serde::protobuf::{ExecutorMetadata, SuccessfulJob}; +use crate::serde::scheduler::ShuffleLayout; use crate::utils::{GrpcClientConfig, create_grpc_client_endpoint}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; @@ -877,7 +878,11 @@ async fn fetch_partition( &metadata.id, &partition_id.into(), location.file_id, - location.is_sort_shuffle, + if location.is_sort_shuffle { + ShuffleLayout::Sort + } else { + ShuffleLayout::Passthrough + }, host, port, flight_transport, diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index d8fc21701c..eec5f3f31d 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -29,6 +29,7 @@ pub mod plan_algebra; mod prefix_merge; mod range_filter; mod range_repartition_common; +pub mod range_shuffle; mod range_shuffle_reader; mod runtime_stats; mod shuffle_reader; @@ -52,6 +53,7 @@ pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicate pub use plan_algebra::{preserves_distribution, preserves_partitioning}; pub use prefix_merge::{FinalizedPartitionState, PrefixMergeExec, ScalarOp, WindowApply}; pub use range_filter::{InputOrder, RangeBound, RangeFilterExec, WidenedBound}; +pub use range_shuffle::RangeShuffleWriterExec; pub use range_shuffle_reader::RangeShuffleReaderExec; pub use runtime_stats::{ MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, diff --git a/ballista/core/src/execution_plans/plan_algebra.rs b/ballista/core/src/execution_plans/plan_algebra.rs index 1f033dcb67..2b0a10711e 100644 --- a/ballista/core/src/execution_plans/plan_algebra.rs +++ b/ballista/core/src/execution_plans/plan_algebra.rs @@ -35,8 +35,8 @@ use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use crate::execution_plans::{ - BufferExec, PrefixMergeExec, RangeFilterExec, RuntimeStatsExec, ShuffleWriterExec, - SortShuffleWriterExec, + BufferExec, PrefixMergeExec, RangeFilterExec, RangeShuffleWriterExec, + RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, }; /// Whitelisted ops preserve the routing key's row set, values, and @@ -52,6 +52,7 @@ pub fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool { .is_some_and(|sort| sort.preserve_partitioning()) // Stage-boundary writers: batches to disk unchanged. || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() || plan.downcast_ref::().is_some() // Pure row-annotation: one input row → one output row with an // added column (window fn result); values, partitioning, count preserved. diff --git a/ballista/core/src/execution_plans/range_shuffle/index.rs b/ballista/core/src/execution_plans/range_shuffle/index.rs new file mode 100644 index 0000000000..c34d427d12 --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle/index.rs @@ -0,0 +1,745 @@ +// 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. + +//! The value index sitting beside a range shuffle file. +//! +//! One row per IPC message in the data file, in file order: +//! +//! ```text +//! key types, nullable first row's key in this batch +//! is_dict Boolean this row describes a dictionary +//! byte_offset UInt64 absolute, from the data file's footer +//! byte_len UInt64 +//! num_rows UInt64, nullable rows in the batch +//! ``` +//! +//! A consumer assigned the value range `[lo, hi)` downloads this file — KB +//! against a data file's MB — binary-searches the key columns for the batches +//! covering its range, and fetches only those byte ranges. +//! +//! # Why the keys are typed columns +//! +//! One column per ORDER BY expression, in lexicographic order, carrying that +//! expression's own type. Multi-column keys, `DESC`, and mixed types all work +//! by construction rather than through a `ScalarValue` encoding, and the +//! index's schema names what it is indexed on. +//! +//! Nullable for two independent reasons: a dictionary row has no key, and a +//! record batch's first key can genuinely be NULL under `NULLS FIRST`. Since +//! those are indistinguishable from the key alone, `is_dict` carries the +//! distinction rather than nullability implying it. +//! +//! # Why dictionaries get rows +//! +//! A record batch referencing a dictionary cannot be decoded from its own +//! bytes alone — the `DictionaryBatch` lives at its own offset earlier in the +//! file. A consumer range-fetching record batches has to fetch those too, so +//! they are in the index rather than requiring a second trip to the data +//! file's footer. +//! +//! The IPC footer does not record which dictionary a block carries (the id is +//! inside the message header), so the index cannot say which record batch +//! needs which dictionary. The reader's rule is therefore to take every +//! dictionary row preceding its selected range — conservative, and correct +//! under both replacement and delta dictionaries. +//! +//! # Why `num_rows` +//! +//! A ROWS-frame window needs N rows of context before its range starts. +//! Summing `num_rows` across the batches below a boundary — over every file +//! covering that band — says how far to widen the value range to reach N rows, +//! computed from indexes alone before any data is fetched. Without it, +//! widening means fetching and decoding to find out, one round trip per guess. +//! +//! Summation assumes each row lives in exactly one file. See +//! `project_range_shuffle_index_open_questions` on the stage-0 write +//! amplification that has to be understood before a consumer relies on it. + +use std::fs::File; +use std::io::{BufReader, BufWriter}; +use std::path::{Path, PathBuf}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::{ + Array, ArrayRef, BooleanArray, BooleanBuilder, RecordBatch, UInt64Array, + UInt64Builder, new_null_array, +}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::arrow::ipc::reader::StreamReader; +use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::common::ScalarValue; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream}; +use futures::{Stream, StreamExt, ready}; + +use crate::error::{BallistaError, Result}; + +use super::ipc_file::FileLayout; + +/// Column holding whether a row describes a dictionary rather than a batch. +pub const IS_DICT_COLUMN: &str = "is_dict"; +/// Column holding a message's absolute byte offset in the data file. +pub const BYTE_OFFSET_COLUMN: &str = "byte_offset"; +/// Column holding a message's length in bytes. +pub const BYTE_LEN_COLUMN: &str = "byte_len"; +/// Column holding a record batch's row count. +pub const NUM_ROWS_COLUMN: &str = "num_rows"; + +/// How many fixed layout columns follow the key columns. +pub const FIXED_COLUMN_COUNT: usize = 4; + +/// Offsets of the fixed columns past the last key column, in the order +/// [`fixed_fields`] emits them. +/// +/// These columns are addressed by **position, never by name**. Key columns are +/// named for the user's sort expressions and sit ahead of these, Arrow permits +/// duplicate field names, and `column_by_name` resolves to the first match — so +/// a query ordering by a column called `byte_offset` would otherwise shadow the +/// real one and be read as whatever type the key happens to be. +const IS_DICT_POSITION: usize = 0; +const BYTE_OFFSET_POSITION: usize = 1; +const BYTE_LEN_POSITION: usize = 2; +const NUM_ROWS_POSITION: usize = 3; + +/// The fixed layout columns, in the order they follow the keys. +/// +/// Defined once so [`index_schema`] and the positional accessors cannot drift: +/// an entry's position here *is* its offset past the last key column. +pub(crate) fn fixed_fields() -> [Field; FIXED_COLUMN_COUNT] { + [ + Field::new(IS_DICT_COLUMN, DataType::Boolean, false), + Field::new(BYTE_OFFSET_COLUMN, DataType::UInt64, false), + Field::new(BYTE_LEN_COLUMN, DataType::UInt64, false), + Field::new(NUM_ROWS_COLUMN, DataType::UInt64, true), + ] +} + +/// How many leading columns of `index` carry sort-key values. +/// +/// Derived from the column count, so a reader holding only the batch can still +/// address the fixed columns without knowing the ordering that produced it. +pub fn key_column_count(index: &RecordBatch) -> Result { + index + .num_columns() + .checked_sub(FIXED_COLUMN_COUNT) + .ok_or_else(|| { + BallistaError::General(format!( + "range shuffle index has {} columns, fewer than the \ + {FIXED_COLUMN_COUNT} fixed layout columns", + index.num_columns() + )) + }) +} + +/// Fixed column at `position` past the keys, downcast to the layout type +/// [`fixed_fields`] declares for it. `name` names it in the error only. +fn fixed_column<'a, T: Array + 'static>( + index: &'a RecordBatch, + position: usize, + name: &str, +) -> Result<&'a T> { + let column = key_column_count(index)? + position; + index + .column(column) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + BallistaError::General(format!( + "range shuffle index column {column} (`{name}`) has type {}, not the \ + layout type", + index.column(column).data_type() + )) + }) +} + +/// Whether each row describes a dictionary rather than a record batch. +pub fn is_dict_column(index: &RecordBatch) -> Result<&BooleanArray> { + fixed_column(index, IS_DICT_POSITION, IS_DICT_COLUMN) +} + +/// Each message's absolute byte offset in the data file. +pub fn byte_offset_column(index: &RecordBatch) -> Result<&UInt64Array> { + fixed_column(index, BYTE_OFFSET_POSITION, BYTE_OFFSET_COLUMN) +} + +/// Each message's length in bytes. +pub fn byte_len_column(index: &RecordBatch) -> Result<&UInt64Array> { + fixed_column(index, BYTE_LEN_POSITION, BYTE_LEN_COLUMN) +} + +/// Each record batch's row count; null on dictionary rows. +pub fn num_rows_column(index: &RecordBatch) -> Result<&UInt64Array> { + fixed_column(index, NUM_ROWS_POSITION, NUM_ROWS_COLUMN) +} + +/// Schema metadata key recording the sort options each key column was written +/// under, so a reader can check its own ordering agrees rather than assume it. +pub const SORT_OPTIONS_METADATA: &str = "ballista.range_shuffle.sort_options"; + +/// The sort key of one record batch's first row, one value per ORDER BY +/// expression. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchKey { + /// Lexicographic key values, positionally matching the index's key columns. + pub values: Vec, + /// Rows the batch carries. + pub num_rows: u64, +} + +/// Path of the value index sitting beside a range shuffle data file: +/// `data-{file_id}.rangeidx.arrow` beside `data-{file_id}.arrow`. +/// +/// The index is itself an Arrow IPC stream, so `.arrow` is its extension and +/// `rangeidx` says which role it plays — the sort shuffle's `.arrow.index` +/// names the two the other way round. +pub fn index_path(data_path: &Path) -> PathBuf { + data_path.with_extension("rangeidx.arrow") +} + +/// True when `data_path` has a value index beside it. +pub fn has_range_index(data_path: &Path) -> bool { + index_path(data_path).exists() +} + +/// Write the value index for one data file. +/// +/// Arrow IPC stream format: the index is read whole, so it has no use for a +/// footer, and the stream framing is what every existing shuffle reader +/// already speaks. +pub fn write_index_file( + path: &Path, + schema: SchemaRef, + layout: &FileLayout, + keys: &[BatchKey], +) -> Result<()> { + let batch = build_index_batch(schema.clone(), layout, keys)?; + let file = File::create(path).map_err(|e| { + BallistaError::General(format!( + "range shuffle index: failed to create {path:?}: {e:?}" + )) + })?; + let mut writer = StreamWriter::try_new(BufWriter::new(file), schema.as_ref())?; + writer.write(&batch)?; + writer.finish()?; + Ok(()) +} + +/// Read the value index beside a range shuffle data file. +pub fn read_index_file(path: &Path) -> Result { + let file = File::open(path).map_err(|e| { + BallistaError::General(format!( + "range shuffle index: failed to open {path:?}: {e:?}" + )) + })?; + let mut reader = StreamReader::try_new(BufReader::new(file), None)?; + let batch = reader.next().transpose()?.ok_or_else(|| { + BallistaError::General(format!("range shuffle index {path:?} holds no rows")) + })?; + if reader.next().is_some() { + return Err(BallistaError::General(format!( + "range shuffle index {path:?} holds more than one batch" + ))); + } + Ok(batch) +} + +/// How many record batches an index describes, ignoring its dictionary rows. +pub fn count_record_batches(index: &RecordBatch) -> Result { + let is_dict = is_dict_column(index)?; + Ok((0..index.num_rows()) + .filter(|&row| !is_dict.value(row)) + .count()) +} + +/// Which record batches of an indexed file can hold rows in `[lo, hi)`. +/// +/// Returns ordinals into the file's record batches — the order +/// `FileReader::set_index` addresses them by — or `None` when the index cannot +/// decide and the caller should read the whole file. +/// +/// `None` on either bound means unbounded on that side. Selection is at batch +/// granularity and deliberately inclusive at the edges: a batch is taken when +/// it *may* hold a row in range, and the exact trim belongs to the +/// `RangeFilterExec` above the reader. Taking a batch too many costs bytes; +/// taking one too few loses rows. +/// +/// # When it declines +/// +/// A NULL key means the batch's first row has no value under the ordering, so +/// where it sits relative to `lo` and `hi` depends on null placement rather +/// than on comparison. Rather than encode that here, the index declines and +/// the caller reads everything — the answer stays right and only the saving is +/// lost. +pub fn select_record_batches( + index: &RecordBatch, + lo: Option<&ScalarValue>, + hi: Option<&ScalarValue>, + descending: bool, +) -> Result>> { + let is_dict = is_dict_column(index)?; + let keys = index.column(0); + + // Record batches in file order, which is the order `set_index` counts in. + let mut first_keys = Vec::with_capacity(index.num_rows()); + for row in 0..index.num_rows() { + if is_dict.value(row) { + continue; + } + if keys.is_null(row) { + return Ok(None); + } + first_keys.push(ScalarValue::try_from_array(keys, row)?); + } + + // `precedes(a, b)` is "a comes before b under this ordering", so the same + // walk serves ASC and DESC. + let precedes = |a: &ScalarValue, b: &ScalarValue| match a.partial_cmp(b) { + Some(std::cmp::Ordering::Less) => !descending, + Some(std::cmp::Ordering::Greater) => descending, + _ => false, + }; + + // Batch `i` covers from its own first key up to the next batch's, so the + // first batch that can hold `lo` is the last one starting at or before it. + let start = match lo { + None => 0, + Some(lo) => first_keys + .iter() + .rposition(|key| !precedes(lo, key)) + .unwrap_or(0), + }; + let end = match hi { + None => first_keys.len(), + Some(hi) => first_keys.iter().filter(|key| precedes(key, hi)).count(), + }; + + Ok(Some((start..end.max(start)).collect())) +} + +/// Collect each batch's first-row sort key as batches pass through to the +/// writer. +/// +/// A stream adapter rather than a hook inside the write path: the keys are the +/// index's business and the IPC framing is the writer's, and keeping them apart +/// means the write path never has to know what the data is sorted on. +/// +/// Every batch yields an entry, including empty ones — the IPC writer emits a +/// block for an empty batch too, and the index pairs keys with blocks +/// positionally, so skipping one would shift every key onto the wrong byte +/// range. +pub struct KeyCollector { + inner: SendableRecordBatchStream, + ordering: LexOrdering, + keys: Vec, +} + +impl KeyCollector { + /// Wrap `inner`, evaluating `ordering` against each batch's first row. + pub fn new(inner: SendableRecordBatchStream, ordering: LexOrdering) -> Self { + Self { + inner, + ordering, + keys: Vec::new(), + } + } + + /// The keys collected so far, in the order the batches were yielded. + pub fn keys(&self) -> &[BatchKey] { + &self.keys + } + + /// Take ownership of the collected keys. + pub fn into_keys(self) -> Vec { + self.keys + } + + /// Evaluate the ordering against `batch`'s first row. + /// + /// An empty batch has no first row, so its key is all NULL — it selects + /// nothing on a range search, which is the truth about a block holding no + /// rows. + fn key_for(&self, batch: &RecordBatch) -> Result { + let num_rows = batch.num_rows() as u64; + let mut values = Vec::with_capacity(self.ordering.len()); + for sort_expr in self.ordering.iter() { + if num_rows == 0 { + let data_type = + sort_expr.expr.data_type(batch.schema_ref()).map_err(|e| { + BallistaError::General(format!( + "range shuffle index: sort expression `{}` has no type: {e}", + sort_expr.expr + )) + })?; + values.push(ScalarValue::try_from(&data_type)?); + continue; + } + let first_row = batch.slice(0, 1); + let evaluated = sort_expr.expr.evaluate(&first_row).map_err(|e| { + BallistaError::General(format!( + "range shuffle index: sort expression `{}` failed to evaluate: {e}", + sort_expr.expr + )) + })?; + let array = evaluated.into_array(1)?; + values.push(ScalarValue::try_from_array(&array, 0)?); + } + Ok(BatchKey { values, num_rows }) + } +} + +impl Stream for KeyCollector { + type Item = datafusion::error::Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + match ready!(self.inner.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + let key = match self.key_for(&batch) { + Ok(key) => key, + Err(e) => { + return Poll::Ready(Some(Err(e.into_datafusion()))); + } + }; + self.keys.push(key); + Poll::Ready(Some(Ok(batch))) + } + other => Poll::Ready(other), + } + } +} + +impl RecordBatchStream for KeyCollector { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } +} + +/// Build the index's schema for a given ORDER BY. +/// +/// Key columns come first, named for the expression they carry so the file +/// says what it is indexed on, followed by the fixed layout columns. +pub fn index_schema(ordering: &LexOrdering, data_schema: &Schema) -> Result { + let mut fields = Vec::with_capacity(ordering.len() + 4); + for sort_expr in ordering.iter() { + let data_type = sort_expr.expr.data_type(data_schema).map_err(|e| { + BallistaError::General(format!( + "range shuffle index: sort expression `{}` has no type against the \ + shuffle schema: {e}", + sort_expr.expr + )) + })?; + fields.push(Field::new(key_column_name(sort_expr), data_type, true)); + } + fields.extend(fixed_fields()); + + let metadata = std::collections::HashMap::from([( + SORT_OPTIONS_METADATA.to_string(), + encode_sort_options(ordering), + )]); + Ok(Arc::new(Schema::new_with_metadata(fields, metadata))) +} + +/// Name of the index column carrying `sort_expr`'s values. +fn key_column_name(sort_expr: &PhysicalSortExpr) -> String { + sort_expr.expr.to_string() +} + +/// Render each key's `ASC`/`DESC` and null placement, comma separated and in +/// lexicographic order. +fn encode_sort_options(ordering: &LexOrdering) -> String { + ordering + .iter() + .map(|sort_expr| { + let direction = if sort_expr.options.descending { + "desc" + } else { + "asc" + }; + let nulls = if sort_expr.options.nulls_first { + "nulls_first" + } else { + "nulls_last" + }; + format!("{direction} {nulls}") + }) + .collect::>() + .join(",") +} + +/// Build the index batch for one data file. +/// +/// `keys` are the record batches' first-row keys in write order, which is the +/// order `layout.record_batches` is in — the data file's own footer is what +/// pairs a key with a byte range, so the two are positional and must be the +/// same length. +/// +/// Rows come out sorted by byte offset, putting dictionaries ahead of the +/// batches referencing them, so a reader walking the index in order sees the +/// file's layout as it is. +pub fn build_index_batch( + schema: SchemaRef, + layout: &FileLayout, + keys: &[BatchKey], +) -> Result { + if layout.record_batches.len() != keys.len() { + return Err(BallistaError::General(format!( + "range shuffle index: {} record blocks against {} keys — the index \ + would pair keys with the wrong byte ranges", + layout.record_batches.len(), + keys.len(), + ))); + } + + let key_count = schema.fields().len() - 4; + + // (offset, len, Some(key index)) for batches, None for dictionaries. + let mut rows: Vec<(u64, u64, Option)> = + Vec::with_capacity(layout.dictionaries.len() + layout.record_batches.len()); + rows.extend( + layout + .dictionaries + .iter() + .map(|block| (block.offset, block.len, None)), + ); + rows.extend( + layout + .record_batches + .iter() + .enumerate() + .map(|(idx, block)| (block.offset, block.len, Some(idx))), + ); + rows.sort_by_key(|(offset, _, _)| *offset); + + let mut key_columns: Vec> = + vec![Vec::with_capacity(rows.len()); key_count]; + let mut is_dict = BooleanBuilder::with_capacity(rows.len()); + let mut byte_offset = UInt64Builder::with_capacity(rows.len()); + let mut byte_len = UInt64Builder::with_capacity(rows.len()); + let mut num_rows = UInt64Builder::with_capacity(rows.len()); + + for (offset, len, key_idx) in &rows { + byte_offset.append_value(*offset); + byte_len.append_value(*len); + match key_idx { + Some(idx) => { + let key = &keys[*idx]; + if key.values.len() != key_count { + return Err(BallistaError::General(format!( + "range shuffle index: batch key has {} values against {} \ + key columns", + key.values.len(), + key_count, + ))); + } + for (column, value) in key_columns.iter_mut().zip(&key.values) { + column.push(value.clone()); + } + is_dict.append_value(false); + num_rows.append_value(key.num_rows); + } + None => { + for (column, field) in key_columns.iter_mut().zip(schema.fields()) { + column.push(ScalarValue::try_from(field.data_type())?); + } + is_dict.append_value(true); + num_rows.append_null(); + } + } + } + + let mut columns: Vec = Vec::with_capacity(schema.fields().len()); + for (values, field) in key_columns.into_iter().zip(schema.fields()) { + columns.push(if values.is_empty() { + new_null_array(field.data_type(), 0) + } else { + ScalarValue::iter_to_array(values)? + }); + } + columns.push(Arc::new(is_dict.finish())); + columns.push(Arc::new(byte_offset.finish())); + columns.push(Arc::new(byte_len.finish())); + columns.push(Arc::new(num_rows.finish())); + + RecordBatch::try_new(schema, columns) + .map_err(|e| BallistaError::General(format!("range shuffle index: {e}"))) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution_plans::sort_shuffle::is_sort_shuffle_output; + use datafusion::arrow::array::{Float64Array, UInt64Array}; + use datafusion::arrow::datatypes::Field; + + /// The positional accessors address the fixed columns by offset past the + /// keys, so their constants have to agree with the order `fixed_fields` + /// emits. Reordering one without the other would read the wrong column and + /// only fail where the types happen to differ. + #[test] + fn fixed_column_positions_match_the_field_order() { + let fields = fixed_fields(); + assert_eq!(fields[IS_DICT_POSITION].name(), IS_DICT_COLUMN); + assert_eq!(fields[BYTE_OFFSET_POSITION].name(), BYTE_OFFSET_COLUMN); + assert_eq!(fields[BYTE_LEN_POSITION].name(), BYTE_LEN_COLUMN); + assert_eq!(fields[NUM_ROWS_POSITION].name(), NUM_ROWS_COLUMN); + } + + /// A sort key named after a fixed column is legal — Arrow permits + /// duplicate field names and resolves a lookup to the *first* match, and + /// key columns sit first. Reading the fixed columns by name would return + /// the key here; by position it stays correct. + #[test] + fn a_key_named_like_a_fixed_column_does_not_shadow_it() { + let mut fields = vec![Field::new(BYTE_OFFSET_COLUMN, DataType::Float64, true)]; + fields.extend(fixed_fields()); + let index = RecordBatch::try_new( + Arc::new(Schema::new(fields)), + vec![ + Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0)])), + Arc::new(BooleanArray::from(vec![false, false])), + Arc::new(UInt64Array::from(vec![100u64, 200])), + Arc::new(UInt64Array::from(vec![10u64, 20])), + Arc::new(UInt64Array::from(vec![Some(4u64), Some(5)])), + ], + ) + .unwrap(); + + assert_eq!(key_column_count(&index).unwrap(), 1); + assert_eq!(byte_offset_column(&index).unwrap().values(), &[100, 200]); + assert_eq!(byte_len_column(&index).unwrap().values(), &[10, 20]); + assert_eq!(num_rows_column(&index).unwrap().values(), &[4, 5]); + assert!(!is_dict_column(&index).unwrap().value(0)); + assert_eq!(count_record_batches(&index).unwrap(), 2); + + // The by-name lookup this replaced would have found the Float64 key. + assert_eq!( + index + .column_by_name(BYTE_OFFSET_COLUMN) + .unwrap() + .data_type(), + &DataType::Float64, + ); + } + + /// An index over batches whose first keys are `keys`, laid out as the + /// writer lays them out. + fn index_over(keys: &[Option]) -> RecordBatch { + let mut fields = vec![Field::new("v", DataType::Float64, true)]; + fields.extend(fixed_fields()); + let schema = Arc::new(Schema::new(fields)); + let rows = keys.len(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Float64Array::from(keys.to_vec())), + Arc::new(BooleanArray::from(vec![false; rows])), + Arc::new(UInt64Array::from( + (0..rows).map(|r| r as u64 * 100).collect::>(), + )), + Arc::new(UInt64Array::from(vec![100u64; rows])), + Arc::new(UInt64Array::from(vec![10u64; rows])), + ], + ) + .unwrap() + } + + fn select( + keys: &[Option], + lo: Option, + hi: Option, + ) -> Option> { + select_record_batches( + &index_over(keys), + lo.map(ScalarValue::from).as_ref(), + hi.map(ScalarValue::from).as_ref(), + false, + ) + .unwrap() + } + + /// Batch `i` holds rows from its own first key up to the next batch's, so + /// a range starting mid-batch has to take that batch, not the one after. + /// Dropping it would silently lose the rows between `lo` and the next key. + #[test] + fn selects_every_batch_that_can_hold_a_row_in_range() { + // First keys 0, 10, 20, 30 — batch 1 holds [10, 20). + let keys = &[Some(0.0), Some(10.0), Some(20.0), Some(30.0)]; + + assert_eq!(select(keys, Some(12.0), Some(19.0)), Some(vec![1])); + assert_eq!(select(keys, Some(12.0), Some(21.0)), Some(vec![1, 2])); + // A bound landing exactly on a boundary takes the batch starting there. + assert_eq!(select(keys, Some(10.0), Some(20.0)), Some(vec![1])); + // Unbounded below starts at the first batch, above runs to the last. + assert_eq!(select(keys, None, Some(15.0)), Some(vec![0, 1])); + assert_eq!(select(keys, Some(25.0), None), Some(vec![2, 3])); + assert_eq!(select(keys, None, None), Some(vec![0, 1, 2, 3])); + } + + /// A range below everything or above everything selects nothing rather + /// than wrapping around or panicking on an empty walk. + #[test] + fn selects_nothing_outside_the_files_range() { + let keys = &[Some(10.0), Some(20.0)]; + assert_eq!(select(keys, Some(0.0), Some(5.0)), Some(vec![])); + assert_eq!(select(keys, Some(30.0), Some(40.0)), Some(vec![1])); + } + + /// A NULL first key has no position under value comparison, so the index + /// declines rather than guessing — the caller reads everything and the + /// answer stays right. + #[test] + fn declines_when_a_key_is_null() { + let keys = &[Some(10.0), None, Some(30.0)]; + assert_eq!(select(keys, Some(12.0), Some(19.0)), None); + } + + /// Under DESC the keys walk downward, so "before" flips; selection has to + /// follow the declared ordering rather than assume ascending. + #[test] + fn follows_a_descending_ordering() { + let index = index_over(&[Some(30.0), Some(20.0), Some(10.0)]); + let selected = select_record_batches( + &index, + Some(&ScalarValue::from(25.0)), + Some(&ScalarValue::from(15.0)), + true, + ) + .unwrap(); + // Descending: `lo` is the high end. 25 sits inside batch 0's [30, 20). + assert_eq!(selected, Some(vec![0, 1])); + } + + /// The index sits in the same directory as the data files and ends in + /// `.arrow` like they do, so its name has to be one no data file can take + /// and one the sort shuffle's probe does not claim. + #[test] + fn index_path_cannot_collide_with_a_data_file() { + let data = Path::new("/work/job/1/0/data-7.arrow"); + let index = index_path(data); + + assert_eq!(index, Path::new("/work/job/1/0/data-7.rangeidx.arrow")); + assert_ne!(index, data.to_path_buf()); + assert!( + !is_sort_shuffle_output(&index), + "the sort shuffle probes for `.arrow.index`, which this must not be", + ); + } +} diff --git a/ballista/core/src/execution_plans/range_shuffle/ipc_file.rs b/ballista/core/src/execution_plans/range_shuffle/ipc_file.rs new file mode 100644 index 0000000000..3177889eb0 --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle/ipc_file.rs @@ -0,0 +1,494 @@ +// 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. + +//! Arrow IPC **file** format write and open for the range shuffle. +//! +//! The passthrough shuffle writes the IPC **stream** format, which has no +//! index: finding the byte range that holds a given value means walking every +//! message from the head. The file format ends with a footer listing one +//! `Block { offset, metadata_length, body_length }` per record batch, so a +//! reader that knows which batches it wants can seek straight to them. +//! +//! The seeking itself sits above this module: `reader` turns a consumer's +//! value range into batch ordinals and opens the file at them, and `remote` +//! turns those ordinals into byte ranges to fetch. What this module provides is +//! the substrate both stand on — a footer to seek against, and a reader for it. +//! +//! # Where the offsets come from +//! +//! The footer, read back after the file is closed — not byte accounting at +//! the writer's boundary. One `FileWriter::write` call emits a batch's +//! dictionary blocks *and* its record block, so a byte counter wrapped around +//! that call cannot say where one ends and the next begins: it would hand out +//! an offset pointing at a dictionary and a length spanning both. Data with no +//! dictionary columns hides that completely, since the two coincide. +//! +//! The footer already separates them, exactly, as computed by the writer that +//! emitted them. It costs one tail read per shuffle file and removes a whole +//! class of silent misaddressing. +//! +//! [`RangeShuffleWriterExec`]: super::RangeShuffleWriterExec +//! [`RangeShuffleReaderExec`]: crate::execution_plans::RangeShuffleReaderExec + +use std::fs::File; +use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom}; +use std::path::Path; + +use datafusion::arrow::ipc::reader::{FileReader, read_footer_length}; +use datafusion::arrow::ipc::writer::FileWriter; +use datafusion::arrow::ipc::{Block, CompressionType, root_as_footer}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::physical_plan::{RecordBatchStream, metrics}; +use futures::StreamExt; +use log::{debug, error}; + +use crate::error::{BallistaError, Result}; +use crate::serde::scheduler::PartitionStats; +use crate::utils::create_write_options; + +/// Leading bytes of an Arrow IPC file: the `ARROW1` magic plus two padding +/// bytes. An IPC stream opens with a continuation marker instead, so the two +/// formats are distinguishable from the head of the file. +const ARROW_FILE_MAGIC: &[u8; 6] = b"ARROW1"; + +/// Where one IPC message sits in the file, copied verbatim from the footer's +/// `Block`. +/// +/// `offset` is absolute from the file start and `[offset, offset + len)` is +/// the byte range a reader fetches to decode this message on its own — which +/// is what a value index turns a cut range into. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MessageBlock { + /// Byte offset of the IPC message, absolute from the file start. + pub offset: u64, + /// Bytes the message occupies: metadata plus body. + pub len: u64, +} + +/// A finished IPC file's message layout, as its footer records it. +/// +/// The two kinds are kept apart because they are consumed differently: a +/// consumer selects the record batches covering its value range, but must +/// also take the dictionaries those batches reference, which sit at their own +/// offsets earlier in the file. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct FileLayout { + /// Dictionary batches, in file order. + pub dictionaries: Vec, + /// Record batches, in file order — positionally the batches as written. + pub record_batches: Vec, +} + +/// True when `path` holds an Arrow IPC **file**, false for an IPC stream. +/// +/// The on-disk layout is authoritative — same principle as +/// `is_sort_shuffle_output`. Deciding by magic rather than by a flag on +/// `PartitionLocation` keeps the format out of the wire protocol, so a +/// consumer, a Flight server, and a debugging human all reach the same +/// answer from the file alone. +/// +/// An unreadable or too-short file reads as "not a file format", leaving the +/// error to whichever reader opens it next and can report it in context. +pub fn is_ipc_file(path: &Path) -> bool { + let Ok(mut file) = File::open(path) else { + return false; + }; + let mut magic = [0u8; ARROW_FILE_MAGIC.len()]; + match file.read_exact(&mut magic) { + Ok(()) => &magic == ARROW_FILE_MAGIC, + Err(_) => false, + } +} + +/// Read a finished IPC file's footer and return where every message sits. +/// +/// The trailer's last 10 bytes give the footer's length, the footer gives the +/// blocks. Called once per shuffle file after the writer closes it, so the +/// index records what the writer actually emitted rather than a reconstruction +/// of it. +pub fn read_file_layout(path: &Path) -> Result { + let mut file = File::open(path).map_err(|e| { + BallistaError::General(format!( + "Failed to open range shuffle file {path:?}: {e:?}" + )) + })?; + + let mut trailer = [0u8; 10]; + file.seek(SeekFrom::End(-10)).map_err(|e| { + BallistaError::General(format!("Failed to seek {path:?} trailer: {e:?}")) + })?; + file.read_exact(&mut trailer).map_err(|e| { + BallistaError::General(format!("Failed to read {path:?} trailer: {e:?}")) + })?; + let footer_len = read_footer_length(trailer).map_err(|e| { + BallistaError::General(format!("{path:?} has no readable IPC footer: {e:?}")) + })?; + + let mut footer_buf = vec![0u8; footer_len]; + file.seek(SeekFrom::End(-10 - footer_len as i64)) + .map_err(|e| { + BallistaError::General(format!("Failed to seek {path:?} footer: {e:?}")) + })?; + file.read_exact(&mut footer_buf).map_err(|e| { + BallistaError::General(format!("Failed to read {path:?} footer: {e:?}")) + })?; + + let footer = root_as_footer(&footer_buf).map_err(|e| { + BallistaError::General(format!("Failed to parse {path:?} footer: {e:?}")) + })?; + + // Taken by bound rather than by the footer vector's concrete type, which + // would pull `flatbuffers` in as a direct dependency just to name it. + fn blocks<'a>( + source: Option>, + ) -> Vec { + source + .map(|blocks| { + blocks + .into_iter() + .map(|block| MessageBlock { + offset: block.offset() as u64, + len: block.metaDataLength() as u64 + block.bodyLength() as u64, + }) + .collect() + }) + .unwrap_or_default() + } + + Ok(FileLayout { + dictionaries: blocks(footer.dictionaries()), + record_batches: blocks(footer.recordBatches()), + }) +} + +/// Open an Arrow IPC file for a whole-file read. +pub fn open_ipc_file(path: &Path) -> Result>> { + let file = File::open(path).map_err(|e| { + BallistaError::General(format!( + "Failed to open range shuffle file {path:?}: {e:?}" + )) + })?; + let file = BufReader::with_capacity(256 * 1024, file); + // Safety: setting `skip_validation` requires `unsafe`, user assures data is valid + let reader = unsafe { + FileReader::try_new(file, None) + .map_err(|e| { + BallistaError::General(format!( + "Failed to create arrow FileReader at {path:?}: {e:?}" + )) + })? + .with_skip_validation(cfg!(feature = "arrow-ipc-optimizations")) + }; + Ok(reader) +} + +/// Stream data to disk in Arrow IPC file format, returning where each batch +/// landed alongside the usual stats. +/// +/// Structured like [`crate::utils::write_stream_to_disk`]: batches cross a +/// bounded channel into a `spawn_blocking` task that owns every synchronous +/// file operation, so no tokio worker blocks on I/O. +pub async fn write_stream_to_ipc_file( + stream: &mut S, + path: &Path, + disk_write_metric: &metrics::Time, + channel_capacity: usize, + compression_type: Option, +) -> Result<(PartitionStats, FileLayout)> +where + S: RecordBatchStream + Unpin + ?Sized, +{ + let schema = stream.schema(); + let path_owned = path.to_owned(); + let write_metric = disk_write_metric.clone(); + + let (tx, mut rx) = tokio::sync::mpsc::channel::(channel_capacity); + + let handle = + tokio::task::spawn_blocking(move || -> Result<(u64, usize, FileLayout)> { + let file = File::create(&path_owned).map_err(|e| { + error!("Failed to create partition file at {:?}: {e:?}", path_owned); + BallistaError::IoError(e) + })?; + + let options = create_write_options(compression_type)?; + let mut writer = FileWriter::try_new_with_options( + BufWriter::new(file), + schema.as_ref(), + options, + )?; + + let mut batches_written = 0; + while let Some(batch) = rx.blocking_recv() { + let timer = write_metric.timer(); + writer.write(&batch)?; + batches_written += 1; + timer.done(); + } + let timer = write_metric.timer(); + writer.finish()?; + timer.done(); + + // Only readable once the footer is on disk, which `finish` above + // is what guarantees. + let layout = read_file_layout(&path_owned)?; + Ok(( + std::fs::metadata(&path_owned).map(|m| m.len()).unwrap_or(0), + batches_written, + layout, + )) + }); + + let mut num_rows = 0; + let mut num_batches = 0; + + let stream_err = loop { + match stream.next().await { + Some(Ok(batch)) => { + num_batches += 1; + num_rows += batch.num_rows(); + if tx.send(batch).await.is_err() { + break None; + } + } + Some(Err(e)) => break Some(e), + None => break None, + } + }; + drop(tx); + + let write_result = handle + .await + .map_err(|e| BallistaError::General(format!("Disk writer task failed: {e}")))?; + + if let Some(e) = stream_err { + if let Err(write_err) = &write_result { + error!("Disk writer also failed: {write_err}"); + } + return Err(e.into()); + } + let (num_bytes, batches_written, layout) = write_result?; + + // Record blocks are addressed positionally by the index that follows, so + // one block per batch written is the invariant that makes that mapping + // valid. A mismatch means the writer re-chunked and every offset would + // name the wrong batch. + if layout.record_batches.len() != batches_written { + return Err(BallistaError::General(format!( + "range shuffle {path:?} footer records {} record batches but {} were \ + written — offsets cannot be mapped to batches", + layout.record_batches.len(), + batches_written, + ))); + } + + debug!( + "range shuffle wrote {path:?}: {} record blocks, {} dictionary blocks: {layout:?}", + layout.record_batches.len(), + layout.dictionaries.len(), + ); + + Ok(( + PartitionStats::new(Some(num_rows as u64), Some(num_batches), Some(num_bytes)), + layout, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{DictionaryArray, Int32Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use datafusion::arrow::ipc::writer::StreamWriter; + use datafusion::arrow::ipc::{Message, root_as_message}; + use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use std::sync::Arc; + use tempfile::tempdir; + + fn batch(values: &[i32]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + ])); + let payload: Vec = values.iter().map(|v| format!("row-{v}")).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(values.to_vec())), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap() + } + + /// Decode the IPC message occupying `[start, end)`, the way a reader + /// issuing a range request over just those bytes would have to. + fn read_message_at(bytes: &[u8], start: usize, end: usize) -> Message<'_> { + // An IPC message opens with a continuation marker and a little-endian + // metadata length, then the flatbuffer header. + let slice = &bytes[start..end]; + assert_eq!( + &slice[0..4], + &[0xff, 0xff, 0xff, 0xff], + "continuation marker" + ); + let meta_len = i32::from_le_bytes(slice[4..8].try_into().unwrap()) as usize; + root_as_message(&slice[8..8 + meta_len]).unwrap() + } + + async fn write_batches( + path: &Path, + batches: Vec, + ) -> (PartitionStats, FileLayout) { + let schema = batches[0].schema(); + let metrics = ExecutionPlanMetricsSet::new(); + let write_time = MetricBuilder::new(&metrics).subset_time("write_time", 0); + let mut stream = RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches.into_iter().map(Ok)), + ); + write_stream_to_ipc_file(&mut stream, path, &write_time, 8, None) + .await + .unwrap() + } + + /// The offsets are what a value index will hand to a range read, so each + /// one has to address a message that decodes on its own. Slicing the file + /// at `[offset, offset + len)` and decoding is the check that catches an + /// offset naming the wrong thing; arithmetic self-consistency would not. + #[tokio::test] + async fn record_offsets_address_decodable_batches() { + let dir = tempdir().unwrap(); + let path = dir.path().join("data-0.arrow"); + let batches = vec![batch(&[1, 2, 3]), batch(&[4, 5]), batch(&[6, 7, 8, 9])]; + + let (_, layout) = write_batches(&path, batches.clone()).await; + + assert_eq!(layout.record_batches.len(), 3, "one block per batch"); + assert!( + layout.dictionaries.is_empty(), + "no dictionary columns in these batches", + ); + + let bytes = std::fs::read(&path).unwrap(); + for (block, expected) in layout.record_batches.iter().zip(&batches) { + let start = block.offset as usize; + let message = read_message_at(&bytes, start, start + block.len as usize); + assert_eq!( + message.header_as_record_batch().map(|rb| rb.length()), + Some(expected.num_rows() as i64), + "block at {start} must hold this batch's record message", + ); + } + } + + /// Dictionary-encoded data is where byte accounting silently broke: one + /// `write` call emits the dictionary and the record batch together, so a + /// counter wrapped around that call hands out an offset pointing at the + /// dictionary. The footer keeps the two apart, and only data with a + /// dictionary column can tell the difference. + #[tokio::test] + async fn separates_dictionary_blocks_from_record_blocks() { + let dir = tempdir().unwrap(); + let path = dir.path().join("dict.arrow"); + + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )])); + let dict: DictionaryArray = + vec!["a", "b", "a", "c"].into_iter().collect(); + let batch = RecordBatch::try_new(schema, vec![Arc::new(dict)]).unwrap(); + + let (_, layout) = write_batches(&path, vec![batch.clone()]).await; + + assert_eq!(layout.record_batches.len(), 1); + assert_eq!( + layout.dictionaries.len(), + 1, + "the dictionary must be recorded as its own block", + ); + + let dictionary = layout.dictionaries[0]; + let record = layout.record_batches[0]; + // The distinction byte counting could not make: the record block + // starts past the dictionary rather than at it. + assert!( + record.offset >= dictionary.offset + dictionary.len, + "record block must start after the dictionary ends, got {record:?} \ + against {dictionary:?}", + ); + + let bytes = std::fs::read(&path).unwrap(); + let start = record.offset as usize; + let message = read_message_at(&bytes, start, start + record.len as usize); + assert_eq!( + message.header_as_record_batch().map(|rb| rb.length()), + Some(batch.num_rows() as i64), + ); + let start = dictionary.offset as usize; + let message = read_message_at(&bytes, start, start + dictionary.len as usize); + assert!( + message.header_as_dictionary_batch().is_some(), + "the dictionary block must hold a dictionary message", + ); + } + + /// The layout has to describe the same file the whole-file reader sees, or + /// a seeking consumer and a streaming one disagree about the contents. + #[tokio::test] + async fn layout_matches_what_the_reader_sees() { + let dir = tempdir().unwrap(); + let path = dir.path().join("data-0.arrow"); + let batches = vec![batch(&[1, 2]), batch(&[3]), batch(&[4, 5, 6])]; + + let (_, layout) = write_batches(&path, batches.clone()).await; + + let reader = open_ipc_file(&path).unwrap(); + assert_eq!(reader.num_batches(), layout.record_batches.len()); + let read_back: Vec = reader.map(|b| b.unwrap()).collect(); + assert_eq!(read_back, batches, "round trip must preserve the batches"); + } + + /// Format detection decides which reader opens a shuffle file, so it has + /// to separate the two IPC formats and not merely detect "some arrow". + #[test] + fn detects_file_format_against_stream_format() { + let dir = tempdir().unwrap(); + let as_file = dir.path().join("file.arrow"); + let as_stream = dir.path().join("stream.arrow"); + let batches = vec![batch(&[1, 2, 3])]; + + tokio::runtime::Runtime::new() + .unwrap() + .block_on(write_batches(&as_file, batches.clone())); + + let out = File::create(&as_stream).unwrap(); + let mut writer = StreamWriter::try_new(out, &batches[0].schema()).unwrap(); + writer.write(&batches[0]).unwrap(); + writer.finish().unwrap(); + + assert!(is_ipc_file(&as_file), "IPC file must be detected"); + assert!(!is_ipc_file(&as_stream), "IPC stream must not be"); + assert!( + !is_ipc_file(&dir.path().join("absent.arrow")), + "a missing file is not an IPC file", + ); + } +} diff --git a/ballista/core/src/execution_plans/range_shuffle/mod.rs b/ballista/core/src/execution_plans/range_shuffle/mod.rs new file mode 100644 index 0000000000..875dae1a58 --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle/mod.rs @@ -0,0 +1,59 @@ +// 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. + +//! Range shuffle: a passthrough shuffle whose output can be seeked into. +//! +//! A stage that routes rows by value range (`OrderedRangeRepartitionExec`) +//! hands each consumer a set of producer files, of which the consumer wants +//! only the rows inside its own cut range. It reads them whole and drops the +//! rest — measured on h2o Q8 at 1e8 rows, the consumer fetches 8.2 GB to feed +//! a window that needs 3.6 GB. +//! +//! Skipping the rest means seeking, and seeking means the data file has to +//! carry a chunk index. The passthrough shuffle's Arrow IPC **stream** format +//! has none. This module writes the IPC **file** format instead, whose footer +//! lists a `Block { offset, metadata_length, body_length }` per record batch. +//! +//! Reads are still whole-file: the point of this layer is the substrate, not +//! the saving. What it establishes is a footer to seek against, read back once +//! per file so the byte range of every record batch — and of every dictionary +//! those batches reference — is known exactly. That layout is what the value +//! index turns a consumer's cut range into. + +mod index; +mod ipc_file; +mod reader; +mod remote; +mod writer; + +// The fixed layout columns are reached through the positional accessors, not +// by name: key columns are named for the user's sort expressions and sit ahead +// of them, and Arrow resolves a duplicate name to the first match. The name +// constants stay internal so no caller outside can reintroduce that lookup. +pub use index::{ + BatchKey, KeyCollector, SORT_OPTIONS_METADATA, build_index_batch, byte_len_column, + byte_offset_column, count_record_batches, has_range_index, index_path, index_schema, + is_dict_column, key_column_count, num_rows_column, read_index_file, + select_record_batches, write_index_file, +}; +pub use ipc_file::{ + FileLayout, MessageBlock, is_ipc_file, open_ipc_file, read_file_layout, + write_stream_to_ipc_file, +}; +pub use reader::{SelectedBatches, open_ipc_file_range, select_local_batches}; +pub use remote::{byte_ranges_for, covers, schema_message}; +pub use writer::RangeShuffleWriterExec; diff --git a/ballista/core/src/execution_plans/range_shuffle/reader.rs b/ballista/core/src/execution_plans/range_shuffle/reader.rs new file mode 100644 index 0000000000..96c87bf093 --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle/reader.rs @@ -0,0 +1,111 @@ +// 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. + +//! Reading only the part of a range shuffle file a consumer's range covers. + +use std::fs::File; +use std::io::BufReader; +use std::path::Path; + +use datafusion::arrow::error::ArrowError; +use datafusion::arrow::ipc::reader::FileReader; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::ScalarValue; +use log::debug; + +use crate::error::Result; + +use super::index::{ + SORT_OPTIONS_METADATA, has_range_index, index_path, read_index_file, + select_record_batches, +}; +use super::ipc_file::open_ipc_file; + +/// Yields the record batches at `ordinals`, seeking to each. +/// +/// The seek is `FileReader::set_index`, which resolves an ordinal against the +/// file's own footer — read once when the file was opened. The index's byte +/// offsets are what a remote reader needs, where paying for a footer read +/// defeats the point; a local reader already has it. +pub struct SelectedBatches { + reader: FileReader>, + ordinals: std::vec::IntoIter, +} + +impl Iterator for SelectedBatches { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + let ordinal = self.ordinals.next()?; + if let Err(e) = self.reader.set_index(ordinal) { + return Some(Err(e)); + } + self.reader.next() + } +} + +/// What a data file's index says about the batches covering `[lo, hi)`. +/// +/// `None` means read the whole file: either there is no index beside it, or +/// the index declined to narrow the range. Both are correctness-preserving — +/// only the saving is lost. +pub fn select_local_batches( + data_path: &Path, + lo: Option<&ScalarValue>, + hi: Option<&ScalarValue>, +) -> Result>> { + if !has_range_index(data_path) { + return Ok(None); + } + let index = read_index_file(&index_path(data_path))?; + let descending = index + .schema() + .metadata() + .get(SORT_OPTIONS_METADATA) + .and_then(|options| options.split(',').next().map(str::to_owned)) + .is_some_and(|options| options.starts_with("desc")); + + select_record_batches(&index, lo, hi, descending) +} + +/// Open a range shuffle file, reading only the batches covering `[lo, hi)`. +/// +/// Selection is at batch granularity, so the stream still carries rows outside +/// the range at either edge — the `RangeFilterExec` above the reader does the +/// exact trim, as it does for a whole-file read. +pub fn open_ipc_file_range( + data_path: &Path, + lo: Option<&ScalarValue>, + hi: Option<&ScalarValue>, +) -> Result { + let reader = open_ipc_file(data_path)?; + let ordinals = match select_local_batches(data_path, lo, hi)? { + Some(ordinals) => { + debug!( + "range shuffle reading {} of {} batches from {data_path:?}", + ordinals.len(), + reader.num_batches(), + ); + ordinals + } + None => (0..reader.num_batches()).collect(), + }; + Ok(SelectedBatches { + reader, + ordinals: ordinals.into_iter(), + }) +} diff --git a/ballista/core/src/execution_plans/range_shuffle/remote.rs b/ballista/core/src/execution_plans/range_shuffle/remote.rs new file mode 100644 index 0000000000..b422a3dc9c --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle/remote.rs @@ -0,0 +1,362 @@ +// 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. + +//! Reading a remote range shuffle file down to the bytes a range covers. +//! +//! Two fetches per source: the index, then the bytes it points at. The +//! consumer does the searching, so the executor never reads an index or +//! compares a value — which is what makes the same two steps work against +//! object storage, where there is no executor to ask. +//! +//! The extra round trip is paid once per source and, batched across sources, +//! once per read. What it buys is the difference between a whole file and the +//! batches that can hold a row the consumer wants. + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::Schema; +use datafusion::arrow::ipc::writer::{ + DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message, +}; +use datafusion::common::ScalarValue; + +use crate::error::{BallistaError, Result}; +use crate::serde::scheduler::ByteRange; + +use super::index::{byte_len_column, byte_offset_column, is_dict_column}; + +/// Encode `schema` as an IPC schema message, the header a stream decoder needs +/// before any batch. +/// +/// Synthesized rather than sliced out of the data file. The file's own schema +/// message sits after the magic padded to the writer's alignment — 64 bytes in +/// arrow-rs, 8 in some other writers — so slicing it means encoding another +/// implementation's padding rule into a byte offset. The consumer already knows +/// the schema it asked for, so it can write the header itself and fetch only +/// batches. +pub fn schema_message(schema: &Schema) -> Result> { + let options = IpcWriteOptions::default(); + let generator = IpcDataGenerator::default(); + let mut tracker = DictionaryTracker::new(false); + let encoded = + generator.schema_to_bytes_with_dictionary_tracker(schema, &mut tracker, &options); + + let mut bytes = Vec::new(); + write_message(&mut bytes, encoded, &options).map_err(|e| { + BallistaError::General(format!("range shuffle: cannot encode schema: {e}")) + })?; + Ok(bytes) +} + +/// The byte ranges holding the batches covering `[lo, hi)`, and the +/// dictionaries they need. +/// +/// Two parts, in file order: +/// +/// - every dictionary block before the selected batches, since a batch +/// referencing a dictionary cannot be decoded without it and the index +/// cannot say which batch needs which +/// - the selected record batches, coalesced, so a contiguous range costs one +/// request +/// +/// No schema range: see [`schema_message`]. `None` when nothing was selected — +/// the consumer's range misses this file entirely and there is nothing to +/// fetch. +pub fn byte_ranges_for( + index: &RecordBatch, + selected: &[usize], +) -> Result>> { + if selected.is_empty() { + return Ok(None); + } + + let is_dict = is_dict_column(index)?; + let offsets = byte_offset_column(index)?; + let lengths = byte_len_column(index)?; + + // `selected` counts record batches; the index counts every message. Walk + // once, mapping one onto the other. + let mut record_ordinal = 0; + let mut dictionaries = Vec::new(); + let mut records = Vec::new(); + let last_selected = selected.last().copied().unwrap_or(0); + + for row in 0..index.num_rows() { + let range = ByteRange { + offset: offsets.value(row), + length: lengths.value(row), + }; + + if is_dict.value(row) { + // Conservative: every dictionary preceding the last batch we want. + // Replacement and delta dictionaries both stay correct, at the cost + // of dictionaries a narrower rule could have skipped. + if records.is_empty() || record_ordinal <= last_selected { + dictionaries.push(range); + } + continue; + } + if selected.contains(&record_ordinal) { + records.push(range); + } + record_ordinal += 1; + } + + if records.is_empty() { + return Ok(None); + } + + let mut ranges = Vec::with_capacity(dictionaries.len() + records.len()); + ranges.extend(dictionaries); + ranges.extend(coalesce(records)); + Ok(Some(ranges)) +} + +/// Merge ranges that touch, so a contiguous run of batches costs one request +/// rather than one per batch. A value range selects a contiguous run, so this +/// is the common case and the one where per-request latency would otherwise +/// undo the saving. +fn coalesce(ranges: Vec) -> Vec { + let mut merged: Vec = Vec::with_capacity(ranges.len()); + for range in ranges { + match merged.last_mut() { + Some(last) if last.offset + last.length == range.offset => { + last.length += range.length; + } + _ => merged.push(range), + } + } + merged +} + +/// Whether an index says anything in this file can fall inside `[lo, hi)`. +/// +/// Used to skip a source outright rather than fetch bytes from it. +pub fn covers( + index: &RecordBatch, + lo: Option<&ScalarValue>, + hi: Option<&ScalarValue>, + descending: bool, +) -> Result { + Ok( + super::index::select_record_batches(index, lo, hi, descending)? + .is_none_or(|selected| !selected.is_empty()), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution_plans::range_shuffle::index::{ + BYTE_OFFSET_COLUMN, fixed_fields, + }; + use datafusion::arrow::array::{BooleanArray, Float64Array, UInt64Array}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use std::sync::Arc; + + /// Build an index whose rows are `(is_dict, offset, len)` in file order. + /// + /// The key column is deliberately named `byte_offset`, colliding with a + /// fixed column: Arrow permits duplicate field names and resolves lookups + /// to the first match, so this fixture reads correctly only while the + /// fixed columns are addressed by position. + fn index(rows: &[(bool, u64, u64)]) -> RecordBatch { + let mut fields = vec![Field::new(BYTE_OFFSET_COLUMN, DataType::Float64, true)]; + fields.extend(fixed_fields()); + RecordBatch::try_new( + Arc::new(Schema::new(fields)), + vec![ + Arc::new(Float64Array::from(vec![Some(1.0); rows.len()])), + Arc::new(BooleanArray::from( + rows.iter().map(|r| r.0).collect::>(), + )), + Arc::new(UInt64Array::from( + rows.iter().map(|r| r.1).collect::>(), + )), + Arc::new(UInt64Array::from( + rows.iter().map(|r| r.2).collect::>(), + )), + Arc::new(UInt64Array::from(vec![Some(1u64); rows.len()])), + ], + ) + .unwrap() + } + + /// A contiguous run of batches has to cost one range, not one per batch, + /// or per-request latency eats the bandwidth saving. + #[test] + fn coalesces_a_contiguous_run_into_one_range() { + // Blocks at 100, 200, 300, 400, each 100 bytes. + let index = index(&[ + (false, 100, 100), + (false, 200, 100), + (false, 300, 100), + (false, 400, 100), + ]); + + let ranges = byte_ranges_for(&index, &[1, 2]).unwrap().unwrap(); + + assert_eq!( + ranges, + vec![ByteRange { + offset: 200, + length: 200 + }], + "two adjacent batches are one request, not two", + ); + } + + /// A batch referencing a dictionary cannot be decoded without it, and the + /// index cannot say which batch needs which — so preceding dictionaries + /// come along. + #[test] + fn includes_dictionaries_preceding_the_selection() { + let index = index(&[ + (true, 100, 50), // dictionary + (false, 150, 100), // batch 0 + (false, 250, 100), // batch 1 + ]); + + let ranges = byte_ranges_for(&index, &[1]).unwrap().unwrap(); + + assert!( + ranges.contains(&ByteRange { + offset: 100, + length: 50 + }), + "the dictionary must be fetched with the batch: {ranges:?}", + ); + assert!( + ranges.contains(&ByteRange { + offset: 250, + length: 100 + }), + "the selected batch must be fetched: {ranges:?}", + ); + } + + /// Selecting nothing means this file holds nothing the consumer wants, so + /// there is no request to make — not a request for zero bytes. + #[test] + fn selecting_nothing_fetches_nothing() { + let index = index(&[(false, 100, 100)]); + assert_eq!(byte_ranges_for(&index, &[]).unwrap(), None); + } + + /// The ranges are handed to an executor that resolves nothing, so what + /// comes back is exactly these bytes concatenated — and it has to decode as + /// an IPC stream. Slicing a real file and decoding the result is the only + /// check that covers the schema prefix, the offsets, and the framing + /// together. + #[tokio::test] + async fn sliced_ranges_decode_as_an_ipc_stream() { + use crate::execution_plans::range_shuffle::{ + index_path, read_index_file, select_record_batches, write_index_file, + }; + use crate::execution_plans::range_shuffle::{ + index_schema, write_stream_to_ipc_file, + }; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::buffer::Buffer; + use datafusion::arrow::ipc::reader::StreamDecoder; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, + }; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use tempfile::tempdir; + + let dir = tempdir().unwrap(); + let path = dir.path().join("data-0.arrow"); + + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let batches: Vec = [vec![0, 1], vec![2, 3], vec![4, 5]] + .iter() + .map(|keys| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(keys.clone()))], + ) + .unwrap() + }) + .collect(); + + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("k", 0), + ))]) + .unwrap(); + + let metrics = ExecutionPlanMetricsSet::new(); + let write_time = MetricBuilder::new(&metrics).subset_time("write_time", 0); + let mut keyed = crate::execution_plans::range_shuffle::KeyCollector::new( + Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + futures::stream::iter(batches.clone().into_iter().map(Ok)), + )), + ordering.clone(), + ); + let (_, layout) = + write_stream_to_ipc_file(&mut keyed, &path, &write_time, 8, None) + .await + .unwrap(); + let keys = keyed.into_keys(); + let index_file = index_path(&path); + write_index_file( + &index_file, + index_schema(&ordering, schema.as_ref()).unwrap(), + &layout, + &keys, + ) + .unwrap(); + + // Ask for the middle batch only, the way a consumer covering [2, 4) + // would. + let index = read_index_file(&index_file).unwrap(); + let selected = select_record_batches( + &index, + Some(&ScalarValue::Int32(Some(2))), + Some(&ScalarValue::Int32(Some(4))), + false, + ) + .unwrap() + .unwrap(); + assert_eq!(selected, vec![1], "only the middle batch holds [2, 4)"); + + let ranges = byte_ranges_for(&index, &selected).unwrap().unwrap(); + + // What the client assembles: the header it writes itself, then the + // bytes the executor concatenates. + let file = std::fs::read(&path).unwrap(); + let mut body = schema_message(schema.as_ref()).unwrap(); + for range in &ranges { + let start = range.offset as usize; + body.extend_from_slice(&file[start..start + range.length as usize]); + } + + let mut decoder = StreamDecoder::new(); + let mut buffer = Buffer::from(body); + let mut decoded = Vec::new(); + while let Some(batch) = decoder.decode(&mut buffer).unwrap() { + decoded.push(batch); + } + assert_eq!( + decoded, + vec![batches[1].clone()], + "the sliced bytes must decode to exactly the selected batch", + ); + } +} diff --git a/ballista/core/src/execution_plans/range_shuffle/writer.rs b/ballista/core/src/execution_plans/range_shuffle/writer.rs new file mode 100644 index 0000000000..11e00e4624 --- /dev/null +++ b/ballista/core/src/execution_plans/range_shuffle/writer.rs @@ -0,0 +1,849 @@ +// 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. + +//! Passthrough shuffle writer whose output can be seeked into. + +use std::fmt::Debug; +use std::future::Future; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::TaskContext; +use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; +use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, SendableRecordBatchStream, Statistics, StatisticsArgs, + statistics::ChildStats, +}; +use futures::TryStreamExt; +use log::debug; +use tokio::sync::oneshot; +use tokio::task::JoinSet; + +use crate::JobId; +use crate::error::BallistaError; +use crate::execution_plans::ObservedWindowState; +use crate::execution_plans::create_shuffle_path; +use crate::execution_plans::shuffle_writer::{ + ShuffleWriteMetrics, WriterState, collect_window_state_against_slice, result_schema, + run_coordinator, summaries_to_batch, walk_child_partition_mapping, +}; +use crate::execution_plans::shuffle_writer_trait::ShuffleWriter; +use crate::extension::SessionConfigExt; +use crate::serde::protobuf::ShuffleWritePartition; + +use super::index::{KeyCollector, index_path, index_schema, write_index_file}; +use super::ipc_file::write_stream_to_ipc_file; + +/// A passthrough shuffle writer that emits Arrow IPC **file** format. +/// +/// Identical to [`ShuffleWriterExec`] in shape — one file per output +/// partition, K concurrent drains, the same K-summary contract to the +/// framework — and differs only in the framing of what it writes. The stream +/// format that [`ShuffleWriterExec`] emits has no index, so a consumer wanting +/// one value range out of a producer file has to read from the head; the file +/// format's footer lists every record batch's byte range, which is what makes +/// a seek possible. +/// +/// # Why a separate writer rather than a mode on the existing one +/// +/// The two formats are not interchangeable on the read side: an IPC stream +/// decoder rejects a file (the leading magic is not a continuation marker) and +/// vice versa. A flag on [`ShuffleWriterExec`] would put every existing reader +/// — local, Flight `do_get`, and the raw block transport — one config change +/// away from failing to decode what it is handed. As a distinct operator, the +/// existing shuffle keeps its format untouched and range shuffle is a new file +/// type that only the readers taught to recognise it will open. +/// +/// This is planted only where the consumer is a +/// [`RangeShuffleReaderExec`](crate::execution_plans::RangeShuffleReaderExec), +/// which is the same condition the reader is chosen under: the stage's child +/// declares an output ordering. +/// +/// # Message layout +/// +/// Each drain reads back its file's footer, so it knows the byte range of +/// every record batch and every dictionary. That layout is what the sidecar +/// index is written from: each block is paired with the first key the +/// `KeyCollector` saw in it, giving a consumer the two halves it needs to turn +/// a value range into byte ranges without opening the data file. +/// +/// [`ShuffleWriterExec`]: crate::execution_plans::ShuffleWriterExec +#[derive(Debug)] +pub struct RangeShuffleWriterExec { + /// Unique ID for the job (query) that this stage is a part of + job_id: JobId, + /// Unique query stage ID within the job + stage_id: usize, + /// Physical execution plan for this query stage + plan: Arc, + /// Ordering the child declares, which the index is keyed on. Captured at + /// construction so the write path and the index agree on the sort key by + /// construction rather than by both consulting the plan separately. + ordering: LexOrdering, + /// Path to write output streams to + work_dir: String, + /// Task id used as `file_id` in shuffle paths so files from different + /// tasks (including retries) don't collide. Stamped by the executor's + /// `create_query_stage_exec`. + task_id: usize, + /// Global partition ids this task's restricted plan covers, in slice + /// order. Position `i` in the child plan corresponds to + /// `global_output_partition_ids[i]` globally. + global_output_partition_ids: Vec, + metrics: ExecutionPlanMetricsSet, + properties: Arc, + /// Shared coordinator handoff state. Clones share the same Arc, so a + /// clone produced by `with_new_children` participates in the same + /// coordinator. + state: Arc>, +} + +impl Clone for RangeShuffleWriterExec { + fn clone(&self) -> Self { + Self { + job_id: self.job_id.clone(), + stage_id: self.stage_id, + plan: self.plan.clone(), + ordering: self.ordering.clone(), + work_dir: self.work_dir.clone(), + task_id: self.task_id, + global_output_partition_ids: self.global_output_partition_ids.clone(), + metrics: self.metrics.clone(), + properties: self.properties.clone(), + state: self.state.clone(), + } + } +} + +impl std::fmt::Display for RangeShuffleWriterExec { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let printable_plan = DisplayableExecutionPlan::with_metrics(self.plan.as_ref()) + .set_show_statistics(true) + .indent(false); + write!( + f, + "RangeShuffleWriterExec: job={} stage={} work_dir={} plan: \n {}", + self.job_id, self.stage_id, self.work_dir, printable_plan + ) + } +} + +impl RangeShuffleWriterExec { + /// Create a new range shuffle writer. `task_id` defaults to 0; the + /// executor stamps the real value at `create_query_stage_exec` time. + pub fn try_new( + job_id: JobId, + stage_id: usize, + plan: Arc, + work_dir: String, + ) -> Result { + // The index is keyed on the child's declared ordering, so a child that + // declares none has nothing to index and this writer is the wrong one. + // The adapter only plants it where an ordering exists, so reaching here + // without one is a planting bug rather than a user error. + let ordering = plan.output_ordering().cloned().ok_or_else(|| { + DataFusionError::Internal(format!( + "RangeShuffleWriterExec requires an ordered input, but `{}` \ + declares none", + plan.name() + )) + })?; + + // This writer never repartitions, so its output partitioning is + // exactly its input plan's. + let partitioning = plan.properties().output_partitioning().clone(); + let output_partition_count = partitioning.partition_count(); + let properties = Arc::new(PlanProperties::new( + datafusion::physical_expr::EquivalenceProperties::new(plan.schema()), + partitioning, + datafusion::physical_plan::execution_plan::EmissionType::Incremental, + datafusion::physical_plan::execution_plan::Boundedness::Bounded, + )); + Ok(Self { + job_id, + stage_id, + plan, + ordering, + work_dir, + task_id: 0, + global_output_partition_ids: (0..output_partition_count).collect(), + metrics: ExecutionPlanMetricsSet::new(), + properties, + state: Arc::new(Mutex::new(WriterState { + initialized: false, + handoffs: (0..output_partition_count).map(|_| None).collect(), + })), + }) + } + + /// Bind this writer to a specific task_id, so shuffle files from + /// different tasks in the same stage don't collide on file_id. + pub fn with_task_id(mut self, task_id: usize) -> Self { + self.task_id = task_id; + self + } + + /// Task id (append-order slot within the stage) this writer is bound to. + pub fn task_id(&self) -> usize { + self.task_id + } + + /// Bind this writer to the task's assigned global partition slice. + pub fn with_global_output_partition_ids( + mut self, + global_output_partition_ids: Vec, + ) -> Self { + self.global_output_partition_ids = global_output_partition_ids; + self + } + + /// Global partition slice this writer instance is bound to. + pub fn global_output_partition_ids(&self) -> &[usize] { + &self.global_output_partition_ids + } + + /// Drain every window-state collector in this stage, translating each + /// capture's task-local partition index to its global one. + /// + /// Shares `ShuffleWriterExec::collect_window_state`'s body: this writer + /// preserves its input partitioning too, so a window can sit under it and + /// its captures must reach the downstream prefix merge. Silently returning + /// none would make that merge arithmetically wrong with nothing later to + /// catch it. + pub fn collect_window_state(&self) -> Result> { + collect_window_state_against_slice( + &self.plan, + &self.global_output_partition_ids, + "RangeShuffleWriterExec", + ) + } + + /// Get the Job ID for this query stage + pub fn job_id(&self) -> &JobId { + &self.job_id + } + + /// Get the Stage ID for this query stage + pub fn stage_id(&self) -> usize { + self.stage_id + } + + /// Work directory shuffle files are written under. + pub fn work_dir(&self) -> &str { + &self.work_dir + } + + /// Executes the shuffle write for this task, draining all K output + /// partitions concurrently. + /// + /// Returns `(handoff_idx, summary)` pairs, where `summary.partition_id` is + /// the **global** output partition id downstream will address. + /// + /// All K must drain concurrently rather than one at a time: + /// `OrderedRangeRepartitionExec` below pushes to all K senders from shared + /// scatter tasks, so draining one to EOF first fills the undrained + /// channels and deadlocks the scatter side. + pub fn execute_shuffle_write( + self, + context: Arc, + ) -> impl Future>> { + let task_id = self.task_id; + let plan = self.plan.clone(); + let partition_map = + walk_child_partition_mapping(&plan, &self.global_output_partition_ids); + let metrics = self.metrics.clone(); + + async move { + let now = Instant::now(); + let config = context.session_config().ballista_config(); + let compression_type = config.shuffle_compression_codec()?; + let channel_capacity = config.shuffle_writer_channel_capacity(); + + let num_partitions = + plan.properties().output_partitioning().partition_count(); + // One schema for every partition this task writes: same ordering, + // same data schema, so building it per drain would just repeat the + // same type resolution. + let index_schema = index_schema(&self.ordering, plan.schema().as_ref()) + .map_err(BallistaError::into_datafusion)?; + let mut handles = JoinSet::new(); + for local_input_partition in 0..num_partitions { + let write_metrics = + ShuffleWriteMetrics::new(local_input_partition, &metrics); + let global_partition = + partition_map.resolve(local_input_partition) as usize; + let path = create_shuffle_path( + &self.work_dir, + &self.job_id, + self.stage_id, + global_partition, + Some(task_id as u64), + false, + )?; + + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + debug!("Writing range shuffle results to {path:?}"); + + let stream = plan.execute(local_input_partition, context.clone())?; + let index_schema = index_schema.clone(); + let ordering = self.ordering.clone(); + handles.spawn(async move { + // The keys are collected from the batches on their way to + // the writer, so a key and the block it describes come from + // the same batch rather than from two passes that could + // disagree. + let mut keyed = KeyCollector::new(stream, ordering); + let (stats, layout) = write_stream_to_ipc_file( + &mut keyed, + path.as_path(), + &write_metrics.write_time, + channel_capacity, + compression_type, + ) + .await + .map_err(BallistaError::into_datafusion)?; + let keys = keyed.into_keys(); + + write_index_file( + index_path(path.as_path()).as_path(), + index_schema, + &layout, + &keys, + ) + .map_err(BallistaError::into_datafusion)?; + + let rows = stats.num_rows.unwrap_or(0) as usize; + write_metrics.input_rows.add(rows); + write_metrics.output_rows.add(rows); + debug!( + "range shuffle partition {global_partition} indexed {} record \ + blocks and {} dictionary blocks", + layout.record_batches.len(), + layout.dictionaries.len(), + ); + Ok::<_, DataFusionError>((local_input_partition, stats)) + }); + } + + let mut results = Vec::with_capacity(num_partitions); + while let Some(joined) = handles.join_next().await { + let (local_input_partition, stats) = joined.map_err(|e| { + DataFusionError::Execution(format!( + "range shuffle-write drain task panicked: {e}" + )) + })??; + results.push(( + local_input_partition, + ShuffleWritePartition { + partition_id: partition_map.resolve(local_input_partition), + num_batches: stats.num_batches.unwrap_or(0), + num_rows: stats.num_rows.unwrap_or(0), + num_bytes: stats.num_bytes.unwrap_or(0), + file_id: Some(task_id as u64), + is_sort_shuffle: false, + }, + )); + } + debug!( + "range shuffle task_id {} drained {} partitions in {}s", + task_id, + num_partitions, + now.elapsed().as_secs() + ); + Ok(results) + } + } +} + +impl DisplayAs for RangeShuffleWriterExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + let partitioning = self.properties().output_partitioning(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "RangeShuffleWriterExec: partitioning: {partitioning}, format: arrow ipc file" + ) + } + DisplayFormatType::TreeRender => { + write!(f, "partitioning={partitioning}") + } + } + } +} + +impl ExecutionPlan for RangeShuffleWriterExec { + fn name(&self) -> &str { + "RangeShuffleWriterExec" + } + + fn schema(&self) -> SchemaRef { + self.plan.schema() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.plan] + } + + /// Owns no expressions — this writer preserves its input partitioning, so + /// any partitioning expressions belong to the child plan. + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + let [_] = children.as_slice() else { + return Err(DataFusionError::Plan( + "Ballista RangeShuffleWriterExec expects single child".to_owned(), + )); + }; + let input = children.pop().expect("single child checked above"); + Ok(Arc::new( + RangeShuffleWriterExec::try_new( + self.job_id.clone(), + self.stage_id, + input, + self.work_dir.clone(), + )? + .with_task_id(self.task_id) + .with_global_output_partition_ids(self.global_output_partition_ids.clone()), + )) + } + + /// Return the stream for output partition `partition`. + /// + /// The first call initializes K oneshot channels and spawns one + /// coordinator that drives every output partition's write, sending each + /// summary to its matching partition's oneshot. Callers should spawn all K + /// `execute(N, ctx)` calls concurrently so the drains don't serialize. + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let schema = result_schema(); + + let mut state = self.state.lock().map_err(|_| { + DataFusionError::Internal( + "RangeShuffleWriterExec state mutex poisoned".to_owned(), + ) + })?; + + if !state.initialized { + state.initialized = true; + let k = state.handoffs.len(); + let mut senders: Vec>>> = + Vec::with_capacity(k); + for slot in state.handoffs.iter_mut() { + let (tx, rx) = oneshot::channel(); + senders.push(tx); + *slot = Some(rx); + } + let writer = self.clone(); + let ctx = context.clone(); + tokio::spawn(async move { + run_coordinator(writer.execute_shuffle_write(ctx), senders).await; + }); + } + + let rx = state + .handoffs + .get_mut(partition) + .and_then(Option::take) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "RangeShuffleWriterExec: execute({partition}) called twice or out of range (have {} partitions)", + state.handoffs.len() + )) + })?; + drop(state); + + let work_dir = self.work_dir.clone(); + let job_id = self.job_id.clone(); + let stage_id = self.stage_id; + let schema_captured = schema.clone(); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { + let summaries = rx.await.map_err(|_| { + DataFusionError::Internal( + "RangeShuffleWriterExec coordinator dropped without sending" + .to_owned(), + ) + })??; + summaries_to_batch( + summaries, + schema_captured, + &work_dir, + &job_id, + stage_id, + ) + }) + .try_flatten(), + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } +} + +impl ShuffleWriter for RangeShuffleWriterExec { + fn job_id(&self) -> &JobId { + &self.job_id + } + + fn stage_id(&self) -> usize { + self.stage_id + } + + /// Always `None`: this writer preserves its input partitioning. + fn shuffle_output_partitioning(&self) -> Option<&Partitioning> { + None + } + + fn input_partition_count(&self) -> usize { + self.plan + .properties() + .output_partitioning() + .partition_count() + } + + fn clone_box(&self) -> Arc { + Arc::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::super::index::{ + SORT_OPTIONS_METADATA, byte_len_column, byte_offset_column, is_dict_column, + num_rows_column, read_index_file, + }; + use super::super::ipc_file::read_file_layout; + use super::*; + use crate::execution_plans::shuffle_reader::fetch_partition_local; + use crate::serde::scheduler::{ + ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, + PartitionId, PartitionLocation, PartitionStats, + }; + use crate::utils::collect_stream; + use datafusion::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_plan::sorts::sort::SortExec; + use datafusion::prelude::SessionContext; + use tempfile::TempDir; + + fn test_batch(keys: &[i32]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + ])); + let payload: Vec = keys.iter().map(|k| format!("row-{k}")).collect(); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(keys.to_vec())), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap() + } + + fn key_ordering() -> LexOrdering { + LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "k", 0, + )))]) + .unwrap() + } + + /// A two-partition input, each partition holding two batches, sorted on + /// the key — this writer indexes on its child's declared ordering, so an + /// unordered child is refused at construction. + fn input_plan() -> Arc { + let partitions = vec![ + vec![test_batch(&[1, 2]), test_batch(&[3, 4])], + vec![test_batch(&[5, 6]), test_batch(&[7, 8])], + ]; + let schema = partitions[0][0].schema(); + let source = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&partitions, schema, None).unwrap(), + ))); + Arc::new(SortExec::new(key_ordering(), source).with_preserve_partitioning(true)) + } + + fn location(job: &str, stage_id: usize, partition: usize) -> PartitionLocation { + PartitionLocation { + map_partition_id: partition, + partition_id: PartitionId { + job_id: job.into(), + stage_id, + partition_id: partition, + }, + executor_meta: ExecutorMetadata { + id: "test-executor".to_string(), + host: "127.0.0.1".to_string(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default(), + os_info: ExecutorOperatingSystemSpecification::default(), + }, + partition_stats: PartitionStats::default(), + file_id: Some(0), + is_sort_shuffle: false, + } + } + + /// Drive every output partition concurrently, the way the executor does — + /// the coordinator only runs once every oneshot receiver has been taken. + async fn drive(writer: Arc, ctx: Arc) { + let k = writer.properties().output_partitioning().partition_count(); + let mut handles = Vec::with_capacity(k); + for n in 0..k { + let writer = writer.clone(); + let ctx = ctx.clone(); + handles.push(tokio::spawn(async move { + let mut stream = writer.execute(n, ctx).unwrap(); + collect_stream(&mut stream).await.unwrap(); + })); + } + for handle in handles { + handle.await.unwrap(); + } + } + + /// The writer's output has to come back through the reader that will + /// actually open it in production, not through a decoder chosen by the + /// test. This is what catches a writer and reader disagreeing on format. + #[tokio::test] + async fn round_trips_through_the_local_reader() { + let work_dir = TempDir::new().unwrap(); + let work_dir_path = work_dir.path().to_str().unwrap().to_owned(); + let job = "job-range-round-trip"; + let stage_id = 1; + + let writer = Arc::new( + RangeShuffleWriterExec::try_new( + job.into(), + stage_id, + input_plan(), + work_dir_path.clone(), + ) + .unwrap(), + ); + drive(writer, SessionContext::new().task_ctx()).await; + + // Partition k's file holds exactly the batches of input partition k, + // in order — this writer passes its input partitioning through. + for (partition, keys) in [(0, vec![1, 2, 3, 4]), (1, vec![5, 6, 7, 8])] { + let loc = location(job, stage_id, partition); + let mut stream = fetch_partition_local(&work_dir_path, &loc).unwrap(); + let batches = collect_stream(&mut stream).await.unwrap(); + + let read_keys: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(read_keys, keys, "partition {partition} round trip"); + } + } + + /// The index is what a consumer reads instead of the data file, so it has + /// to describe that file: a row per message, keys that match the batches' + /// first rows, and byte ranges that decode. + #[tokio::test] + async fn writes_an_index_describing_the_data_file() { + let work_dir = TempDir::new().unwrap(); + let work_dir_path = work_dir.path().to_str().unwrap().to_owned(); + let job = "job-range-index"; + let stage_id = 5; + + let writer = Arc::new( + RangeShuffleWriterExec::try_new( + job.into(), + stage_id, + input_plan(), + work_dir_path.clone(), + ) + .unwrap(), + ); + drive(writer, SessionContext::new().task_ctx()).await; + + // Input partition 0 holds keys 1..4, which the sort coalesces into a + // single output batch, so the index has one row keyed on its first row. + let data = + create_shuffle_path(&work_dir_path, &job.into(), stage_id, 0, Some(0), false) + .unwrap(); + let index = index_path(&data); + assert!(index.exists(), "index must sit beside the data file"); + + let batch = read_index_file(&index).unwrap(); + assert_eq!(batch.num_rows(), 1, "one row per record batch written"); + + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("key column carries the sort expression's own type"); + assert_eq!(keys.values().to_vec(), vec![1], "each batch's first key"); + + let is_dict = is_dict_column(&batch).unwrap(); + assert!( + (0..is_dict.len()).all(|row| !is_dict.value(row)), + "no dictionary columns in this data", + ); + + let num_rows = num_rows_column(&batch).unwrap(); + assert_eq!(num_rows.values().to_vec(), vec![4]); + + // The offsets have to address the data file, not merely be plausible. + let offsets = byte_offset_column(&batch).unwrap(); + let lens = byte_len_column(&batch).unwrap(); + let layout = read_file_layout(&data).unwrap(); + for (row, block) in layout.record_batches.iter().enumerate() { + assert_eq!(offsets.value(row), block.offset, "offset row {row}"); + assert_eq!(lens.value(row), block.len, "len row {row}"); + } + } + + /// The index names what it is indexed on, so a reader can check the file + /// agrees with the ordering it plans to search by rather than assume it. + #[tokio::test] + async fn index_schema_records_the_sort_key_and_options() { + let work_dir = TempDir::new().unwrap(); + let work_dir_path = work_dir.path().to_str().unwrap().to_owned(); + let job = "job-range-index-schema"; + let stage_id = 6; + + let writer = Arc::new( + RangeShuffleWriterExec::try_new( + job.into(), + stage_id, + input_plan(), + work_dir_path.clone(), + ) + .unwrap(), + ); + drive(writer, SessionContext::new().task_ctx()).await; + + let data = + create_shuffle_path(&work_dir_path, &job.into(), stage_id, 0, Some(0), false) + .unwrap(); + let batch = read_index_file(&index_path(&data)).unwrap(); + let schema = batch.schema(); + + assert_eq!( + schema.field(0).name(), + "k@0", + "the key column is named for the expression it carries", + ); + assert!(schema.field(0).is_nullable(), "keys are nullable"); + assert_eq!( + schema + .metadata() + .get(SORT_OPTIONS_METADATA) + .map(String::as_str), + Some("asc nulls_first"), + ); + } + + /// The point of this writer is the seekable format, so the file it leaves + /// behind must actually be one — a regression here is silent until a + /// consumer tries to seek. + #[tokio::test] + async fn writes_the_seekable_ipc_file_format() { + let work_dir = TempDir::new().unwrap(); + let work_dir_path = work_dir.path().to_str().unwrap().to_owned(); + let job = "job-range-format"; + let stage_id = 3; + + let writer = Arc::new( + RangeShuffleWriterExec::try_new( + job.into(), + stage_id, + input_plan(), + work_dir_path.clone(), + ) + .unwrap(), + ); + drive(writer, SessionContext::new().task_ctx()).await; + + for partition in 0..2 { + let path = create_shuffle_path( + &work_dir_path, + &job.into(), + stage_id, + partition, + Some(0), + false, + ) + .unwrap(); + assert!( + super::super::is_ipc_file(&path), + "partition {partition} must be an Arrow IPC file, not a stream", + ); + } + } +} diff --git a/ballista/core/src/execution_plans/range_shuffle_reader.rs b/ballista/core/src/execution_plans/range_shuffle_reader.rs index ed43ed271f..a948120638 100644 --- a/ballista/core/src/execution_plans/range_shuffle_reader.rs +++ b/ballista/core/src/execution_plans/range_shuffle_reader.rs @@ -42,28 +42,31 @@ //! `OrderedRangeRepartitionExec`, which is one-to-one and never fanned //! into a broadcast. //! -//! # Future work +//! # Narrowing what gets read //! -//! Today the reader pulls whole upstream files and lets `StreamingMerge` -//! do the work. Once value-indexed shuffle files land end-to-end (writer -//! side in [PR #2204]), the reader will consult per-file ValueIndex -//! offsets and fetch only the byte ranges covering its target output -//! partition — dropping the "read the whole file to throw most of it -//! away" cost that dominates when K is large. +//! A source written by [`RangeShuffleWriterExec`] carries a value index, and +//! when the scheduler supplies this reader's per-partition value range, only +//! the batches that can hold a row in that range are read. Selection is at +//! batch granularity; the `RangeFilterExec` above still does the exact trim, +//! so this changes the volume read and not the rows produced. //! -//! [PR #2204]: https://github.com/apache/datafusion-ballista/pull/2204 +//! Sources without an index, and partitions with no bounds, are read whole. +//! +//! [`RangeShuffleWriterExec`]: super::RangeShuffleWriterExec use crate::client_pool::BallistaClientPool; +use crate::execution_plans::range_filter::WidenedBound; +use crate::execution_plans::range_shuffle::{is_ipc_file, open_ipc_file_range}; use crate::execution_plans::shuffle_reader::{ - fetch_partition_local, fetch_partition_remote, local_remote_read_split, - stats_for_partition, + LocalShuffleStream, fetch_partition_local, fetch_partition_remote, + fetch_range_remote, local_remote_read_split, stats_for_partition, }; use crate::extension::SessionConfigExt; use crate::serde::scheduler::PartitionLocation; use crate::utils::GrpcClientConfig; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{Result, Statistics}; +use datafusion::common::{Result, Statistics, internal_err}; use datafusion::error::DataFusionError; use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::MemoryConsumer; @@ -81,6 +84,20 @@ use futures::TryStreamExt; use log::debug; use std::sync::Arc; +/// Render per-partition bounds the way `RangeFilterExec` renders its own, so +/// a plan shows whether the reader narrowed its reads and to what. +fn render_bounds(bounds: &[WidenedBound]) -> String { + let rendered = bounds + .iter() + .map(|(lo, hi)| { + let lo = lo.as_ref().map(|v| v.to_string()).unwrap_or("-inf".into()); + let hi = hi.as_ref().map(|v| v.to_string()).unwrap_or("+inf".into()); + format!("[{lo}, {hi})") + }) + .collect::>(); + format!("[{}]", rendered.join(", ")) +} + /// Ordering-preserving shuffle reader. See module docs. #[derive(Debug, Clone)] pub struct RangeShuffleReaderExec { @@ -95,6 +112,15 @@ pub struct RangeShuffleReaderExec { merge_ordering: LexOrdering, /// Row limit pushed down by a consumer. `None` means read everything. fetch: Option, + /// Half-open value range each output partition covers, halo already + /// applied, or `None` when the scheduler had no cuts to give. + /// + /// Used to skip the parts of an indexed source that cannot hold a row this + /// partition wants. Selection is at batch granularity and the + /// `RangeFilterExec` above still does the exact trim, so a bound that is + /// too wide costs bytes and a missing one costs the saving — neither + /// changes the answer. + bounds: Option>, metrics: ExecutionPlanMetricsSet, properties: Arc, work_dir: Option, @@ -129,6 +155,7 @@ impl RangeShuffleReaderExec { partition, merge_ordering, fetch: None, + bounds: None, metrics: ExecutionPlanMetricsSet::new(), properties, work_dir: None, @@ -142,6 +169,25 @@ impl RangeShuffleReaderExec { self } + /// Bind the per-output-partition value ranges this reader may narrow its + /// reads to. One entry per output partition. + pub fn with_bounds(mut self, bounds: Vec) -> Result { + if bounds.len() != self.partition.len() { + return internal_err!( + "RangeShuffleReaderExec got {} bounds for {} output partitions", + bounds.len(), + self.partition.len() + ); + } + self.bounds = Some(bounds); + Ok(self) + } + + /// Per-output-partition value ranges, halo already applied. + pub fn bounds(&self) -> Option<&[WidenedBound]> { + self.bounds.as_deref() + } + /// Late-bound by the executor. pub fn with_work_dir(&self, work_dir: String) -> Self { Self { @@ -182,6 +228,10 @@ impl DisplayAs for RangeShuffleReaderExec { if let Some(fetch) = self.fetch { write!(f, ", fetch: {fetch}")?; } + match &self.bounds { + Some(bounds) => write!(f, ", bounds: {}", render_bounds(bounds))?, + None => write!(f, ", bounds: none")?, + } Ok(()) } DisplayFormatType::TreeRender => { @@ -269,9 +319,34 @@ impl ExecutionPlan for RangeShuffleReaderExec { let mut sub_streams: Vec = Vec::with_capacity(local_locations.len() + remote_locations.len()); + // The range this output partition covers, when the scheduler had cuts + // to give. Sources carrying an index are read down to the batches that + // can hold a row in it; everything else is read whole. + let bound = self + .bounds + .as_ref() + .and_then(|bounds| bounds.get(output_partition)); + for loc in local_locations { - let stream = fetch_partition_local(work_dir, &loc) + let path = loc + .path(work_dir) .map_err(|e| DataFusionError::External(Box::new(e)))?; + // Narrowing needs a range to narrow to and a source that can be + // seeked, and the source itself is what says whether it can: an IPC + // stream has no footer, so asking one to seek fails rather than + // degrading. Nothing configures this — the file answers it. + let stream = match bound.filter(|_| is_ipc_file(path.as_path())) { + Some((lo, hi)) => { + let reader = + open_ipc_file_range(path.as_path(), lo.as_ref(), hi.as_ref()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let schema = self.schema.clone(); + Box::pin(LocalShuffleStream::new(reader, schema)) + as SendableRecordBatchStream + } + None => fetch_partition_local(work_dir, &loc) + .map_err(|e| DataFusionError::External(Box::new(e)))?, + }; sub_streams.push(stream); } @@ -280,7 +355,14 @@ impl ExecutionPlan for RangeShuffleReaderExec { Arc::new((&config.ballista_config()).into()); let customize_endpoint = config.ballista_override_create_grpc_client_endpoint(); - let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight(); + // The `IO_BLOCK_TRANSPORT` action ships raw file bytes for the + // client to decode with a `StreamDecoder`, which only speaks the + // IPC stream format. This reader's sources come from + // `RangeShuffleWriterExec`, which writes the IPC file format, so + // they have to come back decoded over `do_get`. Teaching block + // transport a byte range is what lets these reads use it again, + // and is the same mechanism a remote seek needs. + let prefer_flight = true; let client_pool = self.client_pool.clone(); for loc in remote_locations { @@ -293,15 +375,36 @@ impl ExecutionPlan for RangeShuffleReaderExec { // remote fetch; subsequent polls stream batches directly. No // buffering, no retry — task-level retry covers transport // failures. + // + // With a range to narrow to, the fetch is two round trips — + // index, then the bytes it points at — and the searching + // happens here so the executor stays a byte server. + let bound = bound.cloned(); + let fetch_schema = schema.clone(); let lazy = futures::stream::once(async move { - fetch_partition_remote( - &loc, - grpc_config, - prefer_flight, - customize_endpoint, - client_pool, - ) - .await + match bound { + Some(bound) => { + fetch_range_remote( + &loc, + bound, + fetch_schema, + grpc_config, + customize_endpoint, + client_pool, + ) + .await + } + None => { + fetch_partition_remote( + &loc, + grpc_config, + prefer_flight, + customize_endpoint, + client_pool, + ) + .await + } + } .map_err(|e| DataFusionError::External(Box::new(e))) }) .try_flatten(); diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index 3de1304640..8f1443f937 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/core/src/execution_plans/shuffle_reader.rs @@ -18,13 +18,19 @@ use crate::client::BallistaClient; use crate::client_pool::BallistaClientPool; use crate::error::BallistaError; +use crate::execution_plans::range_filter::WidenedBound; +use crate::execution_plans::range_shuffle::{ + SORT_OPTIONS_METADATA, byte_ranges_for, count_record_batches, is_ipc_file, + open_ipc_file, schema_message, select_record_batches, +}; use crate::execution_plans::sort_shuffle::{ get_index_path, is_sort_shuffle_output, stream_sort_shuffle_partition, }; use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt}; -use crate::serde::scheduler::{PartitionLocation, PartitionStats}; +use crate::serde::scheduler::{PartitionLocation, PartitionStats, ShuffleFileKind}; use crate::utils::GrpcClientConfig; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::error::ArrowError; use datafusion::arrow::ipc::reader::StreamReader; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::runtime::SpawnedTask; @@ -560,17 +566,31 @@ pub fn stats_for_partitions( } } -struct LocalShuffleStream { - reader: StreamReader>, +/// Streams batches off a node-local shuffle file. +/// +/// Generic over the decoder because the two shuffle formats need different +/// ones — `StreamReader` for the IPC stream the passthrough shuffle writes, +/// `FileReader` for the IPC file the range shuffle writes — and they share no +/// arrow trait beyond `Iterator`. The schema is captured at construction +/// rather than delegated for the same reason. +pub(crate) struct LocalShuffleStream { + reader: R, + schema: SchemaRef, } -impl LocalShuffleStream { - pub fn new(reader: StreamReader>) -> Self { - LocalShuffleStream { reader } +impl LocalShuffleStream +where + R: Iterator>, +{ + pub(crate) fn new(reader: R, schema: SchemaRef) -> Self { + LocalShuffleStream { reader, schema } } } -impl Stream for LocalShuffleStream { +impl Stream for LocalShuffleStream +where + R: Iterator> + Unpin, +{ type Item = Result; fn poll_next( @@ -584,9 +604,12 @@ impl Stream for LocalShuffleStream { } } -impl RecordBatchStream for LocalShuffleStream { +impl RecordBatchStream for LocalShuffleStream +where + R: Iterator> + Unpin, +{ fn schema(&self) -> SchemaRef { - self.reader.schema() + self.schema.clone() } } @@ -1056,7 +1079,7 @@ pub(crate) async fn fetch_partition_remote( let metadata = &location.executor_meta; let partition_id = &location.partition_id; let file_id = location.file_id; - let is_sort_shuffle = location.is_sort_shuffle; + let layout = location.layout(); let host = metadata.host.as_str(); let port = metadata.port; @@ -1075,13 +1098,7 @@ pub(crate) async fn fetch_partition_remote( })?; let result = pooled - .fetch_partition( - &metadata.id, - partition_id, - file_id, - is_sort_shuffle, - prefer_flight, - ) + .fetch_partition(&metadata.id, partition_id, file_id, layout, prefer_flight) .await; if result.is_err() { pooled.discard(); @@ -1106,14 +1123,145 @@ pub(crate) async fn fetch_partition_remote( })?; ballista_client - .fetch_partition( - &metadata.id, - partition_id, - file_id, - is_sort_shuffle, - prefer_flight, + .fetch_partition(&metadata.id, partition_id, file_id, layout, prefer_flight) + .await + } +} + +/// Fetch a remote range shuffle source down to the bytes `[lo, hi)` covers. +/// +/// Two round trips: the index, then the ranges it points at. The searching +/// happens here rather than on the executor, so the executor stays a byte +/// server and the same two steps work against object storage. +/// +/// Falls back to fetching the whole partition when the source has no index to +/// search — the answer is the same, only the volume differs. +pub(crate) async fn fetch_range_remote( + location: &PartitionLocation, + bound: WidenedBound, + schema: SchemaRef, + config: Arc, + customize_endpoint: Option>, + client_pool: Option>, +) -> result::Result { + let metadata = &location.executor_meta; + let partition_id = &location.partition_id; + + if let Some(pool) = client_pool { + let mut pooled = pool + .acquire( + metadata.host.as_str(), + metadata.port, + &config, + customize_endpoint, ) .await + .map_err(|error| as_fetch_failed(metadata, partition_id, error))?; + let result = fetch_ranges_with(&mut pooled, location, bound, schema).await; + if result.is_err() { + pooled.discard(); + } + result + } else { + let mut client = new_ballista_client( + metadata.host.as_str(), + metadata.port, + &config, + customize_endpoint, + ) + .await + .map_err(|error| as_fetch_failed(metadata, partition_id, error))?; + fetch_ranges_with(&mut client, location, bound, schema).await + } +} + +/// The two-step fetch itself, over an already-connected client. +async fn fetch_ranges_with( + client: &mut BallistaClient, + location: &PartitionLocation, + bound: WidenedBound, + schema: SchemaRef, +) -> result::Result { + let metadata = &location.executor_meta; + let partition_id = &location.partition_id; + let (lo, hi) = bound; + + let mut index_stream = client + .fetch_byte_ranges( + &metadata.id, + partition_id, + location.file_id, + location.layout(), + ShuffleFileKind::Index, + vec![], + None, + ) + .await?; + let index = crate::utils::collect_stream(&mut index_stream).await?; + let [index] = index.as_slice() else { + return Err(BallistaError::General(format!( + "range shuffle index for {partition_id:?} came back as {} batches", + index.len() + ))); + }; + + let descending = index + .schema() + .metadata() + .get(SORT_OPTIONS_METADATA) + .and_then(|options| options.split(',').next()) + .is_some_and(|options| options.starts_with("desc")); + + let selected = + match select_record_batches(index, lo.as_ref(), hi.as_ref(), descending)? { + Some(selected) => selected, + // The index declined to narrow, so take everything it lists. + None => (0..count_record_batches(index)?).collect(), + }; + + let Some(ranges) = byte_ranges_for(index, &selected)? else { + // Nothing in this file falls inside the range. + return Ok(Box::pin(RecordBatchStreamAdapter::new( + index.schema(), + futures::stream::empty(), + ))); + }; + + debug!( + "range shuffle fetching {} ranges of {:?} from {}", + ranges.len(), + partition_id, + metadata.id, + ); + + client + .fetch_byte_ranges( + &metadata.id, + partition_id, + location.file_id, + location.layout(), + ShuffleFileKind::Data, + ranges, + Some(schema_message(schema.as_ref())?), + ) + .await +} + +/// Report a transport failure as a fetch failure, which is what lets the +/// scheduler retry the task rather than fail the query. +fn as_fetch_failed( + metadata: &crate::serde::scheduler::ExecutorMetadata, + partition_id: &crate::serde::scheduler::PartitionId, + error: BallistaError, +) -> BallistaError { + match error { + BallistaError::GrpcConnectionError(msg) => BallistaError::FetchFailed( + metadata.id.clone(), + partition_id.stage_id, + partition_id.partition_id, + msg, + ), + other => other, } } @@ -1153,6 +1301,24 @@ pub(crate) fn fetch_partition_local( ) }); } + if is_ipc_file(data_path) { + // Range shuffle output. The two IPC framings are not + // interchangeable — a stream decoder rejects a file's leading magic — + // so the format is read off the file itself rather than carried on + // `PartitionLocation`, keeping it out of the wire protocol. + debug!("fetch local range shuffle file: {data_path:?}"); + let reader = open_ipc_file(data_path).map_err(|e| { + BallistaError::FetchFailed( + metadata.id.clone(), + partition_id.stage_id, + partition_id.partition_id, + e.to_string(), + ) + })?; + let schema = reader.schema(); + return Ok(Box::pin(LocalShuffleStream::new(reader, schema))); + } + debug!("fetch local partition file: {data_path:?} "); // Standard single-file shuffle output - read the file directly let reader = fetch_partition_local_inner(path).map_err(|e| { @@ -1164,7 +1330,8 @@ pub(crate) fn fetch_partition_local( e.to_string(), ) })?; - Ok(Box::pin(LocalShuffleStream::new(reader))) + let schema = reader.schema(); + Ok(Box::pin(LocalShuffleStream::new(reader, schema))) } fn fetch_partition_local_inner( @@ -1777,9 +1944,10 @@ mod tests { // from to input partitions test the first one with two batches let file_path = Path::new(path.value(0)); let reader = fetch_partition_local_inner(file_path).unwrap(); + let schema = reader.schema(); let mut stream: Pin> = - async { Box::pin(LocalShuffleStream::new(reader)) }.await; + async { Box::pin(LocalShuffleStream::new(reader, schema)) }.await; let result = utils::collect_stream(&mut stream) .await diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 99f9a7c5f5..338194cd7e 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -74,7 +74,7 @@ use super::shuffle_writer_trait::ShuffleWriter; /// id — the thing that ends up in `ShuffleWritePartition.partition_id` and /// keys `PartitionLocation`s downstream. #[derive(Debug, Clone)] -enum GlobalPartitionMap { +pub(crate) enum GlobalPartitionMap { /// The plan collapses to a single output partition (e.g. /// `SortPreservingMergeExec`). Every local index → global partition 0. Collapsed, @@ -94,7 +94,7 @@ enum GlobalPartitionMap { } impl GlobalPartitionMap { - fn resolve(&self, local: usize) -> u64 { + pub(crate) fn resolve(&self, local: usize) -> u64 { match self { GlobalPartitionMap::Collapsed => 0, GlobalPartitionMap::KSpace => local as u64, @@ -121,7 +121,7 @@ impl GlobalPartitionMap { /// partitioning-preserving passthroughs). /// - If we hit a leaf or a fan-in without recognising it, treat it as /// passthrough (the caller passes `global_output_partition_ids`). -fn walk_child_partition_mapping( +pub(crate) fn walk_child_partition_mapping( plan: &Arc, global_output_partition_ids: &[usize], ) -> GlobalPartitionMap { @@ -193,10 +193,14 @@ pub fn compute_global_output_partition_ids( }; return (0..*k).collect(); } - if stage_plan.is::() { + // Both passthrough writers derive their ids the same way: they never + // repartition, so the child plan's shape decides. + if stage_plan.is::() + || stage_plan.is::() + { let children = stage_plan.children(); let [child] = children.as_slice() else { - unreachable!("ShuffleWriterExec always has exactly one child"); + unreachable!("a passthrough shuffle writer always has exactly one child"); }; return match walk_child_partition_mapping(child, global_input_partition_ids) { GlobalPartitionMap::Collapsed => vec![0], @@ -223,12 +227,13 @@ pub const DEFAULT_SHUFFLE_CHANNEL_CAPACITY: usize = 8; /// task, where each slice member yields one file each containing K logical /// partitions) can hand the full set to a single `execute(N)` stream. /// Passthrough / hash-repart writers produce at most one summary per slot. -struct WriterState { - initialized: bool, +pub(crate) struct WriterState { + pub(crate) initialized: bool, /// One receiver per output partition. `execute(N)` takes `handoffs[N]`; /// the coordinator holds the matching sender and pushes the summaries /// once partition N's files are closed. - handoffs: Vec>>>>, + pub(crate) handoffs: + Vec>>>>, } impl Debug for WriterState { @@ -396,11 +401,11 @@ impl std::fmt::Display for ShuffleWriterExec { } #[derive(Debug, Clone)] -struct ShuffleWriteMetrics { +pub(crate) struct ShuffleWriteMetrics { /// Time spend writing batches to shuffle files - write_time: metrics::Time, - input_rows: metrics::Count, - output_rows: metrics::Count, + pub(crate) write_time: metrics::Time, + pub(crate) input_rows: metrics::Count, + pub(crate) output_rows: metrics::Count, } impl ShuffleWriteMetrics { @@ -411,7 +416,7 @@ impl ShuffleWriteMetrics { /// exactly as it did before K-drain (K tasks × 1 bucket = 1 task × K /// buckets). The scheduler maps `input_partition` to a stage-global /// input partition id via `TaskDescription.global_input_partition_ids`. - fn new(input_partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + pub(crate) fn new(input_partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { let write_time = MetricBuilder::new(metrics).subset_time("write_time", input_partition); @@ -484,6 +489,12 @@ impl ShuffleWriterExec { self.task_id } + /// Work directory shuffle files are written under. Empty until the + /// executor stamps it at `create_query_stage_exec` time. + pub fn work_dir(&self) -> &str { + &self.work_dir + } + /// Bind this writer to the task's assigned global partition slice. pub fn with_global_output_partition_ids( mut self, @@ -507,26 +518,11 @@ impl ShuffleWriterExec { /// catches. Failing the task surfaces it while it is still a failure /// rather than a wrong answer. pub fn collect_window_state(&self) -> Result> { - let mut found: Vec<&PartitionedBoundedWindowAggExec> = Vec::new(); - collect_window_state_operators(&self.plan, &mut found); - found - .into_iter() - .flat_map(|op| op.observed_window_state()) - .map(|observation| { - let global = self - .global_output_partition_ids - .get(observation.partition_idx) - .copied() - .ok_or_else(|| { - DataFusionError::Internal(format!( - "ShuffleWriterExec: window state for local partition {} \ - has no global id (slice covers {:?})", - observation.partition_idx, self.global_output_partition_ids - )) - })?; - Ok((global, observation)) - }) - .collect() + collect_window_state_against_slice( + &self.plan, + &self.global_output_partition_ids, + "ShuffleWriterExec", + ) } /// Get the Job ID for this query stage @@ -765,7 +761,7 @@ impl ExecutionPlan for ShuffleWriterExec { let writer = self.clone(); let ctx = context.clone(); tokio::spawn(async move { - run_coordinator(writer, ctx, senders).await; + run_coordinator(writer.execute_shuffle_write(ctx), senders).await; }); } @@ -865,23 +861,27 @@ pub(crate) fn result_schema() -> SchemaRef { ])) } -/// Drives the shared write work for all K output partitions. Runs -/// `execute_shuffle_write` once, then routes each `(handoff_idx, summary)` -/// pair to the matching sender. Slots that never receive a summary (partition -/// produced no rows) are filled with an empty sentinel so their `execute(N)` -/// stream terminates cleanly. +/// Drives the shared write work for all K output partitions. Awaits +/// `write_work` once — the writer's own `execute_shuffle_write` — then routes +/// each `(handoff_idx, summary)` pair to the matching sender. Slots that never +/// receive a summary (partition produced no rows) are filled with an empty +/// sentinel so their `execute(N)` stream terminates cleanly. +/// +/// Generic over the work future so every passthrough-shaped writer shares this +/// routing; only how a file gets written differs between them. /// /// On failure, sends the error to every sender so no waiting stream hangs. -async fn run_coordinator( - writer: ShuffleWriterExec, - ctx: Arc, +pub(crate) async fn run_coordinator( + write_work: F, senders: Vec>>>, -) { +) where + F: Future>>, +{ let k = senders.len(); let mut senders: Vec>>>> = senders.into_iter().map(Some).collect(); - match writer.execute_shuffle_write(ctx).await { + match write_work.await { Ok(summaries) => { // Bucket per-file summaries by their handoff slot. Passthrough and // hash writers put at most one summary per slot; sort-based @@ -1002,6 +1002,38 @@ fn collect_window_state_operators<'a>( } } +/// Shared body of every writer's `collect_window_state`: walk `plan` for +/// window-state collectors and translate each capture's task-local partition +/// index against `global_output_partition_ids`. +/// +/// `writer` names the caller in the error, since the failure is a plan/slice +/// mismatch and which writer produced it is the first thing worth knowing. +pub(crate) fn collect_window_state_against_slice( + plan: &Arc, + global_output_partition_ids: &[usize], + writer: &str, +) -> Result> { + let mut found: Vec<&PartitionedBoundedWindowAggExec> = Vec::new(); + collect_window_state_operators(plan, &mut found); + found + .into_iter() + .flat_map(|op| op.observed_window_state()) + .map(|observation| { + let global = global_output_partition_ids + .get(observation.partition_idx) + .copied() + .ok_or_else(|| { + DataFusionError::Internal(format!( + "{writer}: window state for local partition {} \ + has no global id (slice covers {global_output_partition_ids:?})", + observation.partition_idx + )) + })?; + Ok((global, observation)) + }) + .collect() +} + #[cfg(test)] #[allow(dead_code, unused_imports)] // clippy false positive with local imports mod tests { diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs index e57a2997ac..d988a260b1 100644 --- a/ballista/core/src/lib.rs +++ b/ballista/core/src/lib.rs @@ -39,7 +39,7 @@ pub const BALLISTA_VERSION: &str = env!("CARGO_PKG_VERSION"); /// /// Zero is reserved as the proto-default "unset" value, produced by executors /// that predate this field — it never matches a real scheduler version. -pub const BALLISTA_PROTOCOL_VERSION: u32 = 2; +pub const BALLISTA_PROTOCOL_VERSION: u32 = 3; /// Prints the current Ballista version to stdout. pub fn print_version() { diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index dbff74c762..313c566324 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -69,6 +69,8 @@ pub mod ballista_physical_plan_node { RangeFilter(super::RangeFilterExecNode), #[prost(message, tag = "14")] PrefixMerge(super::PrefixMergeExecNode), + #[prost(message, tag = "15")] + RangeShuffleWriter(super::RangeShuffleWriterExecNode), } } /// Value-range router over N locally-sorted overlapping input partitions. @@ -398,6 +400,18 @@ pub struct SortShuffleWriterExecNode { #[prost(uint64, optional, tag = "9")] pub memory_limit_per_task_bytes: ::core::option::Option, } +/// Passthrough shuffle writer that emits the seekable Arrow IPC file format. +/// Carries no output partitioning: like `ShuffleWriterExecNode`'s passthrough +/// case it writes its input's partitioning through, one file per partition. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeShuffleWriterExecNode { + #[prost(string, tag = "1")] + pub job_id: ::prost::alloc::string::String, + #[prost(uint32, tag = "2")] + pub stage_id: u32, + #[prost(message, optional, tag = "3")] + pub input: ::core::option::Option<::datafusion_proto::protobuf::PhysicalPlanNode>, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct UnresolvedShuffleExecNode { #[prost(uint32, tag = "1")] @@ -459,6 +473,12 @@ pub struct RangeShuffleReaderExecNode { /// Row limit pushed down by a consuming merge. Absent means read everything. #[prost(uint64, optional, tag = "5")] pub fetch: ::core::option::Option, + /// Half-open value range each output partition covers, halo already applied. + /// One per output partition, or empty when the scheduler had no cuts. Lets + /// the reader skip the parts of an indexed source that cannot hold a row the + /// partition wants; the consuming RangeFilterExec still does the exact trim. + #[prost(message, repeated, tag = "6")] + pub bounds: ::prost::alloc::vec::Vec, } /// CoalescePartitionsRule output: groups upstream partitions into coalesced output partitions. /// Empty when no coalesce is applied (the optional field on the parent message is absent). @@ -677,7 +697,7 @@ pub struct Action { } /// Nested message and enum types in `Action`. pub mod action { - #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] + #[derive(Clone, PartialEq, ::prost::Oneof)] pub enum ActionType { /// Fetch a partition from an executor #[prost(message, tag = "3")] @@ -703,7 +723,15 @@ pub struct ExecutePartition { ::datafusion_proto::protobuf::PhysicalHashRepartition, >, } -#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +/// A half-open byte range of a file: `[offset, offset + length)`. +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ByteRange { + #[prost(uint64, tag = "1")] + pub offset: u64, + #[prost(uint64, tag = "2")] + pub length: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct FetchPartition { #[prost(string, tag = "1")] pub job_id: ::prost::alloc::string::String, @@ -717,8 +745,27 @@ pub struct FetchPartition { pub port: u32, #[prost(uint64, optional, tag = "7")] pub file_id: ::core::option::Option, - #[prost(bool, tag = "8")] - pub is_sort_shuffle: bool, + /// How the producing writer laid its output out, which is what turns the + /// identifiers above into a path. + #[prost(enumeration = "ShuffleLayout", tag = "9")] + pub layout: i32, + /// Which file beside that partition to serve. + #[prost(enumeration = "ShuffleFileKind", tag = "10")] + pub file_kind: i32, + /// Ranges to return, concatenated in request order, as absolute offsets into + /// the file. + /// + /// Empty asks for whatever the identifiers above address. Under the sort + /// layout that is the named partition's slice, which the executor resolves + /// through its own index — one round trip, as it has always been. Under the + /// passthrough layout it is the whole file. + /// + /// With ranges given the executor serves bytes and nothing else: no index + /// read, no comparison, no notion of what a range means. That is what a + /// consumer needs against object storage, where there is no executor to ask, + /// and it is why the range shuffle's index carries byte offsets at all. + #[prost(message, repeated, tag = "11")] + pub byte_ranges: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PartitionLocation { @@ -1794,6 +1841,68 @@ impl ScalarOpNode { } } } +/// How a shuffle writer laid its output out on disk. Selects the path shape, +/// and with it how a client parses the index beside the data. It says nothing +/// about the framing of the bytes inside the data file. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum ShuffleLayout { + /// {stage_id}/{partition_id}/data-{file_id}.arrow — one file per output + /// partition, written by the passthrough and range writers. + Passthrough = 0, + /// {stage_id}/{file_id}/data.arrow — one file per task holding every + /// partition, written by the sort-based writer. + Sort = 1, +} +impl ShuffleLayout { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Passthrough => "SHUFFLE_LAYOUT_PASSTHROUGH", + Self::Sort => "SHUFFLE_LAYOUT_SORT", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "SHUFFLE_LAYOUT_PASSTHROUGH" => Some(Self::Passthrough), + "SHUFFLE_LAYOUT_SORT" => Some(Self::Sort), + _ => None, + } + } +} +/// Which of the files making up a shuffle output is being asked for. Which +/// index sits beside the data, and how to read it, follows from the layout, so +/// it is not spelled out here. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum ShuffleFileKind { + Data = 0, + Index = 1, +} +impl ShuffleFileKind { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Data => "SHUFFLE_FILE_KIND_DATA", + Self::Index => "SHUFFLE_FILE_KIND_INDEX", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "SHUFFLE_FILE_KIND_DATA" => Some(Self::Data), + "SHUFFLE_FILE_KIND_INDEX" => Some(Self::Index), + _ => None, + } + } +} /// Generated client implementations. pub mod scheduler_grpc_client { #![allow( diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 662d8f0f3f..5ba5413ac2 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -59,9 +59,9 @@ use crate::execution_plans::{ BufferExec, BufferMode, ChaosExec, CoalescePlan, FinalizedPartitionState, InputOrder, OrderedRangeRepartitionExec, PartitionGroup, PartitionedBoundedWindowAggExec, PerPartitionFilterExec, PrefixMergeExec, RangeFilterExec, RangeShuffleReaderExec, - RuntimeStatsExec, ScalarOp, ShuffleReaderExec, ShuffleWriterExec, - SortShuffleWriterExec, UnorderedRangeRepartitionExec, UnresolvedShuffleExec, - WindowApply, + RangeShuffleWriterExec, RuntimeStatsExec, ScalarOp, ShuffleReaderExec, + ShuffleWriterExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, + UnresolvedShuffleExec, WindowApply, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -554,6 +554,24 @@ fn decode_scalar_op(op: protobuf::ScalarOpNode) -> ScalarOp { } } +/// Decode a `ScalarValue` carried on a Ballista physical plan node. +fn scalar_from_proto( + proto: &datafusion_proto_common::ScalarValue, +) -> Result { + datafusion::scalar::ScalarValue::try_from(proto).map_err(|e| { + DataFusionError::Internal(format!("failed to decode ScalarValue: {e:?}")) + }) +} + +/// Encode a `ScalarValue` for a Ballista physical plan node. +fn scalar_to_proto( + value: &datafusion::scalar::ScalarValue, +) -> Result { + datafusion_proto_common::ScalarValue::try_from(value).map_err(|e| { + DataFusionError::Internal(format!("failed to encode ScalarValue: {e:?}")) + }) +} + impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { fn try_decode( &self, @@ -599,6 +617,16 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { "".to_string(), // this is intentional but hacky - the executor will fill this in )?)) } + PhysicalPlanType::RangeShuffleWriter(range_shuffle_writer) => { + let input = inputs[0].clone(); + + Ok(Arc::new(RangeShuffleWriterExec::try_new( + range_shuffle_writer.job_id.clone().into(), + range_shuffle_writer.stage_id as usize, + input, + "".to_string(), // the executor fills the work dir in + )?)) + } PhysicalPlanType::SortShuffleWriter(sort_shuffle_writer) => { let input = inputs[0].clone(); @@ -735,10 +763,29 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { partition_location, schema, merge_ordering, - )?; - Ok(Arc::new( - reader.with_fetch_limit(range_reader.fetch.map(|f| f as usize)), - )) + )? + .with_fetch_limit(range_reader.fetch.map(|f| f as usize)); + + // Empty means the scheduler had no cuts to give, which is a + // reader that reads its sources whole — not a reader whose + // every partition covers nothing. + let reader = if range_reader.bounds.is_empty() { + reader + } else { + let bounds = range_reader + .bounds + .iter() + .map(|bound| { + let lo = + bound.lo.as_ref().map(scalar_from_proto).transpose()?; + let hi = + bound.hi.as_ref().map(scalar_from_proto).transpose()?; + Ok::<_, DataFusionError>((lo, hi)) + }) + .collect::, DataFusionError>>()?; + reader.with_bounds(bounds)? + }; + Ok(Arc::new(reader)) } PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => { let schema: SchemaRef = @@ -942,31 +989,24 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { }) } }); - let sv_from_proto = |p: &datafusion_proto_common::ScalarValue| { - datafusion::scalar::ScalarValue::try_from(p).map_err(|e| { - DataFusionError::Internal(format!( - "RangeFilterExec: failed to decode ScalarValue: {e:?}" - )) - }) - }; let halo_lo_proto = node.halo_lo.as_ref().ok_or_else(|| { DataFusionError::Internal( "RangeFilterExecNode missing halo_lo".into(), ) })?; - let halo_lo = sv_from_proto(halo_lo_proto)?; + let halo_lo = scalar_from_proto(halo_lo_proto)?; let halo_hi_proto = node.halo_hi.as_ref().ok_or_else(|| { DataFusionError::Internal( "RangeFilterExecNode missing halo_hi".into(), ) })?; - let halo_hi = sv_from_proto(halo_hi_proto)?; + let halo_hi = scalar_from_proto(halo_hi_proto)?; let raw_bounds = node .raw_bounds .iter() .map(|b| { - let lo = b.lo.as_ref().map(sv_from_proto).transpose()?; - let hi = b.hi.as_ref().map(sv_from_proto).transpose()?; + let lo = b.lo.as_ref().map(scalar_from_proto).transpose()?; + let hi = b.hi.as_ref().map(scalar_from_proto).transpose()?; Ok::<_, DataFusionError>((lo, hi)) }) .collect::, DataFusionError>>()?; @@ -1033,6 +1073,24 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; + Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::RangeShuffleWriter( + protobuf::RangeShuffleWriterExecNode { + job_id: exec.job_id().to_string(), + stage_id: exec.stage_id() as u32, + input: None, + }, + )), + }; + + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode range shuffle writer execution plan: {e:?}" + )) + })?; + Ok(()) } else if let Some(exec) = node.downcast_ref::() { let output_partitioning = match exec.shuffle_output_partitioning() { @@ -1154,6 +1212,16 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { schema: Some(exec.schema().as_ref().try_into()?), merge_ordering, fetch: exec.fetch().map(|f| f as u64), + bounds: exec + .bounds() + .unwrap_or_default() + .iter() + .map(|(lo, hi)| { + let lo = lo.as_ref().map(scalar_to_proto).transpose()?; + let hi = hi.as_ref().map(scalar_to_proto).transpose()?; + Ok::<_, DataFusionError>(protobuf::RangeBound { lo, hi }) + }) + .collect::, DataFusionError>>()?, }, )), }; @@ -1343,20 +1411,13 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { exec.filter_expr(), self.default_codec.as_ref(), )?; - let encode_sv = |sv: &datafusion::scalar::ScalarValue| { - datafusion_proto_common::ScalarValue::try_from(sv).map_err(|e| { - DataFusionError::Internal(format!( - "failed to encode RangeFilterExec ScalarValue: {e:?}" - )) - }) - }; - let halo_lo = encode_sv(exec.halo_lo())?; - let halo_hi = encode_sv(exec.halo_hi())?; + let halo_lo = scalar_to_proto(exec.halo_lo())?; + let halo_hi = scalar_to_proto(exec.halo_hi())?; let raw_bounds_proto = raw_bounds .iter() .map(|(lo, hi)| { - let lo = lo.as_ref().map(&encode_sv).transpose()?; - let hi = hi.as_ref().map(&encode_sv).transpose()?; + let lo = lo.as_ref().map(&scalar_to_proto).transpose()?; + let hi = hi.as_ref().map(&scalar_to_proto).transpose()?; Ok::<_, DataFusionError>(protobuf::RangeBound { lo, hi }) }) .collect::, DataFusionError>>()?; diff --git a/ballista/core/src/serde/scheduler/from_proto.rs b/ballista/core/src/serde/scheduler/from_proto.rs index a9be760c3c..aea773fb2f 100644 --- a/ballista/core/src/serde/scheduler/from_proto.rs +++ b/ballista/core/src/serde/scheduler/from_proto.rs @@ -36,7 +36,7 @@ use crate::error::BallistaError; use crate::extension::SessionConfigHelperExt; use crate::serde::protobuf::{NamedPruningMetrics, NamedRatio}; use crate::serde::scheduler::{ - Action, BallistaFunctionRegistry, ExecutorData, ExecutorMetadata, + Action, BallistaFunctionRegistry, ByteRange, ExecutorData, ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, PartitionId, PartitionLocation, PartitionStats, TaskDefinition, }; @@ -58,7 +58,30 @@ impl TryInto for protobuf::Action { file_id: fetch.file_id, host: fetch.host, port: fetch.port as u16, - is_sort_shuffle: fetch.is_sort_shuffle, + layout: protobuf::ShuffleLayout::try_from(fetch.layout) + .map_err(|_| { + BallistaError::General(format!( + "unknown shuffle layout {} on FetchPartition", + fetch.layout + )) + })? + .into(), + file_kind: protobuf::ShuffleFileKind::try_from(fetch.file_kind) + .map_err(|_| { + BallistaError::General(format!( + "unknown shuffle file kind {} on FetchPartition", + fetch.file_kind + )) + })? + .into(), + byte_ranges: fetch + .byte_ranges + .into_iter() + .map(|range| ByteRange { + offset: range.offset, + length: range.length, + }) + .collect(), }) } _ => Err(BallistaError::General( diff --git a/ballista/core/src/serde/scheduler/mod.rs b/ballista/core/src/serde/scheduler/mod.rs index 970d3d2d07..f5a9c65ca9 100644 --- a/ballista/core/src/serde/scheduler/mod.rs +++ b/ballista/core/src/serde/scheduler/mod.rs @@ -19,6 +19,7 @@ use crate::JobId; use crate::error::BallistaError; use crate::execution_plans::create_shuffle_path; use crate::registry::BallistaFunctionRegistry; +use crate::serde::protobuf; use datafusion::arrow::array::{ ArrayBuilder, StructArray, StructBuilder, UInt64Array, UInt64Builder, }; @@ -53,11 +54,66 @@ pub enum Action { port: u16, /// shuffle file block id file_id: Option, - /// whether this partition uses sort shuffle - is_sort_shuffle: bool, + /// how the producing writer laid its output out on disk + layout: ShuffleLayout, + /// which file making up that output to fetch + file_kind: ShuffleFileKind, + /// byte ranges of that file to return, concatenated in order. Empty + /// asks for whatever the identifiers above address. + byte_ranges: Vec, }, } +impl PartitionLocation { + /// How the writer that produced this partition laid its output out. + /// + /// Derived from `is_sort_shuffle`, which is the same question asked in the + /// vocabulary of one writer. The stored field keeps that name until the + /// construction sites are swept; the wire protocol already speaks layouts. + pub fn layout(&self) -> ShuffleLayout { + if self.is_sort_shuffle { + ShuffleLayout::Sort + } else { + ShuffleLayout::Passthrough + } + } +} + +/// How a shuffle writer laid its output out on disk. +/// +/// Selects the path shape, and with it how a client parses the index beside +/// the data. Says nothing about the framing of the data file's own bytes — +/// that is the file's business, and a reader establishes it by opening it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ShuffleLayout { + /// `{stage_id}/{partition_id}/data-{file_id}.arrow`, one file per output + /// partition. Written by the passthrough and range writers. + #[default] + Passthrough, + /// `{stage_id}/{file_id}/data.arrow`, one file per task holding every + /// partition. Written by the sort-based writer. + Sort, +} + +/// Which file making up a shuffle output is wanted. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ShuffleFileKind { + /// The data itself. + #[default] + Data, + /// The index beside it, whose name and encoding follow from the layout. + Index, +} + +/// A half-open byte range of a file: `[offset, offset + length)`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ByteRange { + /// Where the range starts, from the beginning of the file. + pub offset: u64, + /// How many bytes it spans. + pub length: u64, +} + /// Unique identifier for the output partition of an operator. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct PartitionId { @@ -553,3 +609,39 @@ pub struct TaskDefinition { /// Function registry for UDFs. pub function_registry: Arc, } + +impl From for protobuf::ShuffleLayout { + fn from(layout: ShuffleLayout) -> Self { + match layout { + ShuffleLayout::Passthrough => protobuf::ShuffleLayout::Passthrough, + ShuffleLayout::Sort => protobuf::ShuffleLayout::Sort, + } + } +} + +impl From for ShuffleLayout { + fn from(layout: protobuf::ShuffleLayout) -> Self { + match layout { + protobuf::ShuffleLayout::Passthrough => ShuffleLayout::Passthrough, + protobuf::ShuffleLayout::Sort => ShuffleLayout::Sort, + } + } +} + +impl From for protobuf::ShuffleFileKind { + fn from(kind: ShuffleFileKind) -> Self { + match kind { + ShuffleFileKind::Data => protobuf::ShuffleFileKind::Data, + ShuffleFileKind::Index => protobuf::ShuffleFileKind::Index, + } + } +} + +impl From for ShuffleFileKind { + fn from(kind: protobuf::ShuffleFileKind) -> Self { + match kind { + protobuf::ShuffleFileKind::Data => ShuffleFileKind::Data, + protobuf::ShuffleFileKind::Index => ShuffleFileKind::Index, + } + } +} diff --git a/ballista/core/src/serde/scheduler/to_proto.rs b/ballista/core/src/serde/scheduler/to_proto.rs index 4931f95680..c29f7a4110 100644 --- a/ballista/core/src/serde/scheduler/to_proto.rs +++ b/ballista/core/src/serde/scheduler/to_proto.rs @@ -43,7 +43,9 @@ impl TryInto for Action { file_id, host, port, - is_sort_shuffle, + layout, + file_kind, + byte_ranges, } => Ok(protobuf::Action { action_type: Some(ActionType::FetchPartition(protobuf::FetchPartition { job_id: job_id.into(), @@ -52,7 +54,15 @@ impl TryInto for Action { host, port: port as u32, file_id, - is_sort_shuffle, + layout: protobuf::ShuffleLayout::from(layout) as i32, + file_kind: protobuf::ShuffleFileKind::from(file_kind) as i32, + byte_ranges: byte_ranges + .into_iter() + .map(|range| protobuf::ByteRange { + offset: range.offset, + length: range.length, + }) + .collect(), })), settings: vec![], }), diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 914dc85d81..e8299213aa 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -24,7 +24,7 @@ use ballista_core::client_pool::BallistaClientPool; use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec; use ballista_core::execution_plans::{ - RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec, + RangeShuffleReaderExec, RangeShuffleWriterExec, ShuffleReaderExec, ShuffleWriterExec, }; use ballista_core::serde::protobuf::ShuffleWritePartition; use ballista_core::serde::scheduler::PartitionStats; @@ -191,8 +191,9 @@ impl ExecutionEngine for DefaultExecutionEngine { })? .data; - // the query plan created by the scheduler always starts with a shuffle writer - // (either ShuffleWriterExec or SortShuffleWriterExec) + // the query plan created by the scheduler always starts with a shuffle + // writer (ShuffleWriterExec, RangeShuffleWriterExec, or + // SortShuffleWriterExec) if plan.downcast_ref::().is_some() { let exec = ShuffleWriterExec::try_new( job_id, @@ -205,6 +206,18 @@ impl ExecutionEngine for DefaultExecutionEngine { Ok(Arc::new(DefaultQueryStageExec::new( ShuffleWriterVariant::Passthrough(exec), ))) + } else if plan.downcast_ref::().is_some() { + let exec = RangeShuffleWriterExec::try_new( + job_id, + stage_id, + plan.children()[0].clone(), + work_dir.to_string(), + )? + .with_task_id(task_id) + .with_global_output_partition_ids(global_output_partition_ids); + Ok(Arc::new(DefaultQueryStageExec::new( + ShuffleWriterVariant::Range(exec), + ))) } else if let Some(sort_shuffle_writer) = plan.downcast_ref::() { @@ -223,7 +236,8 @@ impl ExecutionEngine for DefaultExecutionEngine { ))) } else { Err(DataFusionError::Internal( - "Plan passed to new_query_stage_exec is not a ShuffleWriterExec or SortShuffleWriterExec" + "Plan passed to new_query_stage_exec is not a ShuffleWriterExec, \ + RangeShuffleWriterExec, or SortShuffleWriterExec" .to_string(), )) } @@ -236,6 +250,9 @@ pub enum ShuffleWriterVariant { /// Passthrough shuffle writer: preserves its input partitioning, /// one file per output partition. Passthrough(ShuffleWriterExec), + /// Passthrough shuffle writer emitting the seekable Arrow IPC file + /// format, for stages read back in value-range order. + Range(RangeShuffleWriterExec), /// Sort-based shuffle writer. Sort(SortShuffleWriterExec), } @@ -274,6 +291,20 @@ impl Display for DefaultQueryStageExec { writer ) } + ShuffleWriterVariant::Range(writer) => { + let stage_metrics: Vec = writer + .metrics() + .unwrap_or_default() + .iter() + .map(|m| m.to_string()) + .collect(); + write!( + f, + "DefaultQueryStageExec(Range): ({})\n{}", + stage_metrics.join(", "), + writer + ) + } ShuffleWriterVariant::Sort(writer) => { let stage_metrics: Vec = writer .metrics() @@ -304,6 +335,7 @@ impl QueryStageExecutor for DefaultQueryStageExec { ShuffleWriterVariant::Passthrough(writer) => { (Arc::new(writer.clone()), false) } + ShuffleWriterVariant::Range(writer) => (Arc::new(writer.clone()), false), ShuffleWriterVariant::Sort(writer) => (Arc::new(writer.clone()), true), }; debug!( @@ -331,6 +363,7 @@ impl QueryStageExecutor for DefaultQueryStageExec { ShuffleWriterVariant::Passthrough(writer) => { utils::collect_plan_metrics(writer) } + ShuffleWriterVariant::Range(writer) => utils::collect_plan_metrics(writer), ShuffleWriterVariant::Sort(writer) => utils::collect_plan_metrics(writer), } } @@ -347,6 +380,7 @@ impl QueryStageExecutor for DefaultQueryStageExec { // the query. let plan: Arc = match &self.shuffle_writer { ShuffleWriterVariant::Passthrough(writer) => Arc::new(writer.clone()), + ShuffleWriterVariant::Range(writer) => Arc::new(writer.clone()), ShuffleWriterVariant::Sort(writer) => Arc::new(writer.clone()), }; match ballista_core::execution_plans::collect_runtime_stats_reports(&plan) { @@ -363,11 +397,12 @@ impl QueryStageExecutor for DefaultQueryStageExec { fn collect_window_state_reports( &self, ) -> Result> { - // Only the passthrough writer can sit over a window: the sort writer - // is Hash-partitioned by construction, which the prefix rewrite never - // plants. + // Both partitioning-preserving writers can sit over a window; the sort + // writer is Hash-partitioned by construction, which the prefix rewrite + // never plants. let captured = match &self.shuffle_writer { ShuffleWriterVariant::Passthrough(writer) => writer.collect_window_state()?, + ShuffleWriterVariant::Range(writer) => writer.collect_window_state()?, ShuffleWriterVariant::Sort(_) => return Ok(Vec::new()), }; captured diff --git a/ballista/executor/src/flight_service.rs b/ballista/executor/src/flight_service.rs index 3e3c8b079b..9ea40e2d1e 100644 --- a/ballista/executor/src/flight_service.rs +++ b/ballista/executor/src/flight_service.rs @@ -17,7 +17,12 @@ //! Implementation of the Apache Arrow Flight protocol that wraps an executor. +use ballista_core::JobId; use ballista_core::execution_plans::create_shuffle_path; +use ballista_core::execution_plans::range_shuffle::{ + index_path as range_index_path, is_ipc_file, open_ipc_file, +}; +use ballista_core::serde::scheduler::{ByteRange, ShuffleFileKind, ShuffleLayout}; use datafusion::arrow::ipc::reader::StreamReader; use std::convert::TryFrom; use std::fs::File; @@ -43,7 +48,7 @@ use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::arrow::{error::ArrowError, record_batch::RecordBatch}; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, info}; -use std::io::{BufReader, Read, Seek}; +use std::io::BufReader; use tokio::sync::mpsc::channel; use tokio::sync::mpsc::error::SendError; use tokio::{sync::mpsc::Sender, task}; @@ -100,20 +105,30 @@ impl FlightService for BallistaFlightService { stage_id, partition_id, file_id, - is_sort_shuffle, + layout, + file_kind, + byte_ranges, .. } => { - let path = create_shuffle_path( + if !byte_ranges.is_empty() { + // Byte ranges come back as bytes, which `do_get` cannot + // express — it returns decoded FlightData. Serving them + // here would mean decoding on the executor, which is the + // work a ranged read exists to avoid. + return Err(Status::invalid_argument( + "byte ranges are served by the IO_BLOCK_TRANSPORT action, \ + not do_get", + )); + } + let path = resolve_fetch_path( &self.work_dir, job_id, *stage_id, *partition_id, *file_id, - *is_sort_shuffle, - ) - .map_err(|e| { - Status::internal(format!("I/O error, can't create shuffle path: {e}")) - })?; + *layout, + *file_kind, + )?; debug!("FetchPartition reading partition {partition_id} from {path:?}"); // Check if this is a sort-based shuffle output @@ -143,29 +158,47 @@ impl FlightService for BallistaFlightService { )); } - // Standard single-file shuffle output - read the entire file - let file = File::open(&path) - .map_err(|e| { - BallistaError::General(format!( - "Failed to open partition file at {path:?}: {e:?}" - )) - }) - .map_err(|e| from_ballista_err(&e))?; - let file = BufReader::new(file); - // Safety: setting `skip_validation` requires `unsafe`, user assures data is valid - let reader = unsafe { - StreamReader::try_new(file, None) - .map_err(|e| from_arrow_err(&e))? - .with_skip_validation(cfg!(feature = "arrow-ipc-optimizations")) - }; - let (tx, rx) = channel(2); - let schema = reader.schema(); - task::spawn_blocking(move || { - if let Err(e) = read_partition(reader, tx) { - log::warn!("error streaming shuffle partition: {e}"); - } - }); + // Range shuffle writes the IPC file format, whose leading + // magic an IPC stream decoder rejects. The file says which it + // is, so neither the wire protocol nor the caller has to. + let schema = if is_ipc_file(&path) { + debug!("Detected range shuffle format for {path:?}"); + let reader = + open_ipc_file(&path).map_err(|e| from_ballista_err(&e))?; + let schema = reader.schema(); + task::spawn_blocking(move || { + if let Err(e) = read_partition(reader, tx) { + log::warn!("error streaming range shuffle partition: {e}"); + } + }); + schema + } else { + // Standard single-file shuffle output - read the entire file + let file = File::open(&path) + .map_err(|e| { + BallistaError::General(format!( + "Failed to open partition file at {path:?}: {e:?}" + )) + }) + .map_err(|e| from_ballista_err(&e))?; + let file = BufReader::new(file); + // Safety: setting `skip_validation` requires `unsafe`, user assures data is valid + let reader = unsafe { + StreamReader::try_new(file, None) + .map_err(|e| from_arrow_err(&e))? + .with_skip_validation(cfg!( + feature = "arrow-ipc-optimizations" + )) + }; + let schema = reader.schema(); + task::spawn_blocking(move || { + if let Err(e) = read_partition(reader, tx) { + log::warn!("error streaming shuffle partition: {e}"); + } + }); + schema + }; let write_options: IpcWriteOptions = IpcWriteOptions::default() .try_with_compression(Some(CompressionType::LZ4_FRAME)) @@ -264,33 +297,38 @@ impl FlightService for BallistaFlightService { stage_id, partition_id, file_id, - is_sort_shuffle, + layout, + file_kind, + byte_ranges, .. } => { - let path = create_shuffle_path( + let path = resolve_fetch_path( &self.work_dir, job_id, *stage_id, *partition_id, *file_id, - *is_sort_shuffle, - ) - .map_err(|e| { - Status::internal(format!( - "I/O error, can't create shuffle path: {e}" - )) - })?; + *layout, + *file_kind, + )?; debug!("FetchPartition reading {path:?}"); - let stream = if is_sort_shuffle_output(&path) { - // Sort-shuffle: stream the leading schema-header - // bytes followed by the requested partition's - // byte range. The receiver (BlockDataStream) walks - // the resulting concatenated IPC streams. + let stream = if !byte_ranges.is_empty() { + // The caller has read an index and knows which + // bytes it wants. Hand them over and resolve + // nothing — the same request an object store + // serves with a Range header. + stream_byte_ranges(&path, byte_ranges).await? + } else if is_sort_shuffle_output(&path) { + // Sort-shuffle asked for a partition and no ranges, + // so the executor resolves it through the index + // beside the data: one round trip, as ever. When + // that read moves to the caller this arm goes with + // it. stream_sort_shuffle_block(&path, *partition_id).await? } else { - // Hash-shuffle: file contains exactly one partition. + // One partition per file, so the file is the answer. stream_whole_file(&path).await? }; @@ -334,6 +372,110 @@ impl FlightService for BallistaFlightService { } } +/// Resolve a fetch request's identifiers to the file it names. +/// +/// The index's name follows from the layout, which is why the wire does not +/// spell it: a sort-shuffle output's index is its offset table, a passthrough +/// output's is the range shuffle's value index. A layout with no index beside +/// it is a request for something that does not exist, and says so. +fn resolve_fetch_path( + work_dir: &str, + job_id: &JobId, + stage_id: usize, + partition_id: usize, + file_id: Option, + layout: ShuffleLayout, + file_kind: ShuffleFileKind, +) -> Result { + let data = create_shuffle_path( + work_dir, + job_id, + stage_id, + partition_id, + file_id, + matches!(layout, ShuffleLayout::Sort), + ) + .map_err(|e| { + Status::internal(format!("I/O error, can't create shuffle path: {e}")) + })?; + + Ok(match file_kind { + ShuffleFileKind::Data => data, + ShuffleFileKind::Index => match layout { + ShuffleLayout::Sort => get_index_path(data.as_path()), + ShuffleLayout::Passthrough => range_index_path(data.as_path()), + }, + }) +} + +/// Stream the requested byte ranges of `path`, concatenated in request order. +/// +/// The serving side resolves nothing here: a consumer that has read an index +/// knows which bytes it wants, and this hands them over. Ranges are clamped to +/// the file so a stale index cannot walk off the end. +/// +/// Chunked at [`BLOCK_BUFFER_CAPACITY`] like a whole-file read, for two +/// reasons: a range can be far larger than a gRPC message may be, and buffering +/// it whole would hold a consumer's entire share of a partition in the serving +/// executor's memory. +async fn stream_byte_ranges( + path: &std::path::Path, + ranges: &[ByteRange], +) -> Result<::DoActionStream, Status> { + use tokio::io::{AsyncReadExt, AsyncSeekExt}; + + let len = tokio::fs::metadata(path) + .await + .map_err(|e| Status::internal(format!("Failed to stat {path:?}: {e}")))? + .len(); + + for range in ranges { + if range.offset >= len { + return Err(Status::out_of_range(format!( + "range at {} is past the end of {path:?} ({len} bytes)", + range.offset + ))); + } + } + + debug!( + "serving {} byte ranges ({} bytes) from {path:?}", + ranges.len(), + ranges + .iter() + .map(|r| r.length.min(len - r.offset)) + .sum::(), + ); + + let path = path.to_owned(); + let ranges: Vec = ranges.to_vec(); + let stream = futures::stream::iter(ranges) + .then(move |range| { + let path = path.clone(); + async move { + let mut file = tokio::fs::File::open(&path).await.map_err(|e| { + Status::internal(format!("Failed to open {path:?}: {e}")) + })?; + file.seek(std::io::SeekFrom::Start(range.offset)) + .await + .map_err(|e| Status::internal(format!("seek {path:?}: {e}")))?; + let take = range.length.min(len - range.offset); + // Both levels carry `Status` so the flatten has one error type. + Ok::<_, Status>( + ReaderStream::with_capacity(file.take(take), BLOCK_BUFFER_CAPACITY) + .map(|chunk| { + chunk + .map(|bytes| arrow_flight::Result { body: bytes }) + .map_err(|e| Status::internal(format!("I/O error: {e}"))) + }), + ) + } + }) + .try_flatten(); + + Ok(Box::pin(stream)) +} + async fn stream_whole_file( path: &std::path::Path, ) -> Result<::DoActionStream, Status> { @@ -404,12 +546,17 @@ async fn stream_sort_shuffle_block( }))) } -fn read_partition( - reader: StreamReader>, +/// Pump every batch a shuffle-file decoder yields into `tx`. +/// +/// Generic over the decoder: the two shuffle formats need different ones +/// (`StreamReader` for IPC stream, `FileReader` for IPC file) and share no +/// arrow trait past `Iterator`. +fn read_partition( + reader: R, tx: Sender>, ) -> Result<(), FlightError> where - T: Read + Seek, + R: Iterator>, { if tx.is_closed() { return Err(FlightError::Tonic(Box::new(Status::internal( diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index c0a742be6d..fcabee08cb 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -23,7 +23,8 @@ use crate::state::aqe::planner::AdaptiveStageInfo; use crate::state::execution_graph::StageOutput; use ballista_core::JobId; use ballista_core::execution_plans::{ - RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, + RangeFilterExec, RangeShuffleReaderExec, RangeShuffleWriterExec, ShuffleReaderExec, + ShuffleWriter, ShuffleWriterExec, }; use datafusion::common::exec_err; use datafusion::config::ConfigOptions; @@ -40,6 +41,47 @@ use datafusion::{ use std::collections::HashMap; use std::sync::Arc; +/// Swap a passthrough writer for one that emits the seekable Arrow IPC file +/// format, when this stage's output is read back in value-range order. +/// +/// Gated on the same condition as [`BallistaAdapter::build_reader`]'s choice of +/// `RangeShuffleReaderExec` — the stage's child declares an output ordering — +/// because the two have to agree: the range reader is the only consumer taught +/// to open the file format. This is the only place that sees both sides, which +/// is also why the static planner never produces the pair. It plants the +/// arrival-order reader, which would be handed a format it cannot decode. +/// +/// No configuration decides this. What a reader may do with a source is a fact +/// about the source, and the reader establishes it by opening the file; a flag +/// would be a second answer to the same question, free to disagree. +fn use_range_shuffle_writer( + writer: Arc, + exchange: &ExchangeExec, +) -> datafusion::error::Result> { + if exchange.broadcast || exchange.input().output_ordering().is_none() { + return Ok(writer); + } + let as_plan: &dyn ExecutionPlan = writer.as_ref(); + let Some(passthrough) = as_plan.downcast_ref::() else { + // A hash-repartition stage writes its own consolidated format and is + // never read back by the range reader. + return Ok(writer); + }; + let children = passthrough.children(); + let [input] = children.as_slice() else { + return Err(DataFusionError::Internal(format!( + "ShuffleWriterExec must have exactly 1 child, got {}", + children.len() + ))); + }; + Ok(Arc::new(RangeShuffleWriterExec::try_new( + passthrough.job_id().clone(), + passthrough.stage_id(), + Arc::clone(input), + passthrough.work_dir().to_owned(), + )?)) +} + #[derive(Debug, Clone, Default)] pub(crate) struct BallistaAdapter { inputs: HashMap, @@ -183,6 +225,7 @@ impl BallistaAdapter { .clone() .transform_down(|e| adapter.transform_children(e))? .data; + let plan = attach_reader_bounds(plan)?; let stage_id = root.stage_id().ok_or_else(|| { DataFusionError::Execution( "shuffle partitions have to be resolved at this point".to_string(), @@ -198,6 +241,7 @@ impl BallistaAdapter { config, ) .map_err(|e| DataFusionError::External(Box::new(e)))?; + let writer = use_range_shuffle_writer(writer, root)?; Ok(AdaptiveStageInfo { plan: writer, @@ -211,6 +255,7 @@ impl BallistaAdapter { .clone() .transform_down(|e| adapter.transform_children(e))? .data; + let plan = attach_reader_bounds(plan)?; let stage_id = root.stage_id().ok_or_else(|| { DataFusionError::Execution( "shuffle partitions have to be resolved at this point".to_string(), @@ -233,6 +278,48 @@ impl BallistaAdapter { } } +/// Hand every [`RangeShuffleReaderExec`] the value range its output +/// partitions cover, taken from the [`RangeFilterExec`] above it. +/// +/// The filter has already widened the cuts by its own halos, and those widened +/// ranges are exactly what the reader may narrow its reads to: anything inside +/// them the filter might keep, anything outside it would drop anyway. Reading +/// them off the filter rather than recomputing from cuts is what stops the two +/// from disagreeing about halo width. +/// +/// Runs after [`resolve_range_filter_cuts`], which is what puts the bounds on +/// the filter in the first place. A filter with no reader beneath it, or a +/// reader with no filter above it, is left alone — the reader then reads its +/// sources whole, which is the answer it gave before it could narrow at all. +fn attach_reader_bounds( + plan: Arc, +) -> Result, DataFusionError> { + plan.transform_down(|node| { + let Some(range_filter) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(bounds) = range_filter.widened_bounds() else { + return Ok(Transformed::no(node)); + }; + let children = node.children(); + let [child] = children.as_slice() else { + return Ok(Transformed::no(node)); + }; + let Some(reader) = child.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + if bounds.len() != reader.partition.len() { + return Ok(Transformed::no(node)); + } + let reader = Arc::new(reader.clone().with_bounds(bounds)?); + Ok(Transformed::yes(replace_children_if_necessary( + node, + vec![reader], + )?)) + }) + .map(|transformed| transformed.data) +} + /// Walk `plan` and resolve every pending [`RangeFilterExec`]'s bounds from /// its own descendant boundary `ExchangeExec`'s stored routing. Called at /// `adapt_to_ballista` time, once the upstream stage's sketches have diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 277634c48a..ead4ef53d6 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -30,7 +30,8 @@ use log::{debug, error, info, warn}; use ballista_core::JobId; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ - ShuffleWriter, ShuffleWriterExec, SortShuffleWriterExec, UnresolvedShuffleExec, + RangeShuffleWriterExec, ShuffleWriter, ShuffleWriterExec, SortShuffleWriterExec, + UnresolvedShuffleExec, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; @@ -1690,6 +1691,9 @@ impl ExecutionPlanVisitor for ExecutionStageBuilder { // Handle both ShuffleWriterExec and SortShuffleWriterExec if let Some(shuffle_write) = plan.downcast_ref::() { self.current_stage_id = shuffle_write.stage_id(); + } else if let Some(shuffle_write) = plan.downcast_ref::() + { + self.current_stage_id = shuffle_write.stage_id(); } else if let Some(shuffle_write) = plan.downcast_ref::() { self.current_stage_id = shuffle_write.stage_id(); } else if let Some(unresolved_shuffle) = diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 1ee1e191ab..68fd51feb4 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -33,7 +33,8 @@ use log::{debug, warn}; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ - ShuffleWriterExec, SortShuffleWriterExec, TaskRuntimeStats, TaskWindowState, + RangeShuffleWriterExec, ShuffleWriterExec, SortShuffleWriterExec, TaskRuntimeStats, + TaskWindowState, }; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::{ @@ -1331,6 +1332,7 @@ impl Debug for FailedStage { /// scheduling unit. fn stage_input_partitions(plan: &Arc) -> usize { if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() || plan.downcast_ref::().is_some() { plan.children()[0] diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index eeac4f039f..28f398b372 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -314,7 +314,24 @@ fn select_output_partitions( ) else { return Ok(None); }; - return Ok(Some(Arc::new(restricted.with_fetch_limit(reader.fetch())))); + let restricted = restricted.with_fetch_limit(reader.fetch()); + // Bounds are per output partition, so they follow the same slice the + // locations did. Dropping them here would leave the task reading its + // sources whole — right answer, none of the saving. + let restricted = match reader.bounds() { + Some(bounds) => { + let kept = indices + .iter() + .filter_map(|&p| bounds.get(p).cloned()) + .collect(); + let Ok(restricted) = restricted.with_bounds(kept) else { + return Ok(None); + }; + restricted + } + None => restricted, + }; + return Ok(Some(Arc::new(restricted))); } // DataSourceExec: file-backed or in-memory scans.