feat: collect HyperLogLog distinct counts on shuffle writes for AQE

ballista/core/src/config.rs

@@ -203,6 +203,10 @@ pub const BALLISTA_CHAOS_EXECUTION_SEED: &str = "ballista.testing.chaos_executio
 /// Configuration key for the compression codec used in the shuffle write process
 /// Valid values are: none, lz4, zstd
 pub const BALLISTA_SHUFFLE_COMPRESSION_CODEC: &str = "ballista.shuffle.compression.codec";
+/// Configuration key controlling whether shuffle writers sketch the distinct
+/// counts of their hash-partitioning key columns for the adaptive planner.
+pub const BALLISTA_SHUFFLE_COLLECT_DISTINCT_COUNTS: &str =
+    "ballista.shuffle.collect_distinct_counts";
 
 /// Configuration key for the scheduler's per-task partition-slice cap.
 pub const BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK: &str =
@@ -491,6 +495,16 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
             DataType::Utf8,
             Some("lz4".to_string()),
         ),
+        ConfigEntry::new(
+            BALLISTA_SHUFFLE_COLLECT_DISTINCT_COUNTS.to_string(),
+            "Sketch the approximate distinct count of each hash-partitioning key \
+             column with HyperLogLog while writing shuffle output, so the adaptive \
+             planner can estimate join cardinalities downstream of the shuffle. \
+             Costs hashing every key value on the write path and 4 KiB per key \
+             column in each task's status update.".to_string(),
+            DataType::Boolean,
+            Some(true.to_string()),
+        ),
         ConfigEntry::new(
             BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK.to_string(),
             "Upper bound on the number of input partitions packed into a single \
@@ -767,6 +781,12 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE)
     }
 
+    /// Returns whether shuffle writers sketch the distinct counts of their
+    /// hash-partitioning key columns.
+    pub fn shuffle_collect_distinct_counts(&self) -> bool {
+        self.get_bool_setting(BALLISTA_SHUFFLE_COLLECT_DISTINCT_COUNTS)
+    }
+
     /// Returns the bounded-channel capacity for the shuffle writer I/O bridge.
     pub fn shuffle_writer_channel_capacity(&self) -> usize {
         self.get_usize_setting(BALLISTA_SHUFFLE_WRITER_CHANNEL_CAPACITY)

ballista/core/src/execution_plans/range_shuffle/writer.rs

@@ -44,13 +44,13 @@ use crate::JobId;
 use crate::error::BallistaError;
 use crate::execution_plans::create_shuffle_path;
 use crate::execution_plans::shuffle_writer::{
-    ShuffleWriteMetrics, WriterState, collect_window_state_operators, result_schema,
-    run_coordinator, summaries_to_batch,
+    ColumnStatsCollector, ShuffleWriteMetrics, WriterState,
+    collect_window_state_operators, result_schema, run_coordinator, summaries_to_batch,
 };
 use crate::execution_plans::shuffle_writer_trait::ShuffleWriter;
 use crate::execution_plans::{ObservedWindowState, PartitionedBoundedWindowAggExec};
 use crate::extension::SessionConfigExt;
-use crate::serde::protobuf::ShuffleWritePartition;
+use crate::serde::protobuf::{ShuffleWritePartition, TaskColumnStats};
 
 use super::index::{KeyCollector, index_path, index_schema, write_index_file};
 use super::ipc_file::write_stream_to_ipc_file;
@@ -112,6 +112,8 @@ pub struct RangeShuffleWriterExec {
     /// clone produced by `with_new_children` participates in the same
     /// coordinator.
     state: Arc<Mutex<WriterState>>,
+    /// Statistics of the rows this task wrote, shared with clones.
+    column_stats: ColumnStatsCollector,
 }
 
 impl Clone for RangeShuffleWriterExec {
@@ -126,6 +128,7 @@ impl Clone for RangeShuffleWriterExec {
             metrics: self.metrics.clone(),
             properties: self.properties.clone(),
             state: self.state.clone(),
+            column_stats: self.column_stats.clone(),
         }
     }
 }
@@ -187,6 +190,7 @@ impl RangeShuffleWriterExec {
                 initialized: false,
                 handoffs: (0..output_partition_count).map(|_| None).collect(),
             })),
+            column_stats: ColumnStatsCollector::default(),
         })
     }
 
@@ -202,6 +206,20 @@ impl RangeShuffleWriterExec {
         self.task_id
     }
 
+    /// Per-column null counts over all rows this task wrote, one per column
+    /// of the output schema. Empty until the task has written.
+    pub fn column_null_counts(&self) -> Vec<u64> {
+        self.column_stats.null_counts()
+    }
+
+    /// Per-column statistics over all rows this task wrote, as reported to
+    /// the scheduler: a null count for every column. Range partitioning has
+    /// no hash keys, so no column is sketched. Empty until the task has
+    /// written.
+    pub fn column_stats(&self) -> Vec<TaskColumnStats> {
+        self.column_stats.task_column_stats()
+    }
+
     /// Drain every window-state collector in this stage, keyed by the global
     /// output partition each capture belongs to.
     ///
@@ -254,12 +272,15 @@ impl RangeShuffleWriterExec {
         let task_id = self.task_id;
         let plan = self.plan.clone();
         let metrics = self.metrics.clone();
+        let column_stats = self.column_stats.clone();
 
         async move {
             let now = Instant::now();
             let config = context.session_config().ballista_config();
             let compression_type = config.shuffle_compression_codec()?;
             let channel_capacity = config.shuffle_writer_channel_capacity();
+            // Range partitioning has no hash keys, so nothing is sketched.
+            column_stats.reset(plan.schema().fields().len(), &[]);
 
             let num_partitions =
                 plan.properties().output_partitioning().partition_count();
@@ -287,7 +308,8 @@ impl RangeShuffleWriterExec {
 
                 debug!("Writing range shuffle results to {path:?}");
 
-                let stream = plan.execute(global_partition, context.clone())?;
+                let stream =
+                    column_stats.count(plan.execute(global_partition, context.clone())?);
                 let index_schema = index_schema.clone();
                 let ordering = self.ordering.clone();
                 handles.spawn(async move {
@@ -645,6 +667,27 @@ mod tests {
         }
     }
 
+    // A written column with no nulls reports an exact zero, so the scheduler
+    // can tell it from a column that was never counted.
+    #[tokio::test]
+    async fn reports_a_null_count_per_written_column() {
+        let work_dir = TempDir::new().unwrap();
+        let writer = Arc::new(
+            RangeShuffleWriterExec::try_new(
+                JobId::new("job"),
+                1,
+                input_plan(),
+                work_dir.path().to_str().unwrap().to_owned(),
+            )
+            .unwrap(),
+        );
+        assert!(writer.column_null_counts().is_empty());
+
+        drive(writer.clone(), SessionContext::new().task_ctx()).await;
+
+        assert_eq!(writer.column_null_counts(), vec![0, 0]);
+    }
+
     /// The writer's output has to come back through the reader that will
     /// actually open it in production, not through a decoder chosen by the
     /// test. This is what catches a writer and reader disagreeing on format.