Handle dicts for distinct count (#15871)

* Handle dicts for distinct count

* Fix sqllogictests

* Add bench

* Fix no fix the bench

* Do not panic if error type is bad

* Add full bench query

* Set the bench

* Add dict of dict test

* Fix tests

* Rename method

* Increase the grouping test

* Increase the grouping test a bit more :)

* Fix flakiness

---------

Co-authored-by: Dmitrii Blaginin <[email protected]>

Author: blaginin

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

@@ -16,9 +16,11 @@
 // under the License.
 
 mod bytes;
+mod dict;
 mod native;
 
 pub use bytes::BytesDistinctCountAccumulator;
 pub use bytes::BytesViewDistinctCountAccumulator;
+pub use dict::DictionaryCountAccumulator;
 pub use native::FloatDistinctCountAccumulator;
 pub use native::PrimitiveDistinctCountAccumulator;

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

@@ -0,0 +1,70 @@
+// 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.
+
+use arrow::array::{ArrayRef, BooleanArray};
+use arrow::downcast_dictionary_array;
+use datafusion_common::{arrow_datafusion_err, ScalarValue};
+use datafusion_common::{internal_err, DataFusionError};
+use datafusion_expr_common::accumulator::Accumulator;
+
+#[derive(Debug)]
+pub struct DictionaryCountAccumulator {
+    inner: Box<dyn Accumulator>,
+}
+
+impl DictionaryCountAccumulator {
+    pub fn new(inner: Box<dyn Accumulator>) -> Self {
+        Self { inner }
+    }
+}
+
+impl Accumulator for DictionaryCountAccumulator {
+    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+        let values: Vec<_> = values
+            .iter()
+            .map(|dict| {
+                downcast_dictionary_array! {
+                    dict => {
+                        let buff: BooleanArray = dict.occupancy().into();
+                        arrow::compute::filter(
+                            dict.values(),
+                            &buff
+                        ).map_err(|e| arrow_datafusion_err!(e))
+                    },
+                    _ => internal_err!("DictionaryCountAccumulator only supports dictionary arrays")
+                }
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+        self.inner.update_batch(values.as_slice())
+    }
+
+    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+        self.inner.evaluate()
+    }
+
+    fn size(&self) -> usize {
+        self.inner.size()
+    }
+
+    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        self.inner.state()
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+        self.inner.merge_batch(states)
+    }
+}

datafusion/functions-aggregate/benches/count.rs

@@ -17,15 +17,20 @@
 
 use arrow::array::{ArrayRef, BooleanArray};
 use arrow::datatypes::{DataType, Field, Int32Type, Schema};
-use arrow::util::bench_util::{create_boolean_array, create_primitive_array};
+use arrow::util::bench_util::{
+    create_boolean_array, create_dict_from_values, create_primitive_array,
+    create_string_array_with_len,
+};
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use datafusion_expr::{function::AccumulatorArgs, AggregateUDFImpl, GroupsAccumulator};
+use datafusion_expr::{
+    function::AccumulatorArgs, Accumulator, AggregateUDFImpl, GroupsAccumulator,
+};
 use datafusion_functions_aggregate::count::Count;
 use datafusion_physical_expr::expressions::col;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use std::sync::Arc;
 
-fn prepare_accumulator() -> Box<dyn GroupsAccumulator> {
+fn prepare_group_accumulator() -> Box<dyn GroupsAccumulator> {
     let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Int32, true)]));
     let accumulator_args = AccumulatorArgs {
         return_field: Field::new("f", DataType::Int64, true).into(),
@@ -44,13 +49,34 @@ fn prepare_accumulator() -> Box<dyn GroupsAccumulator> {
         .unwrap()
 }
 
+fn prepare_accumulator() -> Box<dyn Accumulator> {
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "f",
+        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
+        true,
+    )]));
+    let accumulator_args = AccumulatorArgs {
+        return_field: Arc::new(Field::new_list_field(DataType::Int64, true)),
+        schema: &schema,
+        ignore_nulls: false,
+        ordering_req: &LexOrdering::default(),
+        is_reversed: false,
+        name: "COUNT(f)",
+        is_distinct: true,
+        exprs: &[col("f", &schema).unwrap()],
+    };
+    let count_fn = Count::new();
+
+    count_fn.accumulator(accumulator_args).unwrap()
+}
+
 fn convert_to_state_bench(
     c: &mut Criterion,
     name: &str,
     values: ArrayRef,
     opt_filter: Option<&BooleanArray>,
 ) {
-    let accumulator = prepare_accumulator();
+    let accumulator = prepare_group_accumulator();
     c.bench_function(name, |b| {
         b.iter(|| {
             black_box(
@@ -89,6 +115,18 @@ fn count_benchmark(c: &mut Criterion) {
         values,
         Some(&filter),
     );
+
+    let arr = create_string_array_with_len::<i32>(20, 0.0, 50);
+    let values =
+        Arc::new(create_dict_from_values::<Int32Type>(200_000, 0.8, &arr)) as ArrayRef;
+
+    let mut accumulator = prepare_accumulator();
+    c.bench_function("count low cardinality dict 20% nulls, no filter", |b| {
+        b.iter(|| {
+            #[allow(clippy::unit_arg)]
+            black_box(accumulator.update_batch(&[values.clone()]).unwrap())
+        })
+    });
 }
 
 criterion_group!(benches, count_benchmark);