feat(core): RuntimeStatsExec — passthrough tap for row counts + quantile sketches

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"

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"] }

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;

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;