feat(core): Replace TDigest with KLL

ballista/client/tests/parallel_window.rs

@@ -0,0 +1,208 @@
+// 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.
+
+//! End-to-end coverage for `ParallelWindowRule`, which rewrites a bounded
+//! RANGE-frame window into `RSE → ORRE → RSE → RangeFilterExec → PBWAG`.
+//!
+//! The rewrite is a pure optimization, so single-node DataFusion is the
+//! oracle: every case runs there, then on Ballista with the rule off, then on
+//! Ballista with the rule on, and all three must agree. A disagreement names
+//! which side moved.
+//!
+//! What the cases vary is what the rewrite has to get right: which end of the
+//! order the NULL run occupies, whether the frame widens the lower bound, the
+//! upper, both, or neither, and whether a DESC key stays on the serial path.
+
+mod common;
+
+#[cfg(test)]
+#[cfg(feature = "standalone")]
+mod parallel_window {
+    use std::fs;
+    use std::path::Path;
+
+    use ballista::prelude::{SessionConfigExt, SessionContextExt};
+    use ballista_core::config::{
+        BALLISTA_ADAPTIVE_PLANNER_ENABLED, BALLISTA_PARALLEL_WINDOW_ENABLED,
+    };
+    use datafusion::arrow::datatypes::{DataType, Field, Schema};
+    use datafusion::arrow::util::pretty::pretty_format_batches;
+    use datafusion::execution::SessionStateBuilder;
+    use datafusion::prelude::*;
+    use rstest::rstest;
+
+    /// Rows per file. Four files of these interleave over the whole key range,
+    /// so no file is range-disjoint from another and ORRE has real routing to
+    /// do rather than passing each input partition through whole.
+    const ROWS_PER_FILE: i64 = 12;
+    const FILES: i64 = 4;
+
+    /// Every fourth key is NULL, so the run is long enough to span more than
+    /// one output partition's share and cannot be mistaken for a rounding
+    /// artifact at whichever end it lands.
+    fn key_is_null(id: i64) -> bool {
+        id % 4 == 3
+    }
+
+    /// One CSV per file, rows dealt round-robin by `id` so file `f` holds keys
+    /// `f, f + FILES, f + 2 * FILES, ...` — every file covers the full range.
+    fn write_table(dir: &Path, nullable: bool) {
+        let table = dir.join("w");
+        fs::create_dir_all(&table).unwrap();
+        for file in 0..FILES {
+            let mut csv = String::from("id,k,v\n");
+            for row in 0..ROWS_PER_FILE {
+                let id = file + row * FILES;
+                let key = if nullable && key_is_null(id) {
+                    String::new()
+                } else {
+                    id.to_string()
+                };
+                csv.push_str(&format!("{id},{key},{}\n", id * 10));
+            }
+            fs::write(table.join(format!("p{file}.csv")), csv).unwrap();
+        }
+    }
+
+    async fn register(ctx: &SessionContext, dir: &Path, nullable: bool) {
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int64, false),
+            Field::new("k", DataType::Int64, nullable),
+            Field::new("v", DataType::Int64, false),
+        ]);
+        ctx.register_csv(
+            "w",
+            dir.join("w").to_str().unwrap(),
+            CsvReadOptions::new()
+                .has_header(true)
+                .schema(&schema)
+                .file_extension(".csv"),
+        )
+        .await
+        .unwrap();
+    }
+
+    async fn run(ctx: &SessionContext, sql: &str) -> String {
+        let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+        pretty_format_batches(&batches).unwrap().to_string()
+    }
+
+    /// `ORDER BY id` on the outside so the comparison is over rows, not over
+    /// whichever order the K output partitions happened to be drained in.
+    fn query(ordering: &str, frame: &str) -> String {
+        format!(
+            "SELECT id, k, SUM(v) OVER (ORDER BY k {ordering} RANGE BETWEEN {frame}) AS w \
+             FROM w ORDER BY id"
+        )
+    }
+
+    /// A Ballista standalone context with AQE on, the rule at `parallel`, and
+    /// K = 4 output partitions.
+    async fn ballista(dir: &Path, nullable: bool, parallel: bool) -> SessionContext {
+        let config = SessionConfig::new_with_ballista()
+            .with_target_partitions(4)
+            .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, true)
+            .set_bool(BALLISTA_PARALLEL_WINDOW_ENABLED, parallel);
+        let state = SessionStateBuilder::new()
+            .with_config(config)
+            .with_default_features()
+            .build();
+        let ctx = SessionContext::standalone_with_state(state).await.unwrap();
+        register(&ctx, dir, nullable).await;
+        ctx
+    }
+
+    #[rstest]
+    #[case::halo_both("ASC NULLS LAST", "5 PRECEDING AND 5 FOLLOWING")]
+    #[case::halo_lower("ASC NULLS LAST", "5 PRECEDING AND CURRENT ROW")]
+    #[case::halo_upper("ASC NULLS LAST", "CURRENT ROW AND 5 FOLLOWING")]
+    #[case::halo_none("ASC NULLS LAST", "CURRENT ROW AND CURRENT ROW")]
+    #[case::nulls_first("ASC NULLS FIRST", "5 PRECEDING AND 5 FOLLOWING")]
+    #[case::nulls_first_halo_none("ASC NULLS FIRST", "CURRENT ROW AND CURRENT ROW")]
+    // DESC is gated to the serial path by `ParallelWindowRule` — RANGE frame
+    // semantics invert with the sort direction. These cases assert the gate
+    // leaves the answer alone rather than that the rewrite fired.
+    #[case::desc_nulls_last("DESC NULLS LAST", "5 PRECEDING AND 5 FOLLOWING")]
+    #[case::desc_nulls_first("DESC NULLS FIRST", "5 PRECEDING AND 5 FOLLOWING")]
+    #[tokio::test]
+    async fn parallel_window_agrees_with_datafusion(
+        #[case] ordering: &str,
+        #[case] frame: &str,
+        // A nullable key is the case the rewrite has to place the NULL run for;
+        // a non-nullable one has no run and exercises the plain value paths.
+        #[values(false, true)] nullable: bool,
+    ) {
+        let case = format!("{ordering}_{frame}_{nullable}").replace(' ', "_");
+        let dir = std::env::temp_dir().join(format!("ballista_parallel_window_{case}"));
+        let _ = fs::remove_dir_all(&dir);
+        fs::create_dir_all(&dir).unwrap();
+        write_table(&dir, nullable);
+
+        let sql = query(ordering, frame);
+
+        let oracle_ctx = SessionContext::new_with_config(
+            SessionConfig::new().with_target_partitions(4),
+        );
+        register(&oracle_ctx, &dir, nullable).await;
+        let oracle = run(&oracle_ctx, &sql).await;
+
+        let serial = run(&ballista(&dir, nullable, false).await, &sql).await;
+        assert_eq!(
+            oracle, serial,
+            "Ballista disagrees with DataFusion before the rule is even on\n{sql}"
+        );
+
+        let parallel = run(&ballista(&dir, nullable, true).await, &sql).await;
+        assert_eq!(
+            oracle, parallel,
+            "the parallel-window rewrite changed the answer\n{sql}"
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    /// The rewrite has to fire for the cases above to mean anything. A plan
+    /// still holding `BoundedWindowAggExec` over a `SortPreservingMergeExec`
+    /// is the serial shape, and an agreeing answer from it proves nothing.
+    #[tokio::test]
+    async fn the_rewrite_actually_fires() {
+        let dir = std::env::temp_dir().join("ballista_parallel_window_explain");
+        let _ = fs::remove_dir_all(&dir);
+        fs::create_dir_all(&dir).unwrap();
+        write_table(&dir, true);
+
+        let sql = query("ASC NULLS LAST", "5 PRECEDING AND 5 FOLLOWING");
+        let ctx = ballista(&dir, true, true).await;
+        // `EXPLAIN` alone shows only the first stage split. The rewrite lands
+        // when AQE re-plans stage 1 against stage 0's reported sketch, which
+        // is reached by executing.
+        let plan = run(&ctx, &format!("EXPLAIN ANALYZE {sql}")).await;
+
+        for operator in [
+            "OrderedRangeRepartitionExec",
+            "RangeFilterExec",
+            "RuntimeStatsExec",
+        ] {
+            assert!(
+                plan.contains(operator),
+                "rewrite did not plant {operator}:\n{plan}"
+            );
+        }
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+}

ballista/core/Cargo.toml

@@ -51,7 +51,6 @@ 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

@@ -92,11 +92,46 @@ message RuntimeStatsExecNode {
   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;
+// A `SortKeySketch` is a wrapper around a KLL, that also includes NULL counts
+//
+// `levels` is an Arrow IPC stream for now. Binary blob encoding was considered
+// so that encoding tricks could be used like delta+varint. These showed a 4x
+// improvement for u64, but only 25% on f64 due to mantissas being effectively
+// random. Ultimately it was decided that a more complex encoder is not worth
+// the added complexity in this PR. We can revisit the decision in the future.
+message SortKeySketchState {
+  // KLL's nominal top-level compactor capacity
+  uint32 k = 1;
+  // Rows whose key was NULL. Not in `levels`, and not recoverable from it.
+  uint64 null_count = 2;
+  // The exact minimum and maximum over every value observed, one element per
+  // ORDER BY expression. Tracked outside the compactor stack so no coin flip
+  // can move them, which is what keeps `quantile(0.0)` and `quantile(1.0)`
+  // exact. Empty when no value was observed.
+  repeated datafusion_common.ScalarValue key_min = 3;
+  repeated datafusion_common.ScalarValue key_max = 4;
+  // The compactor stack: one row per level in level order, single column
+  // `levels: List<Struct<...>>` holding that level's retained keys ascending,
+  // one struct field per ORDER BY expression. An item's weight is
+  // `2^level`, so the row index carries the weight and the total count is
+  // the weighted sum — neither ships.
+  //
+  // Levels are sorted before serialization, so a decoder takes ascending
+  // order as given rather than being told per level.
+  //
+  // A three-level stack over an `Int64` key decodes to one batch of one
+  // column, where the row index is the level:
+  //
+  //   levels: List<item: Struct<expr_0: Int64>>
+  //
+  //   +-------------------------------------------+
+  //   | levels                                    |
+  //   +-------------------------------------------+
+  //   | [{expr_0: 4}, {expr_0: 17}, {expr_0: 23}] |  level 0, weight 1
+  //   | [{expr_0: 9}, {expr_0: 31}]               |  level 1, weight 2
+  //   | [{expr_0: 12}]                            |  level 2, weight 4
+  //   +-------------------------------------------+
+  bytes levels = 5;
 }
 
 // Flow-control operator with an operator-mode enum. The child plan is
@@ -145,10 +180,30 @@ message PerPartitionFilterExecNode {
 // The child plan is plumbed by the framework as `inputs[0]` during decode.
 // Serialization requires bounds to be resolved.
 message RangeFilterExecNode {
-  datafusion.PhysicalExprNode routing_expr = 1;
+  datafusion.PhysicalExprNode filter_expr = 1;
   datafusion_common.ScalarValue halo_lo = 2;
   datafusion_common.ScalarValue halo_hi = 3;
   repeated RangeBound raw_bounds = 4;
+  // What the operator that produced the cuts says about the rows arriving.
+  // Unset claims nothing, which a nullable `filter_expr` has no answer for:
+  // a NULL has a side in the order, not a position among the values.
+  //
+  // `ordered` is also required of the input, so nothing planted between can
+  // reorder the rows the placement was stated against. `unordered_nulls_first`
+  // states the side without demanding an order, which is what leaves an
+  // unordered range-repartition upstream legal.
+  oneof input_order {
+    bool unordered_nulls_first = 5;
+    SortOptions ordered = 6;
+  }
+}
+
+// An arrow `SortOptions`, which has no proto of its own in the DataFusion
+// descriptors — `PhysicalSortExprNode` carries the pair inline beside an
+// expression this message's users already have.
+message SortOptions {
+  bool descending = 1;
+  bool nulls_first = 2;
 }
 
 // Half-open `[lo, hi)` cut range for one input partition. Either side may be
@@ -672,18 +727,32 @@ message RuntimeStatsReport {
   // partition, post-repartition gets one per output sub-partition. The
   // scheduler groups by `order_by` tag and aggregates.
   repeated RuntimeStatsPartitionEntry partitions = 2;
+  // Every partition's observations folded into one sketch, merged on the
+  // executor before the report is sent. Absent in row-count-only mode,
+  // and when no partition observed a value.
+  optional SortKeySketchState sketch = 3;
 }
 
 // One partition's observations from a `RuntimeStatsExec`.
 message RuntimeStatsPartitionEntry {
   uint32 partition_id = 1;
   uint64 row_count = 2;
-  // Present when the `RuntimeStatsExec` was in sketch mode AND this
-  // partition observed at least one non-null routing value.
-  optional QuantileSketchState sketch = 3;
-  // TODO: `optional MinMaxState min_max` — for a lighter post-repartition
-  // mode where the bin-packer just needs (min, max, count) per
-  // sub-partition and a full T-Digest is overkill.
+  reserved 3; // was TDigest
+  reserved "sketch";
+  // This partition's exact key range, one element per ORDER BY expression.
+  // `cut_partitions` routes a whole shuffle file into every downstream
+  // partition whose range overlaps `[key_min, key_max]`, so these are the
+  // only per-partition facts a router needs — never the distribution, which
+  // is why the sketch beside them is merged rather than repeated here.
+  //
+  // Both empty when this partition observed no value, which `null_count`
+  // then distinguishes: some NULLs means a file of NULLs with no value range
+  // to overlap, none means the partition saw nothing at all.
+  repeated datafusion_common.ScalarValue key_min = 4;
+  repeated datafusion_common.ScalarValue key_max = 5;
+  // Rows in this partition whose key was NULL. Folds to the report-level
+  // `SortKeySketchState.null_count`.
+  uint64 null_count = 6;
 }
 
 message ExecutionError {

ballista/core/src/execution_plans/mod.rs

@@ -48,13 +48,13 @@ pub use ordered_range_repartition::OrderedRangeRepartitionExec;
 pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec;
 pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates};
 pub use plan_algebra::{preserves_distribution, preserves_partitioning};
-pub use range_filter::{RangeBound, RangeFilterExec, WidenedBound};
+pub use range_filter::{InputOrder, RangeBound, RangeFilterExec, WidenedBound};
 pub use range_shuffle_reader::RangeShuffleReaderExec;
 pub use runtime_stats::{
     MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats,
     collect_reports as collect_runtime_stats_reports, cut_partitions,
     log_merged_runtime_stats, merge_reports as merge_runtime_stats_reports,
-    repartition_routing_expr, sketch_from_proto, sketch_to_proto,
+    repartition_routing_expr,
 };
 pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec};
 pub use shuffle_reader::{stats_for_partition, stats_for_partitions};