You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2024/02/01 21:04:40 UTC
(arrow-datafusion) branch main updated: Split count_distinct.rs into separate modules (#9087)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8b50774d6a Split count_distinct.rs into separate modules (#9087)
8b50774d6a is described below
commit 8b50774d6afa04f511b56019febc62fd60d26234
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Thu Feb 1 16:04:35 2024 -0500
Split count_distinct.rs into separate modules (#9087)
* Split count_distinct.rs into separate modules
* Remove unecessary typedef
* Rename
* improve module comments
---
.../src/aggregate/count_distinct/mod.rs | 257 ++++-----------------
.../src/aggregate/count_distinct/native.rs | 215 +++++++++++++++++
.../src/aggregate/count_distinct/strings.rs | 6 +-
3 files changed, 261 insertions(+), 217 deletions(-)
diff --git a/datafusion/physical-expr/src/aggregate/count_distinct/mod.rs b/datafusion/physical-expr/src/aggregate/count_distinct/mod.rs
index 41f7c8729e..8baea511c7 100644
--- a/datafusion/physical-expr/src/aggregate/count_distinct/mod.rs
+++ b/datafusion/physical-expr/src/aggregate/count_distinct/mod.rs
@@ -15,39 +15,36 @@
// specific language governing permissions and limitations
// under the License.
+mod native;
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 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,
+ 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 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::native::{
+ FloatDistinctCountAccumulator, PrimitiveDistinctCountAccumulator,
+};
use crate::aggregate::count_distinct::strings::StringDistinctCountAccumulator;
-use crate::aggregate::utils::{down_cast_any_ref, Hashable};
+use crate::aggregate::utils::down_cast_any_ref;
use crate::expressions::format_state_name;
use crate::{AggregateExpr, PhysicalExpr};
-type DistinctScalarValues = ScalarValue;
-
/// Expression for a COUNT(DISTINCT) aggregation.
#[derive(Debug)]
pub struct DistinctCount {
@@ -101,46 +98,46 @@ impl AggregateExpr for DistinctCount {
use TimeUnit::*;
Ok(match &self.state_data_type {
- Int8 => Box::new(NativeDistinctCountAccumulator::<Int8Type>::new()),
- Int16 => Box::new(NativeDistinctCountAccumulator::<Int16Type>::new()),
- Int32 => Box::new(NativeDistinctCountAccumulator::<Int32Type>::new()),
- Int64 => Box::new(NativeDistinctCountAccumulator::<Int64Type>::new()),
- UInt8 => Box::new(NativeDistinctCountAccumulator::<UInt8Type>::new()),
- UInt16 => Box::new(NativeDistinctCountAccumulator::<UInt16Type>::new()),
- UInt32 => Box::new(NativeDistinctCountAccumulator::<UInt32Type>::new()),
- UInt64 => Box::new(NativeDistinctCountAccumulator::<UInt64Type>::new()),
+ Int8 => Box::new(PrimitiveDistinctCountAccumulator::<Int8Type>::new()),
+ Int16 => Box::new(PrimitiveDistinctCountAccumulator::<Int16Type>::new()),
+ Int32 => Box::new(PrimitiveDistinctCountAccumulator::<Int32Type>::new()),
+ Int64 => Box::new(PrimitiveDistinctCountAccumulator::<Int64Type>::new()),
+ UInt8 => Box::new(PrimitiveDistinctCountAccumulator::<UInt8Type>::new()),
+ UInt16 => Box::new(PrimitiveDistinctCountAccumulator::<UInt16Type>::new()),
+ UInt32 => Box::new(PrimitiveDistinctCountAccumulator::<UInt32Type>::new()),
+ UInt64 => Box::new(PrimitiveDistinctCountAccumulator::<UInt64Type>::new()),
Decimal128(_, _) => {
- Box::new(NativeDistinctCountAccumulator::<Decimal128Type>::new())
+ Box::new(PrimitiveDistinctCountAccumulator::<Decimal128Type>::new())
}
Decimal256(_, _) => {
- Box::new(NativeDistinctCountAccumulator::<Decimal256Type>::new())
+ Box::new(PrimitiveDistinctCountAccumulator::<Decimal256Type>::new())
}
- Date32 => Box::new(NativeDistinctCountAccumulator::<Date32Type>::new()),
- Date64 => Box::new(NativeDistinctCountAccumulator::<Date64Type>::new()),
- Time32(Millisecond) => {
- Box::new(NativeDistinctCountAccumulator::<Time32MillisecondType>::new())
- }
+ Date32 => Box::new(PrimitiveDistinctCountAccumulator::<Date32Type>::new()),
+ Date64 => Box::new(PrimitiveDistinctCountAccumulator::<Date64Type>::new()),
+ Time32(Millisecond) => Box::new(PrimitiveDistinctCountAccumulator::<
+ Time32MillisecondType,
+ >::new()),
Time32(Second) => {
- Box::new(NativeDistinctCountAccumulator::<Time32SecondType>::new())
- }
- Time64(Microsecond) => {
- Box::new(NativeDistinctCountAccumulator::<Time64MicrosecondType>::new())
+ Box::new(PrimitiveDistinctCountAccumulator::<Time32SecondType>::new())
}
+ Time64(Microsecond) => Box::new(PrimitiveDistinctCountAccumulator::<
+ Time64MicrosecondType,
+ >::new()),
Time64(Nanosecond) => {
- Box::new(NativeDistinctCountAccumulator::<Time64NanosecondType>::new())
+ Box::new(PrimitiveDistinctCountAccumulator::<Time64NanosecondType>::new())
}
- Timestamp(Microsecond, _) => Box::new(NativeDistinctCountAccumulator::<
+ Timestamp(Microsecond, _) => Box::new(PrimitiveDistinctCountAccumulator::<
TimestampMicrosecondType,
>::new()),
- Timestamp(Millisecond, _) => Box::new(NativeDistinctCountAccumulator::<
+ Timestamp(Millisecond, _) => Box::new(PrimitiveDistinctCountAccumulator::<
TimestampMillisecondType,
>::new()),
- Timestamp(Nanosecond, _) => {
- Box::new(NativeDistinctCountAccumulator::<TimestampNanosecondType>::new())
- }
+ Timestamp(Nanosecond, _) => Box::new(PrimitiveDistinctCountAccumulator::<
+ TimestampNanosecondType,
+ >::new()),
Timestamp(Second, _) => {
- Box::new(NativeDistinctCountAccumulator::<TimestampSecondType>::new())
+ Box::new(PrimitiveDistinctCountAccumulator::<TimestampSecondType>::new())
}
Float16 => Box::new(FloatDistinctCountAccumulator::<Float16Type>::new()),
@@ -175,9 +172,13 @@ impl PartialEq<dyn Any> for DistinctCount {
}
}
+/// General purpose distinct accumulator that works for any DataType by using
+/// [`ScalarValue`]. Some types have specialized accumulators that are (much)
+/// more efficient such as [`PrimitiveDistinctCountAccumulator`] and
+/// [`StringDistinctCountAccumulator`]
#[derive(Debug)]
struct DistinctCountAccumulator {
- values: HashSet<DistinctScalarValues, RandomState>,
+ values: HashSet<ScalarValue, RandomState>,
state_data_type: DataType,
}
@@ -186,7 +187,7 @@ impl DistinctCountAccumulator {
// This method is faster than .full_size(), however it is not suitable for variable length values like strings or complex types
fn fixed_size(&self) -> usize {
std::mem::size_of_val(self)
- + (std::mem::size_of::<DistinctScalarValues>() * self.values.capacity())
+ + (std::mem::size_of::<ScalarValue>() * self.values.capacity())
+ self
.values
.iter()
@@ -199,7 +200,7 @@ impl DistinctCountAccumulator {
// calculates the size as accurate as possible, call to this method is expensive
fn full_size(&self) -> usize {
std::mem::size_of_val(self)
- + (std::mem::size_of::<DistinctScalarValues>() * self.values.capacity())
+ + (std::mem::size_of::<ScalarValue>() * self.values.capacity())
+ self
.values
.iter()
@@ -260,182 +261,6 @@ impl Accumulator for DistinctCountAccumulator {
}
}
-#[derive(Debug)]
-struct NativeDistinctCountAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
- T::Native: Eq + Hash,
-{
- values: HashSet<T::Native, RandomState>,
-}
-
-impl<T> NativeDistinctCountAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
- T::Native: Eq + Hash,
-{
- fn new() -> Self {
- Self {
- values: HashSet::default(),
- }
- }
-}
-
-impl<T> Accumulator for NativeDistinctCountAccumulator<T>
-where
- T: ArrowPrimitiveType + Send + Debug,
- T::Native: Eq + Hash,
-{
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(
- self.values.iter().cloned(),
- )) as ArrayRef;
- let list = Arc::new(array_into_list_array(arr));
- Ok(vec![ScalarValue::List(list)])
- }
-
- fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- if values.is_empty() {
- return Ok(());
- }
-
- let arr = as_primitive_array::<T>(&values[0])?;
- arr.iter().for_each(|value| {
- if let Some(value) = value {
- self.values.insert(value);
- }
- });
-
- Ok(())
- }
-
- fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- if states.is_empty() {
- return Ok(());
- }
- 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| {
- if let Some(list) = maybe_list {
- let list = as_primitive_array::<T>(&list)?;
- self.values.extend(list.values())
- };
- Ok(())
- })
- }
-
- fn evaluate(&mut self) -> Result<ScalarValue> {
- Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
- }
-
- fn size(&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)]
-struct FloatDistinctCountAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
-{
- values: HashSet<Hashable<T::Native>, RandomState>,
-}
-
-impl<T> FloatDistinctCountAccumulator<T>
-where
- T: ArrowPrimitiveType + Send,
-{
- fn new() -> Self {
- Self {
- values: HashSet::default(),
- }
- }
-}
-
-impl<T> Accumulator for FloatDistinctCountAccumulator<T>
-where
- T: ArrowPrimitiveType + Send + Debug,
-{
- fn state(&mut self) -> Result<Vec<ScalarValue>> {
- let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(
- self.values.iter().map(|v| v.0),
- )) as ArrayRef;
- let list = Arc::new(array_into_list_array(arr));
- Ok(vec![ScalarValue::List(list)])
- }
-
- fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- if values.is_empty() {
- return Ok(());
- }
-
- let arr = as_primitive_array::<T>(&values[0])?;
- arr.iter().for_each(|value| {
- if let Some(value) = value {
- self.values.insert(Hashable(value));
- }
- });
-
- Ok(())
- }
-
- fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
- if states.is_empty() {
- return Ok(());
- }
- 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| {
- if let Some(list) = maybe_list {
- let list = as_primitive_array::<T>(&list)?;
- self.values
- .extend(list.values().iter().map(|v| Hashable(*v)));
- };
- Ok(())
- })
- }
-
- fn evaluate(&mut self) -> Result<ScalarValue> {
- Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
- }
-
- fn size(&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)
- }
-}
-
#[cfg(test)]
mod tests {
use arrow::array::{
diff --git a/datafusion/physical-expr/src/aggregate/count_distinct/native.rs b/datafusion/physical-expr/src/aggregate/count_distinct/native.rs
new file mode 100644
index 0000000000..a44e8b772e
--- /dev/null
+++ b/datafusion/physical-expr/src/aggregate/count_distinct/native.rs
@@ -0,0 +1,215 @@
+// 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.
+
+//! Specialized implementation of `COUNT DISTINCT` for "Native" arrays such as
+//! [`Int64Array`] and [`Float64Array`]
+//!
+//! [`Int64Array`]: arrow::array::Int64Array
+//! [`Float64Array`]: arrow::array::Float64Array
+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::ArrayRef;
+use arrow_array::types::ArrowPrimitiveType;
+use arrow_array::PrimitiveArray;
+
+use datafusion_common::cast::{as_list_array, as_primitive_array};
+use datafusion_common::utils::array_into_list_array;
+use datafusion_common::ScalarValue;
+use datafusion_expr::Accumulator;
+
+use crate::aggregate::utils::Hashable;
+
+#[derive(Debug)]
+pub(super) struct PrimitiveDistinctCountAccumulator<T>
+where
+ T: ArrowPrimitiveType + Send,
+ T::Native: Eq + Hash,
+{
+ values: HashSet<T::Native, RandomState>,
+}
+
+impl<T> PrimitiveDistinctCountAccumulator<T>
+where
+ T: ArrowPrimitiveType + Send,
+ T::Native: Eq + Hash,
+{
+ pub(super) fn new() -> Self {
+ Self {
+ values: HashSet::default(),
+ }
+ }
+}
+
+impl<T> Accumulator for PrimitiveDistinctCountAccumulator<T>
+where
+ T: ArrowPrimitiveType + Send + Debug,
+ T::Native: Eq + Hash,
+{
+ fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+ let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(
+ self.values.iter().cloned(),
+ )) as ArrayRef;
+ let list = Arc::new(array_into_list_array(arr));
+ Ok(vec![ScalarValue::List(list)])
+ }
+
+ fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+ if values.is_empty() {
+ return Ok(());
+ }
+
+ let arr = as_primitive_array::<T>(&values[0])?;
+ arr.iter().for_each(|value| {
+ if let Some(value) = value {
+ self.values.insert(value);
+ }
+ });
+
+ Ok(())
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+ if states.is_empty() {
+ return Ok(());
+ }
+ 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| {
+ if let Some(list) = maybe_list {
+ let list = as_primitive_array::<T>(&list)?;
+ self.values.extend(list.values())
+ };
+ Ok(())
+ })
+ }
+
+ fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+ Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
+ }
+
+ fn size(&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)]
+pub(super) struct FloatDistinctCountAccumulator<T>
+where
+ T: ArrowPrimitiveType + Send,
+{
+ values: HashSet<Hashable<T::Native>, RandomState>,
+}
+
+impl<T> FloatDistinctCountAccumulator<T>
+where
+ T: ArrowPrimitiveType + Send,
+{
+ pub(super) fn new() -> Self {
+ Self {
+ values: HashSet::default(),
+ }
+ }
+}
+
+impl<T> Accumulator for FloatDistinctCountAccumulator<T>
+where
+ T: ArrowPrimitiveType + Send + Debug,
+{
+ fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+ let arr = Arc::new(PrimitiveArray::<T>::from_iter_values(
+ self.values.iter().map(|v| v.0),
+ )) as ArrayRef;
+ let list = Arc::new(array_into_list_array(arr));
+ Ok(vec![ScalarValue::List(list)])
+ }
+
+ fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+ if values.is_empty() {
+ return Ok(());
+ }
+
+ let arr = as_primitive_array::<T>(&values[0])?;
+ arr.iter().for_each(|value| {
+ if let Some(value) = value {
+ self.values.insert(Hashable(value));
+ }
+ });
+
+ Ok(())
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+ if states.is_empty() {
+ return Ok(());
+ }
+ 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| {
+ if let Some(list) = maybe_list {
+ let list = as_primitive_array::<T>(&list)?;
+ self.values
+ .extend(list.values().iter().map(|v| Hashable(*v)));
+ };
+ Ok(())
+ })
+ }
+
+ fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+ Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
+ }
+
+ fn size(&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)
+ }
+}
diff --git a/datafusion/physical-expr/src/aggregate/count_distinct/strings.rs b/datafusion/physical-expr/src/aggregate/count_distinct/strings.rs
index d7a9ea5c37..02d30c3506 100644
--- a/datafusion/physical-expr/src/aggregate/count_distinct/strings.rs
+++ b/datafusion/physical-expr/src/aggregate/count_distinct/strings.rs
@@ -15,7 +15,11 @@
// specific language governing permissions and limitations
// under the License.
-//! Specialized implementation of `COUNT DISTINCT` for `StringArray` and `LargeStringArray`
+//! Specialized implementation of `COUNT DISTINCT` for [`StringArray`]
+//! and [`LargeStringArray`]
+//!
+//! [`StringArray`]: arrow::array::StringArray
+//! [`LargeStringArray`]: arrow::array::LargeStringArray
use ahash::RandomState;
use arrow_array::cast::AsArray;