fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group

datafusion/core/tests/memory_limit/mod.rs

@@ -26,7 +26,10 @@ mod nlj_spill_unmatched;
 mod repartition_mem_limit;
 mod union_nullable_spill;
 mod view_spill_compaction;
-use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray};
+use arrow::array::{
+    ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringArray,
+    StringViewArray,
+};
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Int32Type, SchemaRef};
 use arrow_schema::{DataType, Field, Schema};
@@ -125,6 +128,68 @@ async fn group_by_hash() {
         .await
 }
 
+/// A grouped `COUNT(DISTINCT <string>)` gets one accumulator per group, and
+/// each of those owns a hash set of the distinct values it has seen. Those
+/// sets used to be created pre-allocated, which costs far more than the
+/// handful of values a group typically holds, so the query's memory use
+/// tracked the number of groups rather than the amount of data.
+///
+/// With 4,000 groups holding 2 distinct values each, this query needed about
+/// 35.5 MB of budget before the per group pre-allocation was removed and
+/// about 1.9 MB after, so an 8 MB limit is a failure before the change and a
+/// success after it. Spilling is disabled, so completing means the query
+/// genuinely fit in the budget.
+///
+/// The `avg(payload)` is load bearing, and `avg` specifically. Without a
+/// second aggregate, `single_distinct_aggregation_to_group_by` rewrites the
+/// distinct aggregate into a plain two stage `GROUP BY` that does not use
+/// these accumulators at all. That rule tolerates a non-distinct `sum`, `min`
+/// or `max` beside the distinct aggregate, because it re-aggregates its own
+/// partial results over the deduplicated inner group by, and those three
+/// compose with themselves. `avg` does not: averaging per group averages of
+/// different sizes gives the wrong answer, so the rule can never accept it.
+/// That is why ClickBench Q9 keeps its distinct aggregate. Do not replace
+/// this with `count(*)`: `count` is only incidentally rejected today, and
+/// <https://github.com/apache/datafusion/pull/24859> proposes accepting it,
+/// which would rewrite the query and leave this test passing by construction.
+#[tokio::test]
+async fn group_by_count_distinct_utf8() {
+    TestCase::new()
+        .with_query(
+            "select group_key, count(distinct value), avg(payload) from t group by group_key",
+        )
+        .with_scenario(Scenario::GroupedDistinctStrings {
+            groups: 4_000,
+            string_view: false,
+        })
+        .with_config(SessionConfig::new().with_target_partitions(1))
+        .with_memory_limit(8_000_000)
+        .with_expected_success()
+        .run()
+        .await
+}
+
+/// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering
+/// the separate view flavoured hash set. The same query over a `Utf8View`
+/// column needed about 123 MB of budget before the change and about 2.5 MB
+/// after, so 16 MB separates the two.
+#[tokio::test]
+async fn group_by_count_distinct_utf8_view() {
+    TestCase::new()
+        .with_query(
+            "select group_key, count(distinct value), avg(payload) from t group by group_key",
+        )
+        .with_scenario(Scenario::GroupedDistinctStrings {
+            groups: 4_000,
+            string_view: true,
+        })
+        .with_config(SessionConfig::new().with_target_partitions(1))
+        .with_memory_limit(16_000_000)
+        .with_expected_success()
+        .run()
+        .await
+}
+
 #[tokio::test]
 async fn join_by_key_multiple_partitions() {
     let config = SessionConfig::new().with_target_partitions(2);
@@ -982,6 +1047,14 @@ enum Scenario {
         /// If true, splits all input batches into 1 row each
         single_row_batches: bool,
     },
+
+    /// `groups` distinct integer keys paired with a short string value, for
+    /// grouped aggregates that build one accumulator per group.
+    GroupedDistinctStrings {
+        groups: usize,
+        /// If true, the value column is `Utf8View` rather than `Utf8`
+        string_view: bool,
+    },
 }
 
 impl Scenario {
@@ -1056,6 +1129,15 @@ impl Scenario {
                 let table = SortedTableProvider::new(batches, sort_information);
                 Arc::new(table)
             }
+            Self::GroupedDistinctStrings {
+                groups,
+                string_view,
+            } => {
+                let batches = grouped_distinct_string_batches(*groups, *string_view);
+                let table =
+                    MemTable::try_new(batches[0].schema(), vec![batches]).unwrap();
+                Arc::new(table)
+            }
         }
     }
 
@@ -1079,8 +1161,62 @@ impl Scenario {
                 // Use default rules
                 None
             }
+            Self::GroupedDistinctStrings { .. } => {
+                // Disable the rules that would add a repartition, so the test
+                // measures the aggregate's budget rather than a repartition's
+                Some(vec![Arc::new(JoinSelection::new())])
+            }
+        }
+    }
+}
+
+/// Number of distinct string values held by every group produced by
+/// [`grouped_distinct_string_batches`]
+const DISTINCT_VALUES_PER_GROUP: usize = 2;
+
+/// Returns batches of 1024 rows with `groups` distinct keys in `group_key`,
+/// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings
+/// in `value` and an `Int64` `payload` to aggregate over. The values are
+/// `Utf8View` if `string_view` is set, `Utf8` otherwise.
+fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec<RecordBatch> {
+    let value_type = if string_view {
+        DataType::Utf8View
+    } else {
+        DataType::Utf8
+    };
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("group_key", DataType::Int32, false),
+        Field::new("value", value_type, false),
+        Field::new("payload", DataType::Int64, false),
+    ]));
+
+    const ROWS_PER_BATCH: usize = 1024;
+
+    let rows = groups * DISTINCT_VALUES_PER_GROUP;
+    let mut keys = Vec::with_capacity(rows);
+    let mut values = Vec::with_capacity(rows);
+    for value in 0..DISTINCT_VALUES_PER_GROUP {
+        for group in 0..groups {
+            keys.push(group as i32);
+            values.push(format!("value-{value}"));
         }
     }
+
+    keys.chunks(ROWS_PER_BATCH)
+        .zip(values.chunks(ROWS_PER_BATCH))
+        .map(|(keys, values)| {
+            let payloads: ArrayRef =
+                Arc::new(Int64Array::from_iter_values(keys.iter().map(|k| *k as i64)));
+            let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec()));
+            let values: ArrayRef = if string_view {
+                Arc::new(StringViewArray::from_iter_values(values))
+            } else {
+                Arc::new(StringArray::from_iter_values(values))
+            };
+            RecordBatch::try_new(Arc::clone(&schema), vec![keys, values, payloads])
+                .unwrap()
+        })
+        .collect()
 }
 
 fn access_log_batches() -> Vec<RecordBatch> {

datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs

@@ -39,6 +39,18 @@ use std::mem::size_of_val;
 pub struct BytesDistinctCountAccumulator<O: OffsetSizeTrait>(ArrowBytesSet<O>);
 
 impl<O: OffsetSizeTrait> BytesDistinctCountAccumulator<O> {
+    /// The set deliberately does not pre-allocate. `GroupsAccumulatorAdapter`
+    /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a
+    /// high cardinality key holds hundreds of thousands of these at once and
+    /// most of them see only a handful of values.
+    ///
+    /// The ungrouped path builds one of these and grows it to hold every
+    /// distinct value in the input, so it is the caller that had a use for the
+    /// warm up. It loses nothing here: the set grows into exactly the
+    /// capacities a pre-allocated one reaches, which
+    /// `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`
+    /// pins. That is why this constructor needs no signal distinguishing the
+    /// two callers.
     pub fn new(output_type: OutputType) -> Self {
         Self(ArrowBytesSet::new(output_type))
     }
@@ -100,6 +112,8 @@ impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
 pub struct BytesViewDistinctCountAccumulator(ArrowBytesViewSet);
 
 impl BytesViewDistinctCountAccumulator {
+    /// See [`BytesDistinctCountAccumulator::new`] for why the set does not
+    /// pre-allocate.
     pub fn new(output_type: OutputType) -> Self {
         Self(ArrowBytesViewSet::new(output_type))
     }
@@ -151,3 +165,144 @@ impl Accumulator for BytesViewDistinctCountAccumulator {
         size_of_val(self) + self.0.size()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{StringArray, StringViewArray};
+    use datafusion_physical_expr_common::binary_map::INITIAL_MAP_CAPACITY;
+    use datafusion_physical_expr_common::binary_view_map::INITIAL_MAP_CAPACITY as INITIAL_VIEW_MAP_CAPACITY;
+    use std::sync::Arc;
+
+    /// The batch size the aggregate stream drives an ungrouped accumulator with.
+    const BATCH_SIZE: usize = 8192;
+
+    /// Distinct value counts spanning both sides of the warm up capacity, up to
+    /// ones where the two constructors have converged.
+    const CARDINALITIES: [usize; 7] = [0, 1, 100, 1_000, 10_000, 100_000, 500_000];
+
+    /// Cardinalities small enough that the warm up dominates what the set
+    /// holds. This is the per group population, where `GroupsAccumulatorAdapter`
+    /// holds one accumulator per group and most see a handful of values.
+    const PER_GROUP_SCALE: usize = 100;
+
+    /// Cardinalities at which a lazily built set has grown into exactly the
+    /// capacities a pre-allocated one reaches. This is the ungrouped
+    /// population, one set holding every distinct value in the input.
+    const UNGROUPED_SCALE: usize = 10_000;
+
+    /// Longer than the map's inline value length, so the value lands in the
+    /// value buffer rather than inside the hash table entry.
+    fn distinct_value(i: usize) -> String {
+        format!("distinct value number {i}")
+    }
+
+    fn batches(distinct_values: usize, view: bool) -> Vec<ArrayRef> {
+        (0..distinct_values)
+            .step_by(BATCH_SIZE)
+            .map(|start| {
+                let values = (start..(start + BATCH_SIZE).min(distinct_values))
+                    .map(distinct_value);
+                if view {
+                    Arc::new(StringViewArray::from_iter_values(values)) as ArrayRef
+                } else {
+                    Arc::new(StringArray::from_iter_values(values)) as ArrayRef
+                }
+            })
+            .collect()
+    }
+
+    /// The property that decides whether the ungrouped path can afford to share
+    /// the lazy constructor with the per group path.
+    fn assert_lazy_is_not_worse(
+        distinct_values: usize,
+        lazy_size: usize,
+        pre_allocated_size: usize,
+    ) {
+        // The guarantee that lets both paths share one lazy constructor.
+        assert!(
+            lazy_size <= pre_allocated_size,
+            "at {distinct_values} distinct values the lazy set reported \
+             {lazy_size} bytes against the {pre_allocated_size} bytes the \
+             pre-allocated one reported"
+        );
+
+        if distinct_values <= PER_GROUP_SCALE {
+            // The warm up is pure overhead here, and removing it is the whole
+            // point of the change.
+            assert!(
+                lazy_size < pre_allocated_size,
+                "at {distinct_values} distinct values the lazy set should be \
+                 strictly cheaper, but reported {lazy_size} bytes against \
+                 {pre_allocated_size}"
+            );
+        }
+
+        if distinct_values >= UNGROUPED_SCALE {
+            // The hash table's bucket count is a power of two fixed by the
+            // number of entries, and the value buffer grows on a power of two
+            // ladder, so a set that starts empty lands on exactly the
+            // capacities a pre-allocated one reaches. The ungrouped path gives
+            // up nothing by starting empty, which is why these accumulators
+            // need no signal telling them apart from the per group ones.
+            assert_eq!(
+                lazy_size, pre_allocated_size,
+                "at {distinct_values} distinct values the lazy and pre-allocated \
+                 sets should have converged"
+            );
+        }
+    }
+
+    /// An ungrouped `COUNT(DISTINCT <string>)` builds a single accumulator that
+    /// grows to hold every distinct value in its input, which is the population
+    /// the warm up existed for. It must be no more expensive without one.
+    #[test]
+    fn ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set() {
+        for distinct_values in CARDINALITIES {
+            let mut accumulator =
+                BytesDistinctCountAccumulator::<i32>::new(OutputType::Utf8);
+            let mut pre_allocated = ArrowBytesSet::<i32>::with_capacity(
+                OutputType::Utf8,
+                INITIAL_MAP_CAPACITY,
+            );
+
+            for batch in batches(distinct_values, false) {
+                accumulator.update_batch(&[Arc::clone(&batch)]).unwrap();
+                pre_allocated.insert(&batch);
+            }
+
+            assert_eq!(accumulator.0.non_null_len(), distinct_values);
+            assert_lazy_is_not_worse(
+                distinct_values,
+                accumulator.0.size(),
+                pre_allocated.size(),
+            );
+        }
+    }
+
+    /// The `Utf8View` counterpart of
+    /// [`ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`].
+    #[test]
+    fn ungrouped_utf8_view_accumulator_is_never_worse_than_a_pre_allocated_set() {
+        for distinct_values in CARDINALITIES {
+            let mut accumulator =
+                BytesViewDistinctCountAccumulator::new(OutputType::Utf8View);
+            let mut pre_allocated = ArrowBytesViewSet::with_capacity(
+                OutputType::Utf8View,
+                INITIAL_VIEW_MAP_CAPACITY,
+            );
+
+            for batch in batches(distinct_values, true) {
+                accumulator.update_batch(&[Arc::clone(&batch)]).unwrap();
+                pre_allocated.insert(&batch);
+            }
+
+            assert_eq!(accumulator.0.non_null_len(), distinct_values);
+            assert_lazy_is_not_worse(
+                distinct_values,
+                accumulator.0.size(),
+                pre_allocated.size(),
+            );
+        }
+    }
+}

datafusion/physical-expr-common/benches/arrow_bytes_map.rs

@@ -17,7 +17,9 @@
 
 use arrow::array::{ArrayRef, StringArray};
 use criterion::{Criterion, Throughput, criterion_group, criterion_main};
-use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType};
+use datafusion_physical_expr_common::binary_map::{
+    ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType,
+};
 use std::hint::black_box;
 use std::sync::Arc;
 
@@ -57,7 +59,13 @@ fn bench_arrow_bytes_map(c: &mut Criterion) {
     for (name, values) in cases {
         group.bench_function(name, |b| {
             b.iter(|| {
-                let mut map = ArrowBytesMap::<i32, usize>::new(OutputType::Utf8);
+                // The `long_low_cardinality` case is defined by the distinct
+                // values fitting in the pre-allocated buffer, so this benchmark
+                // measures the pre-allocating constructor.
+                let mut map = ArrowBytesMap::<i32, usize>::with_capacity(
+                    OutputType::Utf8,
+                    INITIAL_MAP_CAPACITY,
+                );
                 let mut next_payload = 0;
                 map.insert_if_new(
                     &values,