diff --git a/Cargo.lock b/Cargo.lock index 6922e2d9c..daef5d2a2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1091,6 +1091,7 @@ dependencies = [ "chrono", "clap 4.6.2", "datafusion", + "datafusion-functions-aggregate-common", "datafusion-proto", "datafusion-proto-common", "datafusion-spark", diff --git a/Cargo.toml b/Cargo.toml index 5719b9b16..46f6a3d19 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ clap = { version = "4.5", features = ["derive", "cargo"] } datafusion = "54" datafusion-cli = "54" +datafusion-functions-aggregate-common = "54" datafusion-proto = "54" datafusion-proto-common = "54" datafusion-spark = "54" diff --git a/ballista/core/Cargo.toml b/ballista/core/Cargo.toml index 735918390..fb9c1c985 100644 --- a/ballista/core/Cargo.toml +++ b/ballista/core/Cargo.toml @@ -51,6 +51,7 @@ 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 7bd4191e5..89359aa8c 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -51,9 +51,31 @@ message BallistaPhysicalPlanNode { UnresolvedShuffleExecNode unresolved_shuffle = 3; SortShuffleWriterExecNode sort_shuffle_writer = 4; ChaosExecNode chaos_exec = 5; + RuntimeStatsExecNode runtime_stats = 6; } } +// Runtime-stats tap. Passes batches through unmodified while accumulating: +// * per-partition row counts (always) +// * a quantile sketch over the first ORDER BY expression's Float64 values +// (only when order_by is non-empty) +// Downstream operators inside the same Ballista task read the stats via a +// child-tree walk. The child plan is plumbed by the framework as `inputs[0]`. +message RuntimeStatsExecNode { + // Optional lexicographic ORDER BY. Empty means row-count-only mode + // (no sketch allocated). When non-empty, the first entry drives the + // sketch and the rest are preserved for downstream operators + // (SortExec, BoundedWindowAggExec) that need the full ordering. + 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; +} + message ChaosExecNode { double failure_probability = 1; string fault_type = 2; diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index f863fd477..1cad2faed 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -21,6 +21,7 @@ mod chaos_exec; mod distributed_explain_analyze; mod distributed_query; +mod runtime_stats; mod shuffle_reader; mod shuffle_writer; mod shuffle_writer_trait; @@ -33,6 +34,7 @@ pub use chaos_exec::ChaosExec; use datafusion::common::exec_err; pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; +pub use runtime_stats::{RuntimeStatsExec, sketch_from_proto, sketch_to_proto}; pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec}; pub use shuffle_reader::{stats_for_partition, stats_for_partitions}; pub use shuffle_writer::DEFAULT_SHUFFLE_CHANNEL_CAPACITY; diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs new file mode 100644 index 000000000..43760b55e --- /dev/null +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -0,0 +1,734 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Passthrough operator that accumulates runtime statistics on the data +//! flowing through it. Two accessor families: +//! +//! - **Row count** — always tracked, per input partition. Written via +//! `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. +//! +//! Timing is decoupled from correctness: both accessors are readable at +//! any point. Callers get whatever has flowed through so far — a +//! mid-stream snapshot for callers that decide the sample is accurate +//! enough, a post-drain snapshot for callers that want the full state +//! (typical after a blocking downstream like `SortExec`). +//! +//! The `order_by` field accepts the full `Vec` so +//! multi-key `ORDER BY` survives serde (tie-breakers get preserved for +//! 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. +//! +//! 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 +//! scheduler. Those pieces arrive with the range-repartition operator +//! (which is the first consumer). + +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::common::{Result, Statistics, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalSortExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use datafusion_functions_aggregate_common::tdigest::TDigest; +use futures::stream::StreamExt; +use log::debug; + +/// 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; + +/// Streaming runtime-stats operator. See module-level docs. +pub struct RuntimeStatsExec { + input: Arc, + /// Lexicographic ORDER BY carried through from the wrapping window + /// operator when the caller wants quantile sketching. `None` when + /// only row counting is needed. + order_by: Option>, + /// Always tracked, per input partition. Written via + /// `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]>>, + properties: Arc, +} + +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). + 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 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(|_| { + (0..partition_count) + .map(|_| Mutex::new(TDigest::new(TDIGEST_MAX_SIZE))) + .collect::>() + .into() + }); + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Ok(Self { + input, + order_by, + row_counts, + sketches, + properties, + }) + } + + /// Full ORDER BY carried through, or `None` if the operator was + /// built in row-count-only mode. + pub fn order_by(&self) -> Option<&[PhysicalSortExpr]> { + self.order_by.as_deref() + } + + /// Rows observed on `partition` so far. Cheap `Relaxed` load — the + /// value is a running counter, monotonically non-decreasing. + /// + /// Errors on out-of-range partition. Callers pass a partition id + /// they've already used with `execute`. + pub fn row_count(&self, partition: usize) -> Result { + let counter = self.row_counts.get(partition).ok_or_else(|| { + internal_datafusion_err!( + "RuntimeStatsExec: partition {} out of range (have {})", + partition, + self.row_counts.len() + ) + })?; + Ok(counter.load(Ordering::Relaxed)) + } + + /// Number of partition slots this operator was built with (matches + /// its input's declared partition count). Every slot has its own + /// row counter and — in sketch mode — its own sketch; a given task + /// only fills the slot(s) it actually executes. + pub fn partition_count(&self) -> usize { + self.row_counts.len() + } + + /// Rows observed across all partitions so far. + pub fn total_row_count(&self) -> usize { + self.row_counts + .iter() + .map(|c| c.load(Ordering::Relaxed)) + .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`). + /// + /// 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 { + return Ok(None); + }; + let slot = sketches.get(partition).ok_or_else(|| { + internal_datafusion_err!( + "RuntimeStatsExec: partition {} out of range (have {})", + partition, + sketches.len() + ) + })?; + let guard = slot.lock().map_err(|e| { + internal_datafusion_err!( + "RuntimeStatsExec partition {}: sketch mutex poisoned: {e}", + partition + ) + })?; + Ok(Some(guard.clone())) + } + + /// All partitions merged into one sketch. `Ok(None)` in + /// row-count-only mode. + pub fn merged_quantile_sketch(&self) -> Result> { + let Some(sketches) = self.sketches.as_ref() else { + return Ok(None); + }; + let snapshots: Vec = sketches + .iter() + .enumerate() + .map(|(partition, m)| { + let guard = m.lock().map_err(|e| { + internal_datafusion_err!( + "RuntimeStatsExec partition {}: sketch mutex poisoned: {e}", + partition + ) + })?; + Ok(guard.clone()) + }) + .collect::>>()?; + Ok(Some(TDigest::merge_digests(snapshots.iter()))) + } +} + +impl Debug for RuntimeStatsExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("RuntimeStatsExec") + .field("order_by", &self.order_by) + .finish() + } +} + +impl DisplayAs for RuntimeStatsExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + match &self.order_by { + Some(exprs) => { + let routing = &exprs[0]; + write!( + f, + "RuntimeStatsExec: rows + sketch(routing={} {})", + routing.expr, + if routing.options.descending { + "desc" + } else { + "asc" + } + ) + } + None => write!(f, "RuntimeStatsExec: rows"), + } + } +} + +impl ExecutionPlan for RuntimeStatsExec { + fn name(&self) -> &str { + "RuntimeStatsExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + /// Passthrough: no distribution requirement on the child. + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + /// Passthrough: no ordering requirement on the child. + fn required_input_ordering(&self) -> Vec> { + vec![None] + } + + /// Batches pass through unchanged, so input order is preserved. + /// Overrides default `false`. + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + /// Wrapping this operator doesn't change how the child benefits from + /// its own input partitioning. + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + /// Row count and per-column stats pass through unchanged. + fn partition_statistics(&self, partition: Option) -> Result> { + self.input.partition_statistics(partition) + } + + /// Every input row is emitted exactly once. Overrides default `Unknown`. + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input] = children.as_slice() else { + return internal_err!( + "RuntimeStatsExec expects exactly one child, got {}", + children.len() + ); + }; + // Fresh counters + sketches on rebuild — planning-time + // reshuffles shouldn't carry stale sample state through the + // tree. + Ok(Arc::new(RuntimeStatsExec::try_new( + input.clone(), + self.order_by.clone(), + )?)) + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + if partition >= self.row_counts.len() { + return internal_err!( + "RuntimeStatsExec: partition {} out of range (have {})", + partition, + self.row_counts.len() + ); + } + let input_stream = self.input.execute(partition, ctx)?; + let schema = self.schema(); + // Cloning `Arc`s so downstream operators observing the same + // 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 routing_expr = self + .order_by + .as_ref() + .and_then(|exprs| exprs.first()) + .map(|e| e.expr.clone()); + + let state = StreamState { + input: input_stream, + row_counts, + sketches, + routing_expr, + partition, + }; + let out = futures::stream::unfold(state, |mut state| async move { + let batch = state.input.next().await?; + let forwarded = batch.and_then(|batch| { + state.ingest(&batch)?; + Ok(batch) + }); + Some((forwarded, state)) + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, out))) + } +} + +/// Per-partition streaming state. Owns the input stream and the routing +/// expression; writes to its own row-count slot and (if sketching) its +/// own sketch slot. +struct StreamState { + input: SendableRecordBatchStream, + row_counts: Arc<[AtomicUsize]>, + sketches: Option]>>, + routing_expr: Option>, + partition: usize, +} + +impl StreamState { + /// Update per-partition counters and (if sketching) the digest. + /// Returns `Err` on any failure path — evaluation error, + /// materialisation error, or wrong result type — so the caller + /// propagates to the output stream rather than emitting a batch the + /// stats never observed. + fn ingest( + &mut self, + batch: &datafusion::arrow::record_batch::RecordBatch, + ) -> Result<()> { + let counter = self.row_counts.get(self.partition).ok_or_else(|| { + internal_datafusion_err!( + "RuntimeStatsExec: partition {} out of range (have {}) — \ + execute() should have validated this", + self.partition, + self.row_counts.len() + ) + })?; + + // 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) + { + 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() { + let slot = sketches.get(self.partition).ok_or_else(|| { + internal_datafusion_err!( + "RuntimeStatsExec: partition {} out of range on sketch slot", + self.partition + ) + })?; + let mut sketch = slot.lock().map_err(|e| { + internal_datafusion_err!( + "RuntimeStatsExec partition {}: sketch mutex poisoned: {e}", + self.partition + ) + })?; + *sketch = sketch.merge_unsorted_f64(values); + } + } + + counter.fetch_add(batch.num_rows(), Ordering::Relaxed); + Ok(()) + } +} + +impl Drop for StreamState { + fn drop(&mut self) { + // End-of-stream introspection until the operator learns to + // emit its stats upstream. Log-and-skip if invariants somehow + // broke — panic in Drop is a process-abort footgun. + let Some(counter) = self.row_counts.get(self.partition) else { + log::error!( + "RuntimeStatsExec partition {} missing row-count slot on Drop; \ + skipping end-of-stream log", + self.partition, + ); + return; + }; + let rows = counter.load(Ordering::Relaxed); + if rows == 0 { + return; + } + match self.sketches.as_ref().and_then(|s| s.get(self.partition)) { + Some(slot) => match slot.lock() { + Ok(sketch) => { + debug!( + "RuntimeStatsExec partition {}: rows={} T-Digest count={} min={} max={}", + self.partition, + rows, + sketch.count(), + sketch.min(), + sketch.max(), + ); + } + Err(e) => { + log::error!( + "RuntimeStatsExec partition {}: sketch mutex poisoned on Drop; \ + skipping end-of-stream log: {e}", + self.partition, + ); + } + }, + None => { + debug!( + "RuntimeStatsExec partition {}: rows={}", + self.partition, rows + ); + } + } + } +} + +/// Serialize a T-Digest to the on-wire +/// [`crate::serde::protobuf::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(crate::serde::protobuf::QuantileSketchState { state: proto_state }) +} + +/// Deserialize a [`crate::serde::protobuf::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: &crate::serde::protobuf::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)) +} + +#[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 = crate::serde::protobuf::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 + //! `RuntimeStatsExec`, drain the stream, and verify the accumulated + //! state matches the data that flowed through. + + use super::*; + use datafusion::arrow::array::{Float64Array, Int64Array}; + use datafusion::arrow::compute::SortOptions; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_plan::common; + use datafusion::physical_plan::expressions::col; + use datafusion::prelude::SessionContext; + + fn schema_v_id() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("v", DataType::Float64, false), + Field::new("id", DataType::Int64, false), + ])) + } + + fn batch(schema: &Arc, v: Vec>, id: Vec) -> RecordBatch { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Float64Array::from(v)), + Arc::new(Int64Array::from(id)), + ], + ) + .unwrap() + } + + /// Sketching mode: drain a single-partition input and verify the + /// operator accumulated one row-count per input row and one sketch + /// sample per non-null routing value. + #[tokio::test] + async fn execute_populates_sketch_and_row_count() { + let schema = schema_v_id(); + let b1 = batch( + &schema, + vec![Some(1.0), Some(3.0), Some(5.0)], + vec![10, 11, 12], + ); + let b2 = batch(&schema, vec![Some(2.0), Some(4.0)], vec![20, 21]); + + let memory = Arc::new( + MemorySourceConfig::try_new(&[vec![b1, b2]], 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(); + let output = common::collect(stream).await.unwrap(); + + // Passthrough: same row count out as in. + let out_rows: usize = output.iter().map(|b| b.num_rows()).sum(); + assert_eq!(out_rows, 5); + + // Row-count accessor observed every batch. + 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. + assert_eq!( + stats.merged_quantile_sketch().unwrap().unwrap().count(), + 5.0 + ); + } + + /// Row-count-only mode (`order_by = None`): the operator still + /// counts rows as they stream past, but the sketch accessors stay + /// `None` no matter how much data flows through. + #[tokio::test] + async fn execute_row_count_only_no_sketch() { + let schema = schema_v_id(); + let b1 = batch(&schema, vec![Some(1.0), Some(2.0)], vec![40, 41]); + let b2 = batch(&schema, vec![Some(3.0)], vec![42]); + + let memory = + Arc::new(MemorySourceConfig::try_new(&[vec![b1, b2]], schema, None).unwrap()); + let input: Arc = Arc::new(DataSourceExec::new(memory)); + + let stats = Arc::new(RuntimeStatsExec::try_new(input, None).unwrap()); + + let ctx = SessionContext::new().task_ctx(); + let stream = stats.execute(0, ctx).unwrap(); + let output = common::collect(stream).await.unwrap(); + + 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(), + "row-count-only mode must not allocate a sketch" + ); + assert!(stats.merged_quantile_sketch().unwrap().is_none()); + } +} diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 9f9efa8f1..e118d12f1 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5" + tags = "1, 2, 3, 4, 5, 6" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -51,8 +51,36 @@ pub mod ballista_physical_plan_node { SortShuffleWriter(super::SortShuffleWriterExecNode), #[prost(message, tag = "5")] ChaosExec(super::ChaosExecNode), + #[prost(message, tag = "6")] + RuntimeStats(super::RuntimeStatsExecNode), } } +/// Runtime-stats tap. Passes batches through unmodified while accumulating: +/// +/// * per-partition row counts (always) +/// * a quantile sketch over the first ORDER BY expression's Float64 values +/// (only when order_by is non-empty) +/// Downstream operators inside the same Ballista task read the stats via a +/// child-tree walk. The child plan is plumbed by the framework as `inputs\[0\]`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RuntimeStatsExecNode { + /// Optional lexicographic ORDER BY. Empty means row-count-only mode + /// (no sketch allocated). When non-empty, the first entry drives the + /// sketch and the rest are preserved for downstream operators + /// (SortExec, BoundedWindowAggExec) that need the full ordering. + #[prost(message, repeated, tag = "1")] + pub order_by: ::prost::alloc::vec::Vec< + ::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.... +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct QuantileSketchState { + #[prost(message, repeated, tag = "1")] + pub state: ::prost::alloc::vec::Vec<::datafusion_proto_common::ScalarValue>, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct ChaosExecNode { #[prost(double, tag = "1")] diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 045b303f0..99f92d315 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -31,9 +31,11 @@ use datafusion_proto::logical_plan::file_formats::{ ArrowLogicalExtensionCodec, AvroLogicalExtensionCodec, CsvLogicalExtensionCodec, JsonLogicalExtensionCodec, ParquetLogicalExtensionCodec, }; +use datafusion_proto::physical_plan::from_proto::parse_physical_sort_exprs; use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning; use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; use datafusion_proto::physical_plan::to_proto::serialize_partitioning; +use datafusion_proto::physical_plan::to_proto::serialize_physical_sort_exprs; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext, @@ -54,8 +56,8 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ - ChaosExec, CoalescePlan, PartitionGroup, ShuffleReaderExec, ShuffleWriterExec, - SortShuffleWriterExec, UnresolvedShuffleExec, + ChaosExec, CoalescePlan, PartitionGroup, RuntimeStatsExec, ShuffleReaderExec, + ShuffleWriterExec, SortShuffleWriterExec, UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -543,6 +545,30 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { Some(chaos_exec.seed), )?)) } + PhysicalPlanType::RuntimeStats(node) => { + let [input] = inputs else { + return Err(DataFusionError::Internal(format!( + "RuntimeStatsExec expects exactly 1 input, got {}", + inputs.len() + ))); + }; + // Empty proto vec means row-count-only mode; a non-empty + // vec is the ORDER BY that drives the quantile sketch. + let order_by = if node.order_by.is_empty() { + None + } else { + Some(parse_physical_sort_exprs( + &node.order_by, + &decode_ctx, + input.schema().as_ref(), + &converter, + )?) + }; + Ok(Arc::new(RuntimeStatsExec::try_new( + input.clone(), + order_by, + )?)) + } } } @@ -718,6 +744,29 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + let converter = DefaultPhysicalProtoConverter {}; + // Empty vec on the wire when the operator is in + // row-count-only mode; otherwise serialise the full ORDER BY. + let order_by = match exec.order_by() { + Some(exprs) => serialize_physical_sort_exprs( + exprs.iter().cloned(), + self.default_codec.as_ref(), + &converter, + )?, + None => Vec::new(), + }; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::RuntimeStats( + protobuf::RuntimeStatsExecNode { order_by }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode RuntimeStatsExec: {e:?}" + )) + })?; + Ok(()) } else { Err(DataFusionError::Internal(format!( "Unsupported plan node, name: [{}] ", @@ -1273,4 +1322,186 @@ mod test { assert_eq!(decoded_exec.upstream_partition_count, 4); assert_eq!(decoded_exec.partition.len(), 1); } + + /// `RuntimeStatsExec` in sketching mode round-trips its ORDER BY + /// and keeps the sketch accessor available. Row-count accessor is + /// always available (and empty on a fresh operator). + #[tokio::test] + async fn test_runtime_stats_exec_roundtrip_with_sketch() { + use crate::execution_plans::RuntimeStatsExec; + use datafusion::arrow::compute::SortOptions; + use datafusion::physical_expr::PhysicalSortExpr; + use datafusion::physical_plan::empty::EmptyExec; + + let schema = Arc::new(Schema::new(vec![ + Field::new("v2", DataType::Float64, false), + Field::new("id", DataType::Int64, false), + ])); + let input: Arc = Arc::new(EmptyExec::new(schema.clone())); + let sort_expr = PhysicalSortExpr { + expr: col("v2", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }; + let original = + RuntimeStatsExec::try_new(input.clone(), Some(vec![sort_expr.clone()])) + .unwrap(); + // Fresh operator: row-count zero, sketch present but empty. + // 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.merged_quantile_sketch().unwrap().unwrap().count(), + 0.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()); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec.try_encode(Arc::new(original), &mut buf).unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap(); + + let decoded = decoded_plan + .downcast_ref::() + .expect("Expected RuntimeStatsExec"); + let order_by = decoded + .order_by() + .expect("sketching mode preserves order_by"); + assert_eq!(order_by.len(), 1); + 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.merged_quantile_sketch().unwrap().unwrap().count(), + 0.0 + ); + } + + /// `RuntimeStatsExec` in row-count-only mode — `order_by = None`. + /// Sketch accessors return `None`; row-count accessors work. + #[tokio::test] + async fn test_runtime_stats_exec_roundtrip_row_count_only() { + use crate::execution_plans::RuntimeStatsExec; + use datafusion::physical_plan::empty::EmptyExec; + + let schema = Arc::new(Schema::new(vec![Field::new( + "v2", + DataType::Float64, + false, + )])); + let input: Arc = Arc::new(EmptyExec::new(schema)); + let original = RuntimeStatsExec::try_new(input.clone(), None).unwrap(); + assert!(original.order_by().is_none()); + assert_eq!(original.row_count(0).unwrap(), 0); + assert!( + original.quantile_sketch(0).unwrap().is_none(), + "no sketch was requested at construction" + ); + assert!(original.merged_quantile_sketch().unwrap().is_none()); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec.try_encode(Arc::new(original), &mut buf).unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap(); + let decoded = decoded_plan + .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()); + } + + /// `try_new` refuses an empty ORDER BY — no routing key means the + /// downstream sketcher has nothing to sample. + #[test] + fn test_runtime_stats_exec_rejects_empty_order_by() { + use crate::execution_plans::RuntimeStatsExec; + use datafusion::physical_plan::empty::EmptyExec; + + let schema = Arc::new(Schema::new(vec![Field::new( + "v2", + DataType::Float64, + false, + )])); + let input: Arc = Arc::new(EmptyExec::new(schema)); + let err = RuntimeStatsExec::try_new(input, Some(vec![])) + .expect_err("empty order_by must be rejected"); + assert!( + err.to_string().contains("order_by is Some but empty"), + "got: {err}" + ); + } + + /// `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. + #[test] + fn test_runtime_stats_exec_rejects_non_float64_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; + + let schema = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Float64, true), + Field::new("id", DataType::Int64, false), + ])); + let input: Arc = Arc::new(EmptyExec::new(schema.clone())); + let sort_expr = PhysicalSortExpr { + expr: col("id", 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; + + 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"); + assert!( + err.to_string() + .contains("routing expression must be non-nullable"), + "got: {err}" + ); + } } diff --git a/python/Cargo.lock b/python/Cargo.lock index a0544afc6..0a9be739c 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -511,6 +511,7 @@ dependencies = [ "async-trait", "chrono", "datafusion", + "datafusion-functions-aggregate-common", "datafusion-proto", "datafusion-proto-common", "futures",