andygrove commented on code in PR #1916:
URL:
https://github.com/apache/datafusion-ballista/pull/1916#discussion_r4165161843
##########
ballista/scheduler/src/state/execution_stage.rs:
##########
@@ -1509,6 +1544,88 @@ 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() {
Review Comment:
Nice test. It might be worth one more case next to it, where a task reports
a `column` past the schema width. `get_mut` handles that now, and a test would
keep it from turning back into a panic if this loop gets refactored later.
##########
ballista/scheduler/src/state/execution_stage.rs:
##########
@@ -267,6 +268,8 @@ pub struct SuccessfulStage {
pub stage_metrics: Vec<MetricsSet>,
/// [SessionConfig] used for this stage
pub session_config: Arc<SessionConfig>,
+ /// container to store col stats per stage
Review Comment:
Could this doc say what the vec holds? Something like one entry per column
of the stage's output schema, with `null_count` only `Exact` when every
successful task reported that column. That'll help whoever writes the first
consumer.
##########
ballista/core/src/execution_plans/sort_shuffle/writer.rs:
##########
@@ -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;
Review Comment:
Thanks for switching this. A small test with a `DataType::Null` column would
capture why `logical_null_count()` matters here. A `NullArray::new(3)` column
should report 3, where `null_count()` reported 0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
