perf: try bucket sort

datafusion/common/src/config.rs

@@ -1006,6 +1006,39 @@ config_namespace! {
         /// batches and merged.
         pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024
 
+        /// Sort with a bucket (distribution) sort instead of sorting every batch
+        /// and merging the sorted batches. Rows are scattered on arrival into
+        /// key-range buckets chosen from a sample of the sort keys, every bucket
+        /// is sorted on its own and the buckets are emitted in key order. Falls
+        /// back to the merge sort when the first sort key type is not supported.
+        pub sort_bucket_sort_enabled: bool, default = true
+
+        /// Number of key ranges rows are scattered into on arrival when
+        /// `sort_bucket_sort_enabled` is set (at most 4096). A range whose sort
+        /// keys exceed `sort_bucket_target_bytes` is split when it is sorted, and
+        /// ranges are the unit of spilling.
+        pub sort_bucket_count: usize, default = 64
+
+        /// Target size of the sort keys of a bucket when it is sorted, sized to
+        /// stay in cache. A bucket with larger keys is split by its keys before
+        /// sorting. Only the keys are sorted: the rows are gathered once in sorted
+        /// order, so their size does not count.
+        pub sort_bucket_target_bytes: usize, default = 2 * 1024 * 1024
+
+        /// Number of rows to buffer before choosing the bucket boundaries from
+        /// the sample of keys seen so far.
+        pub sort_bucket_sample_rows: usize, default = 65536
+
+        /// How the bucket sort holds arriving rows until it sorts them:
+        /// - `scatter`: each batch is reordered by key range on arrival, so a range's
+        ///   rows are copied out of a small area when it is sorted (two copies)
+        /// - `gather`: batches are kept as they came and each range keeps the numbers
+        ///   of its rows, copied once when the range is sorted (one copy, read from
+        ///   all over memory)
+        /// - `auto`: times both on the rows held when the boundaries are chosen, and
+        ///   gathers when that is clearly faster (wide rows), else scatters
+        pub sort_bucket_mode: String, default = "auto".to_string()
+
         /// Maximum buffer capacity (in bytes) per partition for BufferExec
         /// inserted during sort pushdown optimization.
         ///

datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs

@@ -92,8 +92,9 @@ async fn test_sort_with_limited_memory() -> Result<()> {
     })
     .await?;
 
-    let total_spill_files_size =
-        metrics.spill_count().unwrap_or_default() * record_batch_size;
+    // Spilled bytes, not spill files: a merge sort spill file is about one run, while
+    // the bucket sort appends many batches to each of its range files
+    let total_spill_files_size = metrics.spilled_bytes().unwrap_or_default();
     assert!(
         total_spill_files_size > pool_size,
         "Total spill files size {total_spill_files_size} should be greater than pool size {pool_size}",
@@ -260,9 +261,11 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
         let memory_pool = Arc::new(FairSpillPool::new(pool_size));
         TaskContext::default()
             .with_session_config(
+                // Checks the merge sort's multi-level merge, not the bucket sort
                 SessionConfig::new()
                     .with_batch_size(record_batch_size)
-                    .with_sort_spill_reservation_bytes(1),
+                    .with_sort_spill_reservation_bytes(1)
+                    .set_bool("datafusion.execution.sort_bucket_sort_enabled", false),
             )
             .with_runtime(Arc::new(
                 RuntimeEnvBuilder::new()
@@ -681,8 +684,9 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory() -> Result<()
         })
         .await?;
 
-    let total_spill_files_size =
-        metrics.spill_count().unwrap_or_default() * record_batch_size;
+    // Spilled bytes, not spill files: a merge sort spill file is about one run, while
+    // the bucket sort appends many batches to each of its range files
+    let total_spill_files_size = metrics.spilled_bytes().unwrap_or_default();
     assert!(
         total_spill_files_size > pool_size,
         "Total spill files size {total_spill_files_size} should be greater than pool size {pool_size}",