Uh oh!
There was an error while loading. Please reload this page.
Optimize COUNT( DISTINCT ...) for strings (up to 9x faster) - #8849
Merged
Conversation
jayzhan211
commented
Jan 13, 2024
Uh oh!
There was an error while loading. Please reload this page.
jayzhan211
commented
Jan 13, 2024
Uh oh!
There was an error while loading. Please reload this page.
jayzhan211
commented
Jan 13, 2024
| const SHORT_STRING_LEN: usize = mem::size_of::<usize>(); | ||
| #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)] |
ContributorAuthor
There was a problem hiding this comment.
Allow Copy since they are all native types
alamb
commented
Jan 13, 2024
Contributor
Thanks @jayzhan211 -- looks basically on the right track. Is there any chance you can run some sort of benchmark on this code? My thinking is that we should get benchmark results showing that the idea actually improves performance before spending too much time polishing I looked at ClickBench and I don't actually think there are any queries that do Q8 looks like it should be helped SELECTCOUNT(DISTINCT "SearchPhrase") FROM hits;however, I am pretty sure datfusion rewrites this query to avoid the distinct with Maybe you could try manually runing a query that can't be rewritten (throw un multiple ❯ SELECTCOUNT(DISTINCT "SearchPhrase"),
COUNT(DISTINCT "MobilePhone"),
COUNT(DISTINCT "MobilePhoneModel")
FROM'hits.parquet';
+-------------------------------------------+------------------------------------------+-----------------------------------------------+
| COUNT(DISTINCT hits.parquet.SearchPhrase) | COUNT(DISTINCT hits.parquet.MobilePhone) | COUNT(DISTINCT hits.parquet.MobilePhoneModel) |
+-------------------------------------------+------------------------------------------+-----------------------------------------------+
| 6019103 | 44 | 166 |
+-------------------------------------------+------------------------------------------+-----------------------------------------------+ |
jayzhan211
commented
Jan 14, 2024
ContributorAuthor
It looks pretty nice.  |
jayzhan211
commented
Jan 14, 2024
ContributorAuthor
Run with simple HashSet has the similar result... I think I need to test with lots of short string data |
ContributorAuthor
I test with injecting string array manually to find out if SSO is better than simple HashSet Compare with n=1e5
Mostly non distinct case
All distinct case
I thought the more the small string is the more performance gains, but it shows that the more the long string is the better TLDR testing file// 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::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,StringArray};use arrow_buffer::BufferBuilder;use chrono::format;use rand::Rng;use std::any::Any;use std::cmp::Eq;use std::fmt::Debug;use std::hash::Hash;use std::mem;use std::sync::Arc;use ahash::RandomState;use arrow::array::{Array,ArrayRef};use std::collections::HashSet;usecrate::aggregate::utils::{down_cast_any_ref,Hashable};usecrate::expressions::format_state_name;usecrate::{AggregateExpr,PhysicalExpr};use datafusion_common::cast::{as_list_array, as_primitive_array, as_string_array};use datafusion_common::utils::array_into_list_array;use datafusion_common::{Result,ScalarValue};use datafusion_execution::memory_pool::proxy::RawTableAllocExt;use datafusion_expr::Accumulator;typeDistinctScalarValues = ScalarValue;/// Expression for a COUNT(DISTINCT) aggregation.#[derive(Debug)]pubstructDistinctCount{/// Column namename:String,/// The DataType used to hold the state for each inputstate_data_type:DataType,/// The input argumentsexpr:Arc<dynPhysicalExpr>,}implDistinctCount{/// Create a new COUNT(DISTINCT) aggregate function.pubfnnew(input_data_type:DataType,expr:Arc<dynPhysicalExpr>,name:String,) -> Self{Self{
name,state_data_type: input_data_type,
expr,}}}macro_rules! native_distinct_count_accumulator {($TYPE:ident) => {{Ok(Box::new(NativeDistinctCountAccumulator::<$TYPE>::new()))}};}macro_rules! float_distinct_count_accumulator {($TYPE:ident) => {{Ok(Box::new(FloatDistinctCountAccumulator::<$TYPE>::new()))}};}implAggregateExprforDistinctCount{/// Return a reference to Any that can be used for downcastingfnas_any(&self) -> &dynAny{self}fnfield(&self) -> Result<Field>{Ok(Field::new(&self.name,DataType::Int64,true))}fnstate_fields(&self) -> Result<Vec<Field>>{Ok(vec![Field::new_list(
format_state_name(&self.name,"count distinct"),Field::new("item",self.state_data_type.clone(),true),false,)])}fnexpressions(&self) -> Vec<Arc<dynPhysicalExpr>>{vec![self.expr.clone()]}fncreate_accumulator(&self) -> Result<Box<dynAccumulator>>{useDataType::*;useTimeUnit::*;match&self.state_data_type{Int8 => native_distinct_count_accumulator!(Int8Type),Int16 => native_distinct_count_accumulator!(Int16Type),Int32 => native_distinct_count_accumulator!(Int32Type),Int64 => native_distinct_count_accumulator!(Int64Type),UInt8 => native_distinct_count_accumulator!(UInt8Type),UInt16 => native_distinct_count_accumulator!(UInt16Type),UInt32 => native_distinct_count_accumulator!(UInt32Type),UInt64 => native_distinct_count_accumulator!(UInt64Type),Decimal128(_, _) => native_distinct_count_accumulator!(Decimal128Type),Decimal256(_, _) => native_distinct_count_accumulator!(Decimal256Type),Date32 => native_distinct_count_accumulator!(Date32Type),Date64 => native_distinct_count_accumulator!(Date64Type),Time32(Millisecond) => {native_distinct_count_accumulator!(Time32MillisecondType)}Time32(Second) => {native_distinct_count_accumulator!(Time32SecondType)}Time64(Microsecond) => {native_distinct_count_accumulator!(Time64MicrosecondType)}Time64(Nanosecond) => {native_distinct_count_accumulator!(Time64NanosecondType)}Timestamp(Microsecond, _) => {native_distinct_count_accumulator!(TimestampMicrosecondType)}Timestamp(Millisecond, _) => {native_distinct_count_accumulator!(TimestampMillisecondType)}Timestamp(Nanosecond, _) => {native_distinct_count_accumulator!(TimestampNanosecondType)}Timestamp(Second, _) => {native_distinct_count_accumulator!(TimestampSecondType)}Float16 => float_distinct_count_accumulator!(Float16Type),Float32 => float_distinct_count_accumulator!(Float32Type),Float64 => float_distinct_count_accumulator!(Float64Type),Utf8 => Ok(Box::new(StringDistinctCountAccumulator::new())),
_ => Ok(Box::new(DistinctCountAccumulator{values:HashSet::default(),state_data_type:self.state_data_type.clone(),})),}}fnname(&self) -> &str{&self.name}}implPartialEq<dynAny>forDistinctCount{fneq(&self,other:&dynAny) -> bool{down_cast_any_ref(other).downcast_ref::<Self>().map(|x| {self.name == x.name
&& self.state_data_type == x.state_data_type
&& self.expr.eq(&x.expr)}).unwrap_or(false)}}#[derive(Debug)]structDistinctCountAccumulator{values:HashSet<DistinctScalarValues,RandomState>,state_data_type:DataType,}implDistinctCountAccumulator{// calculating the size for fixed length values, taking first batch size * number of batches// This method is faster than .full_size(), however it is not suitable for variable length values like strings or complex typesfnfixed_size(&self) -> usize{
std::mem::size_of_val(self)
+ (std::mem::size_of::<DistinctScalarValues>()*self.values.capacity())
+ self.values.iter().next().map(|vals| ScalarValue::size(vals) - std::mem::size_of_val(vals)).unwrap_or(0)
+ std::mem::size_of::<DataType>()}// calculates the size as accurate as possible, call to this method is expensivefnfull_size(&self) -> usize{
std::mem::size_of_val(self)
+ (std::mem::size_of::<DistinctScalarValues>()*self.values.capacity())
+ self.values.iter().map(|vals| ScalarValue::size(vals) - std::mem::size_of_val(vals)).sum::<usize>()
+ std::mem::size_of::<DataType>()}}fnget_vec_string(n:usize) -> Vec<String>{// Create a vector and generate random stringsletmut rng = rand::thread_rng();let random_strings:Vec<String> = (0..n).map(|_| {let random_char = match rng.gen_range(0..3){0 => "a",1 => "b",2 => "cccccccc",
_ => unreachable!(),// This should never happen};
random_char.to_string()}).collect();
random_strings }fnget_distinct_string(n:usize) -> Vec<String>{let distinct_strings:Vec<String> = (1..=n).map(|i| {format!("{}{}", i.to_string(),"aaaaaaaa")// if i < n / 2 {// i.to_string()// } else {// }}).collect();
distinct_strings
// let mut rng = rand::thread_rng();// let random_strings: Vec<String> = (0..n)// .map(|_| {// let random_char = match rng.gen_range(0..usize::MAX) {// x => x.to_string(),// };// random_char.to_string()// })// .collect();// random_strings }implAccumulatorforDistinctCountAccumulator{fnstate(&self) -> Result<Vec<ScalarValue>>{let scalars = self.values.iter().cloned().collect::<Vec<_>>();let arr = ScalarValue::new_list(scalars.as_slice(),&self.state_data_type);Ok(vec![ScalarValue::List(arr)])}fnupdate_batch(&mutself,values:&[ArrayRef]) -> Result<()>{if values.is_empty(){returnOk(());}let arr = &values[0];if arr.data_type() == &DataType::Null{returnOk(());}(0..arr.len()).try_for_each(|index| {if !arr.is_null(index){let scalar = ScalarValue::try_from_array(arr, index)?;self.values.insert(scalar);}Ok(())})}fnmerge_batch(&mutself,states:&[ArrayRef]) -> Result<()>{if states.is_empty(){returnOk(());}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);}Ok(())}fnevaluate(&self) -> Result<ScalarValue>{Ok(ScalarValue::Int64(Some(self.values.len()asi64)))}fnsize(&self) -> usize{match&self.state_data_type{DataType::Boolean | DataType::Null => self.fixed_size(),
d if d.is_primitive() => self.fixed_size(),
_ => self.full_size(),}}}#[derive(Debug)]structNativeDistinctCountAccumulator<T>whereT:ArrowPrimitiveType + Send,T::Native:Eq + Hash,{values:HashSet<T::Native,RandomState>,}impl<T>NativeDistinctCountAccumulator<T>whereT:ArrowPrimitiveType + Send,T::Native:Eq + Hash,{fnnew() -> Self{Self{values:HashSet::default(),}}}impl<T>AccumulatorforNativeDistinctCountAccumulator<T>whereT:ArrowPrimitiveType + Send + Debug,T::Native:Eq + Hash,{fnstate(&self) -> Result<Vec<ScalarValue>>{let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(self.values.iter().cloned(),))asArrayRef;let list = Arc::new(array_into_list_array(arr));Ok(vec![ScalarValue::List(list)])}fnupdate_batch(&mutself,values:&[ArrayRef]) -> Result<()>{if values.is_empty(){returnOk(());}let arr = as_primitive_array::<T>(&values[0])?;
arr.iter().for_each(|value| {ifletSome(value) = value {self.values.insert(value);}});Ok(())}fnmerge_batch(&mutself,states:&[ArrayRef]) -> Result<()>{if states.is_empty(){returnOk(());}assert_eq!(
states.len(),1,"count_distinct states must be single array");let arr = as_list_array(&states[0])?;
arr.iter().try_for_each(|maybe_list| {ifletSome(list) = maybe_list {let list = as_primitive_array::<T>(&list)?;self.values.extend(list.values())};Ok(())})}fnevaluate(&self) -> Result<ScalarValue>{Ok(ScalarValue::Int64(Some(self.values.len()asi64)))}fnsize(&self) -> usize{let estimated_buckets = (self.values.len().checked_mul(8).unwrap_or(usize::MAX)
/ 7).next_power_of_two();// Size of accumulator// + size of entry * number of buckets// + 1 byte for each bucket// + fixed size of HashSet
std::mem::size_of_val(self)
+ std::mem::size_of::<T::Native>()* estimated_buckets
+ estimated_buckets
+ std::mem::size_of_val(&self.values)}}#[derive(Debug)]structFloatDistinctCountAccumulator<T>whereT:ArrowPrimitiveType + Send,{values:HashSet<Hashable<T::Native>,RandomState>,}impl<T>FloatDistinctCountAccumulator<T>whereT:ArrowPrimitiveType + Send,{fnnew() -> Self{Self{values:HashSet::default(),}}}impl<T>AccumulatorforFloatDistinctCountAccumulator<T>whereT:ArrowPrimitiveType + Send + Debug,{fnstate(&self) -> Result<Vec<ScalarValue>>{let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(self.values.iter().map(|v| v.0),))asArrayRef;let list = Arc::new(array_into_list_array(arr));Ok(vec![ScalarValue::List(list)])}fnupdate_batch(&mutself,values:&[ArrayRef]) -> Result<()>{if values.is_empty(){returnOk(());}let arr = as_primitive_array::<T>(&values[0])?;
arr.iter().for_each(|value| {ifletSome(value) = value {self.values.insert(Hashable(value));}});Ok(())}fnmerge_batch(&mutself,states:&[ArrayRef]) -> Result<()>{if states.is_empty(){returnOk(());}assert_eq!(
states.len(),1,"count_distinct states must be single array");let arr = as_list_array(&states[0])?;
arr.iter().try_for_each(|maybe_list| {ifletSome(list) = maybe_list {let list = as_primitive_array::<T>(&list)?;self.values.extend(list.values().iter().map(|v| Hashable(*v)));};Ok(())})}fnevaluate(&self) -> Result<ScalarValue>{Ok(ScalarValue::Int64(Some(self.values.len()asi64)))}fnsize(&self) -> usize{let estimated_buckets = (self.values.len().checked_mul(8).unwrap_or(usize::MAX)
/ 7).next_power_of_two();// Size of accumulator// + size of entry * number of buckets// + 1 byte for each bucket// + fixed size of HashSet
std::mem::size_of_val(self)
+ std::mem::size_of::<T::Native>()* estimated_buckets
+ estimated_buckets
+ std::mem::size_of_val(&self.values)}}#[derive(Debug)]structStringDistinctCountAccumulator2(HashSet<String>);implStringDistinctCountAccumulator2{fnnew() -> Self{Self(HashSet::new())}}implAccumulatorforStringDistinctCountAccumulator2{fnstate(&self) -> Result<Vec<ScalarValue>>{let arr = StringArray::from_iter_values(self.0.iter());let list = Arc::new(array_into_list_array(Arc::new(arr)));Ok(vec![ScalarValue::List(list)])}fnupdate_batch(&mutself,values:&[ArrayRef]) -> Result<()>{if values.is_empty(){returnOk(());}let vs = get_distinct_string(100000);// let vs = get_vec_string(100000);let arr = &StringArray::from_iter_values(vs);// let arr = as_string_array(&values[0])?;
arr.iter().for_each(|value| {ifletSome(value) = value {self.0.insert(value.to_string());}});Ok(())}fnmerge_batch(&mutself,states:&[ArrayRef]) -> Result<()>{if states.is_empty(){returnOk(());}assert_eq!(
states.len(),1,"count_distinct states must be single array");let arr = as_list_array(&states[0])?;
arr.iter().try_for_each(|maybe_list| {ifletSome(list) = maybe_list {let list = as_string_array(&list)?;
list.iter().for_each(|value| {ifletSome(value) = value {self.0.insert(value.to_string());}})};Ok(())})}fnevaluate(&self) -> Result<ScalarValue>{Ok(ScalarValue::Int64(Some(self.0.len()asi64)))}fnsize(&self) -> usize{// Size of accumulator// + SSOStringHashSet size
std::mem::size_of_val(self) + 0}}#[derive(Debug)]structStringDistinctCountAccumulator(SSOStringHashSet);implStringDistinctCountAccumulator{fnnew() -> Self{Self(SSOStringHashSet::new())}}implAccumulatorforStringDistinctCountAccumulator{fnstate(&self) -> Result<Vec<ScalarValue>>{let arr = StringArray::from_iter_values(self.0.iter());let list = Arc::new(array_into_list_array(Arc::new(arr)));Ok(vec![ScalarValue::List(list)])}fnupdate_batch(&mutself,values:&[ArrayRef]) -> Result<()>{if values.is_empty(){returnOk(());}let vs = get_distinct_string(100000);// let vs = get_vec_string(100000);let arr = &StringArray::from_iter_values(vs);// let arr = as_string_array(&values[0])?;
arr.iter().for_each(|value| {ifletSome(value) = value {self.0.insert(value);}});Ok(())}fnmerge_batch(&mutself,states:&[ArrayRef]) -> Result<()>{if states.is_empty(){returnOk(());}assert_eq!(
states.len(),1,"count_distinct states must be single array");let arr = as_list_array(&states[0])?;
arr.iter().try_for_each(|maybe_list| {ifletSome(list) = maybe_list {let list = as_string_array(&list)?;
list.iter().for_each(|value| {ifletSome(value) = value {self.0.insert(value);}})};Ok(())})}fnevaluate(&self) -> Result<ScalarValue>{Ok(ScalarValue::Int64(Some(self.0.len()asi64)))}fnsize(&self) -> usize{// Size of accumulator// + SSOStringHashSet size
std::mem::size_of_val(self) + self.0.size()}}constSHORT_STRING_LEN:usize = mem::size_of::<usize>();#[derive(Debug,PartialEq,Eq,Hash,Clone,Copy)]structSSOStringHeader{/// hash of the string value (used when resizing table)hash:u64,len:usize,offset_or_inline:usize,}implSSOStringHeader{fnevaluate(&self,buffer:&[u8]) -> String{ifself.len <= SHORT_STRING_LEN{self.offset_or_inline.to_string()}else{let offset = self.offset_or_inline;// SAFETY: buffer is only appended to, and we correctly inserted valuesunsafe{
std::str::from_utf8_unchecked(
buffer.get_unchecked(offset..offset + self.len),)}.to_string()}}}// Short String Optimizated HashSet for String// Equivalent to HashSet<String> but with better memory usage#[derive(Default)]structSSOStringHashSet{header_set:HashSet<SSOStringHeader>,long_string_map: hashbrown::raw::RawTable<SSOStringHeader>,map_size:usize,buffer:BufferBuilder<u8>,state:RandomState,}implSSOStringHashSet{fnnew() -> Self{Self::default()}fninsert(&mutself,value:&str){let value_len = value.len();let value_bytes = value.as_bytes();if value_len <= SHORT_STRING_LEN{let inline = value_bytes
.iter().fold(0usize, |acc,&x| acc << 8 | x asusize);let short_string_header = SSOStringHeader{// no need for short string caseshash:0,len: value_len,offset_or_inline: inline,};self.header_set.insert(short_string_header);}else{let hash = self.state.hash_one(value_bytes);let entry = self.long_string_map.get_mut(hash, |header| {// if hash matches, check if the bytes matchlet offset = header.offset_or_inline;let len = header.len;// SAFETY: buffer is only appended to, and we correctly inserted valueslet existing_value =
unsafe{self.buffer.as_slice().get_unchecked(offset..offset + len)};
value_bytes == existing_value
});if entry.is_none(){let offset = self.buffer.len();self.buffer.append_slice(value_bytes);let header = SSOStringHeader{
hash,len: value_len,offset_or_inline: offset,};self.long_string_map.insert_accounted(
header,
|header| header.hash,&mutself.map_size,);self.header_set.insert(header);}}}fniter(&self) -> Vec<String>{self.header_set.iter().map(|header| header.evaluate(self.buffer.as_slice())).collect()}fnlen(&self) -> usize{self.header_set.len()}// NEED HELPEDfnsize(&self) -> usize{self.header_set.len()* mem::size_of::<SSOStringHeader>()
+ self.map_size
+ self.buffer.len()}}implDebugforSSOStringHashSet{fnfmt(&self,f:&mut std::fmt::Formatter<'_>) -> std::fmt::Result{
f.debug_struct("SSOStringHashSet").field("header_set",&self.header_set)// TODO: Print long_string_map.field("map_size",&self.map_size).field("buffer",&self.buffer).field("state",&self.state).finish()}}#[cfg(test)]mod tests {usecrate::expressions::NoOp;usesuper::*;use arrow::array::{ArrayRef,BooleanArray,Float32Array,Float64Array,Int16Array,Int32Array,Int64Array,Int8Array,UInt16Array,UInt32Array,UInt64Array,UInt8Array,};use arrow::datatypes::DataType;use arrow::datatypes::{Float32Type,Float64Type,Int16Type,Int32Type,Int64Type,Int8Type,UInt16Type,UInt32Type,UInt64Type,UInt8Type,};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;macro_rules! state_to_vec_primitive {($LIST:expr, $DATA_TYPE:ident) => {{let arr = ScalarValue::raw_data($LIST).unwrap();let list_arr = as_list_array(&arr).unwrap();let arr = list_arr.values();let arr = as_primitive_array::<$DATA_TYPE>(arr)?;
arr.values().iter().cloned().collect::<Vec<_>>()}};}macro_rules! test_count_distinct_update_batch_numeric {($ARRAY_TYPE:ident, $DATA_TYPE:ident, $PRIM_TYPE:ty) => {{let values:Vec<Option<$PRIM_TYPE>> = vec![Some(1),Some(1),None,Some(3),Some(2),None,Some(2),Some(3),Some(1),];let arrays = vec![Arc::new($ARRAY_TYPE::from(values))asArrayRef];let(states, result) = run_update_batch(&arrays)?;letmut state_vec = state_to_vec_primitive!(&states[0], $DATA_TYPE);
state_vec.sort();
assert_eq!(states.len(),1);
assert_eq!(state_vec, vec![1,2,3]);
assert_eq!(result,ScalarValue::Int64(Some(3)));Ok(())}};}fnstate_to_vec_bool(sv:&ScalarValue) -> Result<Vec<bool>>{let arr = ScalarValue::raw_data(sv)?;let list_arr = as_list_array(&arr)?;let arr = list_arr.values();let bool_arr = as_boolean_array(arr)?;Ok(bool_arr.iter().flatten().collect())}fnrun_update_batch(arrays:&[ArrayRef]) -> Result<(Vec<ScalarValue>,ScalarValue)>{let agg = DistinctCount::new(
arrays[0].data_type().clone(),Arc::new(NoOp::new()),String::from("__col_name__"),);letmut accum = agg.create_accumulator()?;
accum.update_batch(arrays)?;Ok((accum.state()?, accum.evaluate()?))}fnrun_update(data_types:&[DataType],rows:&[Vec<ScalarValue>],) -> Result<(Vec<ScalarValue>,ScalarValue)>{let agg = DistinctCount::new(
data_types[0].clone(),Arc::new(NoOp::new()),String::from("__col_name__"),);letmut accum = agg.create_accumulator()?;let cols = (0..rows[0].len()).map(|i| {
rows.iter().map(|inner| inner[i].clone()).collect::<Vec<ScalarValue>>()}).collect::<Vec<_>>();let arrays:Vec<ArrayRef> = cols
.iter().map(|c| ScalarValue::iter_to_array(c.clone())).collect::<Result<Vec<ArrayRef>>>()?;
accum.update_batch(&arrays)?;Ok((accum.state()?, accum.evaluate()?))}// Used trait to create associated constant for f32 and f64traitSubNormal:'static{constSUBNORMAL:Self;}implSubNormalforf64{constSUBNORMAL:Self = 1.0e-308_f64;}implSubNormalforf32{constSUBNORMAL:Self = 1.0e-38_f32;}macro_rules! test_count_distinct_update_batch_floating_point {($ARRAY_TYPE:ident, $DATA_TYPE:ident, $PRIM_TYPE:ty) => {{let values:Vec<Option<$PRIM_TYPE>> = vec![Some(<$PRIM_TYPE>::INFINITY),Some(<$PRIM_TYPE>::NAN),Some(1.0),Some(<$PRIM_TYPE asSubNormal>::SUBNORMAL),Some(1.0),Some(<$PRIM_TYPE>::INFINITY),None,Some(3.0),Some(-4.5),Some(2.0),None,Some(2.0),Some(3.0),Some(<$PRIM_TYPE>::NEG_INFINITY),Some(1.0),Some(<$PRIM_TYPE>::NAN),Some(<$PRIM_TYPE>::NEG_INFINITY),];let arrays = vec![Arc::new($ARRAY_TYPE::from(values))asArrayRef];let(states, result) = run_update_batch(&arrays)?;letmut state_vec = state_to_vec_primitive!(&states[0], $DATA_TYPE);
dbg!(&state_vec);
state_vec.sort_by(|a, b| match(a, b){(lhs, rhs) => lhs.total_cmp(rhs),});let nan_idx = state_vec.len() - 1;
assert_eq!(states.len(),1);
assert_eq!(&state_vec[..nan_idx],
vec![
<$PRIM_TYPE>::NEG_INFINITY,
-4.5,
<$PRIM_TYPE asSubNormal>::SUBNORMAL,1.0,2.0,3.0,
<$PRIM_TYPE>::INFINITY]);
assert!(state_vec[nan_idx].is_nan());
assert_eq!(result,ScalarValue::Int64(Some(8)));Ok(())}};}macro_rules! test_count_distinct_update_batch_bigint {($ARRAY_TYPE:ident, $DATA_TYPE:ident, $PRIM_TYPE:ty) => {{let values:Vec<Option<$PRIM_TYPE>> = vec![Some(i256::from(1)),Some(i256::from(1)),None,Some(i256::from(3)),Some(i256::from(2)),None,Some(i256::from(2)),Some(i256::from(3)),Some(i256::from(1)),];let arrays = vec![Arc::new($ARRAY_TYPE::from(values))asArrayRef];let(states, result) = run_update_batch(&arrays)?;letmut state_vec = state_to_vec_primitive!(&states[0], $DATA_TYPE);
state_vec.sort();
assert_eq!(states.len(),1);
assert_eq!(state_vec, vec![i256::from(1), i256::from(2), i256::from(3)]);
assert_eq!(result,ScalarValue::Int64(Some(3)));Ok(())}};}#[test]fncount_distinct_update_batch_i8() -> Result<()>{test_count_distinct_update_batch_numeric!(Int8Array,Int8Type,i8)}#[test]fncount_distinct_update_batch_i16() -> Result<()>{test_count_distinct_update_batch_numeric!(Int16Array,Int16Type,i16)}#[test]fncount_distinct_update_batch_i32() -> Result<()>{test_count_distinct_update_batch_numeric!(Int32Array,Int32Type,i32)}#[test]fncount_distinct_update_batch_i64() -> Result<()>{test_count_distinct_update_batch_numeric!(Int64Array,Int64Type,i64)}#[test]fncount_distinct_update_batch_u8() -> Result<()>{test_count_distinct_update_batch_numeric!(UInt8Array,UInt8Type,u8)}#[test]fncount_distinct_update_batch_u16() -> Result<()>{test_count_distinct_update_batch_numeric!(UInt16Array,UInt16Type,u16)}#[test]fncount_distinct_update_batch_u32() -> Result<()>{test_count_distinct_update_batch_numeric!(UInt32Array,UInt32Type,u32)}#[test]fncount_distinct_update_batch_u64() -> Result<()>{test_count_distinct_update_batch_numeric!(UInt64Array,UInt64Type,u64)}#[test]fncount_distinct_update_batch_f32() -> Result<()>{test_count_distinct_update_batch_floating_point!(Float32Array,Float32Type,f32)}#[test]fncount_distinct_update_batch_f64() -> Result<()>{test_count_distinct_update_batch_floating_point!(Float64Array,Float64Type,f64)}#[test]fncount_distinct_update_batch_i256() -> Result<()>{test_count_distinct_update_batch_bigint!(Decimal256Array,Decimal256Type, i256)}#[test]fncount_distinct_update_batch_boolean() -> Result<()>{let get_count = |data:BooleanArray| -> Result<(Vec<bool>,i64)>{let arrays = vec![Arc::new(data)asArrayRef];let(states, result) = run_update_batch(&arrays)?;letmut state_vec = state_to_vec_bool(&states[0])?;
state_vec.sort();let count = match result {ScalarValue::Int64(c) => c.ok_or_else(|| {DataFusionError::Internal("Found None count".to_string())}),
scalar => {internal_err!("Found non int64 scalar value from count: {scalar}")}}?;Ok((state_vec, count))};let zero_count_values = BooleanArray::from(Vec::<bool>::new());let one_count_values = BooleanArray::from(vec![false,false]);let one_count_values_with_null =
BooleanArray::from(vec![Some(true),Some(true),None,None]);let two_count_values = BooleanArray::from(vec![true,false,true,false,true]);let two_count_values_with_null = BooleanArray::from(vec![Some(true),Some(false),None,None,Some(true),Some(false),]);assert_eq!(get_count(zero_count_values)?,(Vec::<bool>::new(),0));assert_eq!(get_count(one_count_values)?,(vec![false],1));assert_eq!(get_count(one_count_values_with_null)?,(vec![true],1));assert_eq!(get_count(two_count_values)?,(vec![false,true],2));assert_eq!(
get_count(two_count_values_with_null)?,(vec![false,true],2));Ok(())}#[test]fncount_distinct_update_batch_all_nulls() -> Result<()>{let arrays = vec![Arc::new(Int32Array::from(
vec![None,None,None,None]asVec<Option<i32>>
))asArrayRef];let(states, result) = run_update_batch(&arrays)?;let state_vec = state_to_vec_primitive!(&states[0],Int32Type);assert_eq!(states.len(),1);assert!(state_vec.is_empty());assert_eq!(result,ScalarValue::Int64(Some(0)));Ok(())}#[test]fncount_distinct_update_batch_empty() -> Result<()>{let arrays = vec![Arc::new(Int32Array::from(vec![0_i32;0]))asArrayRef];let(states, result) = run_update_batch(&arrays)?;let state_vec = state_to_vec_primitive!(&states[0],Int32Type);assert_eq!(states.len(),1);assert!(state_vec.is_empty());assert_eq!(result,ScalarValue::Int64(Some(0)));Ok(())}#[test]fncount_distinct_update() -> Result<()>{let(states, result) = run_update(&[DataType::Int32],&[vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(Some(5))],vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(Some(5))],vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(Some(2))],],)?;assert_eq!(states.len(),1);assert_eq!(result,ScalarValue::Int64(Some(3)));let(states, result) = run_update(&[DataType::UInt64],&[vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(Some(5))],vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(Some(5))],vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(Some(2))],],)?;assert_eq!(states.len(),1);assert_eq!(result,ScalarValue::Int64(Some(3)));Ok(())}#[test]fncount_distinct_update_with_nulls() -> Result<()>{let(states, result) = run_update(&[DataType::Int32],&[// None of these updates contains a None, so these are accumulated.vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(Some(-2))],// Each of these updates contains at least one None, so these// won't be accumulated.vec![ScalarValue::Int32(Some(-1))],vec![ScalarValue::Int32(None)],vec![ScalarValue::Int32(None)],],)?;assert_eq!(states.len(),1);assert_eq!(result,ScalarValue::Int64(Some(2)));let(states, result) = run_update(&[DataType::UInt64],&[// None of these updates contains a None, so these are accumulated.vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(Some(2))],// Each of these updates contains at least one None, so these// won't be accumulated.vec![ScalarValue::UInt64(Some(1))],vec![ScalarValue::UInt64(None)],vec![ScalarValue::UInt64(None)],],)?;assert_eq!(states.len(),1);assert_eq!(result,ScalarValue::Int64(Some(2)));Ok(())}} |
alamb
commented
Jan 14, 2024
Contributor
It seems to me this means that maybe the small string optimization is unnecessary at this time given it doesn't seem to make a significant different to performance 🤔 Maybe we could simplify the code ? |
Contributor
ContributorAuthor
If the number of rows is large > 1e6, then the speed gains of short strings is larger than seconds (5s faster for n=1e6) |
ContributorAuthor
hits.parquet data is not large enough where they are either len 1 or 2. This does not show difference. Hashset SSO URL length is mostly > 8. Improve from 11s to 9s HashSet SSO |
alamb
commented
Jan 15, 2024
Contributor
Nice! I'll add this query to the "extended" benchmark I am working on in #8861 |
alamb
commented
Jan 15, 2024
Contributor
@jayzhan211 -- I plan to review this PR more carefully over the next day or two I am hoping that we can then use the same structure for #7064 |
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
alamb
commented
Jan 22, 2024
Contributor
@jayzhan211 -- Thank you for 0e33b12 |
Contributor
I have an idea for this test, let me push it up and see what you think Update: 3e9289a |
alamb
reviewed
Jan 22, 2024
Uh oh!
There was an error while loading. Please reload this page.
alamb
commented
Jan 22, 2024
Contributor
Ok, it think this PR is now blocked on the following two PRs:
Once those are merged I think we can merge up from this branch, run some final benchmark numbers, and get it reviewed Thanks again @jayzhan211 |
jayzhan211
commented
Jan 23, 2024
ContributorAuthor
Thanks, @alamb. It seems I misunderstand both the goal of fuzz test and memory accounting test 😅 |
alamb
commented
Jan 24, 2024
Contributor
I think this was my bad for not explaining it well. |
alamb
commented
Jan 24, 2024
Contributor
Ok, now I am just waiting on #8950 to merge and then I'll run the benchmarks and I think this PR will be ready for review |
COUNT( DISTINCT ...) for stringsCOUNT( DISTINCT ...) for strings (up to 9x faster)alamb
marked this pull request as ready for review
January 25, 2024 14:30
alamb
approved these changes
Jan 25, 2024
alamb
left a comment
Contributor
There was a problem hiding this comment.
I think this PR is now ready for review. The benchmarks are looking very nice
Thanks again for this great teamwork @jayzhan211
Since I wrote a bunch of this PR I think another committer should also approve prior to merging it
| blake3 = { version = "1.0", optional = true } | ||
| chrono = { workspace = true } | ||
| datafusion-common = { workspace = true } | ||
| datafusion-execution = { workspace = true } |
Contributor
There was a problem hiding this comment.
Needed to use RawTableAlloc trait
| Float32 => float_distinct_count_accumulator!(Float32Type), | ||
| Float64 => float_distinct_count_accumulator!(Float64Type), | ||
| Utf8 => Ok(Box::new(StringDistinctCountAccumulator::<i32>::new())), |
Contributor
There was a problem hiding this comment.
The key contribution in this PR is to add these specialized accumulators
| /// Maximum size of a string that can be inlined in the hash table | ||
| const SHORT_STRING_LEN: usize = mem::size_of::<usize>(); | ||
| /// Entry that is stored in a `SSOStringHashSet` that represents a string |
Contributor
There was a problem hiding this comment.
This explains the core change in this PR and how things work
7 tasks
thinkharderdev
approved these changes
Jan 27, 2024
| return Ok(()); | ||
| } | ||
| self.0.insert(values[0].clone()); |
Contributor
There was a problem hiding this comment.
nit: insert should be able to take a reference right?
Contributor
There was a problem hiding this comment.
Yes, you are right -- I did this in f5e268d
Thank you
6 tasks
alamb
commented
Jan 28, 2024
Contributor
Thank you for the review @thinkharderdev 🙏 |
alamb
commented
Jan 29, 2024
Contributor
🚀 |
This was referenced Jan 31, 2024
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note for reviewers this PR has around 200 lines of code, and the rest is testing
This PR is a collaboration between @jayzhan211 and @alamb
Which issue does this PR close?
Part of #5472
Follow up on #8721
Rationale for this change
Speed up queries that include multiple
COUNT DISTINCTs forStringorLargeStringWhat changes are included in this PR?
Implement a specialized Accumulator for
COUNT DISTINCTthat avoids copying string data or allocating individual stringsAre these changes tested?
Benchmark results:
Clickbench Extended
Admittedly these benchmarks were chosen to highlight this particular change but I am still feeling pretty good with 9x faster query
. See Docs for more details about what these tests are
Entire Clickbench (basically the same)
Entire TPCH_1 (basically the same)
Benchmark tpch_mem.json
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ main_base ┃ bytes-distinctcount ┃ Change ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 219.29ms │ 208.36ms │ no change │
│ QQuery 2 │ 47.19ms │ 45.30ms │ no change │
│ QQuery 3 │ 79.84ms │ 78.33ms │ no change │
│ QQuery 4 │ 75.31ms │ 76.34ms │ no change │
│ QQuery 5 │ 125.54ms │ 126.23ms │ no change │
│ QQuery 6 │ 16.68ms │ 16.15ms │ no change │
│ QQuery 7 │ 337.79ms │ 320.61ms │ +1.05x faster │
│ QQuery 8 │ 82.13ms │ 80.68ms │ no change │
│ QQuery 9 │ 127.99ms │ 128.50ms │ no change │
│ QQuery 10 │ 158.88ms │ 155.37ms │ no change │
│ QQuery 11 │ 33.85ms │ 33.83ms │ no change │
│ QQuery 12 │ 72.23ms │ 70.90ms │ no change │
│ QQuery 13 │ 83.32ms │ 85.77ms │ no change │
│ QQuery 14 │ 26.51ms │ 26.13ms │ no change │
│ QQuery 15 │ 62.08ms │ 60.46ms │ no change │
│ QQuery 16 │ 48.43ms │ 45.78ms │ +1.06x faster │
│ QQuery 17 │ 166.20ms │ 161.16ms │ no change │
│ QQuery 18 │ 471.68ms │ 465.24ms │ no change │
│ QQuery 19 │ 65.11ms │ 65.86ms │ no change │
│ QQuery 20 │ 117.70ms │ 117.06ms │ no change │
│ QQuery 21 │ 368.86ms │ 363.66ms │ no change │
│ QQuery 22 │ 29.75ms │ 29.44ms │ no change │
└──────────────┴───────────┴─────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (main_base) │ 2816.38ms │
│ Total Time (bytes-distinctcount) │ 2761.14ms │
│ Average Time (main_base) │ 128.02ms │
│ Average Time (bytes-distinctcount) │ 125.51ms │
│ Queries Faster │ 2 │
│ Queries Slower │ 0 │
│ Queries with No Change │ 20 │
└────────────────────────────────────┴───────────┘