diff --git a/Cargo.lock b/Cargo.lock index bc93015eb..8d2655dc6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1123,7 +1123,6 @@ dependencies = [ "chrono", "clap 4.6.6", "datafusion", - "datafusion-functions-aggregate-common", "datafusion-proto", "datafusion-proto-common", "datafusion-spark", @@ -1234,7 +1233,6 @@ dependencies = [ "clap 4.6.6", "dashmap", "datafusion", - "datafusion-functions-aggregate-common", "datafusion-proto", "datafusion-substrait", "ferroid", diff --git a/ballista/client/tests/parallel_window.rs b/ballista/client/tests/parallel_window.rs new file mode 100644 index 000000000..deb0cd456 --- /dev/null +++ b/ballista/client/tests/parallel_window.rs @@ -0,0 +1,208 @@ +// 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. + +//! End-to-end coverage for `ParallelWindowRule`, which rewrites a bounded +//! RANGE-frame window into `RSE → ORRE → RSE → RangeFilterExec → PBWAG`. +//! +//! The rewrite is a pure optimization, so single-node DataFusion is the +//! oracle: every case runs there, then on Ballista with the rule off, then on +//! Ballista with the rule on, and all three must agree. A disagreement names +//! which side moved. +//! +//! What the cases vary is what the rewrite has to get right: which end of the +//! order the NULL run occupies, whether the frame widens the lower bound, the +//! upper, both, or neither, and whether a DESC key stays on the serial path. + +mod common; + +#[cfg(test)] +#[cfg(feature = "standalone")] +mod parallel_window { + use std::fs; + use std::path::Path; + + use ballista::prelude::{SessionConfigExt, SessionContextExt}; + use ballista_core::config::{ + BALLISTA_ADAPTIVE_PLANNER_ENABLED, BALLISTA_PARALLEL_WINDOW_ENABLED, + }; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::*; + use rstest::rstest; + + /// Rows per file. Four files of these interleave over the whole key range, + /// so no file is range-disjoint from another and ORRE has real routing to + /// do rather than passing each input partition through whole. + const ROWS_PER_FILE: i64 = 12; + const FILES: i64 = 4; + + /// Every fourth key is NULL, so the run is long enough to span more than + /// one output partition's share and cannot be mistaken for a rounding + /// artifact at whichever end it lands. + fn key_is_null(id: i64) -> bool { + id % 4 == 3 + } + + /// One CSV per file, rows dealt round-robin by `id` so file `f` holds keys + /// `f, f + FILES, f + 2 * FILES, ...` — every file covers the full range. + fn write_table(dir: &Path, nullable: bool) { + let table = dir.join("w"); + fs::create_dir_all(&table).unwrap(); + for file in 0..FILES { + let mut csv = String::from("id,k,v\n"); + for row in 0..ROWS_PER_FILE { + let id = file + row * FILES; + let key = if nullable && key_is_null(id) { + String::new() + } else { + id.to_string() + }; + csv.push_str(&format!("{id},{key},{}\n", id * 10)); + } + fs::write(table.join(format!("p{file}.csv")), csv).unwrap(); + } + } + + async fn register(ctx: &SessionContext, dir: &Path, nullable: bool) { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("k", DataType::Int64, nullable), + Field::new("v", DataType::Int64, false), + ]); + ctx.register_csv( + "w", + dir.join("w").to_str().unwrap(), + CsvReadOptions::new() + .has_header(true) + .schema(&schema) + .file_extension(".csv"), + ) + .await + .unwrap(); + } + + async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + pretty_format_batches(&batches).unwrap().to_string() + } + + /// `ORDER BY id` on the outside so the comparison is over rows, not over + /// whichever order the K output partitions happened to be drained in. + fn query(ordering: &str, frame: &str) -> String { + format!( + "SELECT id, k, SUM(v) OVER (ORDER BY k {ordering} RANGE BETWEEN {frame}) AS w \ + FROM w ORDER BY id" + ) + } + + /// A Ballista standalone context with AQE on, the rule at `parallel`, and + /// K = 4 output partitions. + async fn ballista(dir: &Path, nullable: bool, parallel: bool) -> SessionContext { + let config = SessionConfig::new_with_ballista() + .with_target_partitions(4) + .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, true) + .set_bool(BALLISTA_PARALLEL_WINDOW_ENABLED, parallel); + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + let ctx = SessionContext::standalone_with_state(state).await.unwrap(); + register(&ctx, dir, nullable).await; + ctx + } + + #[rstest] + #[case::halo_both("ASC NULLS LAST", "5 PRECEDING AND 5 FOLLOWING")] + #[case::halo_lower("ASC NULLS LAST", "5 PRECEDING AND CURRENT ROW")] + #[case::halo_upper("ASC NULLS LAST", "CURRENT ROW AND 5 FOLLOWING")] + #[case::halo_none("ASC NULLS LAST", "CURRENT ROW AND CURRENT ROW")] + #[case::nulls_first("ASC NULLS FIRST", "5 PRECEDING AND 5 FOLLOWING")] + #[case::nulls_first_halo_none("ASC NULLS FIRST", "CURRENT ROW AND CURRENT ROW")] + // DESC is gated to the serial path by `ParallelWindowRule` — RANGE frame + // semantics invert with the sort direction. These cases assert the gate + // leaves the answer alone rather than that the rewrite fired. + #[case::desc_nulls_last("DESC NULLS LAST", "5 PRECEDING AND 5 FOLLOWING")] + #[case::desc_nulls_first("DESC NULLS FIRST", "5 PRECEDING AND 5 FOLLOWING")] + #[tokio::test] + async fn parallel_window_agrees_with_datafusion( + #[case] ordering: &str, + #[case] frame: &str, + // A nullable key is the case the rewrite has to place the NULL run for; + // a non-nullable one has no run and exercises the plain value paths. + #[values(false, true)] nullable: bool, + ) { + let case = format!("{ordering}_{frame}_{nullable}").replace(' ', "_"); + let dir = std::env::temp_dir().join(format!("ballista_parallel_window_{case}")); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + write_table(&dir, nullable); + + let sql = query(ordering, frame); + + let oracle_ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(4), + ); + register(&oracle_ctx, &dir, nullable).await; + let oracle = run(&oracle_ctx, &sql).await; + + let serial = run(&ballista(&dir, nullable, false).await, &sql).await; + assert_eq!( + oracle, serial, + "Ballista disagrees with DataFusion before the rule is even on\n{sql}" + ); + + let parallel = run(&ballista(&dir, nullable, true).await, &sql).await; + assert_eq!( + oracle, parallel, + "the parallel-window rewrite changed the answer\n{sql}" + ); + + let _ = fs::remove_dir_all(&dir); + } + + /// The rewrite has to fire for the cases above to mean anything. A plan + /// still holding `BoundedWindowAggExec` over a `SortPreservingMergeExec` + /// is the serial shape, and an agreeing answer from it proves nothing. + #[tokio::test] + async fn the_rewrite_actually_fires() { + let dir = std::env::temp_dir().join("ballista_parallel_window_explain"); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + write_table(&dir, true); + + let sql = query("ASC NULLS LAST", "5 PRECEDING AND 5 FOLLOWING"); + let ctx = ballista(&dir, true, true).await; + // `EXPLAIN` alone shows only the first stage split. The rewrite lands + // when AQE re-plans stage 1 against stage 0's reported sketch, which + // is reached by executing. + let plan = run(&ctx, &format!("EXPLAIN ANALYZE {sql}")).await; + + for operator in [ + "OrderedRangeRepartitionExec", + "RangeFilterExec", + "RuntimeStatsExec", + ] { + assert!( + plan.contains(operator), + "rewrite did not plant {operator}:\n{plan}" + ); + } + + let _ = fs::remove_dir_all(&dir); + } +} diff --git a/ballista/core/Cargo.toml b/ballista/core/Cargo.toml index fb9c1c985..735918390 100644 --- a/ballista/core/Cargo.toml +++ b/ballista/core/Cargo.toml @@ -51,7 +51,6 @@ aws-credential-types = { version = "1.2.0", optional = true } chrono = { version = "0.4", default-features = false } clap = { workspace = true, optional = true } datafusion = { workspace = true } -datafusion-functions-aggregate-common = { workspace = true } datafusion-proto = { workspace = true } datafusion-proto-common = { workspace = true } datafusion-spark = { workspace = true, optional = true, features = ["datafusion"] } diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index ae8098bce..4273fde81 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -92,11 +92,46 @@ message RuntimeStatsExecNode { repeated datafusion.PhysicalSortExprNode order_by = 1; } -// Serialized T-Digest as a fixed-layout `Vec` per -// `TDigest::to_scalar_state()`: max_size, sum, count, max, min, -// centroid_means..., centroid_weights.... -message QuantileSketchState { - repeated datafusion_common.ScalarValue state = 1; +// A `SortKeySketch` is a wrapper around a KLL, that also includes NULL counts +// +// `levels` is an Arrow IPC stream for now. Binary blob encoding was considered +// so that encoding tricks could be used like delta+varint. These showed a 4x +// improvement for u64, but only 25% on f64 due to mantissas being effectively +// random. Ultimately it was decided that a more complex encoder is not worth +// the added complexity in this PR. We can revisit the decision in the future. +message SortKeySketchState { + // KLL's nominal top-level compactor capacity + uint32 k = 1; + // Rows whose key was NULL. Not in `levels`, and not recoverable from it. + uint64 null_count = 2; + // The exact minimum and maximum over every value observed, one element per + // ORDER BY expression. Tracked outside the compactor stack so no coin flip + // can move them, which is what keeps `quantile(0.0)` and `quantile(1.0)` + // exact. Empty when no value was observed. + repeated datafusion_common.ScalarValue key_min = 3; + repeated datafusion_common.ScalarValue key_max = 4; + // The compactor stack: one row per level in level order, single column + // `levels: List>` holding that level's retained keys ascending, + // one struct field per ORDER BY expression. An item's weight is + // `2^level`, so the row index carries the weight and the total count is + // the weighted sum — neither ships. + // + // Levels are sorted before serialization, so a decoder takes ascending + // order as given rather than being told per level. + // + // A three-level stack over an `Int64` key decodes to one batch of one + // column, where the row index is the level: + // + // levels: List> + // + // +-------------------------------------------+ + // | levels | + // +-------------------------------------------+ + // | [{expr_0: 4}, {expr_0: 17}, {expr_0: 23}] | level 0, weight 1 + // | [{expr_0: 9}, {expr_0: 31}] | level 1, weight 2 + // | [{expr_0: 12}] | level 2, weight 4 + // +-------------------------------------------+ + bytes levels = 5; } // Flow-control operator with an operator-mode enum. The child plan is @@ -145,10 +180,30 @@ message PerPartitionFilterExecNode { // The child plan is plumbed by the framework as `inputs[0]` during decode. // Serialization requires bounds to be resolved. message RangeFilterExecNode { - datafusion.PhysicalExprNode routing_expr = 1; + datafusion.PhysicalExprNode filter_expr = 1; datafusion_common.ScalarValue halo_lo = 2; datafusion_common.ScalarValue halo_hi = 3; repeated RangeBound raw_bounds = 4; + // What the operator that produced the cuts says about the rows arriving. + // Unset claims nothing, which a nullable `filter_expr` has no answer for: + // a NULL has a side in the order, not a position among the values. + // + // `ordered` is also required of the input, so nothing planted between can + // reorder the rows the placement was stated against. `unordered_nulls_first` + // states the side without demanding an order, which is what leaves an + // unordered range-repartition upstream legal. + oneof input_order { + bool unordered_nulls_first = 5; + SortOptions ordered = 6; + } +} + +// An arrow `SortOptions`, which has no proto of its own in the DataFusion +// descriptors — `PhysicalSortExprNode` carries the pair inline beside an +// expression this message's users already have. +message SortOptions { + bool descending = 1; + bool nulls_first = 2; } // Half-open `[lo, hi)` cut range for one input partition. Either side may be @@ -672,18 +727,32 @@ message RuntimeStatsReport { // partition, post-repartition gets one per output sub-partition. The // scheduler groups by `order_by` tag and aggregates. repeated RuntimeStatsPartitionEntry partitions = 2; + // Every partition's observations folded into one sketch, merged on the + // executor before the report is sent. Absent in row-count-only mode, + // and when no partition observed a value. + optional SortKeySketchState sketch = 3; } // One partition's observations from a `RuntimeStatsExec`. message RuntimeStatsPartitionEntry { uint32 partition_id = 1; uint64 row_count = 2; - // Present when the `RuntimeStatsExec` was in sketch mode AND this - // partition observed at least one non-null routing value. - optional QuantileSketchState sketch = 3; - // TODO: `optional MinMaxState min_max` — for a lighter post-repartition - // mode where the bin-packer just needs (min, max, count) per - // sub-partition and a full T-Digest is overkill. + reserved 3; // was TDigest + reserved "sketch"; + // This partition's exact key range, one element per ORDER BY expression. + // `cut_partitions` routes a whole shuffle file into every downstream + // partition whose range overlaps `[key_min, key_max]`, so these are the + // only per-partition facts a router needs — never the distribution, which + // is why the sketch beside them is merged rather than repeated here. + // + // Both empty when this partition observed no value, which `null_count` + // then distinguishes: some NULLs means a file of NULLs with no value range + // to overlap, none means the partition saw nothing at all. + repeated datafusion_common.ScalarValue key_min = 4; + repeated datafusion_common.ScalarValue key_max = 5; + // Rows in this partition whose key was NULL. Folds to the report-level + // `SortKeySketchState.null_count`. + uint64 null_count = 6; } message ExecutionError { diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 8c83050bd..2c5f0bee5 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -48,13 +48,13 @@ pub use ordered_range_repartition::OrderedRangeRepartitionExec; pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec; pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates}; pub use plan_algebra::{preserves_distribution, preserves_partitioning}; -pub use range_filter::{RangeBound, RangeFilterExec, WidenedBound}; +pub use range_filter::{InputOrder, RangeBound, RangeFilterExec, WidenedBound}; pub use range_shuffle_reader::RangeShuffleReaderExec; pub use runtime_stats::{ MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, collect_reports as collect_runtime_stats_reports, cut_partitions, log_merged_runtime_stats, merge_reports as merge_runtime_stats_reports, - repartition_routing_expr, sketch_from_proto, sketch_to_proto, + repartition_routing_expr, }; pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec}; pub use shuffle_reader::{stats_for_partition, stats_for_partitions}; diff --git a/ballista/core/src/execution_plans/ordered_range_repartition.rs b/ballista/core/src/execution_plans/ordered_range_repartition.rs index 1f6b7400f..bed1eee41 100644 --- a/ballista/core/src/execution_plans/ordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/ordered_range_repartition.rs @@ -64,11 +64,9 @@ //! //! # Ordering claim //! -//! Constructor requires `input.output_ordering()` to lead with the routing -//! expression — otherwise the merger would produce garbled output. -//! [`PlanProperties::eq_properties`] declares each output partition sorted -//! on `order_by`, letting downstream operators (BWAG, HaloDrop) rely on the -//! claim without inserting a redundant `SortExec`. +//! [`PlanProperties::eq_properties`] declares each output partition sorted on +//! `order_by`, letting downstream operators (BWAG, HaloDrop) rely on the claim +//! without inserting a redundant `SortExec`. //! //! [`StreamingMerge`]: datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder @@ -76,10 +74,12 @@ use std::fmt::{self, Debug, Formatter}; use std::sync::{Arc, Mutex, OnceLock}; use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::runtime::SpawnedTask; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; +use datafusion::common::{ + Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, +}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ Distribution, EquivalenceProperties, LexOrdering, OrderingRequirements, Partitioning, @@ -104,6 +104,7 @@ use tokio_stream::wrappers::ReceiverStream; use crate::execution_plans::range_repartition_common::{ discover_cuts, guarded_scatter, split_batch_by_range, }; +use crate::sort_key::SortKeyCodec; /// Per-output-partition channel capacity, per input source. Matches the /// unordered variant's default; see the discussion there. Total buffered @@ -114,8 +115,7 @@ const CHANNEL_CAPACITY: usize = 2; /// module-level docs. pub struct OrderedRangeRepartitionExec { input: Arc, - /// Lexicographic ORDER BY. `try_new` guarantees the first entry evaluates - /// to `Float64` and matches the input's declared output ordering. + /// Lexicographic ORDER BY order_by: Vec, /// K — number of output partitions. output_partitions: usize, @@ -147,14 +147,13 @@ struct DispatchState { } impl OrderedRangeRepartitionExec { - /// Wrap `input`. `order_by` must be non-empty, the first entry must - /// evaluate to `Float64`, and `input.output_ordering()` must lead with - /// the same expression (otherwise the merger produces garbled output). + /// Creates a new `OrderedRangeRepartitionExec` + /// `order_by` must be non-empty, the first entry must evaluate to a type the sort-key codec + /// encodes, and `order_by` must be a prefix of `input.output_ordering()` pub fn try_new( input: Arc, order_by: Vec, output_partitions: usize, - // TODO: support RANGE & ROW halos ) -> Result { let [routing, ..] = order_by.as_slice() else { return internal_err!( @@ -163,52 +162,32 @@ impl OrderedRangeRepartitionExec { }; let schema = input.schema(); let routing_type = routing.expr.data_type(&schema)?; - if !matches!(routing_type, DataType::Float64) { - // TODO: support all continuous primitives + if SortKeyCodec::try_new(&routing_type, routing.options).is_none() { return internal_err!( - "OrderedRangeRepartitionExec routing expression `{}` must be Float64, got {:?}", + "OrderedRangeRepartitionExec routing expression `{}` has no sort-key encoding for {:?}", routing.expr, routing_type ); } - // TODO: fixed by KLL — a NULL-aware sketch lifts this restriction and - // lets `split_batch_by_range` honor SortOptions::nulls_first properly. - if routing.expr.nullable(&schema)? { - return internal_err!( - "OrderedRangeRepartitionExec: routing expression `{}` must be non-nullable", - routing.expr - ); - } - // Input MUST claim to be sorted on our routing expression — otherwise - // the k-way merge produces garbled output. Sortedness of individual - // input partitions is enforced by the operator upstream (`SortExec` - // with `preserve_partitioning=true`); this check verifies the plan - // node declares that property. - let input_first_sort = input.output_ordering().map(|ordering| ordering.first()); - let Some(input_first) = input_first_sort else { + let lex_ordering = LexOrdering::new(order_by.clone()).ok_or_else(|| { + internal_datafusion_err!("order_by is non-empty but LexOrdering rejected it") + })?; + let Some(input_ordering) = input.output_ordering() else { return internal_err!( "OrderedRangeRepartitionExec requires sorted input — child plan claims no ordering" ); }; - if input_first.expr.as_ref() != routing.expr.as_ref() { + if !input_ordering.starts_with(&lex_ordering) { return internal_err!( - "OrderedRangeRepartitionExec: input's first sort key `{}` does not match \ - routing expression `{}`", - input_first.expr, - routing.expr + "OrderedRangeRepartitionExec: ORDER BY [{lex_ordering}] is not a prefix of \ + the child's declared ordering [{input_ordering}]" ); } // Advertise each output partition as sorted on `order_by`. Downstream // operators (BWAG, HaloDrop) rely on this claim to skip redundant // Sort insertions. - let eq_properties = EquivalenceProperties::new_with_orderings( - schema, - vec![LexOrdering::new(order_by.clone()).ok_or_else(|| { - internal_datafusion_err!( - "order_by is non-empty but LexOrdering rejected it" - ) - })?], - ); + let eq_properties = + EquivalenceProperties::new_with_orderings(schema, vec![lex_ordering]); let properties = Arc::new( PlanProperties::new( eq_properties, @@ -429,16 +408,17 @@ impl OrderedRangeRepartitionExec { senders_per_input.push(senders); } - // Empty `Vec` = discovery failed = single-bucket fallback. + // `Ok(vec![])` = no sketch to read = single-bucket fallback. // Populated once, on the first batch, by whichever scatter task // wins the `OnceLock::get_or_init` race. - let cuts_cell: Arc>> = Arc::new(OnceLock::new()); - let routing_expr = self.order_by[0].expr.clone(); + let cuts_cell: Arc>>> = + Arc::new(OnceLock::new()); + let routing_sort = self.order_by[0].clone(); let mut drop_helper = Vec::with_capacity(input_partitions); for (input_partition, senders) in senders_per_input.into_iter().enumerate() { let child = self.input.clone(); let cuts_cell = cuts_cell.clone(); - let routing_expr = routing_expr.clone(); + let routing_sort = routing_sort.clone(); let ctx = ctx.clone(); let output_partitions = self.output_partitions; // Move senders into an `Arc<[_]>` so scatter and guard can share @@ -472,7 +452,7 @@ impl OrderedRangeRepartitionExec { child, input_partition, ctx, - routing_expr, + routing_sort, scatter_senders, cuts_cell, output_partitions, @@ -570,9 +550,9 @@ async fn scatter_input_partition( child: Arc, input_partition: usize, ctx: Arc, - routing_expr: Arc, + routing_sort: PhysicalSortExpr, senders: Arc<[mpsc::Sender>]>, - cuts_cell: Arc>>, + cuts_cell: Arc>>>, output_partitions: usize, metrics: ScatterMetrics, ) -> Result<()> { @@ -587,12 +567,16 @@ async fn scatter_input_partition( let batch = batch_result?; metrics.input_batches.add(1); metrics.input_rows.add(batch.num_rows()); - let cuts = cuts_cell.get_or_init(|| { - let discover_timer = metrics.discover_cuts_time.timer(); - let cuts = discover_cuts(&child, routing_expr.as_ref(), output_partitions); - discover_timer.done(); - cuts - }); + let cuts = cuts_cell + .get_or_init(|| { + let discover_timer = metrics.discover_cuts_time.timer(); + let cuts = + discover_cuts(&child, routing_sort.expr.as_ref(), output_partitions); + discover_timer.done(); + cuts + }) + .as_ref() + .map_err(|e| internal_datafusion_err!("OrderedRangeRepartitionExec: {e}"))?; // TODO(perf): input is sorted — this per-row `split_batch_by_range` // is legal but wasteful. Two follow-ups worth measuring: // 1. Binary-search batch head/tail against cuts to find slice @@ -601,7 +585,8 @@ async fn scatter_input_partition( // `take_arrays`-materialised sub-batches. Arc bumps replace // allocations under skew. let split_timer = metrics.split_time.timer(); - let splits = split_batch_by_range(&batch, &routing_expr, cuts)?; + let splits = + split_batch_by_range(&batch, &routing_sort.expr, cuts, routing_sort.options)?; split_timer.done(); // Stop the compute timer around the send.await so backpressure // waits get billed to `send_time` alone, not double-counted. @@ -631,6 +616,7 @@ mod tests { use super::*; use crate::execution_plans::RuntimeStatsExec; use datafusion::arrow::array::{Float64Array, Int64Array}; + use datafusion::arrow::datatypes::DataType; use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::SessionStateBuilder; @@ -781,22 +767,23 @@ mod tests { ); } + /// A nullable routing key is routable now: the run is counted beside the + /// values, sized into the cuts, and scattered to the end `nulls_first` + /// names, which is where the read-side filter looks for it. #[test] - fn try_new_rejects_nullable_routing_key() { + fn try_new_accepts_a_nullable_routing_key() { let schema = Arc::new(Schema::new(vec![ - Field::new("v2", DataType::Float64, true), // nullable + Field::new("v2", DataType::Float64, true), Field::new("id", DataType::Int64, false), ])); - let err = OrderedRangeRepartitionExec::try_new( - empty_input(&schema), - vec![asc(&schema, "v2")], - 3, - ) - .expect_err("nullable routing key must be rejected"); - assert!( - err.to_string().contains("must be non-nullable"), - "error should name the nullability constraint, got: {err}" - ); + let sort = asc(&schema, "v2"); + let ordering = LexOrdering::new(vec![sort.clone()]).unwrap(); + let sorted = Arc::new( + SortExec::new(ordering, empty_input(&schema)) + .with_preserve_partitioning(true), + ) as Arc; + OrderedRangeRepartitionExec::try_new(sorted, vec![sort], 3) + .expect("a nullable key must be routable"); } #[test] @@ -814,10 +801,7 @@ mod tests { 4, ) .expect_err("mismatched sort key must be rejected"); - assert!( - err.to_string().contains("does not match routing"), - "got: {err}" - ); + assert!(err.to_string().contains("is not a prefix of"), "got: {err}"); } // ---------- End-to-end ----------------------------------------------- diff --git a/ballista/core/src/execution_plans/range_filter.rs b/ballista/core/src/execution_plans/range_filter.rs index 7df912cae..db7145c16 100644 --- a/ballista/core/src/execution_plans/range_filter.rs +++ b/ballista/core/src/execution_plans/range_filter.rs @@ -20,7 +20,7 @@ //! `execute(k)` applies the predicate //! //! ```text -//! raw_bounds[k].0 - halo_lo <= routing_expr < raw_bounds[k].1 + halo_hi +//! raw_bounds[k].0 - halo_lo <= filter_expr < raw_bounds[k].1 + halo_hi //! ``` //! //! `None` on either bound means unbounded on that side (virtual ±∞). Zero halo @@ -47,29 +47,25 @@ //! stage 0's `RuntimeStatsExec` reports have been merged. `execute` refuses //! while bounds are unresolved; serialization refuses too — over-the-wire //! plans always ship with bounds bound. -//! -//! # Type generality -//! -//! `ScalarValue` at the API + serde surface. The internal fast path is -//! Float64-only today (matches URRE/ORRE T-Digest); widening to other -//! numeric primitives is a KLL-migration follow-up that -//! will land without breaking callers. use std::fmt::{self, Debug, Formatter}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use datafusion::arrow::array::{Array, RecordBatch}; +use datafusion::arrow::array::{Array, ArrayRef, RecordBatch}; +use datafusion::arrow::compute::SortOptions; use datafusion::arrow::compute::filter_record_batch; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::cast::{as_boolean_array, as_float64_array}; +use datafusion::common::cast::as_boolean_array; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{Result, Statistics, internal_err}; use datafusion::execution::TaskContext; use datafusion::logical_expr::Operator; -use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; -use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_expr::expressions::{BinaryExpr, IsNullExpr, Literal}; +use datafusion::physical_expr::{ + Distribution, LexOrdering, OrderingRequirements, PhysicalExpr, PhysicalSortExpr, +}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, @@ -89,9 +85,8 @@ use parking_lot::Mutex; /// means unbounded (virtual ±∞). pub type RangeBound = (Option, Option); -/// Bounds after halo widening. Float64-only internally today — see the -/// "Type generality" section in the module doc. -pub type WidenedBound = (Option, Option); +/// Bounds after halo widening +pub type WidenedBound = (Option, Option); /// Both raw and widened bounds. `raw` is preserved for serialization; the /// executor consumes `widened`. @@ -100,13 +95,38 @@ struct BoundsState { widened: Vec, } +/// The order in which rows will arrive +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InputOrder { + /// Rows arrive in no particular order and the NULL run is at this end. + /// Advertises no input ordering requirement, so an unordered + /// range-repartition upstream stays legal. + Unordered { + /// Which end of the order the NULL run occupies. + nulls_first: bool, + }, + /// Rows arrive in this order. Required of the input, so nothing planted + /// between can reorder them, and the run's end is `nulls_first`. + Ordered(SortOptions), +} + +impl InputOrder { + /// Which end the NULL run occupies. + fn nulls_first(&self) -> bool { + match self { + Self::Unordered { nulls_first } => *nulls_first, + Self::Ordered(options) => options.nulls_first, + } + } +} + /// Filter over an ordered input with a per-input-partition half-open range /// predicate, widened by the operator's halo. Range logic (cuts → per-partition /// half-open ranges → task-slice) lives scheduler-side; RFE is the runtime /// filter that applies the resolved bounds. pub struct RangeFilterExec { input: Arc, - routing_expr: Arc, + filter_expr: Arc, /// Lower halo — subtracted from each partition's `lo` at widen time. halo_lo: ScalarValue, /// Upper halo — added from each partition's `hi` at widen time. @@ -114,11 +134,10 @@ pub struct RangeFilterExec { /// Late-bound: `None` until [`RangeFilterExec::resolve_bounds`]; `execute` and serde /// refuse while unresolved. bounds: Arc>>, - /// True when `input.output_ordering()` leads with `routing_expr` in - /// ascending order. Enables the min/max fast path + binary-search slice - /// in [`RangeFilterStream`] — with a sorted input, per-batch first/last - /// values bound the entire batch, so most batches never touch - /// `filter_record_batch`. + /// The order in which rows will arrive + input_order: Option, + /// True when `input.output_ordering()` leads with `filter_expr` in ascending order. + /// Enables the min/max fast path with a sorted input sorted_on_key: bool, properties: Arc, metrics: ExecutionPlanMetricsSet, @@ -133,16 +152,18 @@ impl RangeFilterExec { /// /// * `input` - upstream operator; its partition count fixes the eventual /// `raw_bounds.len()`. - /// * `routing_expr` - numeric physical expression each row is bucketed by. + /// * `filter_expr` - numeric physical expression each row is compared by. /// * `halo_lo`, `halo_hi` - non-negative widening amounts applied by /// [`RangeFilterExec::resolve_bounds`]. Both must be finite Float64 today. + /// * `input_order` - the order in which rows will arrive pub fn try_new_pending( input: Arc, - routing_expr: Arc, + filter_expr: Arc, halo_lo: ScalarValue, halo_hi: ScalarValue, + input_order: Option, ) -> Result { - Self::try_new_inner(input, routing_expr, halo_lo, halo_hi, None) + Self::try_new_inner(input, filter_expr, halo_lo, halo_hi, input_order, None) } /// Construct with bounds already known. Used by wire decode and by @@ -151,48 +172,49 @@ impl RangeFilterExec { /// /// # Arguments /// - /// * `input`, `routing_expr`, `halo_lo`, `halo_hi` - same as + /// * `input`, `filter_expr`, `halo_lo`, `halo_hi`, `input_order` - same as /// [`Self::try_new_pending`]. /// * `raw_bounds` - one half-open cut range per input partition. Widening /// by halos happens internally; caller passes unwidened. pub fn try_new_resolved( input: Arc, - routing_expr: Arc, + filter_expr: Arc, halo_lo: ScalarValue, halo_hi: ScalarValue, + input_order: Option, raw_bounds: Vec, ) -> Result { - Self::try_new_inner(input, routing_expr, halo_lo, halo_hi, Some(raw_bounds)) + Self::try_new_inner( + input, + filter_expr, + halo_lo, + halo_hi, + input_order, + Some(raw_bounds), + ) } + #[allow(clippy::too_many_arguments)] fn try_new_inner( input: Arc, - routing_expr: Arc, + filter_expr: Arc, halo_lo: ScalarValue, halo_hi: ScalarValue, + input_order: Option, raw_bounds: Option>, ) -> Result { let schema = input.schema(); - let expr_type = routing_expr.data_type(&schema)?; - if !expr_type.is_numeric() { + let expr_type = filter_expr.data_type(&schema)?; + // TODO: as long as halos are 0, we should be able to support things like strings + if !expr_type.is_numeric() && !expr_type.is_temporal() { return internal_err!( - "RangeFilterExec: routing_expr must be numeric, got {expr_type}" - ); - } - let halo_lo_f64 = as_f64(&halo_lo)?; - let halo_hi_f64 = as_f64(&halo_hi)?; - if !halo_lo_f64.is_finite() || halo_lo_f64 < 0.0 { - return internal_err!( - "RangeFilterExec: halo_lo must be finite and non-negative, got {halo_lo_f64}" - ); - } - if !halo_hi_f64.is_finite() || halo_hi_f64 < 0.0 { - return internal_err!( - "RangeFilterExec: halo_hi must be finite and non-negative, got {halo_hi_f64}" + "RangeFilterExec: filter_expr must be numeric or temporal, got {expr_type}" ); } + validate_halo("halo_lo", &halo_lo)?; + validate_halo("halo_hi", &halo_hi)?; let bounds_state = raw_bounds - .map(|raw| build_bounds_state(&input, raw, halo_lo_f64, halo_hi_f64)) + .map(|raw| build_bounds_state(&input, raw, &halo_lo, &halo_hi)) .transpose()?; let properties = Arc::new(PlanProperties::new( input.equivalence_properties().clone(), @@ -200,19 +222,25 @@ impl RangeFilterExec { input.pipeline_behavior(), input.boundedness(), )); + if input_order.is_none() && filter_expr.nullable(&schema)? { + return internal_err!( + "RangeFilterExec: filter_expr is nullable but no input order was given, \ + so which partition holds the NULL run is unknown" + ); + } let sorted_on_key = input .output_ordering() - .map(|ord| { - let first = ord.first(); - first.expr.as_ref() == routing_expr.as_ref() && !first.options.descending - }) - .unwrap_or(false); + .map(|ord| ord.first().clone()) + .is_some_and(|first| { + first.expr.as_ref() == filter_expr.as_ref() && !first.options.descending + }); Ok(Self { input, - routing_expr, + filter_expr, halo_lo, halo_hi, bounds: Arc::new(Mutex::new(bounds_state)), + input_order, sorted_on_key, properties, metrics: ExecutionPlanMetricsSet::new(), @@ -223,9 +251,8 @@ impl RangeFilterExec { /// merge into cuts and the adapter has projected those cuts onto per-input /// partition half-open ranges. Widens by RFE's halos before caching. pub fn resolve_bounds(&self, raw_bounds: Vec) -> Result<()> { - let halo_lo = as_f64(&self.halo_lo)?; - let halo_hi = as_f64(&self.halo_hi)?; - let state = build_bounds_state(&self.input, raw_bounds, halo_lo, halo_hi)?; + let state = + build_bounds_state(&self.input, raw_bounds, &self.halo_lo, &self.halo_hi)?; self.bounds.lock().replace(state); Ok(()) } @@ -246,8 +273,13 @@ impl RangeFilterExec { } /// The physical expression whose value each row is bucketed by. - pub fn routing_expr(&self) -> &Arc { - &self.routing_expr + pub fn filter_expr(&self) -> &Arc { + &self.filter_expr + } + + /// What the cuts' producer said about the rows arriving. + pub fn input_order(&self) -> Option { + self.input_order } /// Halo-widening amount applied to each partition's lower bound. @@ -264,7 +296,7 @@ impl RangeFilterExec { impl Debug for RangeFilterExec { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("RangeFilterExec") - .field("routing_expr", &self.routing_expr.to_string()) + .field("filter_expr", &self.filter_expr.to_string()) .field("halo_lo", &self.halo_lo) .field("halo_hi", &self.halo_hi) .field( @@ -287,7 +319,7 @@ impl DisplayAs for RangeFilterExec { write!( f, "RangeFilterExec: routing={}, halo=[{}, {}], bounds={}", - self.routing_expr, self.halo_lo, self.halo_hi, bounds_str + self.filter_expr, self.halo_lo, self.halo_hi, bounds_str ) } DisplayFormatType::TreeRender => write!(f, "RangeFilterExec"), @@ -319,7 +351,7 @@ impl ExecutionPlan for RangeFilterExec { &self, f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - apply_expression_roots([&self.routing_expr], f) + apply_expression_roots([&self.filter_expr], f) } fn with_new_children( @@ -339,9 +371,10 @@ impl ExecutionPlan for RangeFilterExec { let raw_bounds = self.bounds.lock().as_ref().map(|s| s.raw.clone()); Ok(Arc::new(Self::try_new_inner( input.clone(), - self.routing_expr.clone(), + self.filter_expr.clone(), self.halo_lo.clone(), self.halo_hi.clone(), + self.input_order, raw_bounds, )?)) } @@ -350,8 +383,23 @@ impl ExecutionPlan for RangeFilterExec { vec![Distribution::UnspecifiedDistribution] } + /// Only an [`InputOrder::Ordered`] producer demands one. Requiring it is + /// what stops `EnsureRequirements` sinking an unrelated sort beneath this + /// operator, which would replace the ordering the NULL run's placement and + /// the fast path both read. An unordered producer keeps `None` so no + /// `SortExec` is planted over rows that were never meant to be sorted. fn required_input_ordering(&self) -> Vec> { - vec![None] + let requirement = match self.input_order { + Some(InputOrder::Ordered(options)) => { + LexOrdering::new(vec![PhysicalSortExpr { + expr: self.filter_expr.clone(), + options, + }]) + .map(|lex| OrderingRequirements::new(lex.into())) + } + Some(InputOrder::Unordered { .. }) | None => None, + }; + vec![requirement] } fn maintains_input_order(&self) -> Vec { @@ -386,7 +434,7 @@ impl ExecutionPlan for RangeFilterExec { "RangeFilterExec: execute() called before resolve_bounds()".into(), ) })?; - state.widened.get(partition).copied().ok_or_else(|| { + state.widened.get(partition).cloned().ok_or_else(|| { datafusion::common::DataFusionError::Internal(format!( "RangeFilterExec: partition {partition} out of bounds ({} bounds)", state.widened.len() @@ -394,11 +442,26 @@ impl ExecutionPlan for RangeFilterExec { })? }; let (lo, hi) = widened; - let predicate = build_predicate_from_bounds(self.routing_expr.clone(), lo, hi); + // if the lower or upper side of the range is unbounded, it should include the NULLs + let nulls_first = self.input_order.is_some_and(|order| order.nulls_first()); + let takes_nulls = if nulls_first { + lo.is_none() + } else { + hi.is_none() + }; + let mut predicate = + build_predicate_from_bounds(self.filter_expr.clone(), lo.clone(), hi.clone()); + if takes_nulls { + predicate = Arc::new(BinaryExpr::new( + predicate, + Operator::Or, + Arc::new(IsNullExpr::new(self.filter_expr.clone())), + )); + } let schema = self.schema(); let input = self.input.execute(partition, ctx)?; let fast_path = self.sorted_on_key.then(|| FastPathState { - routing_expr: self.routing_expr.clone(), + filter_expr: self.filter_expr.clone(), lo, hi, }); @@ -426,20 +489,49 @@ impl ExecutionPlan for RangeFilterExec { } } -/// Extract an `f64` from a `ScalarValue`. Restricted to `Float64` today -/// because the surrounding operators (URRE/ORRE, T-Digest) only understand -/// Float64. TODO widen with the KLL migration to accept any -/// `arrow::datatypes::ArrowPrimitiveType`. -fn as_f64(sv: &ScalarValue) -> Result { - match sv { - ScalarValue::Float64(Some(v)) => Ok(*v), - ScalarValue::Float64(None) => { - internal_err!("RangeFilterExec: null ScalarValue is not permitted") - } - other => internal_err!( - "RangeFilterExec: only Float64 ScalarValue supported today, got {other:?}" - ), +/// A halo must be a non-negative, non-NULL width. Float halos are also checked +/// for finiteness: an infinite one widens every bound to cover everything and a +/// NaN one compares false against all of it. +fn validate_halo(name: &str, halo: &ScalarValue) -> Result<()> { + if halo.is_null() { + return internal_err!("RangeFilterExec: {name} must not be NULL"); + } + let finite = match halo { + ScalarValue::Float64(Some(v)) => v.is_finite(), + ScalarValue::Float32(Some(v)) => v.is_finite(), + _ => true, + }; + if !finite { + return internal_err!("RangeFilterExec: {name} must be finite, got {halo}"); + } + if halo < &ScalarValue::new_zero(&halo.data_type())? { + return internal_err!("RangeFilterExec: {name} must be non-negative, got {halo}"); + } + Ok(()) +} + +pub(crate) fn is_zero_halo(halo: &ScalarValue) -> Result { + Ok(halo == &ScalarValue::new_zero(&halo.data_type())?) +} + +pub(crate) fn widen_below( + value: &ScalarValue, + halo: &ScalarValue, +) -> Result { + if is_zero_halo(halo)? { + return Ok(value.clone()); } + value.sub(halo) +} + +pub(crate) fn widen_above( + value: &ScalarValue, + halo: &ScalarValue, +) -> Result { + if is_zero_halo(halo)? { + return Ok(value.clone()); + } + value.add(halo) } /// Validate + widen raw bounds. Emits a `BoundsState` with the raw preserved @@ -447,8 +539,8 @@ fn as_f64(sv: &ScalarValue) -> Result { fn build_bounds_state( input: &Arc, raw: Vec, - halo_lo: f64, - halo_hi: f64, + halo_lo: &ScalarValue, + halo_hi: &ScalarValue, ) -> Result { let partition_count = input.output_partitioning().partition_count(); if raw.len() != partition_count { @@ -460,16 +552,20 @@ fn build_bounds_state( let widened = raw .iter() .map(|(lo, hi)| { - let lo_f = lo.as_ref().map(as_f64).transpose()?.map(|v| v - halo_lo); - let hi_f = hi.as_ref().map(as_f64).transpose()?.map(|v| v + halo_hi); - if let (Some(l), Some(h)) = (lo_f, hi_f) - && l > h - { + let lo_w = lo + .as_ref() + .map(|lo| widen_below(lo, halo_lo)) + .transpose()?; + let hi_w = hi + .as_ref() + .map(|hi| widen_above(hi, halo_hi)) + .transpose()?; + if let (Some(l), Some(h)) = (&lo_w, &hi_w) && l > h { return internal_err!( "RangeFilterExec: widened bound produced inverted [{l}, {h}) — check cuts + halo" ); } - Ok((lo_f, hi_f)) + Ok((lo_w, hi_w)) }) .collect::>>()?; Ok(BoundsState { raw, widened }) @@ -479,22 +575,23 @@ fn build_bounds_state( /// Used both by [`RangeFilterStream`]'s slow path and by the tests that inspect /// the generated expression tree. fn build_predicate_from_bounds( - routing_expr: Arc, - lo: Option, - hi: Option, + filter_expr: Arc, + lo: Option, + hi: Option, ) -> Arc { - let lit = |v: f64| -> Arc { - Arc::new(Literal::new(ScalarValue::Float64(Some(v)))) - }; - let ge = |lo: f64| -> Arc { + let ge = |lo: ScalarValue| -> Arc { Arc::new(BinaryExpr::new( - routing_expr.clone(), + filter_expr.clone(), Operator::GtEq, - lit(lo), + Arc::new(Literal::new(lo)), )) }; - let lt = |hi: f64| -> Arc { - Arc::new(BinaryExpr::new(routing_expr.clone(), Operator::Lt, lit(hi))) + let lt = |hi: ScalarValue| -> Arc { + Arc::new(BinaryExpr::new( + filter_expr.clone(), + Operator::Lt, + Arc::new(Literal::new(hi)), + )) }; match (lo, hi) { (None, None) => Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))), @@ -505,7 +602,7 @@ fn build_predicate_from_bounds( } /// Fast-path state. Present only when the input is sorted ascending on -/// `routing_expr` — then a batch's first and last routing values bound the +/// `filter_expr` — then a batch's first and last values bound the /// whole batch's value range, unlocking three shortcuts: /// /// - `last < lo` or `first >= hi` — batch is entirely outside the partition's @@ -520,9 +617,9 @@ fn build_predicate_from_bounds( /// routing column contains nulls (Float64Array binary search would treat null /// slots as garbage values). struct FastPathState { - routing_expr: Arc, - lo: Option, - hi: Option, + filter_expr: Arc, + lo: Option, + hi: Option, } /// Per-execute counters that split "which fast-path branch fired" so we can @@ -556,8 +653,7 @@ impl RangeFilterStream { .predicate .evaluate(batch) .and_then(|v| v.into_array(batch.num_rows()))?; - let mask = as_boolean_array(&mask)?; - Ok(filter_record_batch(batch, mask)?) + Ok(filter_record_batch(batch, as_boolean_array(&mask)?)?) } /// Try the sorted-input shortcuts. Returns `None` iff the batch is @@ -570,35 +666,39 @@ impl RangeFilterStream { ) -> Result> { let n = batch.num_rows(); let arr = state - .routing_expr + .filter_expr .evaluate(&batch) .and_then(|v| v.into_array(n))?; - let col = as_float64_array(&arr)?; - // Nulls in routing column: `values()` returns garbage for null slots, - // and NULL vs bound comparisons must be false. Slow path handles both. - if col.null_count() > 0 { + // TODO: include NULLs in the fast path + if arr.null_count() > 0 { let filtered = self.slow_filter(&batch)?; return Ok((filtered.num_rows() > 0).then_some(filtered)); } - let first = col.value(0); - let last = col.value(n - 1); + let first = ScalarValue::try_from_array(&arr, 0)?; + let last = ScalarValue::try_from_array(&arr, n - 1)?; // Skip: entire batch is outside the window. - if state.hi.is_some_and(|hi| first >= hi) || state.lo.is_some_and(|lo| last < lo) + if state.hi.as_ref().is_some_and(|hi| &first >= hi) + || state.lo.as_ref().is_some_and(|lo| &last < lo) { self.path_metrics.fast_skip.add(1); return Ok(None); } // Pass-through: entire batch is inside the window. - let above_lo = state.lo.is_none_or(|lo| first >= lo); - let below_hi = state.hi.is_none_or(|hi| last < hi); + let above_lo = state.lo.as_ref().is_none_or(|lo| &first >= lo); + let below_hi = state.hi.as_ref().is_none_or(|hi| &last < hi); if above_lo && below_hi { self.path_metrics.fast_pass.add(1); return Ok(Some(batch)); } // Mixed: partition the sorted column and slice. - let values = col.values(); - let start = state.lo.map_or(0, |lo| values.partition_point(|v| *v < lo)); - let end = state.hi.map_or(n, |hi| values.partition_point(|v| *v < hi)); + let start = match &state.lo { + None => 0, + Some(lo) => partition_point(&arr, n, lo)?, + }; + let end = match &state.hi { + None => n, + Some(hi) => partition_point(&arr, n, hi)?, + }; if start >= end { self.path_metrics.fast_skip.add(1); return Ok(None); @@ -608,6 +708,21 @@ impl RangeFilterStream { } } +/// Binary search - same as the rust version, but works on arrow primitives +fn partition_point(arr: &ArrayRef, len: usize, bound: &ScalarValue) -> Result { + let mut below = 0; + let mut above = len; + while below < above { + let probe = below + (above - below) / 2; + if &ScalarValue::try_from_array(arr, probe)? < bound { + below = probe + 1; + } else { + above = probe; + } + } + Ok(below) +} + impl Stream for RangeFilterStream { type Item = Result; @@ -751,24 +866,28 @@ mod tests { ) -> Arc { let ranges = ranges_from_cuts(cuts); let (lo, hi) = &ranges[partition]; - let lo_f = lo.as_ref().map(|s| match s { - ScalarValue::Float64(Some(v)) => *v - halo_lo, - _ => panic!("float64 only in tests"), - }); - let hi_f = hi.as_ref().map(|s| match s { - ScalarValue::Float64(Some(v)) => *v + halo_hi, - _ => panic!("float64 only in tests"), - }); - build_predicate_from_bounds(v_col(), lo_f, hi_f) + let widen = |bound: &Option, halo: f64| { + bound.as_ref().map(|s| match s { + ScalarValue::Float64(Some(v)) => ScalarValue::Float64(Some(*v + halo)), + other => panic!("float64 only in these shape tests, got {other:?}"), + }) + }; + build_predicate_from_bounds(v_col(), widen(lo, -halo_lo), widen(hi, halo_hi)) } #[test] fn raw_bounds_len_must_match_input_partitions() { let src = v_source(3); let bounds = ranges_from_cuts(&[10.0]); // 2 partitions - let err = - RangeFilterExec::try_new_resolved(src, v_col(), sv(0.0), sv(0.0), bounds) - .unwrap_err(); + let err = RangeFilterExec::try_new_resolved( + src, + v_col(), + sv(0.0), + sv(0.0), + None, + bounds, + ) + .unwrap_err(); assert!( err.to_string() .contains("does not match input partition count"), @@ -785,6 +904,7 @@ mod tests { v_col(), sv(-1.0), sv(0.0), + None, bounds.clone(), ) .unwrap_err(); @@ -794,38 +914,230 @@ mod tests { v_col(), sv(0.0), sv(f64::NAN), + None, bounds, ) .unwrap_err(); assert!(err.to_string().contains("halo_hi")); } + /// Bounds of any ordered type now resolve, and both paths select on them. + /// A zero halo widens by nothing, so a `Float64(0.0)` halo does not refuse + /// an `Int64` key — the scheduler passes that pair for every consumer with + /// no halo at all. + #[tokio::test] + async fn non_float64_bounds_filter_on_both_paths() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let column: Arc = + Arc::new(Column::new_with_schema("v", schema.as_ref()).unwrap()); + let rows: Vec = (0..20).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(datafusion::arrow::array::Int64Array::from(rows))], + ) + .unwrap(); + let sorted = PhysicalSortExpr::new( + column.clone(), + SortOptions { + descending: false, + nulls_first: true, + }, + ); + + for declare_sorted in [false, true] { + let source = + MemorySourceConfig::try_new(&[vec![batch.clone()]], schema.clone(), None) + .unwrap(); + let source = if declare_sorted { + source + .try_with_sort_information(vec![[sorted.clone()].into()]) + .unwrap() + } else { + source + }; + let input: Arc = + Arc::new(DataSourceExec::new(Arc::new(source))); + let rf = Arc::new( + RangeFilterExec::try_new_resolved( + input, + column.clone(), + sv(0.0), + sv(0.0), + None, + vec![( + Some(ScalarValue::Int64(Some(5))), + Some(ScalarValue::Int64(Some(12))), + )], + ) + .unwrap(), + ); + // Sorted input takes the binary-search slice, unsorted the mask. + assert_eq!(rf.sorted_on_key, declare_sorted); + + let ctx = SessionContext::new().task_ctx(); + let mut stream = rf.execute(0, ctx).unwrap(); + let mut selected: Vec = Vec::new(); + while let Some(res) = stream.next().await { + let b = res.unwrap(); + let col = b + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + selected.extend(col.values()); + } + assert_eq!( + selected, + (5..12).collect::>(), + "sorted={declare_sorted}" + ); + } + } + + /// A nullable key whose input declares no ordering on it has no answer for + /// where the run belongs, so construction refuses rather than defaulting to + /// an end and quietly handing the run to whichever partition that names. #[test] - fn non_float64_bounds_are_rejected() { - let src = v_source(2); - let bounds = vec![ - (None, Some(ScalarValue::Int64(Some(5)))), - (Some(ScalarValue::Int64(Some(5))), None), - ]; - let err = - RangeFilterExec::try_new_resolved(src, v_col(), sv(0.0), sv(0.0), bounds) - .unwrap_err(); - assert!(err.to_string().contains("only Float64"), "got: {err}"); + fn a_nullable_key_without_a_declared_order_is_refused() { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, true)])); + let column: Arc = + Arc::new(Column::new_with_schema("v", schema.as_ref()).unwrap()); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Float64Array::from(vec![Some(1.0), None]))], + ) + .unwrap(); + let source = + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(); + let input: Arc = + Arc::new(DataSourceExec::new(Arc::new(source))); + + let err = RangeFilterExec::try_new_resolved( + input, + column, + sv(0.0), + sv(0.0), + None, + vec![(None, None)], + ) + .unwrap_err(); + assert!(err.to_string().contains("NULL run"), "got: {err}"); + } + + /// The NULL run belongs to exactly one partition: the one unbounded at the + /// end the run occupies. Every other partition drops it, and no partition + /// ever compares a value against a NULL bound. + /// + /// Both placements, over K=3 partitions with cuts at 10 and 20, so the run's + /// partition is index 0 under `nulls_first` and index 2 under `nulls_last`. + #[tokio::test] + async fn the_null_run_lands_in_exactly_one_partition() { + for nulls_first in [true, false] { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, true)])); + let column: Arc = + Arc::new(Column::new_with_schema("v", schema.as_ref()).unwrap()); + // Three NULLs beside one value per partition-to-be. + let mut values: Vec> = vec![Some(5.0), Some(15.0), Some(25.0)]; + let nulls = vec![None; 3]; + if nulls_first { + values.splice(0..0, nulls); + } else { + values.extend(nulls); + } + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Float64Array::from(values))], + ) + .unwrap(); + + let bounds = vec![ + (None, Some(sv(10.0))), + (Some(sv(10.0)), Some(sv(20.0))), + (Some(sv(20.0)), None), + ]; + // The filter column has NULLs, so every batch takes the mask path; + // the declared sort only decides which partition claims the run. + let options = SortOptions { + descending: false, + nulls_first, + }; + let sorted = PhysicalSortExpr::new(column.clone(), options); + let source = MemorySourceConfig::try_new( + &[vec![batch.clone()], vec![batch.clone()], vec![batch]], + schema.clone(), + None, + ) + .unwrap() + .try_with_sort_information(vec![[sorted].into()]) + .unwrap(); + let input: Arc = + Arc::new(DataSourceExec::new(Arc::new(source))); + let rf = Arc::new( + RangeFilterExec::try_new_resolved( + input, + column.clone(), + sv(0.0), + sv(0.0), + Some(InputOrder::Ordered(options)), + bounds, + ) + .unwrap(), + ); + assert_eq!(rf.input_order, Some(InputOrder::Ordered(options))); + + let mut null_rows_per_partition = Vec::new(); + let mut value_rows_per_partition = Vec::new(); + for partition in 0..3 { + let ctx = SessionContext::new().task_ctx(); + let mut stream = rf.execute(partition, ctx).unwrap(); + let (mut nulls, mut vals) = (0usize, Vec::new()); + while let Some(res) = stream.next().await { + let b = res.unwrap(); + let col = + b.column(0).as_any().downcast_ref::().unwrap(); + for row in 0..col.len() { + if col.is_null(row) { + nulls += 1; + } else { + vals.push(col.value(row)); + } + } + } + null_rows_per_partition.push(nulls); + value_rows_per_partition.push(vals); + } + + let claimant = if nulls_first { 0 } else { 2 }; + let expected: Vec = + (0..3).map(|p| if p == claimant { 3 } else { 0 }).collect(); + assert_eq!( + null_rows_per_partition, expected, + "nulls_first={nulls_first}: the run belongs to one partition" + ); + // And the values are unaffected by the run riding along. + assert_eq!( + value_rows_per_partition, + vec![vec![5.0], vec![15.0], vec![25.0]], + "nulls_first={nulls_first}" + ); + } } #[test] fn pending_construction_defers_check() { let src = v_source(3); - let rf = - RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0)).unwrap(); + let rf = RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0), None) + .unwrap(); assert!(rf.raw_bounds().is_none()); } #[tokio::test] async fn execute_before_resolve_errors() { let src = v_source(2); - let rf = - RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0)).unwrap(); + let rf = RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0), None) + .unwrap(); let ctx = SessionContext::new().task_ctx(); let Err(err) = rf.execute(0, ctx) else { panic!("execute() should error before resolve_bounds") @@ -836,8 +1148,8 @@ mod tests { #[test] fn resolve_bounds_validates() { let src = v_source(3); - let rf = - RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0)).unwrap(); + let rf = RangeFilterExec::try_new_pending(src, v_col(), sv(0.0), sv(0.0), None) + .unwrap(); // 2 bounds don't fit 3-partition input. let bounds_too_short = ranges_from_cuts(&[1.0]); let err = rf.resolve_bounds(bounds_too_short).unwrap_err(); @@ -883,7 +1195,7 @@ mod tests { } #[test] - fn sorted_on_key_detected_when_input_ascending_on_routing_expr() { + fn sorted_on_key_detected_when_input_ascending_on_filter_expr() { // sorted_v_source is a single-partition DataSourceExec — one raw bound. let src = sorted_v_source(vec![batch(&[1.0, 2.0, 3.0])], asc()); let rf = RangeFilterExec::try_new_resolved( @@ -891,6 +1203,7 @@ mod tests { v_col(), sv(0.0), sv(0.0), + None, vec![(None, None)], ) .unwrap(); @@ -898,7 +1211,7 @@ mod tests { } #[test] - fn sorted_on_key_false_when_input_descending_on_routing_expr() { + fn sorted_on_key_false_when_input_descending_on_filter_expr() { let desc = SortOptions { descending: true, nulls_first: false, @@ -909,6 +1222,7 @@ mod tests { v_col(), sv(0.0), sv(0.0), + None, vec![(None, None)], ) .unwrap(); @@ -921,9 +1235,15 @@ mod tests { // RepartitionExec on a single partition drops ordering information. let src = v_source(2); let bounds = ranges_from_cuts(&[2.0]); - let rf = - RangeFilterExec::try_new_resolved(src, v_col(), sv(0.0), sv(0.0), bounds) - .unwrap(); + let rf = RangeFilterExec::try_new_resolved( + src, + v_col(), + sv(0.0), + sv(0.0), + None, + bounds, + ) + .unwrap(); assert!(!rf.sorted_on_key); } @@ -954,8 +1274,15 @@ mod tests { let ranges = ranges_from_cuts(cuts); let bounds = vec![ranges[global_k].clone()]; Arc::new( - RangeFilterExec::try_new_resolved(input, v_col(), sv(0.0), sv(0.0), bounds) - .unwrap(), + RangeFilterExec::try_new_resolved( + input, + v_col(), + sv(0.0), + sv(0.0), + None, + bounds, + ) + .unwrap(), ) } @@ -1030,6 +1357,7 @@ mod tests { v_expr, sv(0.0), sv(0.0), + Some(InputOrder::Ordered(asc())), vec![ranges[1].clone()], ) .unwrap(), @@ -1086,6 +1414,7 @@ mod tests { v_col(), sv(0.0), sv(0.0), + None, ranges_from_cuts(&[10.0, 20.0]), ) .unwrap(), diff --git a/ballista/core/src/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs index 7765f5ff8..243254cbe 100644 --- a/ballista/core/src/execution_plans/range_repartition_common.rs +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -41,9 +41,11 @@ use std::future::Future; use std::panic::AssertUnwindSafe; use std::sync::Arc; -use datafusion::arrow::array::{Array, Float64Array, RecordBatch, UInt32Array}; +use datafusion::arrow::array::{RecordBatch, UInt32Array}; +use datafusion::arrow::compute::SortOptions; use datafusion::arrow::compute::take_arrays; -use datafusion::common::{Result, internal_datafusion_err}; +use datafusion::arrow::row::{RowConverter, SortField}; +use datafusion::common::{Result, ScalarValue, internal_datafusion_err}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; use futures::FutureExt; @@ -54,59 +56,57 @@ use crate::execution_plans::RuntimeStatsExec; use crate::execution_plans::plan_algebra::preserves_distribution; /// Walk `child`'s subtree for a [`RuntimeStatsExec`] that sketches on our -/// routing expression, snapshot its merged T-Digest, and compute `K - 1` -/// quantile cuts. Any failure to find a matching sketch returns an empty -/// `Vec` — the caller's `split_batch_by_range(&[])` produces a single -/// bucket and every row lands in output partition 0. Never crashes. +/// routing expression, snapshot its merged sketch, and compute `K - 1` +/// quantile cuts. No sketch to read returns an empty `Vec` — the caller's +/// `split_batch_by_range(&[])` produces a single bucket and every row lands +/// in output partition 0. +/// +/// Errors only when a sketch was found and contradicted itself, which the +/// single-bucket fallback would turn into a silently one-partition query. pub(super) fn discover_cuts( child: &Arc, routing_expr: &dyn PhysicalExpr, output_partitions: usize, -) -> Vec { +) -> Result> { let Some(stats) = find_runtime_stats(child, routing_expr) else { warn!( "range-repartition: no matching RuntimeStatsExec found in child subtree — \ single-bucket fallback" ); - return Vec::new(); + return Ok(Vec::new()); }; // Walker returned Some → stats.order_by()'s first entry matches our // routing expression → RuntimeStatsExec's construction contract // guarantees sketch is present. Belt-and-braces arms in case that // invariant ever drifts, plus mutex-poisoning is theoretically possible. - let sketch = match stats.merged_quantile_sketch() { + let sketch = match stats.merged_sort_key_sketch() { Ok(Some(sketch)) => sketch, Ok(None) => { warn!( "range-repartition: matching RuntimeStatsExec has no sketch \ (RuntimeStatsExec contract broken?) — single-bucket fallback" ); - return Vec::new(); + return Ok(Vec::new()); } Err(e) => { warn!( "range-repartition: sketch snapshot failed ({e}) — single-bucket fallback" ); - return Vec::new(); + return Ok(Vec::new()); } }; - // `count()` is the sum of centroid weights — total observed row count - // that fed the digest. Zero means no samples arrived before the - // snapshot; degenerate cuts would follow. - if sketch.count() == 0.0 { + // Zero means no samples arrived before the snapshot; degenerate cuts would follow. + if sketch.count() == 0 { warn!( "range-repartition: matching sketch has no samples yet — single-bucket fallback" ); - return Vec::new(); + return Ok(Vec::new()); } - // K-1 cuts at 1/K, 2/K, ..., (K-1)/K. `estimate_quantile` is monotone by - // construction, so cuts are non-decreasing (ties possible on hot-value - // distributions — `split_batch_by_range` handles those correctly, it - // just skews the resulting distribution). - let k = output_partitions as f64; - (1..output_partitions) - .map(|i| sketch.estimate_quantile(i as f64 / k)) - .collect() + // `cuts` sizes by the whole population, keeps every boundary a real value + // so no consumer compares against a NULL, and puts the NULL run wholly in + // the partition at the end `nulls_first` names — which is where + // `split_batch_by_range` sends it and where `RangeFilterExec` looks for it. + sketch.cuts(output_partitions) } /// Walks `plan`'s subtree through single-child chains only, returning the @@ -152,22 +152,14 @@ pub(super) fn find_runtime_stats<'a>( find_runtime_stats(only_child, routing_expr) } -/// Split `batch` into `K = boundaries.len() + 1` sub-batches under the -/// half-open convention: partition `p` receives rows where -/// `boundaries[p-1] <= key < boundaries[p]` (open at `-∞` on partition 0 -/// and at `+∞` on partition K-1). Output vector is always length K; empty -/// buckets produce empty `RecordBatch`es rather than being omitted, so -/// callers can index by partition id. -/// -/// NULL routing keys land in partition 0 today; a follow-up will honor -/// `sort_options.nulls_first`/`nulls_last` to match SQL semantics. -/// -/// Boundaries are pre-extracted `f64`s; widening to other routing key -/// types generalizes this function and the discovery path together. +/// Split `batch` along boundaries under the half-open convention. Empty buckets produce empty +/// `RecordBatch`es rather than being omitted, so callers can index by partition id. +/// Boundaries must be in `options` order, which is what [`discover_cuts`] produces. pub(super) fn split_batch_by_range( batch: &RecordBatch, routing_expr: &Arc, - boundaries: &[f64], + boundaries: &[ScalarValue], + options: SortOptions, ) -> Result> { let output_partitions = boundaries.len() + 1; let schema = batch.schema(); @@ -176,40 +168,48 @@ pub(super) fn split_batch_by_range( .map(|_| RecordBatch::new_empty(schema.clone())) .collect()); } - let evaluated = routing_expr.evaluate(batch)?; - let array = evaluated.into_array(batch.num_rows())?; - let keys = array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!( - "range-repartition: routing expr produced {:?}, expected Float64", - array.data_type() - ) - })?; + let [first_boundary, ..] = boundaries else { + return Ok(vec![batch.clone()]); + }; + let keys = routing_expr.evaluate(batch)?; + let keys = keys.into_array(batch.num_rows())?; + + let converter = RowConverter::new(vec![SortField::new_with_options( + keys.data_type().clone(), + options, + )]) + .map_err(|e| { + internal_datafusion_err!( + "range-repartition: {:?} has no row encoding: {e}", + keys.data_type() + ) + })?; + let boundaries = converter.convert_columns(&[ScalarValue::iter_to_array( + boundaries.iter().cloned(), + ) + .map_err(|e| { + internal_datafusion_err!( + "range-repartition: boundaries starting {first_boundary:?} do not form \ + one array: {e}" + ) + })?])?; + let boundaries: Vec<_> = boundaries.iter().collect(); + let keys = converter.convert_columns(&[keys])?; - // TODO: vectorized arrow computation let mut buckets: Vec> = (0..output_partitions).map(|_| Vec::new()).collect(); - for row in 0..batch.num_rows() { - let target = if keys.is_null(row) { - 0 - } else { - // partition_point returns the count of elements matching the - // predicate — here `<= key` — which is the target partition - // index under the half-open convention. - boundaries.partition_point(|&cut| cut <= keys.value(row)) - }; - buckets[target].push(row as u32); + for (row_idx, key) in keys.iter().enumerate() { + let bucket_idx = boundaries.partition_point(|boundary| *boundary <= key); + buckets[bucket_idx].push(row_idx as u32); } let mut result = Vec::with_capacity(output_partitions); - for indices in buckets { - if indices.is_empty() { + for row_idxs in buckets { + if row_idxs.is_empty() { result.push(RecordBatch::new_empty(schema.clone())); } else { - let idx_array = UInt32Array::from(indices); - let taken = take_arrays(batch.columns(), &idx_array, None)?; - result.push(RecordBatch::try_new(schema.clone(), taken)?); + let row_idxs = UInt32Array::from(row_idxs); + let bucketed = take_arrays(batch.columns(), &row_idxs, None)?; + result.push(RecordBatch::try_new(schema.clone(), bucketed)?); } } Ok(result) @@ -417,6 +417,21 @@ mod tests { use datafusion::physical_expr::expressions::col; use std::sync::Arc; + /// Boundaries as `Float64` scalars, which is what discovery produces. + fn cuts(values: [f64; N]) -> Vec { + values + .into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect() + } + + fn asc(nulls_first: bool) -> SortOptions { + SortOptions { + descending: false, + nulls_first, + } + } + fn f64_col(schema: &Schema, name: &str) -> Arc { col(name, schema).unwrap() } @@ -456,7 +471,9 @@ mod tests { vec![0, 1, 2, 3, 4, 5, 6], ); let routing = f64_col(&schema, "v2"); - let splits = split_batch_by_range(&batch, &routing, &[0.0, 10.0]).unwrap(); + let splits = + split_batch_by_range(&batch, &routing, &cuts([0.0, 10.0]), asc(true)) + .unwrap(); assert_eq!(splits.len(), 3, "K = boundaries.len() + 1"); let total: usize = splits.iter().map(|b| b.num_rows()).sum(); assert_eq!(total, batch.num_rows(), "no row lost or duplicated"); @@ -472,14 +489,16 @@ mod tests { vec![0, 1, 2, 3], ); let routing = f64_col(&schema, "v2"); - let splits = split_batch_by_range(&batch, &routing, &[0.0, 10.0]).unwrap(); + let splits = + split_batch_by_range(&batch, &routing, &cuts([0.0, 10.0]), asc(true)) + .unwrap(); assert_eq!(splits[0].num_rows(), 1); assert_eq!(splits[1].num_rows(), 2); assert_eq!(splits[2].num_rows(), 1); } #[test] - fn split_routes_nulls_to_partition_zero() { + fn split_routes_the_whole_null_run_to_the_end_it_occupies() { let schema = schema_v2_id(); let batch = batch( &schema, @@ -487,9 +506,52 @@ mod tests { vec![0, 1, 2, 3], ); let routing = f64_col(&schema, "v2"); - let splits = split_batch_by_range(&batch, &routing, &[10.0]).unwrap(); + // Two NULLs, 5.0 below the cut, 50.0 above it. + let first = + split_batch_by_range(&batch, &routing, &cuts([10.0]), asc(true)).unwrap(); + assert_eq!(first[0].num_rows(), 3, "NULLs + 5.0"); + assert_eq!(first[1].num_rows(), 1, "50.0"); + + let last = + split_batch_by_range(&batch, &routing, &cuts([10.0]), asc(false)).unwrap(); + assert_eq!(last[0].num_rows(), 1, "5.0"); + assert_eq!(last[1].num_rows(), 3, "50.0 + NULLs"); + } + + /// A DESC key's boundaries arrive descending, because that is the order + /// the sketch that produced them counts in. Routing has to read them in + /// that same order or every row lands in the mirrored partition. + #[test] + fn split_follows_a_descending_key() { + let schema = schema_v2_id(); + let batch = batch( + &schema, + vec![Some(100.0), Some(10.0), Some(5.0), Some(-1.0)], + vec![0, 1, 2, 3], + ); + let routing = f64_col(&schema, "v2"); + let descending = SortOptions { + descending: true, + nulls_first: false, + }; + let splits = + split_batch_by_range(&batch, &routing, &cuts([10.0, 0.0]), descending) + .unwrap(); + assert_eq!(splits[0].num_rows(), 1, "100.0 sorts above the 10.0 cut"); + assert_eq!(splits[1].num_rows(), 2, "10.0 and 5.0 sit between the cuts"); + assert_eq!(splits[2].num_rows(), 1, "-1.0 sorts below the 0.0 cut"); + } + + /// No boundaries is the discovery fallback, and it has to stay the whole + /// batch in one bucket rather than erroring on the empty boundary set. + #[test] + fn split_without_boundaries_keeps_one_bucket() { + let schema = schema_v2_id(); + let batch = batch(&schema, vec![Some(1.0), None, Some(3.0)], vec![0, 1, 2]); + let routing = f64_col(&schema, "v2"); + let splits = split_batch_by_range(&batch, &routing, &[], asc(true)).unwrap(); + assert_eq!(splits.len(), 1); assert_eq!(splits[0].num_rows(), 3); - assert_eq!(splits[1].num_rows(), 1); } #[test] @@ -497,7 +559,9 @@ mod tests { let schema = schema_v2_id(); let batch = batch(&schema, vec![], vec![]); let routing = f64_col(&schema, "v2"); - let splits = split_batch_by_range(&batch, &routing, &[0.0, 10.0]).unwrap(); + let splits = + split_batch_by_range(&batch, &routing, &cuts([0.0, 10.0]), asc(true)) + .unwrap(); assert_eq!(splits.len(), 3); assert!(splits.iter().all(|b| b.num_rows() == 0)); } diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 6e2701626..60e40bc56 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -22,8 +22,8 @@ //! `AtomicUsize::fetch_add` on the hot path (no cross-partition //! contention, no lock overhead). //! - **Quantile sketch** — optional (only when `order_by` is set at -//! construction). Per-partition `Mutex` keeps writes off any -//! shared lock. +//! construction). A per-partition `Mutex` around each sketch keeps writes +//! off any shared lock. //! //! Timing is decoupled from correctness: both accessors are readable at //! any point. Callers get whatever has flowed through so far — a @@ -36,13 +36,8 @@ //! downstream `SortExec` / `BoundedWindowAggExec` even though only the //! first key drives the sketch today). //! -//! TODO: swap `TDigest` for a generic-over-`Ord` KLL sketch. TDigest is -//! `Float64`-only, single-column, and has no representation for NULLs; -//! a KLL implementation would sketch the full `Vec` -//! (composite keys, non-numeric types) and position NULLs per each -//! sort key's `SortOptions::nulls_first`, letting the operator drop -//! both the "first expression, `Float64` only" restriction and the -//! not-null-routing-column requirement enforced at construction today. +//! The sketch is a [`SortKeySketch`], which covers any fixed-width key and +//! gives NULLs a position per each sort key's `SortOptions::nulls_first`. //! //! This PR lands the tap in isolation: nothing wires it into a plan yet, //! and the executor doesn't yet ship the accumulated state back to the @@ -53,10 +48,12 @@ use std::fmt::{self, Debug, Formatter}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; -use datafusion::arrow::array::Float64Array; -use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; +use datafusion::common::{ + Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, +}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ Distribution, OrderingRequirements, PhysicalExpr, PhysicalSortExpr, @@ -71,19 +68,13 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, StatisticsArgs, apply_expression_roots, statistics::ChildStats, }; -use datafusion_functions_aggregate_common::tdigest::TDigest; use futures::stream::StreamExt; use log::debug; -use crate::serde::protobuf::{ - QuantileSketchState, RuntimeStatsPartitionEntry, RuntimeStatsReport, -}; +use crate::execution_plans::range_filter::{widen_above, widen_below}; +use crate::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; use crate::serde::scheduler::PartitionLocation; - -/// T-Digest centroid budget. 100 is DataFusion's default and gives ~1% -/// quantile error, plenty of margin over the sub-partition counts we -/// expect at bin-pack time. -const TDIGEST_MAX_SIZE: usize = 100; +use crate::sort_key::{SortKeyCodec, SortKeySketch}; /// Streaming runtime-stats operator. See module-level docs. pub struct RuntimeStatsExec { @@ -96,55 +87,52 @@ pub struct RuntimeStatsExec { /// `AtomicUsize::fetch_add` on the hot path — no cross-partition /// contention, no lock overhead. row_counts: Arc<[AtomicUsize]>, - /// Only allocated when `order_by` is `Some`. Sketches over the first - /// ORDER BY expression's `Float64` values. `Mutex`-per-partition to - /// keep writes off any shared lock. - sketches: Option]>>, + /// Only allocated when `order_by` is `Some`. Sketches the first ORDER BY + /// expression. `Mutex`-per-partition to keep writes off any shared lock. + sort_key_sketches: Option]>>, properties: Arc, metrics: ExecutionPlanMetricsSet, } impl RuntimeStatsExec { - /// Wrap `input`. If `order_by` is provided, its first entry drives - /// the per-partition T-Digest; the full slice is preserved for serde - /// and for downstream operators (`SortExec`, `BoundedWindowAggExec`) - /// that need it. When `Some`, at least one expression is required — - /// nothing to sketch on with an empty slice — and the first - /// expression must evaluate to a non-nullable `Float64` (T-Digest is - /// `Float64`-only and has no NULL slot; the KLL swap will lift both - /// restrictions). + /// Wrap `input`. If `order_by` is provided, its first entry drives the + /// per-partition sketch; the full slice is preserved for serde and for + /// downstream operators (`SortExec`, `BoundedWindowAggExec`) that need it. + /// When `Some`, at least one expression is required — nothing to sketch on + /// with an empty slice — and its type must be one [`SortKeyCodec`] encodes. pub fn try_new( input: Arc, order_by: Option>, ) -> Result { - if let Some(exprs) = &order_by { - let [first, ..] = exprs.as_slice() else { - return internal_err!( - "RuntimeStatsExec: order_by is Some but empty; pass None to skip sketching" - ); - }; - let schema = input.schema(); - let routing_type = first.expr.data_type(&schema)?; - if routing_type != DataType::Float64 { - return internal_err!( - "RuntimeStatsExec: routing expression must be Float64, got {routing_type:?}" - ); - } - if first.expr.nullable(&schema)? { - return internal_err!( - "RuntimeStatsExec: routing expression must be non-nullable; \ - T-Digest has no NULL slot (lifts with the KLL swap)" - ); + let codec = match &order_by { + Some(exprs) => { + let [first, ..] = exprs.as_slice() else { + return internal_err!( + "RuntimeStatsExec: order_by is Some but empty; pass None to skip sketching" + ); + }; + let schema = input.schema(); + let routing_type = first.expr.data_type(&schema)?; + // What the codec covers is the only restriction: any + // fixed-width type, nullable or not, in either direction. + let Some(codec) = SortKeyCodec::try_new(&routing_type, first.options) + else { + return internal_err!( + "RuntimeStatsExec: no sort-key encoding for {routing_type:?}" + ); + }; + Some(codec) } - } + None => None, + }; let partition_count = input.output_partitioning().partition_count(); let row_counts: Arc<[AtomicUsize]> = (0..partition_count) .map(|_| AtomicUsize::new(0)) .collect::>() .into(); - let sketches: Option]>> = order_by.as_ref().map(|_| { + let sort_key_sketches: Option]>> = codec.map(|codec| { (0..partition_count) - .map(|_| Mutex::new(TDigest::new(TDIGEST_MAX_SIZE))) + .map(|_| Mutex::new(SortKeySketch::new(codec.clone()))) .collect::>() .into() }); @@ -158,7 +146,7 @@ impl RuntimeStatsExec { input, order_by, row_counts, - sketches, + sort_key_sketches, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -202,15 +190,13 @@ impl RuntimeStatsExec { .sum() } - /// Snapshot of one partition's running quantile sketch. Returns - /// `None` when the operator was built in row-count-only mode (no - /// `order_by`). Cheap clone (a `Vec` of size - /// ≤ `TDIGEST_MAX_SIZE`). + /// Snapshot of one partition's running [`SortKeySketch`]. `None` in + /// row-count-only mode. /// /// Errors if `partition` ≥ input's partition count — callers pass a /// partition id they've already used with `execute`. - pub fn quantile_sketch(&self, partition: usize) -> Result> { - let Some(sketches) = &self.sketches else { + pub fn sort_key_sketch(&self, partition: usize) -> Result> { + let Some(sketches) = &self.sort_key_sketches else { return Ok(None); }; let slot = sketches.get(partition).ok_or_else(|| { @@ -222,33 +208,38 @@ impl RuntimeStatsExec { })?; let guard = slot.lock().map_err(|e| { internal_datafusion_err!( - "RuntimeStatsExec partition {}: sketch mutex poisoned: {e}", + "RuntimeStatsExec partition {}: sort-key sketch mutex poisoned: {e}", partition ) })?; Ok(Some(guard.clone())) } - /// All partitions merged into one sketch. `Ok(None)` in + /// All partitions merged into one [`SortKeySketch`]. `Ok(None)` in /// row-count-only mode. - pub fn merged_quantile_sketch(&self) -> Result> { - let Some(sketches) = self.sketches.as_ref() else { + pub fn merged_sort_key_sketch(&self) -> Result> { + let Some(sketches) = self.sort_key_sketches.as_ref() else { return Ok(None); }; - let snapshots: Vec = sketches - .iter() - .enumerate() - .map(|(partition, m)| { - let guard = m.lock().map_err(|e| { + let mut merged: Option = None; + for (partition, slot) in sketches.iter().enumerate() { + let snapshot = slot + .lock() + .map_err(|e| { internal_datafusion_err!( - "RuntimeStatsExec partition {}: sketch mutex poisoned: {e}", + "RuntimeStatsExec partition {}: sort-key sketch mutex poisoned: {e}", partition ) - })?; - Ok(guard.clone()) - }) - .collect::>>()?; - Ok(Some(TDigest::merge_digests(snapshots.iter()))) + })? + .clone(); + match &mut merged { + // Every slot shares the codec built once in `try_new`, so + // the codec-mismatch arm of `merge` can't fire here. + Some(accumulated) => accumulated.merge(snapshot)?, + None => merged = Some(snapshot), + } + } + Ok(merged) } } @@ -397,7 +388,7 @@ impl ExecutionPlan for RuntimeStatsExec { // state share the counters/sketches. First ORDER BY expression, // if any, drives sketching. let row_counts = self.row_counts.clone(); - let sketches = self.sketches.clone(); + let sort_key_sketches = self.sort_key_sketches.clone(); let routing_expr = self .order_by .as_ref() @@ -406,8 +397,7 @@ impl ExecutionPlan for RuntimeStatsExec { let baseline = BaselineMetrics::new(&self.metrics, partition); // Separates sketch cost from the rest so we can tell whether the - // hot path in sketching mode is the T-Digest merge_unsorted_f64 - // or something else (evaluate/downcast/flatten). + // hot path in sketching mode is the sketch or the evaluate above it. let sketch_time = MetricBuilder::new(&self.metrics).subset_time("sketch_time", partition); let sketch_batches = @@ -416,7 +406,7 @@ impl ExecutionPlan for RuntimeStatsExec { let state = StreamState { input: input_stream, row_counts, - sketches, + sort_key_sketches, routing_expr, partition, baseline, @@ -448,7 +438,7 @@ impl ExecutionPlan for RuntimeStatsExec { struct StreamState { input: SendableRecordBatchStream, row_counts: Arc<[AtomicUsize]>, - sketches: Option]>>, + sort_key_sketches: Option]>>, routing_expr: Option>, partition: usize, baseline: BaselineMetrics, @@ -476,49 +466,32 @@ impl StreamState { })?; // Sketch first: any failure returns before we count rows that - // never made it downstream. With Float64 validated at - // construction, the downcast is belt-and-braces — evaluate() - // itself can still fail for expr-internal reasons. - if let (Some(sketches), Some(routing_expr)) = (&self.sketches, &self.routing_expr) - { + // never made it downstream. Both sketches read one evaluation of + // the routing expression, so running them side by side costs a + // second ingest and not a second evaluate. + if let Some(routing_expr) = &self.routing_expr { let evaluated = routing_expr.evaluate(batch)?; let array = evaluated.into_array(batch.num_rows())?; - let f64_arr = - array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!( - "RuntimeStatsExec partition {}: routing expr produced {:?}, \ - expected Float64", - self.partition, - array.data_type() - ) - })?; - // Construction rejects nullable routing exprs, so nulls - // shouldn't reach us; keep the flatten as cheap defense - // against a Field/data mismatch. NULLs are still forwarded - // downstream and counted — they just can't enter the - // sketch until the KLL swap gives us a `nulls_first`-aware - // slot. - let values: Vec = f64_arr.iter().flatten().collect(); - if !values.is_empty() { + if !array.is_empty() { + self.sketch_batches.add(1); + } + + if let Some(sketches) = &self.sort_key_sketches { let slot = sketches.get(self.partition).ok_or_else(|| { internal_datafusion_err!( - "RuntimeStatsExec: partition {} out of range on sketch slot", + "RuntimeStatsExec: partition {} out of range on sort-key sketch slot", self.partition ) })?; let mut sketch = slot.lock().map_err(|e| { internal_datafusion_err!( - "RuntimeStatsExec partition {}: sketch mutex poisoned: {e}", + "RuntimeStatsExec partition {}: sort-key sketch mutex poisoned: {e}", self.partition ) })?; let sketch_timer = self.sketch_time.timer(); - *sketch = sketch.merge_unsorted_f64(values); + sketch.ingest(array.as_ref())?; sketch_timer.done(); - self.sketch_batches.add(1); } } @@ -544,31 +517,30 @@ impl Drop for StreamState { if rows == 0 { return; } - match self.sketches.as_ref().and_then(|s| s.get(self.partition)) { - Some(slot) => match slot.lock() { + if let Some(slot) = self + .sort_key_sketches + .as_ref() + .and_then(|sketches| sketches.get(self.partition)) + { + match slot.lock() { Ok(sketch) => { debug!( - "RuntimeStatsExec partition {}: rows={} T-Digest count={} min={} max={}", + "RuntimeStatsExec partition {}: sort-key sketch count={} nulls={} \ + min={:?} max={:?}", self.partition, - rows, sketch.count(), + sketch.null_count(), sketch.min(), sketch.max(), ); } Err(e) => { log::error!( - "RuntimeStatsExec partition {}: sketch mutex poisoned on Drop; \ - skipping end-of-stream log: {e}", + "RuntimeStatsExec partition {}: sort-key sketch mutex poisoned \ + on Drop; skipping end-of-stream log: {e}", self.partition, ); } - }, - None => { - debug!( - "RuntimeStatsExec partition {}: rows={}", - self.partition, rows - ); } } } @@ -644,65 +616,48 @@ fn stats_to_report( let mut partitions = Vec::with_capacity(partition_count); for partition_id in 0..partition_count { let row_count = stats.row_count(partition_id)? as u64; - let sketch = match stats.quantile_sketch(partition_id)? { - Some(sk) if sk.count() > 0.0 => Some(sketch_to_proto(&sk)?), - _ => None, + // The router needs each partition's value range, never its + // distribution — hence extremes here and one merged sketch below. + let (key_min, key_max, null_count) = match stats.sort_key_sketch(partition_id)? { + Some(sk) => ( + extreme_to_proto(sk.value_min()?)?, + extreme_to_proto(sk.value_max()?)?, + sk.null_count(), + ), + None => (Vec::new(), Vec::new(), 0), }; partitions.push(RuntimeStatsPartitionEntry { partition_id: partition_id as u32, row_count, - sketch, + key_min, + key_max, + null_count, }); } + let sort_key_sketch = match stats.merged_sort_key_sketch()? { + Some(sk) if sk.count() > 0 => Some(sk.to_proto()?), + _ => None, + }; Ok(RuntimeStatsReport { order_by, partitions, + sketch: sort_key_sketch, }) } -/// Serialize a T-Digest to the on-wire -/// [`QuantileSketchState`]. -/// -/// Wraps `TDigest::to_scalar_state()` — the 6-element canonical form -/// `(max_size, sum, count, max, min, centroids_as_list)` — each element -/// encoded via `datafusion_proto_common::ScalarValue::try_from`. -pub fn sketch_to_proto(sketch: &TDigest) -> Result { - let state = sketch.to_scalar_state(); - let proto_state = state - .iter() - .map(datafusion_proto_common::ScalarValue::try_from) - .collect::, _>>() - .map_err(|e| { - internal_datafusion_err!("failed to encode TDigest to proto: {e:?}") - })?; - Ok(QuantileSketchState { state: proto_state }) -} - -/// Deserialize a [`QuantileSketchState`] into a -/// T-Digest. -/// -/// Reverses [`sketch_to_proto`]. Guards against corrupted wire input by -/// checking the element count before calling -/// `TDigest::from_scalar_state`, which would panic on invalid shape. -pub fn sketch_from_proto(proto: &QuantileSketchState) -> Result { - let scalars = proto - .state - .iter() - .map(datafusion::common::ScalarValue::try_from) - .collect::, _>>() - .map_err(|e| { - internal_datafusion_err!( - "failed to decode QuantileSketchState scalars: {e:?}" - ) - })?; - if scalars.len() != 6 { - return internal_err!( - "QuantileSketchState: expected 6 elements per TDigest::to_scalar_state, got {} \ - — likely wire corruption", - scalars.len() - ); - } - Ok(TDigest::from_scalar_state(&scalars)) +/// One extreme as the wire's `repeated ScalarValue`: a tuple with one element +/// per key column, empty when no value was observed. +fn extreme_to_proto( + extreme: Option, +) -> Result> { + extreme + .map(|value| { + datafusion_proto_common::ScalarValue::try_from(&value).map_err(|e| { + internal_datafusion_err!("failed to encode key extreme {value:?}: {e:?}") + }) + }) + .transpose() + .map(|encoded| encoded.into_iter().collect()) } /// One group's merged view: sketches from every report sharing the same @@ -720,16 +675,23 @@ pub struct MergedRuntimeStats { pub task_count: usize, /// Sum of `row_count` across every partition entry in the group. pub total_rows: u64, + /// Rows in the group whose key was NULL. `0` in row-count-only mode + pub null_count: u64, /// `partition_count - 1` cut points at quantiles `i/partition_count` - /// on the merged T-Digest. Empty when `partition_count < 2` or no + /// on the merged sketch. Empty when `partition_count < 2` or no /// non-empty sketches were merged. - pub cuts: Vec, - /// Merged T-Digest's `min()` if at least one non-empty sketch - /// contributed; `None` in row-count-only mode. - pub min: Option, - /// Merged T-Digest's `max()` if at least one non-empty sketch - /// contributed; `None` in row-count-only mode. - pub max: Option, + pub cuts: Vec, + /// Merged sketch's minimum if at least one non-empty sketch contributed; + /// `None` in row-count-only mode. + pub min: Option, + /// Merged sketch's maximum if at least one non-empty sketch contributed; + /// `None` in row-count-only mode. + pub max: Option, + /// Which end of the order the NULL run occupies, from the `order_by` tag. + /// Travels with the cuts because every consumer that routes by them also + /// has to place the run, and re-deriving it invites the two to disagree. + /// Meaningless when `cuts` is empty. + pub nulls_first: bool, } /// Group `RuntimeStatsReport`s by `order_by` wire tag, merge the T-Digests @@ -792,49 +754,60 @@ fn merge_group(group: &[&RuntimeStatsReport]) -> Result { } let mut total_rows: u64 = 0; - let mut sketches: Vec = Vec::new(); + // The key's direction and NULL placement live once per report, in the + // `order_by` tag every consumer already reads to know which expression a + // sketch describes. + let options = first.order_by.first().map(|sort| SortOptions { + descending: !sort.asc, + nulls_first: sort.nulls_first, + }); + let mut merged: Option = None; for report in group { for entry in &report.partitions { total_rows = total_rows.saturating_add(entry.row_count); - if let Some(proto_sketch) = entry.sketch.as_ref() { - let sketch = sketch_from_proto(proto_sketch)?; - if sketch.count() > 0.0 { - sketches.push(sketch); - } - } + } + let Some(state) = report.sketch.as_ref() else { + continue; + }; + let Some(options) = options else { + return internal_err!( + "runtime stats merge: a report carries a sketch with an empty \ + order_by tag, so the key's ordering is unknown" + ); + }; + let sketch = SortKeySketch::try_from_proto(state, options)?; + match &mut merged { + Some(accumulated) => accumulated.merge(sketch)?, + None => merged = Some(sketch), } } - if sketches.is_empty() { + let Some(merged) = merged.filter(|sketch| sketch.count() > 0) else { return Ok(MergedRuntimeStats { order_by_len: first.order_by.len(), partition_count, task_count, total_rows, + null_count: 0, cuts: Vec::new(), min: None, max: None, + nulls_first: false, }); - } - - let merged_sketch = TDigest::merge_digests(sketches.iter()); - let cuts: Vec = if partition_count > 1 { - (1..partition_count) - .map(|cut_index| { - merged_sketch.estimate_quantile(cut_index as f64 / partition_count as f64) - }) - .collect() - } else { - Vec::new() }; + + let cuts = merged.cuts(partition_count)?; + let (min, max) = (merged.min()?, merged.max()?); Ok(MergedRuntimeStats { order_by_len: first.order_by.len(), partition_count, task_count, total_rows, + null_count: merged.null_count(), cuts, - min: Some(merged_sketch.min()), - max: Some(merged_sketch.max()), + min, + max, + nulls_first: options.is_some_and(|options| options.nulls_first), }) } @@ -943,24 +916,28 @@ pub fn repartition_routing_expr( pub fn cut_partitions( original_partitions: Vec>, reports: &[TaskRuntimeStats], - global_cuts: &[f64], - halo_lo: f64, - halo_hi: f64, + global_cuts: &[ScalarValue], + halo_lo: &ScalarValue, + halo_hi: &ScalarValue, + nulls_first: bool, ) -> Result>> { use std::collections::HashMap; // Index sketches by (producer_task_id, sub_part_id). Under // ShuffleWriter(Passthrough) file_id == task_id, so PartitionLocation's // (file_id, partition_id.partition_id) is the same pair. - let sketches: HashMap<(usize, u32), &QuantileSketchState> = reports - .iter() - .flat_map(|stats| { - stats.report.partitions.iter().filter_map(move |entry| { - let sketch = entry.sketch.as_ref()?; - Some(((stats.producer_task_id, entry.partition_id), sketch)) + // Only each file's value range is needed, never its distribution, which + // is why the sketch beside these is merged once per report rather than + // repeated per partition. + let ranges: HashMap<(usize, u32), &RuntimeStatsPartitionEntry> = + reports + .iter() + .flat_map(|stats| { + stats.report.partitions.iter().map(move |entry| { + ((stats.producer_task_id, entry.partition_id), entry) + }) }) - }) - .collect(); + .collect(); debug_assert!( global_cuts.windows(2).all(|w| w[0] <= w[1]), @@ -978,18 +955,27 @@ pub fn cut_partitions( ); }; let sub_part_id = file.partition_id.partition_id as u32; - // Fold "no sketch" and "empty sketch" into one Option — both - // mean "no routing info for this file" - let sketch = sketches - .get(&(task_id as usize, sub_part_id)) - .map(|proto| sketch_from_proto(proto)) - .transpose()? - .filter(|s| s.count() > 0.0); - let Some(sketch) = sketch else { + // A file of NULLs has a null count but no value range, and it + // belongs wholly to the partition holding the run — which the + // `nulls_first` end names, not an overlap check. + let entry = ranges.get(&(task_id as usize, sub_part_id)); + let null_only = entry + .is_some_and(|entry| entry.null_count > 0 && entry.key_min.is_empty()); + if null_only { + let run = if nulls_first { 0 } else { partition_count - 1 }; + remapped[run].push(file); + continue; + } + let range = entry.and_then(|entry| { + let lo = entry.key_min.first()?; + let hi = entry.key_max.first()?; + Some((lo, hi)) + }); + let Some((min_proto, max_proto)) = range else { // No routing info. Safe to skip only if the file has zero rows if file.partition_stats.num_rows != Some(0) { return internal_err!( - "range-repartition remap: file has num_rows={:?} but no usable sketch (task_id={task_id}, sub_part_id={sub_part_id})", + "range-repartition remap: file has num_rows={:?} but no usable key range (task_id={task_id}, sub_part_id={sub_part_id})", file.partition_stats.num_rows ); } @@ -1001,9 +987,22 @@ pub fn cut_partitions( // Monotone cuts → the set of matching buckets is a contiguous // range [b_lo, b_hi], found by two partition_points over // `global_cuts` with the sketch shifted by the halos. - let (sketch_min, sketch_max) = (sketch.min(), sketch.max()); - let b_lo = global_cuts.partition_point(|&c| c <= sketch_min - halo_hi); - let b_hi = global_cuts.partition_point(|&c| c <= sketch_max + halo_lo); + let sketch_min = ScalarValue::try_from(min_proto).map_err(|e| { + internal_datafusion_err!( + "range-repartition remap: undecodable key_min: {e:?}" + ) + })?; + let sketch_max = ScalarValue::try_from(max_proto).map_err(|e| { + internal_datafusion_err!( + "range-repartition remap: undecodable key_max: {e:?}" + ) + })?; + // Typed widening with the same zero shortcut `RangeFilterExec` + // uses, so a zero halo of one type cannot refuse a key of another. + let reach_lo = widen_below(&sketch_min, halo_hi)?; + let reach_hi = widen_above(&sketch_max, halo_lo)?; + let b_lo = global_cuts.partition_point(|cut| cut <= &reach_lo); + let b_hi = global_cuts.partition_point(|cut| cut <= &reach_hi); for bucket in &mut remapped[b_lo..=b_hi] { bucket.push(file.clone()); } @@ -1063,67 +1062,6 @@ pub fn log_merged_runtime_stats( } } -#[cfg(test)] -mod wire_tests { - use super::*; - - /// Round-trip a populated T-Digest through the wire. Count / min / - /// max should survive unchanged; quantile queries should agree to - /// within floating-point equality since serde is lossless. - #[test] - fn tdigest_wire_roundtrip_preserves_populated_sketch() { - let mut original = TDigest::new(100); - original = original - .merge_unsorted_f64(vec![1.0, 5.0, 10.0, 20.0, 30.0, 50.0, 75.0, 100.0]); - let proto = sketch_to_proto(&original).unwrap(); - let decoded = sketch_from_proto(&proto).unwrap(); - assert_eq!(decoded.count(), original.count()); - assert_eq!(decoded.min(), original.min()); - assert_eq!(decoded.max(), original.max()); - // Quantile agreement at the median. - assert_eq!( - decoded.estimate_quantile(0.5), - original.estimate_quantile(0.5), - ); - } - - /// Empty T-Digest — zero centroids, no samples — still survives - /// the round-trip. This is the case an executor hits when a - /// `RuntimeStatsExec` was present in the plan but no batches - /// flowed through (e.g. empty input partition). - #[test] - fn tdigest_wire_roundtrip_preserves_empty_sketch() { - let original = TDigest::new(100); - let proto = sketch_to_proto(&original).unwrap(); - let decoded = sketch_from_proto(&proto).unwrap(); - assert_eq!(decoded.count(), 0.0); - assert_eq!(decoded.max_size(), original.max_size()); - } - - /// Corrupted wire input (wrong element count) is caught before - /// `TDigest::from_scalar_state` gets a chance to panic. - #[test] - fn sketch_from_proto_rejects_wrong_shape() { - use datafusion::common::ScalarValue; - let proto = QuantileSketchState { - state: (0..3) - .map(|_| { - datafusion_proto_common::ScalarValue::try_from(&ScalarValue::Float64( - Some(0.0), - )) - .unwrap() - }) - .collect(), - }; - let err = sketch_from_proto(&proto) - .expect_err("wrong-count wire input must be rejected before decode"); - assert!( - err.to_string().contains("expected 6 elements"), - "got: {err}" - ); - } -} - #[cfg(test)] mod stream_tests { //! End-to-end: build a small in-memory input, wrap it in @@ -1135,6 +1073,7 @@ mod stream_tests { use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; + use datafusion::common::ScalarValue; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::physical_expr::PhysicalSortExpr; @@ -1200,16 +1139,74 @@ mod stream_tests { assert_eq!(stats.row_count(0).unwrap(), 5); assert_eq!(stats.total_row_count(), 5); - // Sketch observed every non-null routing value. - let sketch = stats.quantile_sketch(0).unwrap().unwrap(); - assert_eq!(sketch.count(), 5.0); - assert_eq!(sketch.min(), 1.0); - assert_eq!(sketch.max(), 5.0); - // Merged over the (single) partition matches the per-partition view. + // Sketch observed every routing value. + let sort_key = stats.sort_key_sketch(0).unwrap().unwrap(); + assert_eq!(sort_key.count(), 5); + assert_eq!(sort_key.null_count(), 0); + assert_eq!( + sort_key.min().unwrap(), + Some(ScalarValue::Float64(Some(1.0))) + ); assert_eq!( - stats.merged_quantile_sketch().unwrap().unwrap().count(), - 5.0 + sort_key.max().unwrap(), + Some(ScalarValue::Float64(Some(5.0))) ); + assert_eq!(stats.merged_sort_key_sketch().unwrap().unwrap().count(), 5); + } + + /// Volume through the operator, not just a handful of rows: past KLL's + /// k=800 level-0 capacity so compaction actually runs, and ascending to + /// match the post-`SortExec` position the operator occupies in a + /// range-repartition plan. + #[tokio::test] + async fn the_sketch_sees_every_row_through_the_operator() { + let schema = schema_v_id(); + const ROWS: i64 = 5_000; + let batches: Vec = (0..5) + .map(|chunk| { + let lo = chunk * (ROWS / 5); + let hi = lo + (ROWS / 5); + batch( + &schema, + (lo..hi).map(|v| Some(v as f64)).collect(), + (lo..hi).collect(), + ) + }) + .collect(); + + let memory = Arc::new( + MemorySourceConfig::try_new(&[batches], schema.clone(), None).unwrap(), + ); + let input: Arc = Arc::new(DataSourceExec::new(memory)); + let sort_expr = PhysicalSortExpr { + expr: col("v", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }; + let stats = + Arc::new(RuntimeStatsExec::try_new(input, Some(vec![sort_expr])).unwrap()); + + let ctx = SessionContext::new().task_ctx(); + let stream = stats.execute(0, ctx).unwrap(); + common::collect(stream).await.unwrap(); + + let sketch = stats.merged_sort_key_sketch().unwrap().unwrap(); + assert_eq!(sketch.count(), ROWS as u64); + assert_eq!(sketch.null_count(), 0); + // The extremes are tracked outside the compactor, so compaction + // cannot have moved them however many times it ran. + assert_eq!(sketch.min().unwrap(), Some(ScalarValue::Float64(Some(0.0)))); + assert_eq!( + sketch.max().unwrap(), + Some(ScalarValue::Float64(Some((ROWS - 1) as f64))) + ); + // K-1 real, non-decreasing boundaries over what it saw. + let cuts = sketch.cuts(8).unwrap(); + assert_eq!(cuts.len(), 7); + assert!(cuts.iter().all(|cut| !cut.is_null())); + assert!(cuts.windows(2).all(|pair| pair[0] <= pair[1])); } /// Row-count-only mode (`order_by = None`): the operator still @@ -1234,10 +1231,10 @@ mod stream_tests { assert_eq!(output.iter().map(|b| b.num_rows()).sum::(), 3); assert_eq!(stats.row_count(0).unwrap(), 3); assert!( - stats.quantile_sketch(0).unwrap().is_none(), + stats.sort_key_sketch(0).unwrap().is_none(), "row-count-only mode must not allocate a sketch" ); - assert!(stats.merged_quantile_sketch().unwrap().is_none()); + assert!(stats.merged_sort_key_sketch().unwrap().is_none()); } } @@ -1325,11 +1322,28 @@ mod collect_tests { }; assert_eq!(entry.partition_id, 0); assert_eq!(entry.row_count, 5); - let proto_sketch = entry.sketch.as_ref().expect("sketch present in wire"); - let round_tripped = sketch_from_proto(proto_sketch).unwrap(); - assert_eq!(round_tripped.count(), 5.0); - assert_eq!(round_tripped.min(), 1.0); - assert_eq!(round_tripped.max(), 5.0); + // The report's merged sketch is what routing reads; the entry + // carries only the key range a file router needs. + let state = report.sketch.as_ref().expect("sketch present in wire"); + let round_tripped = SortKeySketch::try_from_proto( + state, + SortOptions { + descending: false, + nulls_first: true, + }, + ) + .unwrap(); + assert_eq!(round_tripped.count(), 5); + assert_eq!( + round_tripped.min().unwrap(), + Some(ScalarValue::Float64(Some(1.0))) + ); + assert_eq!( + round_tripped.max().unwrap(), + Some(ScalarValue::Float64(Some(5.0))) + ); + assert_eq!(entry.key_min.len(), 1, "entry carries its own range"); + assert_eq!(entry.key_max.len(), 1); } /// Row-count-only mode: report emitted, but its partition entry @@ -1351,10 +1365,8 @@ mod collect_tests { }; assert!(report.order_by.is_empty()); assert_eq!(report.partitions.len(), 1); - assert!( - report.partitions[0].sketch.is_none(), - "no sketch in row-count-only mode" - ); + assert!(report.sketch.is_none(), "no sketch in row-count-only mode"); + assert!(report.partitions[0].key_min.is_empty()); assert_eq!(report.partitions[0].row_count, 5); } @@ -1418,6 +1430,9 @@ mod merge_tests { //! sharing an `order_by` tag, verify the merged view (total rows, //! cuts, min/max) reflects the union of the underlying samples. + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::datatypes::DataType; + use super::*; use crate::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; use datafusion_proto::protobuf::PhysicalSortExprNode; @@ -1426,31 +1441,82 @@ mod merge_tests { /// from that slot's `values`. Slot `slot_id` in the resulting /// report has `row_count = values[slot_id].len()` and a sketch /// over those values. + /// The wire tag matching [`sketch_options`]. A report carrying a sketch + /// must have one, since the tag is where the key's ordering lives. + fn sketch_tag() -> Vec { + vec![PhysicalSortExprNode { + expr: None, + asc: true, + nulls_first: true, + }] + } + + /// The ordering a fixture sketch is built under. Must match the tag the + /// merge reads, since that is where direction and NULL placement live. + fn sketch_options() -> SortOptions { + SortOptions { + descending: false, + nulls_first: true, + } + } + fn sketching_report( order_by: Vec, values_per_slot: Vec>, ) -> RuntimeStatsReport { + // One merged sketch per report, as the executor builds it, plus the + // per-partition row counts. + let codec = SortKeyCodec::try_new(&DataType::Float64, sketch_options()).unwrap(); + let mut merged = SortKeySketch::new(codec); let partitions = values_per_slot .into_iter() .enumerate() .map(|(slot_id, slot_values)| { let row_count = slot_values.len() as u64; - let sketch = if slot_values.is_empty() { - None - } else { - let digest = TDigest::new(100).merge_unsorted_f64(slot_values); - Some(sketch_to_proto(&digest).unwrap()) - }; + merged + .ingest(&Float64Array::from(slot_values)) + .expect("Float64 samples into a Float64 sketch"); RuntimeStatsPartitionEntry { partition_id: slot_id as u32, row_count, - sketch, + ..Default::default() } }) .collect(); + let sketch = (merged.count() > 0).then(|| merged.to_proto().unwrap()); RuntimeStatsReport { order_by, partitions, + sketch, + } + } + + /// A report whose key is NULL in every row. The sketch retains nothing, + /// so its NULL count is all that survives the round-trip. + fn all_null_report( + order_by: Vec, + nulls_per_slot: Vec, + ) -> RuntimeStatsReport { + let codec = SortKeyCodec::try_new(&DataType::Float64, sketch_options()).unwrap(); + let mut merged = SortKeySketch::new(codec); + let partitions = nulls_per_slot + .into_iter() + .enumerate() + .map(|(slot_id, nulls)| { + merged + .ingest(&Float64Array::from(vec![None::; nulls])) + .expect("NULL Float64 samples into a Float64 sketch"); + RuntimeStatsPartitionEntry { + partition_id: slot_id as u32, + row_count: nulls as u64, + ..Default::default() + } + }) + .collect(); + RuntimeStatsReport { + order_by, + partitions, + sketch: Some(merged.to_proto().unwrap()), } } @@ -1462,6 +1528,14 @@ mod merge_tests { } } + /// A cut as the number the band assertions are written in terms of. + fn as_f64(cut: &ScalarValue) -> f64 { + match cut { + ScalarValue::Float64(Some(v)) => *v, + other => panic!("expected a Float64 cut, got {other:?}"), + } + } + /// Two reports over disjoint value ranges — merged sketch spans the /// union, total_rows sums, and the partition_count=2 midpoint cut /// falls between the two ranges. @@ -1469,23 +1543,24 @@ mod merge_tests { fn merge_reports_combines_disjoint_ranges() { // Both reports share an empty `order_by` — we just need two // reports that land in the same group. - let low_range = sketching_report(vec![], vec![vec![1.0, 2.0, 3.0], vec![]]); - let high_range = sketching_report(vec![], vec![vec![], vec![10.0, 11.0, 12.0]]); + let low_range = sketching_report(sketch_tag(), vec![vec![1.0, 2.0, 3.0], vec![]]); + let high_range = + sketching_report(sketch_tag(), vec![vec![], vec![10.0, 11.0, 12.0]]); let group = only_group(&[low_range, high_range]); assert_eq!(group.partition_count, 2); assert_eq!(group.task_count, 2); assert_eq!(group.total_rows, 6); let midpoint = match group.cuts.as_slice() { - [midpoint] => *midpoint, + [midpoint] => as_f64(midpoint), other => panic!("expected exactly one cut, got {other:?}"), }; assert!( (3.0..=10.0).contains(&midpoint), "midpoint cut should land between ranges (got {midpoint})" ); - assert_eq!(group.min, Some(1.0)); - assert_eq!(group.max, Some(12.0)); + assert_eq!(group.min.as_ref().map(as_f64), Some(1.0)); + assert_eq!(group.max.as_ref().map(as_f64), Some(12.0)); } /// partition_count=4 cuts on a uniform [0, 100) sample land roughly @@ -1502,12 +1577,12 @@ mod merge_tests { uniform[50..75].to_vec(), uniform[75..100].to_vec(), ]; - let report = sketching_report(vec![], values_per_slot); + let report = sketching_report(sketch_tag(), values_per_slot); let group = only_group(&[report]); assert_eq!(group.partition_count, 4); let (p25, p50, p75) = match group.cuts.as_slice() { - [p25, p50, p75] => (*p25, *p50, *p75), + [p25, p50, p75] => (as_f64(p25), as_f64(p50), as_f64(p75)), other => panic!("expected 3 cuts, got {other:?}"), }; // Loose bounds — T-Digest quantile estimates aren't exact, but @@ -1527,14 +1602,15 @@ mod merge_tests { RuntimeStatsPartitionEntry { partition_id: 0, row_count: row_counts[0], - sketch: None, + ..Default::default() }, RuntimeStatsPartitionEntry { partition_id: 1, row_count: row_counts[1], - sketch: None, + ..Default::default() }, ], + ..Default::default() }; let group = only_group(&[make_report([100, 200]), make_report([300, 400])]); assert_eq!(group.partition_count, 2); @@ -1544,13 +1620,31 @@ mod merge_tests { assert!(group.max.is_none()); } + /// A key that is NULL in every row has no value to cut on, so `cuts` + /// comes back empty — the same answer a report carrying no sketch at all + /// gives. The two have to stay distinguishable: one is a degenerate + /// distribution whose every row belongs in a single partition, the other + /// is a broken invariant a range-repartition stage must refuse to route + /// on. + #[test] + fn merge_reports_distinguishes_an_all_null_key_from_a_missing_sketch() { + let group = only_group(&[ + all_null_report(sketch_tag(), vec![3, 1]), + all_null_report(sketch_tag(), vec![2, 4]), + ]); + + assert_eq!(group.total_rows, 10); + assert!(group.cuts.is_empty(), "no value to cut on"); + assert_eq!(group.null_count, group.total_rows, "every key was NULL"); + } + /// Mismatched partition counts within a group surface as an error — /// the caller (scheduler / slice-D consumer) sees the invariant /// break rather than silently getting a partial merge. #[test] fn merge_reports_errors_on_mismatched_partition_counts() { - let two_partitions = sketching_report(vec![], vec![vec![1.0], vec![2.0]]); - let one_partition = sketching_report(vec![], vec![vec![3.0]]); + let two_partitions = sketching_report(sketch_tag(), vec![vec![1.0], vec![2.0]]); + let one_partition = sketching_report(sketch_tag(), vec![vec![3.0]]); let err = merge_reports(&[two_partitions, one_partition]) .expect_err("mismatched partition counts must error"); let message = err.to_string(); @@ -1560,38 +1654,30 @@ mod merge_tests { ); } - /// A wire-corrupted sketch — one whose scalar-state length is wrong - /// — surfaces the underlying `sketch_from_proto` error rather than - /// getting silently dropped. + /// A sketch the decoder cannot read surfaces as an error rather than + /// getting silently dropped, which would size partitions from a + /// population missing whatever that report observed. #[test] fn merge_reports_propagates_sketch_decode_errors() { - use datafusion::common::ScalarValue; - - // Six scalars is the valid shape; three is a corrupted wire. - let corrupt_sketch = QuantileSketchState { - state: (0..3) - .map(|_| { - datafusion_proto_common::ScalarValue::try_from(&ScalarValue::Float64( - Some(0.0), - )) - .unwrap() - }) - .collect(), + let corrupt = crate::serde::protobuf::SortKeySketchState { + k: 800, + null_count: 0, + key_min: vec![], + key_max: vec![], + levels: b"not an arrow stream".to_vec(), }; let report = RuntimeStatsReport { - order_by: vec![], + order_by: sketch_tag(), partitions: vec![RuntimeStatsPartitionEntry { partition_id: 0, row_count: 1, - sketch: Some(corrupt_sketch), + ..Default::default() }], + sketch: Some(corrupt), }; let err = merge_reports(&[report]) - .expect_err("corrupt sketch must surface as an error"); - assert!( - err.to_string().contains("expected 6 elements"), - "expected shape-error propagation, got: {err}" - ); + .expect_err("an undecodable sketch must surface as an error"); + assert!(!err.to_string().is_empty(), "got: {err}"); } /// Empty input → empty output; ensures no panics or spurious groups. @@ -1767,26 +1853,53 @@ mod overlap_remap_tests { } } - /// Build a report whose sub-parts each carry a T-Digest sketch over - /// `values`. Slot `sub_part_id` gets a sketch of `values[sub_part_id]`. + /// A zero halo of the routing key's type, which is what the scheduler + /// passes for a consumer with no halo at all. + const ZERO_HALO: ScalarValue = ScalarValue::Float64(Some(0.0)); + + /// Cuts as `Float64` scalars, which is what merging produces. + fn cuts_f64(values: [f64; N]) -> Vec { + values + .into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect() + } + + /// Build a report whose sub-parts carry the key range routing reads. + /// Slot `sub_part_id` covers `values[sub_part_id]`. fn sketch_report( producer_task_id: usize, values_per_sub_part: Vec>, ) -> TaskRuntimeStats { + let scalar = |v: f64| { + vec![ + datafusion_proto_common::ScalarValue::try_from(&ScalarValue::Float64( + Some(v), + )) + .unwrap(), + ] + }; let partitions = values_per_sub_part .into_iter() .enumerate() .map(|(sub_part_id, samples)| { - let sketch = if samples.is_empty() { - None - } else { - let digest = TDigest::new(100).merge_unsorted_f64(samples.clone()); - Some(sketch_to_proto(&digest).unwrap()) + let extremes = samples.iter().copied().fold( + None::<(f64, f64)>, + |acc, v| match acc { + None => Some((v, v)), + Some((lo, hi)) => Some((lo.min(v), hi.max(v))), + }, + ); + let (key_min, key_max) = match extremes { + Some((lo, hi)) => (scalar(lo), scalar(hi)), + None => (Vec::new(), Vec::new()), }; RuntimeStatsPartitionEntry { partition_id: sub_part_id as u32, row_count: samples.len() as u64, - sketch, + key_min, + key_max, + ..Default::default() } }) .collect(); @@ -1795,6 +1908,7 @@ mod overlap_remap_tests { report: RuntimeStatsReport { order_by: vec![], partitions, + ..Default::default() }, } } @@ -1809,12 +1923,19 @@ mod overlap_remap_tests { sketch_report(200, vec![vec![20.0, 25.0, 29.0]]), ]; // Cut at 15 → partition 0 = (-∞, 15), partition 1 = [15, +∞). - let cuts = vec![15.0]; + let cuts = cuts_f64([15.0]); // Passthrough map: both producers wrote to sub_part_id=0. let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; - let remapped = - cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .unwrap(); assert_eq!(remapped.len(), 2, "K = cuts.len() + 1"); // Partition 0: only producer 100. assert_eq!(remapped[0].len(), 1); @@ -1831,11 +1952,18 @@ mod overlap_remap_tests { fn overlap_remap_straddling_producer_appears_in_both_partitions() { // Producer 300 covers [5, 25) — straddles the cut at 15. let reports = vec![sketch_report(300, vec![vec![5.0, 15.0, 25.0]])]; - let cuts = vec![15.0]; + let cuts = cuts_f64([15.0]); let original_partitions = vec![vec![location(0, 300)]]; - let remapped = - cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .unwrap(); assert_eq!(remapped.len(), 2); assert_eq!(remapped[0].len(), 1, "straddler in partition 0"); assert_eq!(remapped[0][0].file_id, Some(300)); @@ -1849,14 +1977,21 @@ mod overlap_remap_tests { #[test] fn overlap_remap_missing_file_id_errors() { let reports = vec![sketch_report(100, vec![vec![1.0, 2.0, 3.0]])]; - let cuts = vec![10.0]; + let cuts = cuts_f64([10.0]); // Loc has file_id=None — invalid for URRE/ORRE stages. let mut bad = location(0, 100); bad.file_id = None; let original_partitions = vec![vec![bad]]; - let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) - .expect_err("missing file_id must surface as an error"); + let err = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .expect_err("missing file_id must surface as an error"); assert!( err.to_string().contains("missing file_id"), "unexpected error: {err}" @@ -1874,11 +2009,18 @@ mod overlap_remap_tests { // Report from producer 100, but original_partitions only has // producer 200 — no file to route. let reports = vec![sketch_report(100, vec![vec![1.0, 2.0, 3.0]])]; - let cuts = vec![10.0]; + let cuts = cuts_f64([10.0]); let original_partitions = vec![vec![location(0, 200)]]; - let remapped = - cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1889,11 +2031,18 @@ mod overlap_remap_tests { #[test] fn overlap_remap_empty_sketches_produce_empty_partitions() { let reports = vec![sketch_report(100, vec![vec![]])]; - let cuts = vec![10.0]; + let cuts = cuts_f64([10.0]); let original_partitions = vec![vec![location(0, 100)]]; - let remapped = - cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .unwrap(); assert_eq!(remapped.len(), 2); assert!(remapped[0].is_empty()); assert!(remapped[1].is_empty()); @@ -1904,17 +2053,24 @@ mod overlap_remap_tests { #[test] fn missing_sketch_with_rows_errors() { let reports = vec![sketch_report(100, vec![vec![1.0, 2.0]])]; - let cuts = vec![10.0]; + let cuts = cuts_f64([10.0]); // File 200 has 5 rows but no report entry exists for it. let mut orphan = location(0, 200); orphan.partition_stats = PartitionStats::new(Some(5), None, None); let original_partitions = vec![vec![orphan]]; - let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) - .expect_err("file with rows but no sketch must error"); + let err = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .expect_err("file with rows but no sketch must error"); let msg = err.to_string(); assert!( - msg.contains("num_rows=Some(5)") && msg.contains("no usable sketch"), + msg.contains("num_rows=Some(5)") && msg.contains("no usable key range"), "unexpected error: {msg}" ); } @@ -1940,7 +2096,7 @@ mod overlap_remap_tests { // F: spans the whole range → every bucket sketch_report(6, vec![vec![0.0, 20.0, 100.0]]), ]; - let cuts = vec![10.0, 20.0, 30.0]; + let cuts = cuts_f64([10.0, 20.0, 30.0]); let original_partitions = vec![vec![ location(0, 1), location(0, 2), @@ -1950,8 +2106,15 @@ mod overlap_remap_tests { location(0, 6), ]]; - let remapped = - cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0).unwrap(); + let remapped = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .unwrap(); assert_eq!(remapped.len(), 4); let ids = |b: &[PartitionLocation]| { let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); @@ -1970,16 +2133,23 @@ mod overlap_remap_tests { #[test] fn missing_sketch_with_unknown_rows_errors() { let reports: Vec = vec![]; - let cuts = vec![10.0]; + let cuts = cuts_f64([10.0]); let mut orphan = location(0, 100); orphan.partition_stats = PartitionStats::default(); // num_rows = None let original_partitions = vec![vec![orphan]]; - let err = cut_partitions(original_partitions, &reports, &cuts, 0.0, 0.0) - .expect_err("file with unknown rows but no sketch must error"); + let err = cut_partitions( + original_partitions, + &reports, + &cuts, + &ZERO_HALO, + &ZERO_HALO, + true, + ) + .expect_err("file with unknown rows but no sketch must error"); let msg = err.to_string(); assert!( - msg.contains("num_rows=None") && msg.contains("no usable sketch"), + msg.contains("num_rows=None") && msg.contains("no usable key range"), "unexpected error: {msg}" ); } @@ -1997,9 +2167,9 @@ mod overlap_remap_tests { #[test] fn overlap_remap_halo_band_widens_both_sides_without_bleeding_to_far_partitions() { // K=5, asymmetric halos so we can tell halo_lo and halo_hi apart. - let cuts = vec![10.0, 20.0, 30.0, 40.0]; - let halo_lo = 1.0; - let halo_hi = 2.0; + let cuts = cuts_f64([10.0, 20.0, 30.0, 40.0]); + let halo_lo = ScalarValue::Float64(Some(1.0)); + let halo_hi = ScalarValue::Float64(Some(2.0)); // Effective partition ranges: // P0: (-∞, 12) P1: [9, 22) P2: [19, 32) P3: [29, 42) P4: [39, +∞) let reports = vec![ @@ -2026,9 +2196,15 @@ mod overlap_remap_tests { location(0, 500), ]]; - let remapped = - cut_partitions(original_partitions, &reports, &cuts, halo_lo, halo_hi) - .unwrap(); + let remapped = cut_partitions( + original_partitions, + &reports, + &cuts, + &halo_lo, + &halo_hi, + true, + ) + .unwrap(); let ids = |b: &[PartitionLocation]| { let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); v.sort(); diff --git a/ballista/core/src/execution_plans/unordered_range_repartition.rs b/ballista/core/src/execution_plans/unordered_range_repartition.rs index cdf331875..e86d9db9c 100644 --- a/ballista/core/src/execution_plans/unordered_range_repartition.rs +++ b/ballista/core/src/execution_plans/unordered_range_repartition.rs @@ -54,27 +54,18 @@ //! degraded-but-alive beats the alternative, and downstream sees an empty //! stream on the K-1 partitions that got no data. //! -//! # Type generality -//! -//! The impl hardcodes Float64 downcast internally (that's what DataFusion's -//! T-Digest speaks). The public API and the sibling [`RuntimeStatsExec`] -//! stay type-agnostic; widening to other `Ord` `ScalarValue` types replaces -//! the downcast + boundary computation, no API break. -//! -//! Sibling `OrderedRangeRepartitionExec` (not yet built) handles the sorted -//! case (N sorted → M sorted range-disjoint via k-way merge). See -//! `docs/source/contributors-guide/parallel-window-kll-adaptive.md`. -//! -//! [`RuntimeStatsExec`]: crate::execution_plans::RuntimeStatsExec +//! [`RuntimeStatsExec`]: super::RuntimeStatsExec use std::fmt::{self, Debug, Formatter}; use std::sync::{Arc, Mutex, OnceLock}; use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::runtime::SpawnedTask; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; +use datafusion::common::{ + Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, +}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{ Distribution, EquivalenceProperties, OrderingRequirements, Partitioning, @@ -95,6 +86,7 @@ use tokio_stream::wrappers::ReceiverStream; use crate::execution_plans::range_repartition_common::{ discover_cuts, guarded_scatter, split_batch_by_range, }; +use crate::sort_key::SortKeyCodec; /// Per-output-partition channel capacity. Small = tight backpressure; the /// classic double-buffering shape (one batch in-flight while consumer works @@ -109,7 +101,7 @@ pub struct UnorderedRangeRepartitionExec { input: Arc, /// Lexicographic ORDER BY carried through from the wrapping window /// operator. `try_new` guarantees at least one element; the first entry - /// (a `Float64` column, until we widen) drives routing. + /// drives routing. order_by: Vec, /// K — number of output partitions. K=1 collapses all P inputs to a /// single bucket (the same shape discovery-failure fallback produces); @@ -140,7 +132,7 @@ struct DispatchState { impl UnorderedRangeRepartitionExec { /// Wrap `input`. `order_by` must be non-empty and its first entry must - /// evaluate to `Float64` against `input.schema()`. `output_partitions` + /// evaluate to a type the sort-key codec encodes. `output_partitions` /// is K; any value works, K=1 gives a coalesce-shaped passthrough. pub fn try_new( input: Arc, @@ -155,22 +147,13 @@ impl UnorderedRangeRepartitionExec { }; let schema = input.schema(); let routing_type = routing.expr.data_type(&schema)?; - if !matches!(routing_type, DataType::Float64) { - // TODO: support all continuous primitives + if SortKeyCodec::try_new(&routing_type, routing.options).is_none() { return internal_err!( - "UnorderedRangeRepartitionExec routing expression `{}` must be Float64, got {:?}", + "UnorderedRangeRepartitionExec routing expression `{}` has no sort-key encoding for {:?}", routing.expr, routing_type ); } - // TODO: fixed by KLL — a NULL-aware sketch lifts this restriction and - // lets `split_batch_by_range` honor SortOptions::nulls_first properly. - if routing.expr.nullable(&schema)? { - return internal_err!( - "UnorderedRangeRepartitionExec: routing expression `{}` must be non-nullable", - routing.expr - ); - } let properties = Arc::new( PlanProperties::new( EquivalenceProperties::new(schema), @@ -335,16 +318,17 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { state.receivers = receivers; state.initialized = true; let senders: Arc<[mpsc::Sender>]> = senders.into(); - let cuts_cell: Arc>> = Arc::new(OnceLock::new()); + let cuts_cell: Arc>>> = + Arc::new(OnceLock::new()); let input_partitions = self.input.output_partitioning().partition_count(); - let routing_expr = self.order_by[0].expr.clone(); // TODO: KLL for multi-column? + let routing_sort = self.order_by[0].clone(); let mut drop_helper = Vec::with_capacity(input_partitions); for input_partition in 0..input_partitions { let child = self.input.clone(); let scatter_senders = senders.clone(); let guard_senders = senders.clone(); let cuts_cell = cuts_cell.clone(); - let routing_expr = routing_expr.clone(); + let routing_sort = routing_sort.clone(); let ctx = ctx.clone(); let output_partitions = self.output_partitions; // `guarded_scatter` wraps the body in `catch_unwind`: a @@ -356,7 +340,7 @@ impl ExecutionPlan for UnorderedRangeRepartitionExec { child, input_partition, ctx, - routing_expr, + routing_sort, scatter_senders, cuts_cell, output_partitions, @@ -404,9 +388,9 @@ async fn scatter_input_partition( child: Arc, input_partition: usize, ctx: Arc, - routing_expr: Arc, + routing_sort: PhysicalSortExpr, senders: Arc<[mpsc::Sender>]>, - cuts_cell: Arc>>, + cuts_cell: Arc>>>, output_partitions: usize, ) -> Result<()> { let mut stream = child.execute(input_partition, ctx)?; @@ -419,9 +403,14 @@ async fn scatter_input_partition( return Ok(()); } let batch = batch_result?; - let cuts = cuts_cell.get_or_init(|| { - discover_cuts(&child, routing_expr.as_ref(), output_partitions) - }); + let cuts = cuts_cell + .get_or_init(|| { + discover_cuts(&child, routing_sort.expr.as_ref(), output_partitions) + }) + .as_ref() + .map_err(|e| { + internal_datafusion_err!("UnorderedRangeRepartitionExec: {e}") + })?; // TODO(perf): `split_batch_by_range` materialises K sub-batches per // input batch via `take_arrays` — one copy per row into a fresh // allocation. Unlike the ordered variant we can't slice @@ -429,7 +418,8 @@ async fn scatter_input_partition( // `Arc` broadcast + receiver-side filter would skip // the scatter-side allocations at the cost of duplicating the // filter work K times. Worth measuring under skew. - let splits = split_batch_by_range(&batch, &routing_expr, cuts)?; + let splits = + split_batch_by_range(&batch, &routing_sort.expr, cuts, routing_sort.options)?; for (output, sub) in splits.into_iter().enumerate() { if sub.num_rows() == 0 { continue; @@ -451,6 +441,7 @@ mod tests { use super::*; use crate::execution_plans::RuntimeStatsExec; use datafusion::arrow::array::{Float64Array, Int64Array}; + use datafusion::arrow::datatypes::DataType; use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::SessionStateBuilder; @@ -501,36 +492,35 @@ mod tests { ); } + /// An `Int64` key and a nullable key are both routable now. The sketch + /// encodes any fixed-width type, and a NULL run is counted beside the + /// values, sized into the cuts, and scattered to one end. #[test] - fn try_new_rejects_non_float64_routing_key() { - let schema = schema_v2_id(); - let err = UnorderedRangeRepartitionExec::try_new( - empty_input(&schema), - vec![asc(&schema, "id")], // Int64 - 3, - ) - .expect_err("Int64 routing key must be rejected"); - assert!( - err.to_string().contains("must be Float64"), - "error should name the type mismatch, got: {err}" - ); - } - - #[test] - fn try_new_rejects_nullable_routing_key() { + fn try_new_accepts_widened_routing_keys() { let schema = Arc::new(Schema::new(vec![ - Field::new("v2", DataType::Float64, true), // nullable + Field::new("v2", DataType::Float64, true), Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), ])); + for key in ["v2", "id"] { + UnorderedRangeRepartitionExec::try_new( + empty_input(&schema), + vec![asc(&schema, key)], + 3, + ) + .unwrap_or_else(|e| panic!("{key} must be routable: {e}")); + } + + // A variable-width key has no fixed-width encoding, and says so. let err = UnorderedRangeRepartitionExec::try_new( empty_input(&schema), - vec![asc(&schema, "v2")], + vec![asc(&schema, "name")], 3, ) - .expect_err("nullable routing key must be rejected"); + .expect_err("Utf8 routing key has no sort-key encoding"); assert!( - err.to_string().contains("must be non-nullable"), - "error should name the nullability constraint, got: {err}" + err.to_string().contains("no sort-key encoding"), + "error should name what is missing, got: {err}" ); } diff --git a/ballista/core/src/kll.rs b/ballista/core/src/kll.rs index cc0b9b710..4658a32ed 100644 --- a/ballista/core/src/kll.rs +++ b/ballista/core/src/kll.rs @@ -261,6 +261,62 @@ impl KllSketch { } } + /// Nominal top-level compactor capacity + pub fn k(&self) -> usize { + self.k + } + + /// The compactor stack, level 0 first. `levels()[h]` holds the items + /// retained at height `h`, each standing for `2^h` of the stream, so the + /// stack plus `k` is the whole sketch apart from its extremes. + /// + /// Order within a level is not guaranteed. A caller that needs one + /// sorts, which is what [`Self::from_parts`] assumes on the way back in. + pub fn levels(&self) -> &[Vec] { + &self.levels + } + + /// Rebuild from what [`Self::levels`], [`Self::k`], [`Self::min`] and + /// [`Self::max`] expose. Levels are sorted here, so a caller that + /// serialized them in any order still gets a sketch whose per-level + /// ordering invariant holds. + /// + /// `None` when `levels` describes a stack this sketch could not have + /// produced — no levels at all, or a level holding more than its + /// capacity — rather than building something whose next compaction would + /// misbehave. + /// + /// The PRNG is reseeded from OS entropy, matching [`Clone`]: compaction + /// decisions after a rebuild are deliberately uncorrelated with the ones + /// the original made. + pub fn from_parts( + k: usize, + mut levels: Vec>, + min: Option, + max: Option, + ) -> Option { + if levels.is_empty() || k < MIN_LEVEL_WIDTH { + return None; + } + let num_levels = levels.len(); + for (height, level) in levels.iter().enumerate() { + if level.len() > level_capacity(k, num_levels, height) { + return None; + } + } + for level in &mut levels { + level.sort_unstable(); + } + Some(Self { + levels, + sorted: vec![true; num_levels], + k, + rng: StdRng::seed_from_u64(rand::random::()), + min, + max, + }) + } + /// Consume `other` and fold its content into `self`. /// /// Same-height compactors concatenate: items promoted to level `h` in diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs index f7ff5ed42..41bab4a43 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 = 1; +pub const BALLISTA_PROTOCOL_VERSION: u32 = 2; /// 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 9e05f2ed2..acc22bd55 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -107,13 +107,52 @@ pub struct RuntimeStatsExecNode { ::datafusion_proto::protobuf::PhysicalSortExprNode, >, } -/// Serialized T-Digest as a fixed-layout `Vec` per -/// `TDigest::to_scalar_state()`: max_size, sum, count, max, min, -/// centroid_means..., centroid_weights.... +/// A `SortKeySketch` is a wrapper around a KLL, that also includes NULL counts +/// +/// `levels` is an Arrow IPC stream for now. Binary blob encoding was considered +/// so that encoding tricks could be used like delta+varint. These showed a 4x +/// improvement for u64, but only 25% on f64 due to mantissas being effectively +/// random. Ultimately it was decided that a more complex encoder is not worth +/// the added complexity in this PR. We can revisit the decision in the future. #[derive(Clone, PartialEq, ::prost::Message)] -pub struct QuantileSketchState { - #[prost(message, repeated, tag = "1")] - pub state: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, +pub struct SortKeySketchState { + /// KLL's nominal top-level compactor capacity + #[prost(uint32, tag = "1")] + pub k: u32, + /// Rows whose key was NULL. Not in `levels`, and not recoverable from it. + #[prost(uint64, tag = "2")] + pub null_count: u64, + /// The exact minimum and maximum over every value observed, one element per + /// ORDER BY expression. Tracked outside the compactor stack so no coin flip + /// can move them, which is what keeps `quantile(0.0)` and `quantile(1.0)` + /// exact. Empty when no value was observed. + #[prost(message, repeated, tag = "3")] + pub key_min: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + #[prost(message, repeated, tag = "4")] + pub key_max: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + /// The compactor stack: one row per level in level order, single column + /// `levels: List>` holding that level's retained keys ascending, + /// one struct field per ORDER BY expression. An item's weight is + /// `2^level`, so the row index carries the weight and the total count is + /// the weighted sum — neither ships. + /// + /// Levels are sorted before serialization, so a decoder takes ascending + /// order as given rather than being told per level. + /// + /// A three-level stack over an `Int64` key decodes to one batch of one + /// column, where the row index is the level: + /// + /// levels: List\> + /// + /// +-------------------------------------------+ + /// \| levels | + /// +-------------------------------------------+ + /// \| \[{expr_0: 4}, {expr_0: 17}, {expr_0: 23}\] | level 0, weight 1 + /// \| \[{expr_0: 9}, {expr_0: 31}\] | level 1, weight 2 + /// \| \[{expr_0: 12}\] | level 2, weight 4 + /// +-------------------------------------------+ + #[prost(bytes = "vec", tag = "5")] + pub levels: ::prost::alloc::vec::Vec, } /// Flow-control operator with an operator-mode enum. The child plan is /// plumbed by the framework as `inputs\[0\]` during decode. @@ -163,7 +202,7 @@ pub struct PerPartitionFilterExecNode { #[derive(Clone, PartialEq, ::prost::Message)] pub struct RangeFilterExecNode { #[prost(message, optional, tag = "1")] - pub routing_expr: ::core::option::Option< + pub filter_expr: ::core::option::Option< ::datafusion_proto::protobuf::PhysicalExprNode, >, #[prost(message, optional, tag = "2")] @@ -172,6 +211,44 @@ pub struct RangeFilterExecNode { pub halo_hi: ::core::option::Option<::datafusion_proto_common::ScalarValue>, #[prost(message, repeated, tag = "4")] pub raw_bounds: ::prost::alloc::vec::Vec, + /// What the operator that produced the cuts says about the rows arriving. + /// Unset claims nothing, which a nullable `filter_expr` has no answer for: + /// a NULL has a side in the order, not a position among the values. + /// + /// `ordered` is also required of the input, so nothing planted between can + /// reorder the rows the placement was stated against. `unordered_nulls_first` + /// states the side without demanding an order, which is what leaves an + /// unordered range-repartition upstream legal. + #[prost(oneof = "range_filter_exec_node::InputOrder", tags = "5, 6")] + pub input_order: ::core::option::Option, +} +/// Nested message and enum types in `RangeFilterExecNode`. +pub mod range_filter_exec_node { + /// What the operator that produced the cuts says about the rows arriving. + /// Unset claims nothing, which a nullable `filter_expr` has no answer for: + /// a NULL has a side in the order, not a position among the values. + /// + /// `ordered` is also required of the input, so nothing planted between can + /// reorder the rows the placement was stated against. `unordered_nulls_first` + /// states the side without demanding an order, which is what leaves an + /// unordered range-repartition upstream legal. + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum InputOrder { + #[prost(bool, tag = "5")] + UnorderedNullsFirst(bool), + #[prost(message, tag = "6")] + Ordered(super::SortOptions), + } +} +/// An arrow `SortOptions`, which has no proto of its own in the DataFusion +/// descriptors — `PhysicalSortExprNode` carries the pair inline beside an +/// expression this message's users already have. +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct SortOptions { + #[prost(bool, tag = "1")] + pub descending: bool, + #[prost(bool, tag = "2")] + pub nulls_first: bool, } /// Half-open `[lo, hi)` cut range for one input partition. Either side may be /// unset to signal ±∞. @@ -981,6 +1058,11 @@ pub struct RuntimeStatsReport { /// scheduler groups by `order_by` tag and aggregates. #[prost(message, repeated, tag = "2")] pub partitions: ::prost::alloc::vec::Vec, + /// Every partition's observations folded into one sketch, merged on the + /// executor before the report is sent. Absent in row-count-only mode, + /// and when no partition observed a value. + #[prost(message, optional, tag = "3")] + pub sketch: ::core::option::Option, } /// One partition's observations from a `RuntimeStatsExec`. #[derive(Clone, PartialEq, ::prost::Message)] @@ -989,14 +1071,23 @@ pub struct RuntimeStatsPartitionEntry { pub partition_id: u32, #[prost(uint64, tag = "2")] pub row_count: u64, - /// Present when the `RuntimeStatsExec` was in sketch mode AND this - /// partition observed at least one non-null routing value. + /// This partition's exact key range, one element per ORDER BY expression. + /// `cut_partitions` routes a whole shuffle file into every downstream + /// partition whose range overlaps `\[key_min, key_max\]`, so these are the + /// only per-partition facts a router needs — never the distribution, which + /// is why the sketch beside them is merged rather than repeated here. /// - /// TODO: `optional MinMaxState min_max` — for a lighter post-repartition - /// mode where the bin-packer just needs (min, max, count) per - /// sub-partition and a full T-Digest is overkill. - #[prost(message, optional, tag = "3")] - pub sketch: ::core::option::Option, + /// Both empty when this partition observed no value, which `null_count` + /// then distinguishes: some NULLs means a file of NULLs with no value range + /// to overlap, none means the partition saw nothing at all. + #[prost(message, repeated, tag = "4")] + pub key_min: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + #[prost(message, repeated, tag = "5")] + pub key_max: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, + /// Rows in this partition whose key was NULL. Folds to the report-level + /// `SortKeySketchState.null_count`. + #[prost(uint64, tag = "6")] + pub null_count: u64, } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct ExecutionError {} diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 56d2b8eb8..db28fb077 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -56,17 +56,18 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ - BufferExec, BufferMode, ChaosExec, CoalescePlan, OrderedRangeRepartitionExec, - PartitionGroup, PartitionedBoundedWindowAggExec, PerPartitionFilterExec, - RangeFilterExec, RangeShuffleReaderExec, RuntimeStatsExec, ShuffleReaderExec, - ShuffleWriterExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, - UnresolvedShuffleExec, + BufferExec, BufferMode, ChaosExec, CoalescePlan, InputOrder, + OrderedRangeRepartitionExec, PartitionGroup, PartitionedBoundedWindowAggExec, + PerPartitionFilterExec, RangeFilterExec, RangeShuffleReaderExec, RuntimeStatsExec, + ShuffleReaderExec, ShuffleWriterExec, SortShuffleWriterExec, + UnorderedRangeRepartitionExec, UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, ballista_physical_plan_node::PhysicalPlanType, }; use crate::serde::scheduler::PartitionLocation; +use datafusion::arrow::compute::SortOptions; pub use generated::ballista as protobuf; /// Generated protobuf code from Ballista protocol definitions. @@ -715,18 +716,31 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { ))); }; let schema = input.schema(); - let routing_expr_proto = node.routing_expr.as_ref().ok_or_else(|| { + let filter_expr_proto = node.filter_expr.as_ref().ok_or_else(|| { DataFusionError::Internal( - "RangeFilterExecNode missing routing_expr".into(), + "RangeFilterExecNode missing filter_expr".into(), ) })?; - let routing_expr = + let filter_expr = datafusion_proto::physical_plan::from_proto::parse_physical_expr( - routing_expr_proto, + filter_expr_proto, ctx, schema.as_ref(), self, )?; + let input_order = node.input_order.as_ref().map(|order| match order { + protobuf::range_filter_exec_node::InputOrder::UnorderedNullsFirst( + nulls_first, + ) => InputOrder::Unordered { + nulls_first: *nulls_first, + }, + protobuf::range_filter_exec_node::InputOrder::Ordered(options) => { + InputOrder::Ordered(SortOptions { + descending: options.descending, + nulls_first: options.nulls_first, + }) + } + }); let sv_from_proto = |p: &datafusion_proto_common::ScalarValue| { datafusion::scalar::ScalarValue::try_from(p).map_err(|e| { DataFusionError::Internal(format!( @@ -757,9 +771,10 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { .collect::, DataFusionError>>()?; Ok(Arc::new(RangeFilterExec::try_new_resolved( input.clone(), - routing_expr, + filter_expr, halo_lo, halo_hi, + input_order, raw_bounds, )?)) } @@ -1096,9 +1111,9 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { "RangeFilterExec: cannot serialize before resolve_bounds()".into(), ) })?; - let routing_expr = + let filter_expr = datafusion_proto::physical_plan::to_proto::serialize_physical_expr( - exec.routing_expr(), + exec.filter_expr(), self.default_codec.as_ref(), )?; let encode_sv = |sv: &datafusion::scalar::ScalarValue| { @@ -1121,10 +1136,23 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { let proto = protobuf::BallistaPhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::RangeFilter( protobuf::RangeFilterExecNode { - routing_expr: Some(routing_expr), + filter_expr: Some(filter_expr), halo_lo: Some(halo_lo), halo_hi: Some(halo_hi), raw_bounds: raw_bounds_proto, + input_order: exec.input_order().map(|order| match order { + InputOrder::Unordered { nulls_first } => { + protobuf::range_filter_exec_node::InputOrder::UnorderedNullsFirst(nulls_first) + } + InputOrder::Ordered(options) => { + protobuf::range_filter_exec_node::InputOrder::Ordered( + protobuf::SortOptions { + descending: options.descending, + nulls_first: options.nulls_first, + }, + ) + } + }), }, )), }; @@ -2017,14 +2045,14 @@ mod test { // EmptyExec exposes one partition, so partition-0 is the only slot. assert_eq!(original.row_count(0).unwrap(), 0); assert_eq!(original.total_row_count(), 0); - assert_eq!(original.quantile_sketch(0).unwrap().unwrap().count(), 0.0); + assert_eq!(original.sort_key_sketch(0).unwrap().unwrap().count(), 0); assert_eq!( - original.merged_quantile_sketch().unwrap().unwrap().count(), - 0.0 + original.merged_sort_key_sketch().unwrap().unwrap().count(), + 0 ); // Out-of-range partition surfaces as an internal error, not a panic. assert!(original.row_count(1).is_err()); - assert!(original.quantile_sketch(1).is_err()); + assert!(original.sort_key_sketch(1).is_err()); let codec = BallistaPhysicalExtensionCodec::default(); let mut buf: Vec = vec![]; @@ -2051,10 +2079,10 @@ mod test { assert_eq!(order_by[0].expr.to_string(), sort_expr.expr.to_string()); assert!(!order_by[0].options.descending); assert_eq!(decoded.row_count(0).unwrap(), 0); - assert_eq!(decoded.quantile_sketch(0).unwrap().unwrap().count(), 0.0); + assert_eq!(decoded.sort_key_sketch(0).unwrap().unwrap().count(), 0); assert_eq!( - decoded.merged_quantile_sketch().unwrap().unwrap().count(), - 0.0 + decoded.merged_sort_key_sketch().unwrap().unwrap().count(), + 0 ); } @@ -2075,10 +2103,10 @@ mod test { assert!(original.order_by().is_none()); assert_eq!(original.row_count(0).unwrap(), 0); assert!( - original.quantile_sketch(0).unwrap().is_none(), + original.sort_key_sketch(0).unwrap().is_none(), "no sketch was requested at construction" ); - assert!(original.merged_quantile_sketch().unwrap().is_none()); + assert!(original.merged_sort_key_sketch().unwrap().is_none()); let codec = BallistaPhysicalExtensionCodec::default(); let mut buf: Vec = vec![]; @@ -2098,8 +2126,8 @@ mod test { .downcast_ref::() .expect("Expected RuntimeStatsExec"); assert!(decoded.order_by().is_none()); - assert!(decoded.quantile_sketch(0).unwrap().is_none()); - assert!(decoded.merged_quantile_sketch().unwrap().is_none()); + assert!(decoded.sort_key_sketch(0).unwrap().is_none()); + assert!(decoded.merged_sort_key_sketch().unwrap().is_none()); } /// `try_new` refuses an empty ORDER BY — no routing key means the @@ -2124,10 +2152,12 @@ mod test { } /// `try_new` refuses a routing expression whose evaluated type is - /// not `Float64` — TDigest can't ingest anything else, so failing - /// at construction beats a downcast error mid-stream. + /// not encodable — the sketch needs a fixed-width key, so failing + /// An `Int64` key and a nullable key are both sketchable now: the codec + /// covers every fixed-width type and NULLs are counted beside the values + /// rather than having nowhere to go. #[test] - fn test_runtime_stats_exec_rejects_non_float64_routing_expr() { + fn test_runtime_stats_exec_accepts_widened_routing_exprs() { use crate::execution_plans::RuntimeStatsExec; use datafusion::arrow::compute::SortOptions; use datafusion::physical_expr::PhysicalSortExpr; @@ -2135,53 +2165,29 @@ mod test { use datafusion::physical_plan::expressions::col; let schema = Arc::new(Schema::new(vec![ - Field::new("v", DataType::Float64, true), + Field::new("nullable_float", DataType::Float64, true), Field::new("id", DataType::Int64, false), + Field::new("when", DataType::Utf8, false), ])); let input: Arc = Arc::new(EmptyExec::new(schema.clone())); - let sort_expr = PhysicalSortExpr { - expr: col("id", schema.as_ref()).unwrap(), + let sort_on = |name: &str| PhysicalSortExpr { + expr: col(name, schema.as_ref()).unwrap(), options: SortOptions { descending: false, nulls_first: true, }, }; - let err = RuntimeStatsExec::try_new(input, Some(vec![sort_expr])) - .expect_err("non-Float64 routing expr must be rejected"); - assert!( - err.to_string() - .contains("routing expression must be Float64"), - "got: {err}" - ); - } - /// `try_new` refuses a nullable routing expression — TDigest has no - /// NULL slot, so allowing nulls would silently exclude them from - /// the sketch while `row_count` still saw them. The KLL swap lifts - /// this by positioning nulls per `SortOptions::nulls_first`. - #[test] - fn test_runtime_stats_exec_rejects_nullable_routing_expr() { - use crate::execution_plans::RuntimeStatsExec; - use datafusion::arrow::compute::SortOptions; - use datafusion::physical_expr::PhysicalSortExpr; - use datafusion::physical_plan::empty::EmptyExec; - use datafusion::physical_plan::expressions::col; + for name in ["nullable_float", "id"] { + RuntimeStatsExec::try_new(input.clone(), Some(vec![sort_on(name)])) + .unwrap_or_else(|e| panic!("{name} must be sketchable: {e}")); + } - let schema = - Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, true)])); - let input: Arc = Arc::new(EmptyExec::new(schema.clone())); - let sort_expr = PhysicalSortExpr { - expr: col("v", schema.as_ref()).unwrap(), - options: SortOptions { - descending: false, - nulls_first: true, - }, - }; - let err = RuntimeStatsExec::try_new(input, Some(vec![sort_expr])) - .expect_err("nullable routing expr must be rejected"); + // What the codec does not cover is still refused, and says so. + let err = RuntimeStatsExec::try_new(input, Some(vec![sort_on("when")])) + .expect_err("a variable-width key has no fixed-width encoding"); assert!( - err.to_string() - .contains("routing expression must be non-nullable"), + err.to_string().contains("no sort-key encoding"), "got: {err}" ); } @@ -2269,7 +2275,7 @@ mod test { RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(3)).unwrap(), ); use datafusion::scalar::ScalarValue; - let routing_expr: Arc = Arc::new(Column::new("v", 0)); + let filter_expr: Arc = Arc::new(Column::new("v", 0)); // K=3 raw bounds: (-∞, 10), [10, 20), [20, +∞) let raw_bounds: Vec<(Option, Option)> = vec![ (None, Some(ScalarValue::Float64(Some(10.0)))), @@ -2279,11 +2285,18 @@ mod test { ), (Some(ScalarValue::Float64(Some(20.0))), None), ]; + // Round-tripped alongside the bounds: the placement of the NULL run is + // not recoverable from the child on the far side. + let original_input_order = Some(InputOrder::Ordered(SortOptions { + descending: false, + nulls_first: true, + })); let original = RangeFilterExec::try_new_resolved( input.clone(), - routing_expr.clone(), + filter_expr.clone(), ScalarValue::Float64(Some(3.0)), ScalarValue::Float64(Some(0.0)), + original_input_order, raw_bounds.clone(), ) .unwrap(); @@ -2308,7 +2321,8 @@ mod test { assert_eq!(decoded.raw_bounds().unwrap(), raw_bounds); assert_eq!(decoded.halo_lo(), &ScalarValue::Float64(Some(3.0))); assert_eq!(decoded.halo_hi(), &ScalarValue::Float64(Some(0.0))); - assert_eq!(decoded.routing_expr().to_string(), routing_expr.to_string()); + assert_eq!(decoded.filter_expr().to_string(), filter_expr.to_string()); + assert_eq!(decoded.input_order(), original_input_order); } /// A pending RangeFilterExec (unresolved bounds) refuses to serialize — @@ -2329,12 +2343,13 @@ mod test { RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(2)).unwrap(), ); use datafusion::scalar::ScalarValue; - let routing_expr: Arc = Arc::new(Column::new("v", 0)); + let filter_expr: Arc = Arc::new(Column::new("v", 0)); let pending = RangeFilterExec::try_new_pending( input, - routing_expr, + filter_expr, ScalarValue::Float64(Some(0.0)), ScalarValue::Float64(Some(0.0)), + None, ) .unwrap(); diff --git a/ballista/core/src/sort_key.rs b/ballista/core/src/sort_key.rs index 3b85fea00..f50f228b1 100644 --- a/ballista/core/src/sort_key.rs +++ b/ballista/core/src/sort_key.rs @@ -122,7 +122,13 @@ //! order, and variable-width types have no fixed encoding — leaving those //! to the arrow-row path. -use datafusion::arrow::array::{Array, ArrowPrimitiveType, AsArray, PrimitiveArray}; +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, AsArray, ListArray, PrimitiveArray, RecordBatch, + StructArray, new_empty_array, +}; +use datafusion::arrow::buffer::OffsetBuffer; use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{ DataType, Date32Type, Date64Type, DurationMicrosecondType, DurationMillisecondType, @@ -132,9 +138,25 @@ use datafusion::arrow::datatypes::{ TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; -use datafusion::common::{Result, ScalarValue, internal_datafusion_err}; +use datafusion::arrow::datatypes::{Field, Fields, Schema}; +use datafusion::arrow::ipc::reader::StreamReader; +use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::common::{Result, ScalarValue, internal_datafusion_err, internal_err}; use crate::kll::KllSketch; +use crate::serde::protobuf::SortKeySketchState; + +/// Field name for the one list-of-levels column in a serialized sketch. +const LEVELS_FIELD_NAME: &str = "levels"; +/// Field name of the single key column inside each item struct. +/// +/// The item is a struct so a multi-column key is siblings alongside it rather +/// than a second payload shape. What holds that back is this module's key, +/// not the format: a [`SortKeyCodec`] encodes one column to one `u64`. Adding +/// columns means teaching the encode side to emit them *and* relaxing +/// [`SortKeySketch::try_from_proto`], which refuses anything but one field so +/// a half-widened producer fails loudly instead of silently dropping a key. +const KEY_FIELD_NAME: &str = "expr_0"; /// Bijection between a primitive's native value and a `u64` whose ascending /// order matches the native ascending order. @@ -530,7 +552,13 @@ impl SortKeySketch { /// execution. pub fn ingest(&mut self, array: &dyn Array) -> Result<()> { let keys = self.codec.encode(array)?; - self.sketch.absorb_slice(&keys); + // Sketching a sort key usually means sitting above the `SortExec` + // that produced it, so the sorted path is the common case rather + // than a special one. It verifies sortedness in O(n) comparisons and + // falls through when the input isn't sorted, so it is correct to + // call unconditionally — a `DESC` key, whose encoding inverts every + // bit and so arrives descending, takes that fallback. + self.sketch.absorb_sorted_slice(&keys); self.null_count += array.null_count() as u64; Ok(()) } @@ -584,6 +612,28 @@ impl SortKeySketch { self.extreme(!self.codec.options().nulls_first, self.sketch.max()) } + /// The least value observed, or `None` when no value was. + /// + /// Distinct from [`Self::min`], which answers about the least *row* and so + /// reports a typed NULL when NULLs sort first. A router comparing a key + /// against a range wants this one: a NULL bound makes its comparison NULL + /// and drops every row it was meant to select. + pub fn value_min(&self) -> Result> { + self.sketch + .min() + .map(|key| self.codec.decode(*key)) + .transpose() + } + + /// The greatest value observed, or `None` when no value was. Counterpart + /// to [`Self::value_min`]. + pub fn value_max(&self) -> Result> { + self.sketch + .max() + .map(|key| self.codec.decode(*key)) + .transpose() + } + /// Shared body of [`Self::min`] and [`Self::max`]: the extreme is a /// NULL when the NULL run is on `nulls_are_on_this_end` and non-empty, /// otherwise it is `value_extreme` decoded. @@ -649,10 +699,6 @@ impl SortKeySketch { /// The rank *among the values* that population-quantile `q` asks for, /// or `None` when it lands inside the NULL run. - /// - /// This is the remap [`Self::quantile`] documents, split out so that - /// [`Self::cuts`] can resolve every boundary against one pass over the - /// sketch. Callers must have checked that something was observed. fn value_rank(&self, q: f64) -> Option { let rank = (q.clamp(0.0, 1.0) * self.count() as f64) as u64; // No NULL run to step over, so the population rank *is* the value @@ -672,48 +718,321 @@ impl SortKeySketch { } } - /// The `partitions - 1` boundaries that split everything observed into - /// `partitions` equally-sized runs, in sort order. + /// The `partitions - 1` boundaries splitting everything observed into + /// `partitions` runs of equal size, in sort order. If a cut _would be_ NULL, it is adjusted to + /// include the nearest value, so this function never returns a NULL cut. + /// + /// The run is indivisible, so it takes the partition at its end whole and + /// only the values beside it balance. /// - /// Empty when `partitions < 2` or nothing was observed, and otherwise - /// exactly `partitions - 1` long so a caller can index it by output - /// partition. Entries may repeat where one value dominates, and may be - /// typed NULLs where the NULL run spans a boundary; both are faithful - /// answers about a skewed distribution rather than errors. - pub fn cuts(&self, partitions: usize) -> Result> { - if partitions < 2 || self.count() == 0 { + /// ```text + /// nulls_first nulls_last + /// ┌──────┬─────────────────┐ ┌─────────────────┬──────┐ + /// │NULLs │ values │ │ values │NULLs │ + /// └──────┴─────────────────┘ └─────────────────┴──────┘ + /// 0 n N 0 v N + /// + /// rank = max(pop − n, ...) rank = min(pop, ...) + /// pulls UP from min pulls DOWN from max + /// ``` + /// + /// | nulls_first | NULLs | values | K | cuts | partition sizes | + /// |-------------|------:|-------:|--:|----------------|-------------------| + /// | true | 10 | 1..=90 | 4 | `[15, 40, 65]` | 24 / 25 / 25 / 26 | + /// | true | 60 | 1..=40 | 4 | `[1, 14, 27]` | 60 / 13 / 13 / 14 | + /// | false | 60 | 1..=40 | 4 | `[14, 27, 40]` | 13 / 13 / 13 / 61 | + /// | false | 90 | 1..=10 | 4 | `[4, 7, 10]` | 3 / 3 / 3 / 91 | + /// + /// Empty when `partitions < 2` or no value was observed. + /// + /// Should only error if: + /// 1. invalid sketch: min/max is empty but levels are not - guarded against in proto decode + /// 2. codec.decode() failure - guarded against in try_new + pub fn cuts(&self, partition_cnt: usize) -> Result> { + if partition_cnt < 2 || self.sketch.count() == 0 { return Ok(Vec::new()); } - let targets: Vec> = (1..partitions) - .map(|cut| self.value_rank(cut as f64 / partitions as f64)) - .collect(); - // One sorted pass over the retained items for every boundary. Going - // through `quantile` per cut instead re-sorts the whole retained set - // each time, which at 256 partitions costs more than the ingest that - // built the sketch. - let value_ranks: Vec = targets.iter().flatten().copied().collect(); - let mut value_keys = self.sketch.at_ranks(&value_ranks).into_iter(); - - targets - .iter() - .map(|target| match target { - None => self.codec.null_value(), - // Having observed something, every rank the remap produces - // names a row, so the only way back is a full vector. - // Dropping a missing cut would renumber every partition - // above it, and silently. - Some(rank) => { - let key = value_keys.next().flatten().ok_or_else(|| { - internal_datafusion_err!( - "SortKeySketch: no value at rank {rank} of {} values", - self.sketch.count() - ) - })?; - self.codec.decode(*key) - } + /// NULLs are indivisible, so when they outgrow a partition they take one whole. + const PARTITION_FOR_NULLS: usize = 1; + /// `at_ranks` is 1-based: rank `r` names the `r`-th value, and nothing is rank 0. + const FIRST_RANK: u64 = 1; + + let total_cnt = self.count(); + let sketch_cnt = self.sketch.count(); + // Partitions sharing the not-NULLs: all but the one partition consumed by NULLs + let not_null_parts = (partition_cnt - PARTITION_FOR_NULLS) as u128; + let cut_cnt = partition_cnt - 1; + let last_cut_idx = cut_cnt - 1; + // null_count > total_cnt / partition_cnt (but without error inducing division) + let nulls_outgrow_a_partition = + self.null_count as u128 * partition_cnt as u128 > total_cnt as u128; + let mut sketch_ranks: Vec = Vec::with_capacity(cut_cnt); + for cut_idx in 0..cut_cnt { + let num_parts_closed = cut_idx as u128 + 1; + let total_rank = + (num_parts_closed * total_cnt as u128 / partition_cnt as u128) as u64; + if self.codec.options().nulls_first { + let sketch_rank = if !nulls_outgrow_a_partition { + // no cut ever lands on a NULL value anyway + total_rank.saturating_sub(self.null_count) + } else { + // save a whole partition for the NULLs, divide evenly amongst the rest + (cut_idx as u128 * sketch_cnt as u128 / not_null_parts) as u64 + + FIRST_RANK + }; + let cuts_below = cut_idx as u64; + let lowest_unreserved_rank = cuts_below + FIRST_RANK; + let rank = sketch_rank.max(lowest_unreserved_rank).min(sketch_cnt); + sketch_ranks.push(rank); + } else { + // mirror of above + let sketch_rank = if !nulls_outgrow_a_partition { + total_rank + } else { + // round up, since we're going the other way + (num_parts_closed * sketch_cnt as u128).div_ceil(not_null_parts) + as u64 + }; + let cuts_above = (last_cut_idx - cut_idx) as u64; + let highest_unreserved_rank = sketch_cnt.saturating_sub(cuts_above); + let rank = sketch_rank.min(highest_unreserved_rank).max(FIRST_RANK); + sketch_ranks.push(rank); + } + } + + // Turn ranks into values + self.sketch + .at_ranks(&sketch_ranks) + .into_iter() + .zip(&sketch_ranks) + .map(|(key, rank)| { + let key = key.ok_or_else(|| { + internal_datafusion_err!( + "SortKeySketch: no value at rank {rank} of {} values", + self.sketch.count() + ) + })?; + self.codec.decode(*key) }) .collect() } + + /// Serialize to [`SortKeySketchState`]. See that message for the layout + /// and for what was priced against it. + /// + /// The key's direction and NULL placement are deliberately absent: they + /// live once per report, in the `order_by` tag every consumer already + /// reads to know which expression a sketch describes. + pub fn to_proto(&self) -> Result { + let mut offsets: Vec = Vec::with_capacity(self.sketch.levels().len() + 1); + offsets.push(0); + let mut keys: Vec = Vec::new(); + for level in self.sketch.levels() { + let mut ascending = level.clone(); + ascending.sort_unstable(); + keys.extend_from_slice(&ascending); + offsets.push(i32::try_from(keys.len()).map_err(|_| { + internal_datafusion_err!( + "SortKeySketch: {} retained items overflow an arrow list offset", + keys.len() + ) + })?); + } + + let item = Arc::new(Field::new( + KEY_FIELD_NAME, + self.codec.data_type().clone(), + false, + )); + let items = StructArray::try_new( + Fields::from(vec![item]), + vec![self.decode_all(&keys)?], + None, + )?; + let levels = ListArray::try_new( + Arc::new(Field::new("item", items.data_type().clone(), false)), + OffsetBuffer::new(offsets.into()), + Arc::new(items), + None, + )?; + let schema = Arc::new(Schema::new(vec![Field::new( + LEVELS_FIELD_NAME, + levels.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(levels)])?; + + let mut ipc = Vec::new(); + let mut writer = StreamWriter::try_new(&mut ipc, &schema)?; + writer.write(&batch)?; + writer.finish()?; + drop(writer); + + Ok(SortKeySketchState { + k: u32::try_from(self.sketch.k()).map_err(|_| { + internal_datafusion_err!( + "SortKeySketch: k={} exceeds u32", + self.sketch.k() + ) + })?, + null_count: self.null_count, + key_min: self.extreme_proto(self.sketch.min())?, + key_max: self.extreme_proto(self.sketch.max())?, + levels: ipc, + }) + } + + /// Rebuild what [`Self::to_proto`] wrote. `options` comes from the + /// report's `order_by` tag; the key's type comes from the payload's own + /// arrow schema, so the two together reconstruct the codec. + /// + /// Errors on a payload this sketch could not have produced — a schema + /// that isn't one list of structs, or a compactor stack over capacity — + /// rather than returning a sketch whose answers would be quietly wrong. + pub fn try_from_proto( + proto: &SortKeySketchState, + options: SortOptions, + ) -> Result { + let mut reader = + StreamReader::try_new(std::io::Cursor::new(&proto.levels), None)?; + let batch = reader.next().transpose()?.ok_or_else(|| { + internal_datafusion_err!("SortKeySketchState: levels payload holds no batch") + })?; + let levels = batch + .column_by_name(LEVELS_FIELD_NAME) + .and_then(|column| column.as_list_opt::()) + .ok_or_else(|| { + internal_datafusion_err!( + "SortKeySketchState: expected a `{LEVELS_FIELD_NAME}` list column, got {:?}", + batch.schema() + ) + })?; + + let key_type = match levels.value_type() { + DataType::Struct(fields) => match fields.as_ref() { + [only] => only.data_type().clone(), + other => { + return internal_err!( + "SortKeySketchState: expected one key field per item, got {}", + other.len() + ); + } + }, + other => { + return internal_err!( + "SortKeySketchState: expected items to be structs, got {other:?}" + ); + } + }; + let codec = SortKeyCodec::try_new(&key_type, options).ok_or_else(|| { + internal_datafusion_err!( + "SortKeySketchState: {key_type:?} is not an encodable key" + ) + })?; + + let mut stack: Vec> = Vec::with_capacity(levels.len()); + for level in 0..levels.len() { + let items = levels.value(level); + let keys = items + .as_struct_opt() + .map(|items| codec.encode(items.column(0).as_ref())) + .transpose()? + .ok_or_else(|| { + internal_datafusion_err!( + "SortKeySketchState: level {level} is not a struct array" + ) + })?; + stack.push(keys); + } + + let key_min = Self::extreme_key(&codec, &proto.key_min)?; + let key_max = Self::extreme_key(&codec, &proto.key_max)?; + // The extremes are set on the first insert and never cleared, so a + // stack holding keys has both and an empty one has neither. `cuts` + // reads a rank off an extreme and has no answer when it is missing. + let has_keys = stack.iter().any(|level| !level.is_empty()); + if has_keys != key_min.is_some() || has_keys != key_max.is_some() { + return internal_err!( + "SortKeySketchState: {} retained keys against key_min={} key_max={} — \ + a stack holding keys carries both extremes", + stack.iter().map(Vec::len).sum::(), + proto.key_min.len(), + proto.key_max.len() + ); + } + + let sketch = KllSketch::from_parts(proto.k as usize, stack, key_min, key_max) + .ok_or_else(|| { + internal_datafusion_err!( + "SortKeySketchState: k={} and the level widths describe a stack KLL \ + could not have produced", + proto.k + ) + })?; + Ok(Self { + codec, + sketch, + null_count: proto.null_count, + }) + } + + /// Every key decoded into one array of the codec's type, in the order + /// given. Its own function because `iter_to_array` refuses an empty + /// iterator, and a sketch that observed nothing still serializes. + fn decode_all(&self, keys: &[u64]) -> Result { + if keys.is_empty() { + return Ok(new_empty_array(self.codec.data_type())); + } + ScalarValue::iter_to_array( + keys.iter() + .map(|key| self.codec.decode(*key)) + .collect::>>()?, + ) + } + + /// One extreme as the wire's `repeated ScalarValue`: a tuple with one + /// element per key column, empty when nothing was observed. + fn extreme_proto( + &self, + key: Option<&u64>, + ) -> Result> { + key.map(|key| { + let value = self.codec.decode(*key)?; + datafusion_proto_common::ScalarValue::try_from(&value).map_err(|e| { + internal_datafusion_err!( + "SortKeySketch: failed to encode {value:?}: {e:?}" + ) + }) + }) + .transpose() + .map(|encoded| encoded.into_iter().collect()) + } + + /// Reverses [`Self::extreme_proto`]. + fn extreme_key( + codec: &SortKeyCodec, + proto: &[datafusion_proto_common::ScalarValue], + ) -> Result> { + let [value] = proto else { + return match proto { + [] => Ok(None), + other => internal_err!( + "SortKeySketchState: expected one element per key column in an \ + extreme, got {}", + other.len() + ), + }; + }; + let value = ScalarValue::try_from(value).map_err(|e| { + internal_datafusion_err!("SortKeySketchState: undecodable extreme: {e:?}") + })?; + match codec.encode(value.to_array()?.as_ref())?.as_slice() { + [key] => Ok(Some(*key)), + // A NULL extreme would encode to nothing. The extremes are the + // sketch's *value* bounds, so a NULL there is a producer bug. + _ => internal_err!("SortKeySketchState: extreme encoded to no key"), + } + } } #[cfg(test)] @@ -1191,31 +1510,312 @@ mod tests { ); } - /// `cuts` splits the population, NULL run included, so a mostly-NULL - /// column spends its low cuts inside that run and only crosses into - /// the values once the run is behind it. + /// Boundaries as the `Int64` scalars `cuts` returns for an `Int64` key. + fn i64_cuts(values: &[i64]) -> Vec { + values + .iter() + .map(|v| ScalarValue::Int64(Some(*v))) + .collect() + } + + /// A population of 12 over K=4 gives each partition a share of 3, which is + /// the smallest size where the end effects of the half-open convention stay + /// within one row. At 10 the same cuts spread by 2 and these assertions + /// would be measuring integer rounding rather than the repair. /// - /// 60 NULLs then values 1..=40, NULLs first. Quartile ranks are 25, 50 - /// and 75 of 100; the first two sit inside the 60-row NULL run, and the - /// third is 15 rows past it, which is value 15 of 40. + /// No NULLs, so nothing displaces anything and the ranks land where the + /// population quartiles say: 3, 6, 9. + #[test] + fn cuts_balance_when_nothing_is_null() { + let values = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]; + let column = vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(10), + Some(11), + Some(12), + ]; + let sketch = int_sketch(column, sort_options(false, true)); + + let cuts = sketch.cuts(4).unwrap(); + + assert_eq!(cuts, i64_cuts(&[3, 6, 9])); + assert_eq!(partition_sizes(0, &values, &cuts, true), vec![2, 3, 3, 4]); + } + + /// A run of 2 against a share of 3 fits inside partition 0 beside the + /// values below the first cut, so every boundary keeps its population rank + /// shifted down by the run and nothing needs repairing. + #[test] + fn cuts_balance_when_the_run_fits_inside_one_partition() { + let values = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + let column = vec![ + None, + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(10), + ]; + let sketch = int_sketch(column, sort_options(false, true)); + + let cuts = sketch.cuts(4).unwrap(); + + assert_eq!(cuts, i64_cuts(&[1, 4, 7])); + assert_eq!(partition_sizes(2, &values, &cuts, true), vec![2, 3, 3, 4]); + } + + /// A run of 6 against a share of 3 is indivisible, so partition 0 takes it + /// whole and its size cannot be improved. What the boundaries above it owe + /// is an even split of the six values — `[2, 2, 2]`, not the `[4, 1, 1]` + /// that keeping the raw population ranks would give. + #[test] + fn cuts_balance_when_the_run_outgrows_one_partition() { + let values = vec![1, 2, 3, 4, 5, 6]; + let column = vec![ + None, + None, + None, + None, + None, + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + ]; + let sketch = int_sketch(column, sort_options(false, true)); + + let cuts = sketch.cuts(4).unwrap(); + + assert_eq!(cuts, i64_cuts(&[1, 3, 5])); + assert_eq!(partition_sizes(6, &values, &cuts, true), vec![6, 2, 2, 2]); + } + + /// The run leaves fewer values than partitions to spread them over, so the + /// best the boundaries can do is one value each. Every cut is still a real + /// value: a NULL boundary would silently drop its partition. + #[test] + fn cuts_balance_when_the_run_leaves_one_value_per_partition() { + let values = vec![1, 2, 3]; + let column = vec![ + None, + None, + None, + None, + None, + None, + None, + None, + None, + Some(1), + Some(2), + Some(3), + ]; + let sketch = int_sketch(column, sort_options(false, true)); + + let cuts = sketch.cuts(4).unwrap(); + + assert_eq!(cuts, i64_cuts(&[1, 2, 3])); + assert_eq!(partition_sizes(9, &values, &cuts, true), vec![9, 1, 1, 1]); + } + + /// K=8 over 6 values: there are not enough distinct values to give every + /// partition one, so the top boundary repeats and the partition between the + /// repeated pair gets nothing. Repeating beats emitting a NULL or a + /// decreasing boundary, both of which consumers read as a dropped range. #[test] - fn cuts_split_the_population_including_nulls() { - let mut values: Vec> = vec![None; 60]; - values.extend((1..=40).map(Some)); + fn cuts_balance_when_partitions_outnumber_the_values() { + let values = vec![1, 2, 3, 4, 5, 6]; + let column = vec![ + None, + None, + None, + None, + None, + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + ]; + let sketch = int_sketch(column, sort_options(false, true)); + + let cuts = sketch.cuts(8).unwrap(); + + assert_eq!(cuts, i64_cuts(&[1, 2, 3, 4, 5, 6, 6])); + assert_eq!( + partition_sizes(6, &values, &cuts, true), + vec![6, 1, 1, 1, 1, 1, 0, 1] + ); + } + + /// Partition sizes a router would produce from `cuts`, under the half-open + /// convention every consumer uses: partition 0 takes the NULLs and every + /// value below `cuts[0]`, partition `i` takes `[cuts[i - 1], cuts[i])`, and + /// the last takes everything from the final cut up. + fn partition_sizes( + nulls: usize, + values: &[i64], + cuts: &[ScalarValue], + nulls_first: bool, + ) -> Vec { + let bound = |cut: &ScalarValue| match cut { + ScalarValue::Int64(Some(value)) => *value, + other => panic!("expected a non-NULL Int64 boundary, got {other:?}"), + }; + let count = |predicate: &dyn Fn(i64) -> bool| { + values.iter().filter(|value| predicate(**value)).count() + }; + let first = bound(&cuts[0]); + let last = bound(cuts.last().expect("at least one cut")); + // The run rides the partition at its own end of the sort order. + let (low_nulls, high_nulls) = if nulls_first { (nulls, 0) } else { (0, nulls) }; + let mut sizes = vec![low_nulls + count(&|value| value < first)]; + for pair in cuts.windows(2) { + let (lo, hi) = (bound(&pair[0]), bound(&pair[1])); + sizes.push(count(&|value| lo <= value && value < hi)); + } + sizes.push(high_nulls + count(&|value| value >= last)); + sizes + } + + /// A run shorter than one partition's share costs no balance at all: the + /// boundaries keep their population ranks, so the run simply shares the + /// lowest partition with the values below the first cut. + #[test] + fn cuts_are_unrepaired_when_the_null_run_is_small() { + let mut values: Vec> = vec![None; 10]; + values.extend((1..=90).map(Some)); let sketch = int_sketch(values, sort_options(false, true)); let cuts = sketch.cuts(4).unwrap(); - assert_eq!(cuts.len(), 3, "K-1 cuts for K partitions"); + // Population ranks 25/50/75 minus the 10 NULLs, with no clamping. assert_eq!( - cuts[0], - ScalarValue::Int64(None), - "the first quarter is entirely NULL" + cuts, + vec![ + ScalarValue::Int64(Some(15)), + ScalarValue::Int64(Some(40)), + ScalarValue::Int64(Some(65)), + ] ); - assert_eq!(cuts[1], ScalarValue::Int64(None), "so is the second"); + // Which leaves 10 NULLs + values 1..14, then 25, 25 and 26 rows: the + // NULL run cost one row of balance against a perfect 25 apiece. + } + + /// The `nulls_last` mirror. The run sits above the values, so every + /// adjustment reverses: boundaries pull down toward the maximum and the run + /// takes the *top* partition. Same properties as its `nulls_first` twin, + /// and the shared no-NULL cases must agree between the two. + #[test] + fn cuts_mirror_the_repair_when_nulls_sort_last() { + for (nulls, value_count, partitions) in [ + (60usize, 40i64, 4usize), + (90, 10, 4), + (60, 40, 8), + (10, 90, 4), + (95, 5, 4), + ] { + let values: Vec = (1..=value_count).collect(); + let mut column: Vec> = values.iter().copied().map(Some).collect(); + column.extend(std::iter::repeat_n(None, nulls)); + let sketch = int_sketch(column, sort_options(false, false)); + + let cuts = sketch.cuts(partitions).unwrap(); + let context = + format!("{nulls} NULLs last, {value_count} values, K={partitions}"); + assert_eq!(cuts.len(), partitions - 1, "{context}"); + assert!( + cuts.iter().all(|cut| !cut.is_null()), + "{context}: got {cuts:?}" + ); + assert!( + cuts.windows(2).all(|pair| pair[0] <= pair[1]), + "{context}: boundaries must be non-decreasing, got {cuts:?}" + ); + + let sizes = partition_sizes(nulls, &values, &cuts, false); + assert_eq!( + sizes.iter().sum::(), + nulls + value_count as usize, + "{context}: every row lands somewhere, got {sizes:?}" + ); + // The run takes the last partition, so only the others can balance. + let below_the_run = &sizes[..sizes.len() - 1]; + let spread = + below_the_run.iter().max().unwrap() - below_the_run.iter().min().unwrap(); + assert!( + spread <= 1, + "{context}: partitions below the run should differ by at most \ + one row, got {sizes:?}" + ); + } + } + + /// The population rank sits *below* the even split once the run owns the + /// top partition, so honouring it strands the partitions beneath: cuts + /// `[1, 3]` give sizes `[0, 2, 3]`, an empty partition beside a double one. + /// The even split is what the run leaves room for, and nothing above it can + /// be improved by consulting a rank the run already displaced. + #[test] + fn cuts_fill_the_low_partitions_when_the_run_owns_the_top_one() { + let values = vec![1, 2, 3]; + let column = vec![Some(1), Some(2), Some(3), None, None]; + let sketch = int_sketch(column, sort_options(false, false)); + + let cuts = sketch.cuts(3).unwrap(); + + assert_eq!(cuts, i64_cuts(&[2, 3])); + assert_eq!(partition_sizes(2, &values, &cuts, false), vec![1, 1, 3]); + } + + /// With no NULLs there is no run at either end, so the two layouts have + /// nothing to mirror and must produce identical boundaries. + #[test] + fn cuts_agree_across_null_placement_when_nothing_is_null() { + let column: Vec> = (1..=100).map(Some).collect(); + let first = int_sketch(column.clone(), sort_options(false, true)); + let last = int_sketch(column, sort_options(false, false)); + for partitions in [2usize, 4, 8, 16] { + assert_eq!( + first.cuts(partitions).unwrap(), + last.cuts(partitions).unwrap(), + "K={partitions}" + ); + } + } + + /// With no NULLs there is no run to sit above or below the values, so + /// `nulls_last` needs nothing mirrored and must not be refused. + #[test] + fn cuts_allow_nulls_last_without_nulls() { + let sketch = + int_sketch((1..=100).map(Some).collect(), sort_options(false, false)); + let cuts = sketch.cuts(4).unwrap(); assert_eq!( - cuts[2], - ScalarValue::Int64(Some(15)), - "the third crosses out of the run" + cuts, + vec![ + ScalarValue::Int64(Some(25)), + ScalarValue::Int64(Some(50)), + ScalarValue::Int64(Some(75)), + ] ); } @@ -1597,4 +2197,125 @@ mod tests { .expect_err("type mismatch must not silently reinterpret"); assert!(err.to_string().contains("built for"), "got: {err}"); } + /// A round trip has to preserve every answer the sketch gives, not just + /// its byte count: the extremes exactly, and every rank the compactor + /// stack encodes. Rebuilding the stack wrongly — a dropped level, a + /// weight off by a factor of two — leaves `count` intact while moving + /// the quantiles, so the quantile sweep is the assertion that matters. + #[test] + fn wire_round_trip_preserves_every_answer() { + let options = SortOptions { + descending: false, + nulls_first: true, + }; + let codec = SortKeyCodec::try_new(&DataType::Float64, options).unwrap(); + let mut original = SortKeySketch::new(codec); + // Past k so the stack has several levels with real weights, rather + // than one level where every item weighs 1 and a broken rebuild + // would still answer correctly. + let values: Vec> = (0..5_000) + .map(|row| Some(1.0 + row as f64 * 99.0 / 5_000.0)) + .chain((0..40).map(|_| None)) + .collect(); + original.ingest(&Float64Array::from(values)).unwrap(); + + let decoded = + SortKeySketch::try_from_proto(&original.to_proto().unwrap(), options) + .unwrap(); + + assert_eq!(decoded.count(), original.count()); + assert_eq!(decoded.null_count(), original.null_count()); + assert_eq!(decoded.codec(), original.codec()); + assert_eq!(decoded.min().unwrap(), original.min().unwrap()); + assert_eq!(decoded.max().unwrap(), original.max().unwrap()); + for step in 0..=100 { + let q = step as f64 / 100.0; + assert_eq!( + decoded.quantile(q).unwrap(), + original.quantile(q).unwrap(), + "quantile({q}) diverged across the wire" + ); + } + assert_eq!(decoded.cuts(8).unwrap(), original.cuts(8).unwrap()); + } + + /// An empty sketch still has to survive the wire: a task may hold a + /// partition slot it never executed, and the report emits every slot. + #[test] + fn wire_round_trip_preserves_empty_sketch() { + let options = SortOptions::default(); + let codec = SortKeyCodec::try_new(&DataType::Int64, options).unwrap(); + let original = SortKeySketch::new(codec); + + let decoded = + SortKeySketch::try_from_proto(&original.to_proto().unwrap(), options) + .unwrap(); + + assert_eq!(decoded.count(), 0); + assert_eq!(decoded.null_count(), 0); + assert_eq!(decoded.min().unwrap(), None); + assert_eq!(decoded.max().unwrap(), None); + assert!(decoded.cuts(8).unwrap().is_empty()); + } + + /// A payload carrying keys but no extremes is the one shape that makes + /// `cuts` fallible, so it has to die at decode instead. + #[test] + fn wire_refuses_retained_keys_without_extremes() { + let options = SortOptions::default(); + let codec = SortKeyCodec::try_new(&DataType::Int64, options).unwrap(); + let mut original = SortKeySketch::new(codec); + original.ingest(&Int64Array::from(vec![1, 2, 3])).unwrap(); + + let mut proto = original.to_proto().unwrap(); + proto.key_max.clear(); + + let err = SortKeySketch::try_from_proto(&proto, options) + .expect_err("retained keys without a key_max must be refused"); + assert!( + err.to_string().contains("carries both extremes"), + "got: {err}" + ); + } + + /// NULLs are counted beside the sketch rather than in it, so they have + /// their own way of not surviving serialization. + #[test] + fn wire_round_trip_preserves_a_nulls_only_sketch() { + let options = SortOptions { + descending: false, + nulls_first: false, + }; + let codec = SortKeyCodec::try_new(&DataType::Int64, options).unwrap(); + let mut original = SortKeySketch::new(codec); + original + .ingest(&Int64Array::from(vec![None, None, None])) + .unwrap(); + + let decoded = + SortKeySketch::try_from_proto(&original.to_proto().unwrap(), options) + .unwrap(); + + assert_eq!(decoded.count(), 3); + assert_eq!(decoded.null_count(), 3); + // NULLs sort last here, so both extremes are the typed NULL. + assert_eq!(decoded.min().unwrap(), original.min().unwrap()); + assert_eq!(decoded.max().unwrap(), original.max().unwrap()); + } + + /// A truncated or foreign payload must fail rather than decode into a + /// sketch whose answers are quietly wrong. + #[test] + fn wire_rejects_a_payload_it_did_not_write() { + let proto = SortKeySketchState { + k: 800, + null_count: 0, + key_min: vec![], + key_max: vec![], + levels: b"not an arrow stream".to_vec(), + }; + let err = SortKeySketch::try_from_proto(&proto, SortOptions::default()) + .expect_err("a non-IPC payload must not decode"); + assert!(!err.to_string().is_empty()); + } } diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index 627e34498..2a05099f8 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -89,7 +89,6 @@ name = "tpch_plan_stability" path = "tests/tpch_plan_stability/main.rs" [dev-dependencies] -datafusion-functions-aggregate-common = { workspace = true } regex = "1" rstest = { workspace = true } serde_json = "1" diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 9e39c5ffb..c0a742be6 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -277,12 +277,12 @@ fn resolve_range_filter_cuts( ); }; let routing = descend_to_boundary_routing(child)?; - if !rf.routing_expr().eq(&routing.routing_expr) { + if !rf.filter_expr().eq(&routing.routing_expr) { return datafusion::common::internal_err!( - "RangeFilterExec routing_expr `{}` disagrees with its descendant \ + "RangeFilterExec filter_expr `{}` disagrees with its descendant \ boundary ExchangeExec's routing_expr `{}` — plant-time invariant \ broken", - rf.routing_expr(), + rf.filter_expr(), routing.routing_expr ); } @@ -327,15 +327,14 @@ fn descend_to_boundary_routing( /// Project K-1 cuts to K half-open `(cuts[k-1], cuts[k])` ranges with `None` /// sentinels at ±∞. This is the pure range-partitioning projection — no halo /// arithmetic here (RFE widens internally at resolve time). -fn raw_bounds_from_cuts(cuts: &[f64]) -> Vec<(Option, Option)> { +fn raw_bounds_from_cuts( + cuts: &[ScalarValue], +) -> Vec<(Option, Option)> { let k = cuts.len() + 1; (0..k) .map(|i| { - let lo = i - .checked_sub(1) - .and_then(|j| cuts.get(j).copied()) - .map(|v| ScalarValue::Float64(Some(v))); - let hi = cuts.get(i).copied().map(|v| ScalarValue::Float64(Some(v))); + let lo = i.checked_sub(1).and_then(|j| cuts.get(j).cloned()); + let hi = cuts.get(i).cloned(); (lo, hi) }) .collect() diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index b2aa8d970..628e4d036 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -19,6 +19,7 @@ use ballista_core::execution_plans::{ CoalescePlan, stats_for_partition, stats_for_partitions, }; use ballista_core::serde::scheduler::PartitionLocation; +use datafusion::common::ScalarValue; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::Statistics; use datafusion::{ @@ -41,7 +42,7 @@ use std::sync::{Arc, atomic::AtomicI64}; /// task-specialization time to build per-downstream-partition range filters /// (see `PerPartitionFilterExec`). /// -/// `cuts` are `K - 1` monotone `f64` boundaries expressed in the value space +/// `cuts` are `K - 1` monotone boundaries expressed in the value space /// of `routing_expr`; downstream partition `k` owns `[cuts[k-1], cuts[k])` /// with virtual `-∞`/`+∞` sentinels on the ends (matching the range /// repartition's write-side convention). `routing_expr` is the same @@ -49,7 +50,10 @@ use std::sync::{Arc, atomic::AtomicI64}; /// with the writer's placement decision. #[derive(Clone, Debug)] pub struct RangeRepartitionRouting { - pub cuts: Vec, + pub cuts: Vec, + /// Which end of the order holds the NULL run, carried with the cuts so + /// the file router and the read-side filter cannot disagree about it. + pub nulls_first: bool, pub routing_expr: Arc, } @@ -522,6 +526,14 @@ mod range_repartition_routing_tests { ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, }; + /// Cuts as `Float64` scalars, which is what merging produces. + fn cuts_f64(values: [f64; N]) -> Vec { + values + .into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect() + } + fn v_source() -> Arc { let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); @@ -537,7 +549,8 @@ mod range_repartition_routing_tests { fn sample_routing() -> RangeRepartitionRouting { RangeRepartitionRouting { - cuts: vec![10.0, 20.0, 30.0], + cuts: cuts_f64([10.0, 20.0, 30.0]), + nulls_first: true, routing_expr: v_routing_expr(), } } @@ -555,7 +568,7 @@ mod range_repartition_routing_tests { let recovered = exchange .range_repartition_routing() .expect("routing must be Some after resolve"); - assert_eq!(recovered.cuts, vec![10.0, 20.0, 30.0]); + assert_eq!(recovered.cuts, cuts_f64([10.0, 20.0, 30.0])); } #[test] @@ -564,11 +577,12 @@ mod range_repartition_routing_tests { let exchange = ExchangeExec::new(v_source(), None, 42); exchange.resolve_range_repartition_routing(sample_routing()); exchange.resolve_range_repartition_routing(RangeRepartitionRouting { - cuts: vec![100.0], + cuts: cuts_f64([100.0]), + nulls_first: true, routing_expr: v_routing_expr(), }); let recovered = exchange.range_repartition_routing().unwrap(); - assert_eq!(recovered.cuts, vec![100.0], "second resolve wins"); + assert_eq!(recovered.cuts, cuts_f64([100.0]), "second resolve wins"); } /// `with_new_children` must carry the routing slot through: transform @@ -594,6 +608,6 @@ mod range_repartition_routing_tests { let recovered = rebuilt_exchange .range_repartition_routing() .expect("routing must survive with_new_children"); - assert_eq!(recovered.cuts, vec![10.0, 20.0, 30.0]); + assert_eq!(recovered.cuts, cuts_f64([10.0, 20.0, 30.0])); } } diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index c45236cc4..89a654043 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -294,11 +294,16 @@ impl AdaptiveExecutionGraph { /// have already established via `range_repartition_routing_expr` that /// the stage's plan warrants routing, and passes the recovered expr in. /// - /// `Ok(None)` means the stage produced no rows (nothing to route - /// through — passthrough is safe). `Err` means the stage's plan says - /// it should route but something went wrong recovering the cuts — - /// an invariant break, not a soft fallback (would misroute real data + /// `Ok(None)` means the stage has a single output partition, so there is + /// no boundary to route across. `Err` means the stage's plan says it + /// should route but something went wrong recovering the cuts — an + /// invariant break, not a soft fallback (would misroute real data /// downstream). + /// + /// Every other stage gets routing back even when its cuts are empty. A + /// boundary always has a consuming `RangeFilterExec`, which resolves its + /// bounds from the routing parked on that boundary, so declining to park + /// leaves the filter unresolvable rather than saving anyone work. fn repartition_routing( running_stage: &RunningStage, routing_expr: Arc, @@ -330,10 +335,6 @@ impl AdaptiveExecutionGraph { ))); } }; - if entry.total_rows == 0 { - debug!("range-repartition stage {stage_id}: no rows produced, passthrough"); - return Ok(None); - } if entry.partition_count < 2 { // K=1: single output partition — no cuts needed, no routing to // recover. Everything flows to the one downstream partition. @@ -343,14 +344,17 @@ impl AdaptiveExecutionGraph { ); return Ok(None); } - if entry.cuts.is_empty() { + if entry.cuts.is_empty() && entry.null_count != entry.total_rows { + // entirely NULL or empty are valid but degenerate cases return Err(BallistaError::General(format!( - "range-repartition stage {stage_id}: {} rows, K={}, but no cuts (sketch missed)", - entry.total_rows, entry.partition_count + "range-repartition stage {stage_id}: {} rows ({} NULL), K={}, \ + but no cuts (sketch missed)", + entry.total_rows, entry.null_count, entry.partition_count ))); } Ok(Some(RangeRepartitionRouting { cuts: entry.cuts.clone(), + nulls_first: entry.nulls_first, routing_expr, })) } @@ -400,8 +404,9 @@ impl AdaptiveExecutionGraph { partitions, reports, &routing.cuts, - halo_lo, - halo_hi, + &halo_lo, + &halo_hi, + routing.nulls_first, ) .map_err(|err| { BallistaError::General(format!( @@ -1486,8 +1491,8 @@ impl ExecutionGraph for AdaptiveExecutionGraph { fn downstream_halos( full_plan: &Arc, producer_stage_id: usize, -) -> datafusion::common::Result<(f64, f64)> { - let mut result: Option<(f64, f64)> = None; +) -> datafusion::common::Result<(ScalarValue, ScalarValue)> { + let mut result: Option<(ScalarValue, ScalarValue)> = None; full_plan.apply(|node| { let Some(rf) = node.downcast_ref::() else { return Ok(TreeNodeRecursion::Continue); @@ -1508,7 +1513,7 @@ fn downstream_halos( if exchange.stage_id() != Some(producer_stage_id) { return Ok(TreeNodeRecursion::Continue); } - result = Some((scalar_to_f64(rf.halo_lo())?, scalar_to_f64(rf.halo_hi())?)); + result = Some((rf.halo_lo().clone(), rf.halo_hi().clone())); Ok(TreeNodeRecursion::Stop) })?; result.ok_or_else(|| { @@ -1519,16 +1524,3 @@ fn downstream_halos( )) }) } - -/// Halos travel through the RFE public API as `ScalarValue` for future -/// type widening (Interval, timestamps under KLL). Today the internal -/// routing math is `f64`; any other variant is a shape violation upstream -/// and we fail loud rather than silently zero-widen. -fn scalar_to_f64(sv: &ScalarValue) -> datafusion::common::Result { - match sv { - ScalarValue::Float64(Some(v)) => Ok(*v), - other => datafusion::common::internal_err!( - "only f64 halos are implemented, got: {other:?}" - ), - } -} diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs index 822d82f28..9a6248757 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs @@ -576,6 +576,7 @@ mod tests { Arc::new(Column::new("v", 0)), ScalarValue::Float64(Some(0.0)), ScalarValue::Float64(Some(0.0)), + None, ) .unwrap(), ); diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs index fb0038206..5f76feb01 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/parallel_window.rs @@ -35,8 +35,8 @@ //! computed exprs — separate rewrites) //! - `RANGE` frame with finite `PRECEDING` / `FOLLOWING` / `CurrentRow` //! bounds (UNBOUNDED frames go down a different path) -//! - ORDER BY column is `Float64` today (T-Digest restriction; lifts when -//! the sketch swaps to KLL) +//! - ORDER BY column is a type the sort-key sketch encodes (any fixed-width +//! type, nullable or not) //! //! # Rewrite //! @@ -75,9 +75,10 @@ use std::sync::Arc; use ballista_core::config::BallistaConfig; use ballista_core::execution_plans::{ - OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec, RangeFilterExec, - RuntimeStatsExec, + InputOrder, OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec, + RangeFilterExec, RuntimeStatsExec, }; +use ballista_core::sort_key::SortKeyCodec; use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::DataType; use datafusion::common::config::ConfigOptions; @@ -172,27 +173,33 @@ fn fmt_bound(bound: &WindowFrameBound) -> String { } } -/// Extract the halo width in `f64` from a bound. `CurrentRow` → `Some(0.0)`. -/// Non-numeric scalars (e.g. Interval bounds) return `None` — a shape gate, -/// widened alongside KLL. -fn halo_from_bound(bound: &WindowFrameBound) -> Option { +/// The halo the frame declares, if the key can be widened by it. The bound +/// *is* the halo; this converts its units and checks the widening typechecks. +/// `None` declines the rewrite, leaving the query on the non-parallel path. +fn halo_from_bound(bound: &WindowFrameBound, key_type: &DataType) -> Option { + // key_type Timestamp(ns) | Float64 let scalar = match bound { - WindowFrameBound::CurrentRow => return Some(0.0), + // Zero widens by nothing and `RangeFilterExec` short-circuits before + // the arithmetic, so its type need not pair with the key's. + WindowFrameBound::CurrentRow => return Some(ScalarValue::Float64(Some(0.0))), WindowFrameBound::Preceding(s) | WindowFrameBound::Following(s) => s, }; - match scalar { - ScalarValue::Int8(Some(v)) => Some(*v as f64), - ScalarValue::Int16(Some(v)) => Some(*v as f64), - ScalarValue::Int32(Some(v)) => Some(*v as f64), - ScalarValue::Int64(Some(v)) => Some(*v as f64), - ScalarValue::UInt8(Some(v)) => Some(*v as f64), - ScalarValue::UInt16(Some(v)) => Some(*v as f64), - ScalarValue::UInt32(Some(v)) => Some(*v as f64), - ScalarValue::UInt64(Some(v)) => Some(*v as f64), - ScalarValue::Float32(Some(v)) => Some(*v as f64), - ScalarValue::Float64(Some(v)) => Some(*v), - _ => None, + // scalar Interval(1 day) | Int64(3) + if scalar.is_null() { + return None; } + // Into the key's units where that means something. An interval must not + // become a timestamp: it is already the delta type a halo is. + let halo = scalar.cast_to(key_type).unwrap_or_else(|_| scalar.clone()); + // halo Interval(1 day) | Float64(3.0) + if ScalarValue::new_zero(&halo.data_type()).is_ok_and(|zero| halo == zero) { + return Some(halo); + } + // Not computing the halo — checking the filter will be able to apply it. + let probe = ScalarValue::new_zero(key_type).ok()?; + probe.sub(&halo).ok()?; // Timestamp - Interval -> Timestamp | Float64 + probe.add(&halo).ok()?; + Some(halo) } /// Match the parallel-window shape rooted at `node` and, if it fits, splice @@ -254,13 +261,6 @@ fn maybe_rewrite_bwag( if subtree_contains_our_rewrite(window.children().as_slice()) { return Ok(None); } - let (Some(halo_lo), Some(halo_hi)) = ( - halo_from_bound(&frame.start_bound), - halo_from_bound(&frame.end_bound), - ) else { - return Ok(None); - }; - let node_children = node.children(); let [immediate] = node_children.as_slice() else { return datafusion::common::internal_err!( @@ -282,12 +282,18 @@ fn maybe_rewrite_bwag( } let source_schema = base_source.schema(); - // Route on the ORDER BY column. ORRE requires Float64 today (T-Digest - // restriction; lifts when the sketch swaps to KLL). + // Route on the ORDER BY column. Any key the sketch encodes will do; a + // variable-width one declines the rewrite rather than failing the query. let routing_type = order.expr.data_type(&source_schema)?; - if !matches!(routing_type, DataType::Float64) { + if SortKeyCodec::try_new(&routing_type, order.options).is_none() { return Ok(None); } + let (Some(halo_lo), Some(halo_hi)) = ( + halo_from_bound(&frame.start_bound, &routing_type), + halo_from_bound(&frame.end_bound, &routing_type), + ) else { + return Ok(None); + }; let sort_expr = normalize_sort_expr(order); let rse1: Arc = Arc::new(RuntimeStatsExec::try_new( @@ -318,8 +324,9 @@ fn maybe_rewrite_bwag( let wide_filter: Arc = Arc::new(RangeFilterExec::try_new_pending( rse2, sort_expr.expr.clone(), - ScalarValue::Float64(Some(halo_lo)), - ScalarValue::Float64(Some(halo_hi)), + halo_lo, + halo_hi, + Some(InputOrder::Ordered(sort_expr.options)), )?); // Wrap BWAG in PartitionedBoundedWindowAggExec instead of collapsing @@ -342,6 +349,7 @@ fn maybe_rewrite_bwag( sort_expr.expr.clone(), ScalarValue::Float64(Some(0.0)), ScalarValue::Float64(Some(0.0)), + Some(InputOrder::Ordered(sort_expr.options)), )?); debug!( @@ -371,6 +379,7 @@ fn normalize_sort_expr(expr: &PhysicalSortExpr) -> PhysicalSortExpr { mod tests { use super::*; use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::arrow::datatypes::{IntervalMonthDayNano, TimeUnit}; use datafusion::config::ExtensionOptions; use datafusion::datasource::empty::EmptyTable; use datafusion::physical_plan::displayable; @@ -382,6 +391,7 @@ mod tests { Field::new("id2", DataType::Int64, false), Field::new("id3", DataType::Int64, false), Field::new("v2", DataType::Float64, false), + Field::new("name", DataType::Utf8, false), ])); let ctx = SessionContext::new(); ctx.register_table("large", Arc::new(EmptyTable::new(schema)))?; @@ -526,9 +536,11 @@ mod tests { Ok(()) } + /// An `Int64` ORDER BY is routable now: the sketch encodes it, the cuts + /// come back as `Int64`, and `RANGE 3 PRECEDING` keeps an `Int64` halo + /// rather than rounding through `f64`. #[tokio::test] - async fn no_rewrite_on_non_float64_order_key() -> datafusion::common::Result<()> { - // id3 is Int64; ORRE requires Float64 today (T-Digest restriction). + async fn rewrites_an_int64_order_key() -> datafusion::common::Result<()> { let plan = plan( "SELECT sum(v2) OVER (ORDER BY id3 \ RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ @@ -537,31 +549,82 @@ mod tests { .await?; let rewritten = optimize(plan)?; let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); + assert!( + rendered.contains("OrderedRangeRepartitionExec"), + "an Int64 order key should now be rewritten:\n{rendered}" + ); + Ok(()) + } + + /// A key the sketch cannot encode declines the rewrite rather than + /// failing: the query still runs, just not in parallel. + #[tokio::test] + async fn no_rewrite_on_a_key_without_an_encoding() -> datafusion::common::Result<()> { + let plan = plan( + "SELECT sum(v2) OVER (ORDER BY name \ + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) \ + FROM large", + ) + .await?; + let rewritten = optimize(plan)?; + let rendered = format!("{}", displayable(rewritten.as_ref()).indent(true)); assert!( !rendered.contains("OrderedRangeRepartitionExec"), - "non-Float64 order key should not be rewritten:\n{rendered}" + "a Utf8 order key has no sort-key encoding:\n{rendered}" ); Ok(()) } + /// A halo comes back in a type the key can be widened by: the key's own + /// units for a numeric offset, the delta type for a temporal one. #[test] - fn halo_from_bound_reads_all_numeric_variants() { - assert_eq!(halo_from_bound(&WindowFrameBound::CurrentRow), Some(0.0)); + fn halo_from_bound_lands_in_a_type_the_key_can_widen_by() { + let float = DataType::Float64; + let int = DataType::Int64; + assert_eq!( - halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Int64(Some(3)))), - Some(3.0) + halo_from_bound(&WindowFrameBound::CurrentRow, &float), + Some(ScalarValue::Float64(Some(0.0))) ); + // An integer offset casts into a Float64 key's units... assert_eq!( - halo_from_bound(&WindowFrameBound::Following(ScalarValue::Float64(Some( - 2.5 - )))), - Some(2.5) + halo_from_bound( + &WindowFrameBound::Preceding(ScalarValue::Int64(Some(3))), + &float + ), + Some(ScalarValue::Float64(Some(3.0))) ); + // ...and stays an integer on an Int64 key, where the old `as f64` + // rounded every bound through a float. assert_eq!( - halo_from_bound(&WindowFrameBound::Preceding(ScalarValue::Utf8(Some( - "x".into() - )))), + halo_from_bound( + &WindowFrameBound::Preceding(ScalarValue::Int64(Some(3))), + &int + ), + Some(ScalarValue::Int64(Some(3))) + ); + // A bound that cannot widen the key declines the rewrite rather than + // handing the filter a halo it would refuse at runtime. + assert_eq!( + halo_from_bound( + &WindowFrameBound::Preceding(ScalarValue::Utf8(Some("x".into()))), + &float + ), None ); } + + /// The pairing that motivated typing this: an interval offset on a + /// timestamp key, which is not the key's type and must not be cast to it. + #[test] + fn halo_from_bound_keeps_an_interval_against_a_timestamp_key() { + let key = DataType::Timestamp(TimeUnit::Nanosecond, None); + let day = + ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::new(0, 1, 0))); + assert_eq!( + halo_from_bound(&WindowFrameBound::Preceding(day.clone()), &key), + Some(day), + "kept as the delta type, not cast to the key's" + ); + } } diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index ab132af49..0a68a058c 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -32,6 +32,7 @@ use ballista_core::serde::scheduler::{ ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, PartitionId, PartitionLocation, PartitionStats, }; +use datafusion::common::ScalarValue; use datafusion::datasource::MemTable; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; @@ -240,7 +241,11 @@ async fn should_skip_coalesce_when_leaf_has_range_repartition_routing() planner.set_repartition_routing( 0, RangeRepartitionRouting { - cuts: vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0], + cuts: [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0] + .into_iter() + .map(|v| ScalarValue::Float64(Some(v))) + .collect(), + nulls_first: true, routing_expr: Arc::new(Column::new("c", 0)), }, )?; diff --git a/ballista/scheduler/src/state/aqe/test/range_repartition.rs b/ballista/scheduler/src/state/aqe/test/range_repartition.rs index 129e90d2d..3c17b1d97 100644 --- a/ballista/scheduler/src/state/aqe/test/range_repartition.rs +++ b/ballista/scheduler/src/state/aqe/test/range_repartition.rs @@ -21,10 +21,13 @@ //! cuts on it, and `cut_partitions` duplicates straddlers so //! downstream can inject a `PerPartitionFilterExec` to trim them. +use crate::state::aqe::AdaptiveExecutionGraph; use crate::state::aqe::execution_plan::RangeRepartitionRouting; use crate::state::aqe::planner::AdaptivePlanner; +use crate::state::execution_stage::RunningStage; use ballista_core::execution_plans::{ RuntimeStatsExec, UnorderedRangeRepartitionExec, cut_partitions, + repartition_routing_expr, }; use ballista_core::extension::SessionConfigExt; use ballista_core::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; @@ -32,14 +35,18 @@ use ballista_core::serde::scheduler::{ ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, PartitionId, PartitionLocation, PartitionStats, }; +use ballista_core::sort_key::{SortKeyCodec, SortKeySketch}; +use datafusion::arrow::array::Float64Array; use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::ScalarValue; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::physical_expr::PhysicalSortExpr; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::expressions::col; use datafusion::prelude::SessionConfig; +use datafusion_proto::protobuf::PhysicalSortExprNode; use std::sync::Arc; fn v_schema() -> Arc { @@ -133,7 +140,7 @@ async fn routing_parks_when_range_repartition_is_plan_root() .map(|e| e.plan.stage_id()) .expect("a runnable stage must exist"); - // Producer task 7, sub-part 0, sketched [5, 15, 25], straddles the cut at 15. + // Producer task 7, sub-part 0, covering [5, 25], straddles the cut at 15. let reports = vec![ballista_core::execution_plans::TaskRuntimeStats { producer_task_id: 7, report: RuntimeStatsReport { @@ -141,16 +148,26 @@ async fn routing_parks_when_range_repartition_is_plan_root() partitions: vec![RuntimeStatsPartitionEntry { partition_id: 0, row_count: 3, - sketch: Some(ballista_core::execution_plans::sketch_to_proto( - &datafusion_functions_aggregate_common::tdigest::TDigest::new(100) - .merge_unsorted_f64(vec![5.0, 15.0, 25.0]), - )?), + key_min: vec![datafusion_proto::protobuf::ScalarValue::try_from( + &ScalarValue::Float64(Some(5.0)), + )?], + key_max: vec![datafusion_proto::protobuf::ScalarValue::try_from( + &ScalarValue::Float64(Some(25.0)), + )?], + ..Default::default() }], + ..Default::default() }, }]; - let cuts = vec![15.0]; - let remapped = - cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts, 0.0, 0.0)?; + let cuts = vec![ScalarValue::Float64(Some(15.0))]; + let remapped = cut_partitions( + vec![vec![location(0, 7, 3)]], + &reports, + &cuts, + &ScalarValue::Float64(Some(0.0)), + &ScalarValue::Float64(Some(0.0)), + true, + )?; // `cut_partitions` must duplicate the straddler into both partitions — // the read-side filter is expected to trim on read. @@ -159,6 +176,7 @@ async fn routing_parks_when_range_repartition_is_plan_root() let routing = RangeRepartitionRouting { cuts: cuts.clone(), + nulls_first: true, routing_expr: col("v", v_schema().as_ref()).unwrap(), }; planner.set_repartition_routing(stage_id, routing)?; @@ -202,7 +220,8 @@ async fn set_repartition_routing_errs_when_stage_has_no_exchange() .expect("a runnable stage must exist"); let routing = RangeRepartitionRouting { - cuts: vec![0.0], + cuts: vec![ScalarValue::Float64(Some(0.0))], + nulls_first: true, routing_expr: col("v", v_schema().as_ref()).unwrap(), }; let result = planner.set_repartition_routing(stage_id, routing); @@ -237,7 +256,8 @@ async fn routing_parks_when_range_repartition_has_a_parent() .expect("a runnable stage must exist"); let routing = RangeRepartitionRouting { - cuts: vec![15.0], + cuts: vec![ScalarValue::Float64(Some(15.0))], + nulls_first: true, routing_expr: col("v", v_schema().as_ref()).unwrap(), }; planner.set_repartition_routing(stage_id, routing)?; @@ -253,3 +273,119 @@ async fn routing_parks_when_range_repartition_has_a_parent() Ok(()) } + +/// A producer report whose key was NULL in every row. The sketch retains no +/// values, so its NULL count is all that crosses the wire. +fn all_null_report(nulls_per_sub_part: Vec) -> RuntimeStatsReport { + let options = SortOptions { + descending: false, + nulls_first: false, + }; + let codec = SortKeyCodec::try_new(&DataType::Float64, options) + .expect("Float64 is a sortable key type"); + let mut sketch = SortKeySketch::new(codec); + let partitions = nulls_per_sub_part + .into_iter() + .enumerate() + .map(|(sub_part_id, nulls)| { + sketch + .ingest(&Float64Array::from(vec![None::; nulls])) + .expect("NULL Float64 samples into a Float64 sketch"); + RuntimeStatsPartitionEntry { + partition_id: sub_part_id as u32, + row_count: nulls as u64, + null_count: nulls as u64, + ..Default::default() + } + }) + .collect(); + RuntimeStatsReport { + // The tag is where the key's ordering lives, so it has to agree with + // the sort expression `stats_over_urre_root` sketches on. + order_by: vec![PhysicalSortExprNode { + expr: None, + asc: true, + nulls_first: false, + }], + partitions, + sketch: Some(sketch.to_proto().expect("a NULL-only sketch serializes")), + } +} + +/// A key that is NULL in every row is a real distribution, not a missed +/// sketch. No value exists for a boundary to name, so the cuts come back +/// empty and the whole population belongs in the one partition that leaves — +/// still a valid range partitioning. Refusing to route here fails any query +/// whose routing column happens to hold only NULLs. +#[test] +fn repartition_routing_accepts_a_key_that_is_null_in_every_row() { + let mut stage = RunningStage::new( + 1, + 0, + stats_over_urre_root(), + 2, + vec![], + std::collections::HashMap::new(), + Arc::new(SessionConfig::default()), + ); + stage.append_runtime_stats_reports(7, vec![all_null_report(vec![4, 6])]); + + let routing_expr = repartition_routing_expr(stage.plan.as_ref()) + .expect("the URRE spine is a shape the walker recognizes") + .expect("a range-repartition stage routes on an expression"); + let routing = AdaptiveExecutionGraph::repartition_routing(&stage, routing_expr) + .expect("an all-NULL key is a distribution, not an invariant break") + .expect("rows were observed, so the stage still routes"); + + assert!(routing.cuts.is_empty(), "a NULL is not a value to cut on"); +} + +/// A producer report from a task that read nothing: sub-part entries with no +/// rows and no sketch to carry, which is what `RuntimeStatsExec` emits when +/// no batch ever reached it. +fn empty_report(sub_parts: usize) -> RuntimeStatsReport { + RuntimeStatsReport { + order_by: vec![PhysicalSortExprNode { + expr: None, + asc: true, + nulls_first: false, + }], + partitions: (0..sub_parts) + .map(|sub_part_id| RuntimeStatsPartitionEntry { + partition_id: sub_part_id as u32, + row_count: 0, + ..Default::default() + }) + .collect(), + sketch: None, + } +} + +/// A stage that produced no rows still has to park its routing. The +/// downstream `RangeFilterExec` resolves its bounds from the boundary +/// `ExchangeExec`, and every range-repartition boundary has one — an +/// unparked boundary fails that resolution instead of yielding the empty +/// result. No rows means no cuts, which is the same single valid partition +/// an all-NULL key leaves. +#[test] +fn repartition_routing_parks_a_boundary_for_a_stage_that_produced_no_rows() { + let mut stage = RunningStage::new( + 1, + 0, + stats_over_urre_root(), + 2, + vec![], + std::collections::HashMap::new(), + Arc::new(SessionConfig::default()), + ); + stage.append_runtime_stats_reports(7, vec![empty_report(2)]); + + let routing_expr = repartition_routing_expr(stage.plan.as_ref()) + .expect("the URRE spine is a shape the walker recognizes") + .expect("a range-repartition stage routes on an expression"); + let routing = AdaptiveExecutionGraph::repartition_routing(&stage, routing_expr) + .expect("no rows is not an invariant break") + .expect("the boundary still needs routing for its downstream filter"); + + assert!(routing.cuts.is_empty(), "no rows, so no value to cut on"); +} diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 4643e3efd..220a2b38a 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -1752,9 +1752,10 @@ mod tests { ballista_core::serde::protobuf::RuntimeStatsPartitionEntry { partition_id: marker_partition_id, row_count: 0, - sketch: None, + ..Default::default() }, ], + ..Default::default() } } diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index a9ab5e502..879d62e35 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -117,9 +117,10 @@ fn restrict( .collect::>()?; return Ok(Arc::new(RangeFilterExec::try_new_resolved( new_child, - rf.routing_expr().clone(), + rf.filter_expr().clone(), rf.halo_lo().clone(), rf.halo_hi().clone(), + rf.input_order(), sliced_bounds, )?)); } @@ -727,7 +728,7 @@ mod tests { ) .unwrap(); use datafusion::scalar::ScalarValue; - let routing_expr: Arc = Arc::new(Column::new("v", 0)); + let filter_expr: Arc = Arc::new(Column::new("v", 0)); // K=4 raw bounds derived from cuts [100, 200, 300]. let sv = |v: f64| ScalarValue::Float64(Some(v)); let raw_bounds: Vec<(Option, Option)> = vec![ @@ -739,9 +740,10 @@ mod tests { let plan: Arc = Arc::new( RangeFilterExec::try_new_resolved( Arc::new(reader) as Arc, - routing_expr, + filter_expr, ScalarValue::Float64(Some(0.0)), ScalarValue::Float64(Some(0.0)), + None, raw_bounds.clone(), ) .unwrap(), diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 60e533047..faa9a74de 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -1007,15 +1007,17 @@ fn log_runtime_stats_arrival( let non_empty_partitions = report.partitions.iter().filter(|p| p.row_count > 0).count(); let total_rows: u64 = report.partitions.iter().map(|p| p.row_count).sum(); - let sketch_count = report + let ranged_partitions = report .partitions .iter() - .filter(|p| p.sketch.is_some()) + .filter(|p| !p.key_min.is_empty() && !p.key_max.is_empty()) .count(); + // Counted rather than assumed: an entry whose range never got filled + // in would route no files, and nothing else here would say so. debug!( "RuntimeStats arrival: executor={} job={} stage={} task={} \ report[{}] order_by_len={} partitions={} non_empty={} \ - total_rows={} sketches={}", + total_rows={} key_ranges={}/{} sort_key_sketch={}", executor.id, status.job_id, status.stage_id, @@ -1025,11 +1027,53 @@ fn log_runtime_stats_arrival( report.partitions.len(), non_empty_partitions, total_rows, - sketch_count, + ranged_partitions, + report.partitions.len(), + describe_sort_key_sketch(report), ); } } +/// Decode the report's merged [`SortKeySketch`] far enough to say what +/// arrived. Rebuilding it here is the point: a byte count proves the field +/// crossed, where a decoded count and range prove it survived. +/// +/// Any failure is described rather than propagated — this is a log line, and +/// the query's data was already produced correctly. +/// +/// [`SortKeySketch`]: ballista_core::sort_key::SortKeySketch +fn describe_sort_key_sketch( + report: &ballista_core::serde::protobuf::RuntimeStatsReport, +) -> String { + use ballista_core::sort_key::SortKeySketch; + use datafusion::arrow::compute::SortOptions; + + let Some(state) = report.sketch.as_ref() else { + return "none".to_string(); + }; + // The key's direction and NULL placement are not in the sketch — they + // live once here, in the tag that says which expression it describes. + let Some(first) = report.order_by.first() else { + return "undescribable (sketch present with an empty order_by tag)".to_string(); + }; + let options = SortOptions { + descending: !first.asc, + nulls_first: first.nulls_first, + }; + match SortKeySketch::try_from_proto(state, options) { + Ok(sketch) => format!( + "{{bytes={} k={} count={} nulls={} min={:?} max={:?}}}", + state.levels.len(), + state.k, + sketch.count(), + sketch.null_count(), + sketch.value_min(), + sketch.value_max(), + ), + Err(e) => format!("undecodable ({e})"), + } +} + #[cfg(test)] mod tests { use super::*;