From ef8d0a44c085472d03d02f9e2a0a34a63ded384b Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 24 Aug 2026 16:05:29 -0600 Subject: [PATCH 1/2] feat(core): RangeShuffleWriterExec writing the seekable Arrow IPC file format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A stage routed by value range hands each consumer producer files it reads whole and mostly discards. Measured on h2o Q8 at 1e8 rows, the consumer fetches 8.2 GB to feed a window needing 3.6 GB. Skipping the excess means seeking, and seeking needs a chunk index, which the IPC stream format the passthrough shuffle writes does not have. This adds a writer that emits the IPC file format, whose footer indexes every record batch. Reads stay whole-file, so volume is unchanged — what lands is the substrate: a footer to seek against, plus the per-batch byte offsets captured at write time, which are the content of the value index that follows. A separate operator rather than a flag on ShuffleWriterExec: the two IPC framings are not interchangeable on the read side, so a mode flag would put every existing reader one config change away from failing to decode. As a distinct writer the existing shuffle keeps its format untouched, and only the readers taught to recognise the new one will open it. The pair is planted in the AQE adapter, the only place that sees both sides, under the same condition that plants RangeShuffleReaderExec — the stage's child declares an output ordering. The static planner never plants it, since it plants the arrival-order reader. Off by default behind `ballista.shuffle.range.enabled`. Readers pick a decoder by sniffing the file's magic rather than by a flag on PartitionLocation, keeping the format out of the wire protocol; the local reader and the Flight do_get path both do this. Remote reads route over do_get while the flag is on, because the IO_BLOCK_TRANSPORT action ships raw bytes to a stream-only decoder. Verified against DataFusion on h2o Q8 with the flag on and off, and the files on disk confirm the split: the range-routed stage is ARROW1, the result stage stays a stream. Co-Authored-By: Claude Opus 5 (1M context) fix(docs): regenerate config table so the padding matches update_config_docs Co-Authored-By: Claude Opus 5 (1M context) fix(core): take range shuffle block offsets from the footer, not a byte counter The offsets a value index hands to a range read have to name a record batch. Counting bytes around `FileWriter::write` cannot produce them: one such call emits the batch's dictionary blocks *and* its record block, so the counter reports an offset pointing at a dictionary and a length spanning both. Data with no dictionary columns hides this — the two coincide — which is why the h2o benchmark and the previous tests were green. The footer already holds what the counter was reconstructing, separated by kind and computed by the writer that emitted the bytes. Read it back after `finish`, at the cost of one tail read per shuffle file, and drop the counter. `read_file_layout` returns dictionary and record blocks as `FileLayout`, and the writer asserts the record-block count equals the batches written, so a future re-chunking fails loudly instead of producing an index that addresses the wrong batches. Tests now slice the file at each recorded range and decode it, rather than checking the offsets are self-consistent — arithmetic agreed with itself while being wrong. One covers dictionary-encoded data, which is the only shape that can tell the two schemes apart. Co-Authored-By: Claude Opus 5 (1M context) feat(core): write a value index beside each range shuffle file One row per IPC message in the data file, in file order: key types, nullable first row's key in this batch is_dict Boolean row describes a dictionary byte_offset UInt64 absolute, from the data file footer byte_len UInt64 num_rows UInt64, nullable rows in the batch A consumer assigned `[lo, hi)` reads this instead of the data file, binary searches the key columns, and fetches only the covering byte ranges. Measured on h2o Q8: 5 KB of index against 19 MB of data. Keys are one typed column per ORDER BY expression rather than an encoded blob, so multi-column keys, DESC and mixed types work by construction, and the schema names what the file is indexed on. They are nullable for two independent reasons — a dictionary row has no key, and a first key can genuinely be NULL under NULLS FIRST — so `is_dict` carries that distinction rather than nullability implying it. Dictionaries get rows because a record batch referencing one cannot be decoded from its own bytes alone. The IPC footer does not record which dictionary a block carries, so a reader takes every dictionary row preceding its range: conservative, and correct under replacement and delta alike. `num_rows` is what lets a ROWS-frame halo be computed from indexes before any data is fetched — sum counts below a boundary until N rows are covered, rather than fetching and decoding to find out. Its summation across files assumes each row lives in exactly one file, which the unexplained stage-0 write amplification has to settle before a consumer relies on it. Keys are collected by a stream adapter on the way to the writer, so a key and the block describing it come from the same batch rather than from two passes that could disagree. The writer now requires an ordered input, since a child declaring no ordering has nothing to index. Nothing reads the index yet. Co-Authored-By: Claude Opus 5 (1M context) feat(core): read only the batches a range shuffle consumer's range covers The reader consults each source's value index, binary searches its assigned range, and reads only the batches that can hold a row in it. Selection is at batch granularity and the RangeFilterExec above still does the exact trim, so this changes the volume read and not the rows produced. h2o Q8 at 1e7, read amplification (reader rows / rows the halo trim keeps): K before after 8 1.69 1.41 32 1.38 1.22 128 1.20 1.17 About half the available saving, because only local sources seek today; remote sources still come back whole over `do_get`. Bounds come from the RangeFilterExec above the reader, which has already widened the cuts by its own halos — reading them off the filter rather than recomputing from cuts is what stops the two disagreeing about halo width. They are per output partition, so `restrict_plan_to_partitions` slices them alongside the locations; dropping them there left every task reading whole files, which is how this was first found to be doing nothing. Removes `ballista.shuffle.range.enabled`. Whether a source can be seeked is a fact about the source, and the reader establishes it by opening the file — a flag was a second answer to the same question, free to disagree with the first. It did: with the flag off the writer emitted IPC stream while the reader still received bounds, so it looked for a footer that was not there. The range writer is now planted whenever the range reader is, which is the condition that already had to hold for the two to agree. Co-Authored-By: Claude Opus 5 (1M context) feat: fetch remote range shuffle sources by byte range Read amplification on h2o Q8 at 1e7, reader rows over rows the halo trim keeps: K before local seek only this 8 1.69 1.41 1.01 32 1.38 1.22 1.07 128 1.20 1.17 1.19 At K=8 the consumer fetches 373 MB where it needs 372 MB, against 626 MB before — within half a percent of the floor. K=128 gains nothing because a cut is 0.77 wide against a halo of 3, so one 8192-row batch already spans more than the band a partition wants and selection cannot go below a batch. The consumer fetches the index, searches it, and asks for byte ranges. The executor resolves nothing: no index read, no comparison, no notion of what a range means. That is what makes the same two steps work against object storage, where there is no executor to ask, and it is why the index carries byte offsets rather than leaving the lookup to whoever holds the file. Sort shuffle keeps its single round trip: coordinates with no ranges still ask the executor to resolve a partition through its own index. Moving that lookup to the caller needs no further protocol change — fetch the index, read two i64s, ask for the bytes — and would delete `stream_sort_shuffle_block` rather than add a path beside it. Protocol version 2 -> 3. `FetchPartition.is_sort_shuffle` becomes `layout`, because the bool selected a path shape and reads as data-vs-index once a `file_kind` sits next to it. `byte_ranges` is empty for "whatever these identifiers address" and otherwise absolute file offsets. The schema message is written by the consumer from the schema it already knows, not sliced out of the data file: arrow-rs pads the file magic to its write alignment, 64 bytes, so slicing the header encodes another implementation's padding rule into an offset. That cost a debugging cycle — offset 8 held padding, and the decoder read the zeros as end-of-stream. Co-Authored-By: Claude Opus 5 (1M context) fix(executor): chunk a ranged shuffle response instead of buffering it A ranged read built the whole requested range in memory and sent it as one gRPC message. At 1e8 rows that is ~29 MB against a 16 MB receive limit, so the consumer failed with a transport error wearing an IpcError's clothing: `execute_do_action` maps transport failures through ArrowError, which made a message-size rejection read as a corrupt stream. Chunked at BLOCK_BUFFER_CAPACITY like a whole-file read, and streamed rather than collected — buffering held a consumer's entire share of a partition in the serving executor's memory, which is the same defect measured in bytes instead of in errors. Found by measuring 1e8 to fill in a table, which is the argument for filling in tables. Co-Authored-By: Claude Opus 5 (1M context) fix(docker): base runtime images on ubuntu:26.04 These Dockerfiles are not hermetic: `.local/build.sh` compiles on the host and each image does `COPY target/${RELEASE_FLAG}/ /root/`, so the host's glibc must be no newer than the base image's. Building on Ubuntu 26.04 (glibc 2.43) against `ubuntu:24.04` (2.39) produces images that fail at the loader before `main` runs: /root/ballista-executor: /lib/x86_64-linux-gnu/libm.so.6: version `GLIBC_2.43' not found (required by /root/ballista-executor) `chaos.Dockerfile` is left alone — it builds inside `rust:1-bookworm` and runs on `debian:bookworm-slim`, so it carries its own toolchain and is unaffected by the host. Co-Authored-By: Claude Opus 5 (1M context) fix(core): address the range shuffle index's fixed columns by position Key columns are named for the user's sort expressions and sit ahead of the fixed layout columns. Arrow permits duplicate field names and resolves `column_by_name` to the first match, so a query ordering by a column called `byte_offset` shadowed the real one and was read as whatever type the key happened to be. The fixed columns are now reached through `is_dict_column`, `byte_offset_column`, `byte_len_column` and `num_rows_column`, which index `key_column_count(index) + position` and downcast to the declared layout type. `fixed_fields` defines those columns once, so `index_schema` and the accessors cannot drift; a test asserts each position constant names the field it addresses. The name constants stay internal rather than being re-exported, so no caller outside the module can reintroduce the lookup. Co-Authored-By: Claude Opus 5 (1M context) docs(core): say where range shuffle seeking happens, not that it doesn't `ipc_file`'s module doc still said "Nothing seeks yet — RangeShuffleWriterExec writes whole files and RangeShuffleReaderExec reads them whole". Both halves are now false in the same commit that added the seeking: the reader opens a file at the batches a value range covers, and the remote path fetches only those byte ranges. Describe where the seeking lives instead of what has not landed, so the doc states behaviour that survives rather than progress that rots. Co-Authored-By: Claude Opus 5 (1M context)  Conflicts:  ballista/core/src/diagram.rs --- ballista/core/proto/ballista.proto | 66 +- ballista/core/src/client.rs | 120 ++- .../src/execution_plans/distributed_query.rs | 7 +- ballista/core/src/execution_plans/mod.rs | 2 + .../core/src/execution_plans/plan_algebra.rs | 5 +- .../execution_plans/range_shuffle/index.rs | 745 +++++++++++++++ .../execution_plans/range_shuffle/ipc_file.rs | 494 ++++++++++ .../src/execution_plans/range_shuffle/mod.rs | 59 ++ .../execution_plans/range_shuffle/reader.rs | 111 +++ .../execution_plans/range_shuffle/remote.rs | 362 ++++++++ .../execution_plans/range_shuffle/writer.rs | 849 ++++++++++++++++++ .../execution_plans/range_shuffle_reader.rs | 145 ++- .../src/execution_plans/shuffle_reader.rs | 218 ++++- .../src/execution_plans/shuffle_writer.rs | 120 ++- ballista/core/src/lib.rs | 2 +- ballista/core/src/serde/generated/ballista.rs | 119 ++- ballista/core/src/serde/mod.rs | 119 ++- .../core/src/serde/scheduler/from_proto.rs | 27 +- ballista/core/src/serde/scheduler/mod.rs | 96 +- ballista/core/src/serde/scheduler/to_proto.rs | 14 +- ballista/executor/src/execution_engine.rs | 49 +- ballista/executor/src/flight_service.rs | 243 ++++- ballista/scheduler/src/state/aqe/adapter.rs | 89 +- .../scheduler/src/state/execution_graph.rs | 6 +- .../scheduler/src/state/execution_stage.rs | 4 +- ballista/scheduler/src/state/task_builder.rs | 19 +- 26 files changed, 3868 insertions(+), 222 deletions(-) create mode 100644 ballista/core/src/execution_plans/range_shuffle/index.rs create mode 100644 ballista/core/src/execution_plans/range_shuffle/ipc_file.rs create mode 100644 ballista/core/src/execution_plans/range_shuffle/mod.rs create mode 100644 ballista/core/src/execution_plans/range_shuffle/reader.rs create mode 100644 ballista/core/src/execution_plans/range_shuffle/remote.rs create mode 100644 ballista/core/src/execution_plans/range_shuffle/writer.rs 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..ae44796da5 --- /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. Nothing consumes that layout yet +/// — it is logged — but it is what the value index is built from, so having +/// it here makes the index a matter of writing a sidecar rather than of +/// re-plumbing the writer. +/// +/// [`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. From 96991336284126a2114f3f63f7f984240ff9e0e6 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Tue, 25 Aug 2026 19:34:08 -0600 Subject: [PATCH 2/2] docs(core): say what consumes the range shuffle footer layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The writer's doc said "Nothing consumes that layout yet — it is logged". Both halves are false: the layout is passed straight into `write_index_file`, which pairs each block with the first key the `KeyCollector` saw in it. That sidecar is the whole point of reading the footer back. Co-Authored-By: Claude Opus 5 (1M context) --- ballista/core/src/execution_plans/range_shuffle/writer.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ballista/core/src/execution_plans/range_shuffle/writer.rs b/ballista/core/src/execution_plans/range_shuffle/writer.rs index ae44796da5..11e00e4624 100644 --- a/ballista/core/src/execution_plans/range_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/range_shuffle/writer.rs @@ -83,10 +83,10 @@ use super::ipc_file::write_stream_to_ipc_file; /// # Message layout /// /// Each drain reads back its file's footer, so it knows the byte range of -/// every record batch and every dictionary. Nothing consumes that layout yet -/// — it is logged — but it is what the value index is built from, so having -/// it here makes the index a matter of writing a sidecar rather than of -/// re-plumbing the writer. +/// 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)]