This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 62228ecc2 feat(core): RuntimeStatsExec — passthrough tap for row
counts + quantile sketches (#2094)
62228ecc2 is described below
commit 62228ecc2a9396e91fc07760981ed11e0b8e7e17
Author: Brent Gardner <[email protected]>
AuthorDate: Tue Jul 21 06:47:43 2026 -0600
feat(core): RuntimeStatsExec — passthrough tap for row counts + quantile
sketches (#2094)
* feat(core): add RuntimeStatsExec — passthrough operator that taps row
counts and quantile sketches
Passthrough physical operator that observes input batches as they stream
past and accumulates two families of runtime state, without changing the
data:
* per-partition row count (`AtomicUsize`, always)
* per-partition quantile sketch over the first `ORDER BY` expression's
`Float64` values (`Mutex<TDigest>`, only when `order_by = Some`)
Both accessors are readable at any point, so callers get mid-stream
snapshots when the sample is accurate enough or post-drain snapshots
when they want the fully-observed state.
Ships 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 operators (the first consumer).
Includes:
* `datafusion-functions-aggregate-common` dep for `TDigest`
* proto `RuntimeStatsExecNode` (tag 6 on `BallistaPhysicalPlanNode`)
and `QuantileSketchState` (T-Digest wire format via `to_scalar_state`)
* `sketch_to_proto` / `sketch_from_proto` helpers
* codec encode/decode registration
* nine tests: three wire round-trip, three streaming (populated sketch,
NULL-key skip, row-count-only), and three plan-node serde roundtrips
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* chore: sync python/Cargo.lock with new
datafusion-functions-aggregate-common dep
Adds the transitive entry to ballista-core in python/Cargo.lock so
`cargo check --locked` succeeds in the python workspace.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* feedback
* RuntimeStatsExec: reject nullable routing at construction
TDigest has no NULL slot, so a nullable routing column silently
excludes NULLs from the sketch while `row_count` still counts them —
skewing the sketch away from what actually flowed through. Rejecting
nullable routing at construction closes the gap. KLL (the planned
swap) will lift this by positioning NULLs per each key's
`SortOptions::nulls_first`.
Also:
- Extend module TODO to name KLL's NULL positioning story.
- Rewrite the `ingest` comment about NULLs — it's a TDigest limitation,
not a semantic decision (NULLs are still forwarded and counted).
- Delete `execute_skips_null_routing_keys_from_sketch` — unreachable
under the new construction check; replaced by a rejection test in
serde tests.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
* feedback: RuntimeStatsExec explicit passthrough defaults
Matches the algebraic-method surface on BufferExec — adds explicit
`required_input_distribution`, `required_input_ordering`, and
`benefits_from_input_partitioning` alongside the three
behaviour-overriding methods already there. All six now spell out the
passthrough contract so a reader doesn't have to consult trait
defaults to know how the planner will treat this node.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
---------
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
---
Cargo.lock | 1 +
Cargo.toml | 1 +
ballista/core/Cargo.toml | 1 +
ballista/core/proto/ballista.proto | 22 +
ballista/core/src/execution_plans/mod.rs | 2 +
ballista/core/src/execution_plans/runtime_stats.rs | 734 +++++++++++++++++++++
ballista/core/src/serde/generated/ballista.rs | 30 +-
ballista/core/src/serde/mod.rs | 235 ++++++-
python/Cargo.lock | 1 +
9 files changed, 1024 insertions(+), 3 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 74ae6d104..1d2d0e675 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1091,6 +1091,7 @@ dependencies = [
"chrono",
"clap 4.6.3",
"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 c114a72e9..2728fcc0b 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<ScalarValue>` 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<TDigest>` 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<PhysicalSortExpr>` 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<PhysicalSortExpr>`
+//! (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<dyn ExecutionPlan>,
+ /// 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<Vec<PhysicalSortExpr>>,
+ /// 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<Arc<[Mutex<TDigest>]>>,
+ properties: Arc<PlanProperties>,
+}
+
+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<dyn ExecutionPlan>,
+ order_by: Option<Vec<PhysicalSortExpr>>,
+ ) -> Result<Self> {
+ 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::<Vec<_>>()
+ .into();
+ let sketches: Option<Arc<[Mutex<TDigest>]>> =
order_by.as_ref().map(|_| {
+ (0..partition_count)
+ .map(|_| Mutex::new(TDigest::new(TDIGEST_MAX_SIZE)))
+ .collect::<Vec<_>>()
+ .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<usize> {
+ 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<Centroid>` 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<Option<TDigest>>
{
+ 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<Option<TDigest>> {
+ let Some(sketches) = self.sketches.as_ref() else {
+ return Ok(None);
+ };
+ let snapshots: Vec<TDigest> = 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::<Result<Vec<_>>>()?;
+ 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<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ /// Passthrough: no distribution requirement on the child.
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ vec![Distribution::UnspecifiedDistribution]
+ }
+
+ /// Passthrough: no ordering requirement on the child.
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+ vec![None]
+ }
+
+ /// Batches pass through unchanged, so input order is preserved.
+ /// Overrides default `false`.
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+
+ /// Wrapping this operator doesn't change how the child benefits from
+ /// its own input partitioning.
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ vec![false]
+ }
+
+ /// Row count and per-column stats pass through unchanged.
+ fn partition_statistics(&self, partition: Option<usize>) ->
Result<Arc<Statistics>> {
+ 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<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ 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<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ 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<Arc<[Mutex<TDigest>]>>,
+ routing_expr: Option<Arc<dyn datafusion::physical_plan::PhysicalExpr>>,
+ 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::<Float64Array>()
+ .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> = 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<crate::serde::protobuf::QuantileSketchState> {
+ let state = sketch.to_scalar_state();
+ let proto_state = state
+ .iter()
+ .map(datafusion_proto_common::ScalarValue::try_from)
+ .collect::<std::result::Result<Vec<_>, _>>()
+ .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<TDigest> {
+ let scalars = proto
+ .state
+ .iter()
+ .map(datafusion::common::ScalarValue::try_from)
+ .collect::<std::result::Result<Vec<_>, _>>()
+ .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<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new("v", DataType::Float64, false),
+ Field::new("id", DataType::Int64, false),
+ ]))
+ }
+
+ fn batch(schema: &Arc<Schema>, v: Vec<Option<f64>>, id: Vec<i64>) ->
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<dyn ExecutionPlan> =
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<dyn ExecutionPlan> =
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::<usize>(), 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 31708a701..670f936bb 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<ScalarValue>` 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 213a600a7..a433f0ac5 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,
@@ -550,6 +552,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,
+ )?))
+ }
}
}
@@ -728,6 +754,29 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
))
})?;
Ok(())
+ } else if let Some(exec) = node.downcast_ref::<RuntimeStatsExec>() {
+ 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: [{}] ",
@@ -1327,4 +1376,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<dyn ExecutionPlan> =
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<u8> = 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::<RuntimeStatsExec>()
+ .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<dyn ExecutionPlan> = 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<u8> = 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::<RuntimeStatsExec>()
+ .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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> =
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<dyn ExecutionPlan> =
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 d5d6b2b1a..0dabbe276 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",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
