Optimize COUNT( DISTINCT ...) for strings (up to 9x faster)

benchmarks/queries/clickbench/README.md

@@ -29,7 +29,6 @@ SELECT COUNT(DISTINCT "SearchPhrase"), COUNT(DISTINCT "MobilePhone"), COUNT(DIST
 FROM hits;
 ```
 
-
 ### Q1: Data Exploration
 
 **Question**: "How many distinct "hit color", "browser country" and "language" are there in the dataset?"
@@ -42,7 +41,7 @@ SELECT COUNT(DISTINCT "HitColor"), COUNT(DISTINCT "BrowserCountry"), COUNT(DISTI
 FROM hits;
 ```
 
-### Q2: Top 10 anaylsis
+### Q2: Top 10 analysis
 
 **Question**: "Find the top 10 "browser country" by number of distinct "social network"s, 
 including the distinct counts of "hit color", "browser language",

datafusion/core/tests/fuzz_cases/distinct_count_string_fuzz.rs

@@ -0,0 +1,211 @@
+// 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.
+
+//! Compare DistinctCount for string with naive HashSet and Short String Optimized HashSet
+
+use std::sync::Arc;
+
+use arrow::array::ArrayRef;
+use arrow::record_batch::RecordBatch;
+use arrow_array::{Array, GenericStringArray, OffsetSizeTrait, UInt32Array};
+
+use arrow_array::cast::AsArray;
+use datafusion::datasource::MemTable;
+use rand::rngs::StdRng;
+use rand::{thread_rng, Rng, SeedableRng};
+use std::collections::HashSet;
+use tokio::task::JoinSet;
+
+use datafusion::prelude::{SessionConfig, SessionContext};
+use test_utils::stagger_batch;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn distinct_count_string_test() {
+ // max length of generated strings
+ let mut join_set = JoinSet::new();
+ let mut rng = thread_rng();
+ for null_pct in [0.0, 0.01, 0.1, 0.5] {
+ for _ in 0..100 {
+ let max_len = rng.gen_range(1..50);
+ let num_strings = rng.gen_range(1..100);
+ let num_distinct_strings = if num_strings > 1 {
+ rng.gen_range(1..num_strings)
+ } else {
+ num_strings
+ };
+ let generator = BatchGenerator {
+ max_len,
+ num_strings,
+ num_distinct_strings,
+ null_pct,
+ rng: StdRng::from_seed(rng.gen()),
+ };
+ join_set.spawn(async move { run_distinct_count_test(generator).await });
+ }
+ }
+ while let Some(join_handle) = join_set.join_next().await {
+ // propagate errors
+ join_handle.unwrap();
+ }
+}
+
+/// Run COUNT DISTINCT using SQL and compare the result to computing the
+/// distinct count using HashSet<String>
+async fn run_distinct_count_test(mut generator: BatchGenerator) {
+ let input = generator.make_input_batches();
+
+ let schema = input[0].schema();
+ let session_config = SessionConfig::new().with_batch_size(50);
+ let ctx = SessionContext::new_with_config(session_config);
+
+ // split input into two partitions
+ let partition_len = input.len() / 2;
+ let partitions = vec![
+ input[0..partition_len].to_vec(),
+ input[partition_len..].to_vec(),
+ ];
+
+ let provider = MemTable::try_new(schema, partitions).unwrap();
+ ctx.register_table("t", Arc::new(provider)).unwrap();
+ // input has two columns, a and b. The result is the number of distinct
+ // values in each column.
+ //
+ // Note, we need at least two count distinct aggregates to trigger the
+ // count distinct aggregate. Otherwise, the optimizer will rewrite the
+ // `COUNT(DISTINCT a)` to `COUNT(*) from (SELECT DISTINCT a FROM t)`
+ let results = ctx
+ .sql("SELECT COUNT(DISTINCT a), COUNT(DISTINCT b) FROM t")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+
+ // get all the strings from the first column of the result (distinct a)
+ let expected_a = extract_distinct_strings::<i32>(&input, 0).len();
+ let result_a = extract_i64(&results, 0);
+ assert_eq!(expected_a, result_a);
+
+ // get all the strings from the second column of the result (distinct b(
+ let expected_b = extract_distinct_strings::<i64>(&input, 1).len();
+ let result_b = extract_i64(&results, 1);
+ assert_eq!(expected_b, result_b);
+}
+
+/// Return all (non null) distinct strings from column col_idx
+fn extract_distinct_strings<O: OffsetSizeTrait>(
+ results: &[RecordBatch],
+ col_idx: usize,
+) -> Vec<String> {
+ results
+ .iter()
+ .flat_map(|batch| {
+ let array = batch.column(col_idx).as_string::<O>();
+ // remove nulls via 'flatten'
+ array.iter().flatten().map(|s| s.to_string())
+ })
+ .collect::<HashSet<_>>()
+ .into_iter()
+ .collect()
+}
+
+// extract the value from the Int64 column in col_idx in batch and return
+// it as a usize
+fn extract_i64(results: &[RecordBatch], col_idx: usize) -> usize {
+ assert_eq!(results.len(), 1);
+ let array = results[0]
+ .column(col_idx)
+ .as_any()
+ .downcast_ref::<arrow::array::Int64Array>()
+ .unwrap();
+ assert_eq!(array.len(), 1);
+ assert!(!array.is_null(0));
+ array.value(0).try_into().unwrap()
+}
+
+struct BatchGenerator {
+ //// The maximum length of the strings
+ max_len: usize,
+ /// the total number of strings in the output
+ num_strings: usize,
+ /// The number of distinct strings in the columns
+ num_distinct_strings: usize,
+ /// The percentage of nulls in the columns
+ null_pct: f64,
+ /// Random number generator
+ rng: StdRng,
+}
+
+impl BatchGenerator {
+ /// Make batches of random strings with a random length columns "a" and "b":
+ ///
+ /// * "a" is a StringArray
+ /// * "b" is a LargeStringArray
+ fn make_input_batches(&mut self) -> Vec<RecordBatch> {
+ // use a random number generator to pick a random sized output
+
+ let batch = RecordBatch::try_from_iter(vec![
+ ("a", self.gen_data::<i32>()),
+ ("b", self.gen_data::<i64>()),
+ ])
+ .unwrap();
+
+ stagger_batch(batch)
+ }
+
+ /// Creates a StringArray or LargeStringArray with random strings according
+ /// to the parameters of the BatchGenerator
+ fn gen_data<O: OffsetSizeTrait>(&mut self) -> ArrayRef {
+ // table of strings from which to draw
+ let distinct_strings: GenericStringArray<O> = (0..self.num_distinct_strings)
+ .map(|_| Some(random_string(&mut self.rng, self.max_len)))
+ .collect();
+
+ // pick num_strings randomly from the distinct string table
+ let indicies: UInt32Array = (0..self.num_strings)
+ .map(|_| {
+ if self.rng.gen::<f64>() < self.null_pct {
+ None
+ } else if self.num_distinct_strings > 1 {
+ let range = 1..(self.num_distinct_strings as u32);
+ Some(self.rng.gen_range(range))
+ } else {
+ Some(0)
+ }
+ })
+ .collect();
+
+ let options = None;
+ arrow::compute::take(&distinct_strings, &indicies, options).unwrap()
+ }
+}
+
+/// Return a string of random characters of length 1..=max_len
+fn random_string(rng: &mut StdRng, max_len: usize) -> String {
+ // pick characters at random (not just ascii)
+ match max_len {
+ 0 => "".to_string(),
+ 1 => String::from(rng.gen::<char>()),
+ _ => {
+ let len = rng.gen_range(1..=max_len);
+ rng.sample_iter::<char, _>(rand::distributions::Standard)
+ .take(len)
+ .map(char::from)
+ .collect::<String>()
+ }
+ }
+}

datafusion/core/tests/fuzz_cases/mod.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 mod aggregate_fuzz;
+mod distinct_count_string_fuzz;
 mod join_fuzz;
 mod merge_fuzz;
 mod sort_fuzz;

datafusion/physical-expr/Cargo.toml

@@ -54,6 +54,7 @@ blake2 = { version = "^0.10.2", optional = true }
 blake3 = { version = "1.0", optional = true }
 chrono = { workspace = true }
 datafusion-common = { workspace = true }
+datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
 half = { version = "2.1", default-features = false }
 hashbrown = { version = "0.14", features = ["raw"] }

datafusion/physical-expr/src/aggregate/count_distinct.rs → datafusion/physical-expr/src/aggregate/count_distinct/mod.rs

@@ -15,34 +15,37 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::datatypes::{DataType, Field, TimeUnit};
-use arrow_array::types::{
- ArrowPrimitiveType, Date32Type, Date64Type, Decimal128Type, Decimal256Type,
- Float16Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
- Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
- TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
- TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
-};
-use arrow_array::PrimitiveArray;
+mod strings;
 
 use std::any::Any;
 use std::cmp::Eq;
+use std::collections::HashSet;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::sync::Arc;
 
 use ahash::RandomState;
 use arrow::array::{Array, ArrayRef};
-use std::collections::HashSet;
+use arrow::datatypes::{DataType, Field, TimeUnit};
+use arrow_array::types::{
+ ArrowPrimitiveType, Date32Type, Date64Type, Decimal128Type, Decimal256Type,
+ Float16Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
+ Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
+ TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
+ TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+};
+use arrow_array::PrimitiveArray;
 
-use crate::aggregate::utils::{down_cast_any_ref, Hashable};
-use crate::expressions::format_state_name;
-use crate::{AggregateExpr, PhysicalExpr};
 use datafusion_common::cast::{as_list_array, as_primitive_array};
 use datafusion_common::utils::array_into_list_array;
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::Accumulator;
 
+use crate::aggregate::count_distinct::strings::StringDistinctCountAccumulator;
+use crate::aggregate::utils::{down_cast_any_ref, Hashable};
+use crate::expressions::format_state_name;
+use crate::{AggregateExpr, PhysicalExpr};
+
 type DistinctScalarValues = ScalarValue;
 
 /// Expression for a COUNT(DISTINCT) aggregation.
@@ -61,10 +64,10 @@ impl DistinctCount {
  pub fn new(
  input_data_type: DataType,
  expr: Arc<dyn PhysicalExpr>,
- name: String,
+ name: impl Into<String>,
  ) -> Self {
  Self {
- name,
+ name: name.into(),
  state_data_type: input_data_type,
  expr,
  }
@@ -152,6 +155,9 @@ impl AggregateExpr for DistinctCount {
  Float32 => float_distinct_count_accumulator!(Float32Type),
  Float64 => float_distinct_count_accumulator!(Float64Type),
 
+ Utf8 => Ok(Box::new(StringDistinctCountAccumulator::<i32>::new())),
+ LargeUtf8 => Ok(Box::new(StringDistinctCountAccumulator::<i64>::new())),
+
  _ => Ok(Box::new(DistinctCountAccumulator {
  values: HashSet::default(),
  state_data_type: self.state_data_type.clone(),
@@ -244,7 +250,7 @@ impl Accumulator for DistinctCountAccumulator {
  assert_eq!(states.len(), 1, "array_agg states must be singleton!");
  let scalar_vec = ScalarValue::convert_array_to_scalar_vec(&states[0])?;
  for scalars in scalar_vec.into_iter() {
- self.values.extend(scalars)
+ self.values.extend(scalars);
  }
  Ok(())
  }
@@ -440,9 +446,6 @@ where
 
 #[cfg(test)]
 mod tests {
- use crate::expressions::NoOp;
-
- use super::*;
  use arrow::array::{
  ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
  Int64Array, Int8Array, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
@@ -454,10 +457,15 @@ mod tests {
  };
  use arrow_array::Decimal256Array;
  use arrow_buffer::i256;
+
  use datafusion_common::cast::{as_boolean_array, as_list_array, as_primitive_array};
  use datafusion_common::internal_err;
  use datafusion_common::DataFusionError;
 
+ use crate::expressions::NoOp;
+
+ use super::*;
+
  macro_rules! state_to_vec_primitive {
  ($LIST:expr, $DATA_TYPE:ident) => {{
  let arr = ScalarValue::raw_data($LIST).unwrap();