You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by bh...@apache.org on 2026/10/03 08:18:40 UTC
(datafusion-ballista) branch main updated: feat: Column statistics support v1 (#1916)
This is an automated email from the ASF dual-hosted git repository.
coderfender 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 96c1fb92d feat: Column statistics support v1 (#1916)
96c1fb92d is described below
commit 96c1fb92d411e136e4a46c53d6d3ef2020d2ae68
Author: Bhargava Vadlamani <11...@users.noreply.github.com>
AuthorDate: Sat Oct 3 01:18:34 2026 -0700
feat: Column statistics support v1 (#1916)
* add_col_stats_framework
* add_col_stats_framework
* add_col_stats_framework
* add_col_stats_framework
* consume_col_stats_successful_stage
* consume_col_stats_successful_stage
* add_tests
* add_tests
* fix_lint
* fix_lint
* fix_lint
* fix_lint
* propogate_column_stats_along
* address_review_comments
* make stricter fallbacks and add further tests lint
* make col stats encapsulated
* address review comments and add default col stats method across carriers
* address review feedback
---
ballista/core/proto/ballista.proto | 20 +++
.../src/execution_plans/sort_shuffle/writer.rs | 153 ++++++++++++++++++++-
ballista/core/src/serde/generated/ballista.rs | 47 +++++++
ballista/executor/src/execution_engine.rs | 30 +++-
ballista/executor/src/execution_loop.rs | 2 +
ballista/executor/src/executor.rs | 10 +-
ballista/executor/src/executor_server.rs | 2 +
ballista/executor/src/lib.rs | 12 +-
ballista/scheduler/src/scheduler_server/mod.rs | 1 +
.../src/scheduler_server/query_stage_scheduler.rs | 1 +
ballista/scheduler/src/state/execution_stage.rs | 144 ++++++++++++++++++-
ballista/scheduler/src/test_utils.rs | 2 +
12 files changed, 408 insertions(+), 16 deletions(-)
diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto
index dc90b9228..acee01bd6 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -829,6 +829,25 @@ message FailedTask {
}
}
+// Per-column statistics for one task's shuffle output. An empty list on
+// SuccessfulTask means stats were not collected for that task.
+message TaskColumnStats {
+ // Index into the output schema's fields.
+ uint32 column = 1;
+ uint64 null_count = 2;
+ repeated ColumnSketch sketches = 3;
+}
+
+message ColumnSketch {
+ SketchType type = 1;
+ bytes data = 2;
+}
+
+enum SketchType {
+ SKETCH_TYPE_UNSPECIFIED = 0;
+ SKETCH_TYPE_HLL = 1;
+}
+
message SuccessfulTask {
string executor_id = 1;
// TODO tasks are currently always shuffle writes but this will not always be the case
@@ -859,6 +878,7 @@ message SuccessfulTask {
// (`sort_shuffle::get_index_path`), and send only a reference here. That
// keeps the completion message fixed-size regardless of aggregate.
repeated WindowStateReport window_state = 4;
+ repeated TaskColumnStats task_column_stats = 5;
}
// One finalized window-aggregate state from a task's
diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
index a25737ba5..142d2136a 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
@@ -23,6 +23,7 @@
//! concurrently and then emits one file holding every bucket, laid out
//! partition-major so each output partition stays one contiguous range.
+use datafusion::arrow::array::Array;
use std::fmt::Debug;
use std::fs::File;
use std::io::{BufWriter, Seek, Write};
@@ -92,6 +93,8 @@ struct InputPartitionOutput {
/// Kept alive so the pool keeps accounting for this input until the
/// consolidated write has consumed its buffers.
_reservation: MemoryReservation,
+ // collect null counts for better planning
+ null_counts: Vec<u64>,
}
/// What one output partition's encoded rows amount to.
@@ -170,6 +173,7 @@ fn encode_buffered_partitions(
struct WriterState {
initialized: bool,
handoffs: Vec<Option<oneshot::Receiver<Result<Vec<ShuffleWritePartition>>>>>,
+ column_null_counts: Vec<u64>,
}
impl Debug for WriterState {
@@ -496,6 +500,7 @@ impl SortShuffleWriterExec {
state: Arc::new(Mutex::new(WriterState {
initialized: false,
handoffs: (0..output_partition_count).map(|_| None).collect(),
+ column_null_counts: Vec::new(),
})),
})
}
@@ -508,6 +513,15 @@ impl SortShuffleWriterExec {
self
}
+ /// Compute null counts for each column. This will be enhanced in future to add further
+ /// col stats for better query planning
+ pub fn column_null_counts(&self) -> Vec<u64> {
+ self.state
+ .lock()
+ .map(|s| s.column_null_counts.clone())
+ .unwrap_or_default()
+ }
+
/// Task id (append-order slot within the stage) this writer instance
/// is bound to.
pub fn task_id(&self) -> usize {
@@ -630,10 +644,14 @@ impl SortShuffleWriterExec {
// `MemoryPool` as the sole spill trigger.
let memory_limit = config.memory_limit_per_task_bytes;
let per_task_budget_enabled = memory_limit > 0;
+ let mut null_counts = vec![0u64; schema.fields().len()];
while let Some(result) = stream.next().await {
let input_batch = result?;
metrics.input_rows.add(input_batch.num_rows());
+ for (i, col) in input_batch.columns().iter().enumerate() {
+ null_counts[i] += col.logical_null_count() as u64;
+ }
// Compute partition assignment for every row.
let timer = metrics.repart_time.timer();
@@ -761,6 +779,7 @@ impl SortShuffleWriterExec {
spill_manager,
schema,
_reservation: reservation,
+ null_counts,
})
}
}
@@ -1160,6 +1179,21 @@ async fn run_coordinator(
}
writes.abort_all();
+ // Fold each input partition's per-column null counts into the task-level
+ // total and publish it on the shared state for `execute_query_stage`.
+ let mut column_null_counts: Vec<u64> = Vec::new();
+ for output in outputs.iter() {
+ if column_null_counts.len() < output.null_counts.len() {
+ column_null_counts.resize(output.null_counts.len(), 0);
+ }
+ for (slot, n) in column_null_counts.iter_mut().zip(&output.null_counts) {
+ *slot += n;
+ }
+ }
+ if let Ok(mut state) = writer.state.lock() {
+ state.column_null_counts = column_null_counts;
+ }
+
// Every input partition has finished bucketing; emit the task's single
// data file from all of them, then describe each output partition once.
// The file is keyed by `task_id`, which is unique within (job, stage).
@@ -1385,7 +1419,7 @@ mod tests {
use super::*;
use crate::error::BallistaError;
use crate::execution_plans::ChaosExec;
- use datafusion::arrow::array::{StringArray, UInt32Array};
+ use datafusion::arrow::array::{NullArray, StringArray, UInt32Array};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::datasource::memory::MemorySourceConfig;
@@ -2042,6 +2076,123 @@ mod tests {
Ok(())
}
+ /// Build a sort-shuffle writer over `partitions`, drive every output
+ /// partition to completion (which is what publishes the task-level column
+ /// stats onto the shared state), and return the per-column null counts.
+ async fn collect_null_counts(
+ schema: SchemaRef,
+ partitions: Vec<Vec<RecordBatch>>,
+ hash_col: &str,
+ num_output_partitions: usize,
+ ) -> Result<Vec<u64>> {
+ let source = Arc::new(MemorySourceConfig::try_new(
+ &partitions,
+ schema.clone(),
+ None,
+ )?);
+ let input: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(source));
+ let task_ctx = SessionContext::new().task_ctx();
+ let work_dir = TempDir::new()?;
+ let col_idx = schema.index_of(hash_col)?;
+
+ let writer = Arc::new(SortShuffleWriterExec::try_new(
+ "null_count_job".into(),
+ 1,
+ input,
+ work_dir.path().to_str().unwrap().to_string(),
+ Partitioning::Hash(
+ vec![Arc::new(Column::new(hash_col, col_idx))],
+ num_output_partitions,
+ ),
+ SortShuffleConfig::default(),
+ )?);
+
+ for r in drive_partition_results(writer.clone(), task_ctx).await {
+ r.expect("partition drive should succeed");
+ }
+ Ok(writer.column_null_counts())
+ }
+
+ #[tokio::test]
+ async fn column_null_counts_fold_across_inputs() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, true),
+ Field::new("b", DataType::Utf8, true),
+ ]));
+ // Per batch: "a" has 1 null, "b" has 2.
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(UInt32Array::from(vec![Some(1), None, Some(3)])),
+ Arc::new(StringArray::from(vec![None, None, Some("x")])),
+ ],
+ )?;
+ // Two input partitions of the same batch -> the coordinator must SUM
+ // them: "a" = 1x2 = 2, "b" = 2x2 = 4. Also proves the shared-state
+ // read-back (we assert on the original Arc, not the driven clone) and
+ // that nulls are counted over the whole batch despite hashing on "a".
+ let partitions = vec![vec![batch.clone()], vec![batch]];
+ let counts = collect_null_counts(schema, partitions, "a", 2).await?;
+ assert_eq!(counts, vec![2, 4]);
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn column_null_counts_zero_for_fully_populated_column() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, false),
+ Field::new("b", DataType::Utf8, true),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(UInt32Array::from(vec![Some(1), Some(2), Some(3)])),
+ Arc::new(StringArray::from(vec![Some("p"), None, Some("q")])),
+ ],
+ )?;
+ let counts = collect_null_counts(schema, vec![vec![batch]], "a", 2).await?;
+ assert_eq!(counts, vec![0, 1]);
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn column_null_counts_empty_input_is_schema_width_zeros() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, true),
+ Field::new("b", DataType::Utf8, true),
+ ]));
+ // A zero-row batch: the accumulator is still schema-width and the
+ // coordinator's fold still runs, yielding zeros rather than an empty vec.
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(UInt32Array::from(Vec::<Option<u32>>::new())),
+ Arc::new(StringArray::from(Vec::<Option<&str>>::new())),
+ ],
+ )?;
+ let counts = collect_null_counts(schema, vec![vec![batch]], "a", 2).await?;
+ assert_eq!(counts, vec![0, 0]);
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn column_null_counts_counts_logical_nulls_for_null_type() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, true),
+ Field::new("n", DataType::Null, true),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(UInt32Array::from(vec![Some(1), Some(2), Some(3)])),
+ Arc::new(NullArray::new(3)),
+ ],
+ )?;
+ let counts = collect_null_counts(schema, vec![vec![batch]], "a", 2).await?;
+ assert_eq!(counts, vec![0, 3]);
+ Ok(())
+ }
+
#[tokio::test]
async fn multi_spill_round_trips() -> Result<()> {
// Trigger 2 only, with a larger total payload and a tighter budget than
diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs
index 4b78b7f6c..2e14088d7 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -1149,6 +1149,25 @@ pub mod failed_task {
TaskKilled(super::TaskKilled),
}
}
+/// Per-column statistics for one task's shuffle output. An empty list on
+/// SuccessfulTask means stats were not collected for that task.
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct TaskColumnStats {
+ /// Index into the output schema's fields.
+ #[prost(uint32, tag = "1")]
+ pub column: u32,
+ #[prost(uint64, tag = "2")]
+ pub null_count: u64,
+ #[prost(message, repeated, tag = "3")]
+ pub sketches: ::prost::alloc::vec::Vec<ColumnSketch>,
+}
+#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
+pub struct ColumnSketch {
+ #[prost(enumeration = "SketchType", tag = "1")]
+ pub r#type: i32,
+ #[prost(bytes = "vec", tag = "2")]
+ pub data: ::prost::alloc::vec::Vec<u8>,
+}
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct SuccessfulTask {
#[prost(string, tag = "1")]
@@ -1184,6 +1203,8 @@ pub struct SuccessfulTask {
/// keeps the completion message fixed-size regardless of aggregate.
#[prost(message, repeated, tag = "4")]
pub window_state: ::prost::alloc::vec::Vec<WindowStateReport>,
+ #[prost(message, repeated, tag = "5")]
+ pub task_column_stats: ::prost::alloc::vec::Vec<TaskColumnStats>,
}
/// One finalized window-aggregate state from a task's
/// `BoundedWindowAggExec`.
@@ -1907,6 +1928,32 @@ impl ShuffleFileKind {
}
}
}
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+pub enum SketchType {
+ Unspecified = 0,
+ Hll = 1,
+}
+impl SketchType {
+ /// String value of the enum field names used in the ProtoBuf definition.
+ ///
+ /// The values are not transformed in any way and thus are considered stable
+ /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+ pub fn as_str_name(&self) -> &'static str {
+ match self {
+ Self::Unspecified => "SKETCH_TYPE_UNSPECIFIED",
+ Self::Hll => "SKETCH_TYPE_HLL",
+ }
+ }
+ /// Creates an enum from field names used in the ProtoBuf definition.
+ pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+ match value {
+ "SKETCH_TYPE_UNSPECIFIED" => Some(Self::Unspecified),
+ "SKETCH_TYPE_HLL" => Some(Self::Hll),
+ _ => None,
+ }
+ }
+}
/// Generated client implementations.
pub mod scheduler_grpc_client {
#![allow(
diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs
index 9b4653308..c158b3a38 100644
--- a/ballista/executor/src/execution_engine.rs
+++ b/ballista/executor/src/execution_engine.rs
@@ -26,7 +26,7 @@ use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec;
use ballista_core::execution_plans::{
RangeShuffleReaderExec, RangeShuffleWriterExec, ShuffleReaderExec, ShuffleWriterExec,
};
-use ballista_core::serde::protobuf::ShuffleWritePartition;
+use ballista_core::serde::protobuf::{ShuffleWritePartition, TaskColumnStats};
use ballista_core::serde::scheduler::PartitionStats;
use ballista_core::{JobId, utils};
use datafusion::arrow::array::{
@@ -119,6 +119,15 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display {
) -> Result<Vec<ballista_core::serde::protobuf::WindowStateReport>> {
Ok(Vec::new())
}
+
+ /// Collect per-column statistics folded across this task's shuffle output.
+ /// Called at task completion, like [`Self::collect_runtime_stats_reports`],
+ /// and rides the same `SuccessfulTask` message. An empty list means stats
+ /// weren't collected. Default returns empty — implementers override to drain
+ /// counts stashed during execution.
+ fn collect_column_stats(&self) -> Vec<TaskColumnStats> {
+ Vec::new()
+ }
}
/// Default execution engine using DataFusion's ShuffleWriterExec.
@@ -354,6 +363,7 @@ impl QueryStageExecutor for DefaultQueryStageExec {
result.is_ok(),
DisplayableExecutionPlan::with_metrics(plan_arc.as_ref()).indent(true)
);
+
result
}
@@ -414,6 +424,24 @@ impl QueryStageExecutor for DefaultQueryStageExec {
})
.collect()
}
+
+ fn collect_column_stats(&self) -> Vec<TaskColumnStats> {
+ // Only the sort writer stashes per-column null counts during its drain
+ // loop; the passthrough/range paths report none for now.
+ match &self.shuffle_writer {
+ ShuffleWriterVariant::Sort(writer) => writer
+ .column_null_counts()
+ .into_iter()
+ .enumerate()
+ .map(|(column, null_count)| TaskColumnStats {
+ column: column as u32,
+ null_count,
+ sketches: vec![],
+ })
+ .collect(),
+ _ => vec![],
+ }
+ }
}
/// Spawn K parallel `plan.execute(N, ctx)` calls against a shuffle writer,
diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs
index a439dff03..10369a9e8 100644
--- a/ballista/executor/src/execution_loop.rs
+++ b/ballista/executor/src/execution_loop.rs
@@ -419,6 +419,7 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
.collect::<Result<Vec<_>, BallistaError>>()
.ok();
let runtime_stats = query_stage_exec.collect_runtime_stats_reports();
+ let column_stats = query_stage_exec.collect_column_stats();
// Collect only when the task otherwise succeeded: a failed task's
// partial state is meaningless, and its own error is the useful one.
// A collection failure fails the task — these are load-bearing for the
@@ -453,6 +454,7 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
operator_metrics,
runtime_stats,
window_state,
+ column_stats,
},
));
diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs
index 123d9a4aa..b6121c2b1 100644
--- a/ballista/executor/src/executor.rs
+++ b/ballista/executor/src/executor.rs
@@ -29,8 +29,7 @@ use ballista_core::JobId;
use ballista_core::RuntimeProducer;
use ballista_core::error::BallistaError;
use ballista_core::registry::BallistaFunctionRegistry;
-use ballista_core::serde::protobuf;
-use ballista_core::serde::protobuf::ExecutorRegistration;
+use ballista_core::serde::protobuf::{ExecutorRegistration, ShuffleWritePartition};
use ballista_core::serde::scheduler::TaskKey;
use dashmap::DashMap;
use datafusion::execution::context::TaskContext;
@@ -240,7 +239,7 @@ impl Executor {
key: TaskKey,
query_stage_exec: Arc<dyn QueryStageExecutor>,
task_ctx: Arc<TaskContext>,
- ) -> Result<Vec<protobuf::ShuffleWritePartition>, BallistaError> {
+ ) -> Result<Vec<ShuffleWritePartition>, BallistaError> {
let (task, abort_handle) = futures::future::abortable(
query_stage_exec.execute_query_stage(key.task_id, task_ctx),
);
@@ -317,8 +316,7 @@ mod test {
use ballista_core::RuntimeProducer;
use ballista_core::error::BallistaError;
use ballista_core::execution_plans::ShuffleWriterExec;
- use ballista_core::serde::protobuf;
- use ballista_core::serde::protobuf::ExecutorRegistration;
+ use ballista_core::serde::protobuf::{ExecutorRegistration, ShuffleWritePartition};
use ballista_core::serde::scheduler::TaskKey;
use ballista_core::utils::default_config_producer;
use datafusion::arrow::datatypes::{Schema, SchemaRef};
@@ -452,7 +450,7 @@ mod test {
}
/// The result `execute_query_stage` hands back once a spawned task unwinds.
- type TaskOutcome = Result<Vec<protobuf::ShuffleWritePartition>, BallistaError>;
+ type TaskOutcome = Result<Vec<ShuffleWritePartition>, BallistaError>;
/// Builds an executor over `work_dir`, along with the session context whose
/// runtime its tasks run on.
diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs
index 3313cbe54..53274c4e2 100644
--- a/ballista/executor/src/executor_server.rs
+++ b/ballista/executor/src/executor_server.rs
@@ -465,6 +465,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
.collect::<Result<Vec<_>, BallistaError>>()
.ok();
let runtime_stats = exec.collect_runtime_stats_reports();
+ let column_stats = exec.collect_column_stats();
// Collect only when the task otherwise succeeded: a failed task's
// partial state is meaningless, and its own error is the useful one.
// A collection failure fails the task — these are load-bearing for the
@@ -499,6 +500,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
operator_metrics,
runtime_stats,
window_state,
+ column_stats,
},
);
diff --git a/ballista/executor/src/lib.rs b/ballista/executor/src/lib.rs
index 9499005a0..624134a58 100644
--- a/ballista/executor/src/lib.rs
+++ b/ballista/executor/src/lib.rs
@@ -59,15 +59,14 @@ pub use standalone::new_standalone_executor;
pub use standalone::new_standalone_executor_from_builder;
pub use standalone::new_standalone_executor_from_state;
-use log::info;
-
use crate::shutdown::Shutdown;
use ballista_core::serde::protobuf::{
FailedTask, OperatorMetricsSet, RuntimeStatsReport, ShuffleWritePartition,
- SuccessfulTask, TaskStatus, WindowStateReport, task_status,
+ SuccessfulTask, TaskColumnStats, TaskStatus, WindowStateReport, task_status,
};
use ballista_core::serde::scheduler::TaskKey;
use ballista_core::utils::GrpcServerConfig;
+use log::info;
/// [ArrowFlightServerProvider] provides a function which creates a new Arrow Flight server.
///
@@ -115,6 +114,9 @@ pub struct TaskCompletionExtras {
/// window, already stamped with the global partition each entry belongs
/// to by the stage's `ShuffleWriterExec`.
pub window_state: Vec<WindowStateReport>,
+ /// Per-column statistics folded across this task's shuffle output. Empty
+ /// when the executed plan collects none (e.g. non-sort shuffle paths).
+ pub column_stats: Vec<TaskColumnStats>,
}
/// Converts a task execution result into a [`TaskStatus`] protobuf message.
@@ -123,7 +125,7 @@ pub struct TaskCompletionExtras {
/// along with timing and metrics information into a status message that
/// can be sent back to the scheduler.
pub fn as_task_status(
- execution_result: ballista_core::error::Result<Vec<ShuffleWritePartition>>,
+ execution_result: Result<Vec<ShuffleWritePartition>, BallistaError>,
executor_id: String,
stage_attempt_num: usize,
key: TaskKey,
@@ -134,6 +136,7 @@ pub fn as_task_status(
operator_metrics,
runtime_stats,
window_state,
+ column_stats,
} = extras;
let metrics = operator_metrics.unwrap_or_default();
let task_id = key.task_id;
@@ -159,6 +162,7 @@ pub fn as_task_status(
executor_id,
partitions,
runtime_stats,
+ task_column_stats: column_stats,
window_state,
})),
}
diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs
index b606ec6c6..9d3ee8809 100644
--- a/ballista/scheduler/src/scheduler_server/mod.rs
+++ b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -687,6 +687,7 @@ mod test {
executor_id: "executor-1".to_owned(),
partitions,
runtime_stats: vec![],
+ task_column_stats: vec![],
window_state: vec![],
})),
};
diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
index 6f4e0be22..f3964de17 100644
--- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
+++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs
@@ -896,6 +896,7 @@ mod tests {
})
.collect(),
runtime_stats: vec![],
+ task_column_stats: vec![],
window_state: vec![],
})),
})
diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs
index 68fd51feb..64988abe4 100644
--- a/ballista/scheduler/src/state/execution_stage.rs
+++ b/ballista/scheduler/src/state/execution_stage.rs
@@ -15,14 +15,15 @@
// specific language governing permissions and limitations
// under the License.
+use datafusion::common::ColumnStatistics;
+use datafusion::common::stats::Precision;
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::aggregate_statistics::AggregateStatistics;
use std::collections::{HashMap, HashSet, VecDeque};
use std::convert::TryInto;
use std::fmt::{Debug, Formatter};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
-
-use datafusion::config::ConfigOptions;
-use datafusion::physical_optimizer::aggregate_statistics::AggregateStatistics;
//use datafusion::physical_optimizer::join_selection::JoinSelection;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
@@ -267,6 +268,9 @@ pub struct SuccessfulStage {
pub stage_metrics: Vec<MetricsSet>,
/// [SessionConfig] used for this stage
pub session_config: Arc<SessionConfig>,
+ /// Per-column stats, one entry per output-schema column. `null_count` is
+ /// `Exact` only when every successful task reported that column.
+ pub output_column_stats: Vec<ColumnStatistics>,
}
/// If a stage fails, it will be with an error message
@@ -665,6 +669,35 @@ impl RunningStage {
warn!("The metrics for stage {} should not be none", self.stage_id);
vec![]
});
+ let num_columns = self.plan.schema().fields().len();
+ let mut output_column_stats = vec![
+ ColumnStatistics {
+ null_count: Precision::Exact(0),
+ ..ColumnStatistics::new_unknown()
+ };
+ num_columns
+ ];
+ for info in self.task_infos.iter() {
+ let task_status::Status::Successful(task_status) = &info.task_status else {
+ continue;
+ };
+ let mut reported = vec![false; num_columns];
+ for col_stats in &task_status.task_column_stats {
+ let column = col_stats.column as usize;
+ if let Some(slot) = output_column_stats.get_mut(column) {
+ slot.null_count = slot
+ .null_count
+ .add(&Precision::Exact(col_stats.null_count as usize));
+ reported[column] = true;
+ }
+ }
+ for (slot, reported) in output_column_stats.iter_mut().zip(reported) {
+ if !reported {
+ slot.null_count = Precision::Absent;
+ }
+ }
+ }
+
SuccessfulStage {
stage_id: self.stage_id,
stage_attempt_num: self.stage_attempt_num,
@@ -675,6 +708,7 @@ impl RunningStage {
task_infos,
stage_metrics,
session_config: self.session_config.clone(),
+ output_column_stats,
}
}
@@ -1416,7 +1450,8 @@ impl StageOutput {
mod tests {
use super::*;
use ballista_core::serde::protobuf::{
- OperatorMetric, SuccessfulTask, TaskStatus, operator_metric, task_status,
+ OperatorMetric, SuccessfulTask, TaskColumnStats, TaskStatus, operator_metric,
+ task_status,
};
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::prelude::SessionConfig;
@@ -1449,6 +1484,7 @@ mod tests {
executor_id: "executor-1".to_string(),
partitions: vec![],
runtime_stats: vec![],
+ task_column_stats: vec![],
window_state: vec![],
})),
metrics: vec![],
@@ -1509,6 +1545,105 @@ mod tests {
});
}
+ /// Build a running stage over a two-column schema plus a successful task
+ /// reporting the given per-column null counts.
+ fn stage_with_col_stats(tasks: Vec<Vec<TaskColumnStats>>) -> RunningStage {
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, true),
+ Field::new("b", DataType::UInt32, true),
+ ]));
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(schema));
+ let mut stage = RunningStage::new(
+ 1,
+ 0,
+ plan,
+ tasks.len(),
+ vec![],
+ HashMap::new(),
+ Arc::new(SessionConfig::default()),
+ );
+ for (task_id, task_column_stats) in tasks.into_iter().enumerate() {
+ stage.task_infos.push(TaskInfo {
+ task_id,
+ scheduled_time: 50,
+ launch_time: 100,
+ start_exec_time: 200,
+ end_exec_time: 300,
+ finish_time: 300,
+ task_status: task_status::Status::Successful(SuccessfulTask {
+ executor_id: "executor-1".to_string(),
+ partitions: vec![],
+ runtime_stats: vec![],
+ task_column_stats,
+ window_state: vec![],
+ }),
+ global_input_partition_ids: vec![task_id],
+ vcores_consumed: 1,
+ });
+ }
+ stage
+ }
+
+ fn col(column: u32, null_count: u64) -> TaskColumnStats {
+ TaskColumnStats {
+ column,
+ null_count,
+ sketches: vec![],
+ }
+ }
+
+ #[test]
+ fn test_to_successful_folds_column_null_counts_across_tasks() {
+ let stage = stage_with_col_stats(vec![
+ vec![col(0, 1), col(1, 2)],
+ vec![col(0, 3), col(1, 4)],
+ ]);
+ let successful = stage.to_successful();
+ assert_eq!(
+ successful.output_column_stats[0].null_count,
+ Precision::Exact(4)
+ );
+ assert_eq!(
+ successful.output_column_stats[1].null_count,
+ Precision::Exact(6)
+ );
+ }
+
+ #[test]
+ fn test_to_successful_marks_unreported_column_absent() {
+ // Second task reports only column 0 (e.g. an older executor); column 1
+ // must degrade to Absent rather than silently undercounting as Exact.
+ let stage =
+ stage_with_col_stats(vec![vec![col(0, 1), col(1, 2)], vec![col(0, 3)]]);
+ let successful = stage.to_successful();
+ assert_eq!(
+ successful.output_column_stats[0].null_count,
+ Precision::Exact(4)
+ );
+ assert_eq!(
+ successful.output_column_stats[1].null_count,
+ Precision::Absent
+ );
+ }
+
+ #[test]
+ fn test_to_successful_skips_column_past_schema_width() {
+ // Column index past the 2-wide output schema must be skipped by
+ // get_mut, not panic (the fold runs on the stage event loop).
+ let stage = stage_with_col_stats(vec![vec![col(0, 1), col(5, 9)]]);
+ let successful = stage.to_successful();
+ assert_eq!(successful.output_column_stats.len(), 2);
+ assert_eq!(
+ successful.output_column_stats[0].null_count,
+ Precision::Exact(1)
+ );
+ assert_eq!(
+ successful.output_column_stats[1].null_count,
+ Precision::Absent
+ );
+ }
+
/// Verify that a normal status update transitions the task to Successful.
#[test]
fn test_update_task_info_normal_update_succeeds() {
@@ -1885,6 +2020,7 @@ mod tests {
executor_id: executor.to_string(),
partitions: vec![],
runtime_stats: vec![],
+ task_column_stats: vec![],
window_state: vec![],
});
stage.append_runtime_stats_reports(
diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs
index d6209df92..5545b975e 100644
--- a/ballista/scheduler/src/test_utils.rs
+++ b/ballista/scheduler/src/test_utils.rs
@@ -309,6 +309,7 @@ pub fn default_task_runner() -> impl TaskRunner {
executor_id: executor_id.clone(),
partitions: partitions.clone(),
runtime_stats: vec![],
+ task_column_stats: vec![],
window_state: vec![],
})),
});
@@ -1296,6 +1297,7 @@ pub fn mock_completed_task(task: TaskDescription, executor_id: &str) -> TaskStat
executor_id: executor_id.to_owned(),
partitions,
runtime_stats: vec![],
+ task_column_stats: vec![],
window_state: vec![],
})),
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org