Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement special min/max accumulator for Strings and Binary (10% faster for Clickbench Q28) #12792

Merged
merged 20 commits into from
Oct 13, 2024
Merged

Implement special min/max accumulator for Strings and Binary (10% faster for Clickbench Q28) #12792

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
8bc00f8
Implement special min/max accumulator for Strings: `MinMaxBytesAccumu…
alamb Oct 7, 2024
19b3297
fix bug
alamb Oct 8, 2024
4b3e625
fix msrv
alamb Oct 8, 2024
64e9861
move code, handle filters
alamb Oct 8, 2024
c4f6271
simplify
alamb Oct 8, 2024
c3145cd
Merge remote-tracking branch 'apache/main' into alamb/min_max_strings
alamb Oct 9, 2024
126a9a8
Add functional tests
alamb Oct 10, 2024
a7ebb56
remove unecessary test
alamb Oct 10, 2024
2d5957d
improve docs
alamb Oct 10, 2024
2671b2d
improve docs
alamb Oct 10, 2024
a43288e
Merge remote-tracking branch 'apache/main' into alamb/min_max_strings
alamb Oct 10, 2024
043ac35
cleanup
alamb Oct 10, 2024
e454838
improve comments
alamb Oct 10, 2024
8618673
Merge remote-tracking branch 'apache/main' into alamb/min_max_strings
alamb Oct 10, 2024
c7aa11f
fix diagram
alamb Oct 10, 2024
8e876bf
fix accounting
alamb Oct 10, 2024
1389fb7
Use correct type in memory accounting
alamb Oct 11, 2024
558e1ba
Merge remote-tracking branch 'apache/main' into alamb/min_max_strings
alamb Oct 11, 2024
7853d8f
Add TODO comment
alamb Oct 13, 2024
81431a0
Merge remote-tracking branch 'apache/main' into alamb/min_max_strings
alamb Oct 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ impl NullState {
///
/// When value_fn is called it also sets
///
/// 1. `self.seen_values[group_index]` to true for all rows that had a non null vale
/// 1. `self.seen_values[group_index]` to true for all rows that had a non null value
pub fn accumulate<T, F>(
&mut self,
group_indices: &[usize],
Expand Down
115 changes: 113 additions & 2 deletions datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/nulls.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,22 @@
// specific language governing permissions and limitations
// under the License.

//! [`set_nulls`], and [`filtered_null_mask`], utilities for working with nulls
//! [`set_nulls`], other utilities for working with nulls

use arrow::array::{Array, ArrowNumericType, BooleanArray, PrimitiveArray};
use arrow::array::{
Array, ArrayRef, ArrowNumericType, AsArray, BinaryArray, BinaryViewArray,
BooleanArray, LargeBinaryArray, LargeStringArray, PrimitiveArray, StringArray,
StringViewArray,
};
use arrow::buffer::NullBuffer;
use arrow::datatypes::DataType;
use datafusion_common::{not_impl_err, Result};
use std::sync::Arc;

/// Sets the validity mask for a `PrimitiveArray` to `nulls`
/// replacing any existing null mask
///
/// See [`set_nulls_dyn`] for a version that works with `Array`
pub fn set_nulls<T: ArrowNumericType + Send>(
array: PrimitiveArray<T>,
nulls: Option<NullBuffer>,
Expand Down Expand Up @@ -91,3 +100,105 @@ pub fn filtered_null_mask(
let opt_filter = opt_filter.and_then(filter_to_nulls);
NullBuffer::union(opt_filter.as_ref(), input.nulls())
}

/// Applies optional filter to input, returning a new array of the same type
/// with the same data, but with any values that were filtered out set to null
pub fn apply_filter_as_nulls(
input: &dyn Array,
opt_filter: Option<&BooleanArray>,
) -> Result<ArrayRef> {
let nulls = filtered_null_mask(opt_filter, input);
set_nulls_dyn(input, nulls)
}

/// Replaces the nulls in the input array with the given `NullBuffer`
///
/// TODO: replace when upstreamed in arrow-rs: <https://github.com/apache/arrow-rs/issues/6528>
pub fn set_nulls_dyn(input: &dyn Array, nulls: Option<NullBuffer>) -> Result<ArrayRef> {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is supporting code for replacing the null buffers in arrays

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this could be private

if let Some(nulls) = nulls.as_ref() {
assert_eq!(nulls.len(), input.len());
}

let output: ArrayRef = match input.data_type() {
DataType::Utf8 => {
let input = input.as_string::<i32>();
// safety: values / offsets came from a valid string array, so are valid utf8
// and we checked nulls has the same length as values
unsafe {
Arc::new(StringArray::new_unchecked(
input.offsets().clone(),
input.values().clone(),
nulls,
))
}
}
DataType::LargeUtf8 => {
let input = input.as_string::<i64>();
// safety: values / offsets came from a valid string array, so are valid utf8
// and we checked nulls has the same length as values
unsafe {
Arc::new(LargeStringArray::new_unchecked(
input.offsets().clone(),
input.values().clone(),
nulls,
))
}
}
DataType::Utf8View => {
let input = input.as_string_view();
// safety: values / views came from a valid string view array, so are valid utf8
// and we checked nulls has the same length as values
unsafe {
Arc::new(StringViewArray::new_unchecked(
input.views().clone(),
input.data_buffers().to_vec(),
nulls,
))
}
}

DataType::Binary => {
let input = input.as_binary::<i32>();
// safety: values / offsets came from a valid binary array
// and we checked nulls has the same length as values
unsafe {
Arc::new(BinaryArray::new_unchecked(
input.offsets().clone(),
input.values().clone(),
nulls,
))
}
}
DataType::LargeBinary => {
let input = input.as_binary::<i64>();
// safety: values / offsets came from a valid large binary array
// and we checked nulls has the same length as values
unsafe {
Arc::new(LargeBinaryArray::new_unchecked(
input.offsets().clone(),
input.values().clone(),
nulls,
))
}
}
DataType::BinaryView => {
let input = input.as_binary_view();
// safety: values / views came from a valid binary view array
// and we checked nulls has the same length as values
unsafe {
Arc::new(BinaryViewArray::new_unchecked(
input.views().clone(),
input.data_buffers().to_vec(),
nulls,
))
}
}
_ => {
return not_impl_err!("Applying nulls {:?}", input.data_type());
Comment on lines +196 to +197
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have to support this for any other data types?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not necessairly -- I am hoping we put this code upstream in arrow-rs and can remove it entirely from datafusion eventually

}
};
assert_eq!(input.len(), output.len());
assert_eq!(input.data_type(), output.data_type());

Ok(output)
}
123 changes: 69 additions & 54 deletions datafusion/functions-aggregate/src/min_max.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
//! [`Max`] and [`MaxAccumulator`] accumulator for the `max` function
//! [`Min`] and [`MinAccumulator`] accumulator for the `min` function

mod min_max_bytes;

use arrow::array::{
ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array,
Decimal128Array, Decimal256Array, Float16Array, Float32Array, Float64Array,
Expand Down Expand Up @@ -50,6 +52,7 @@ use arrow::datatypes::{
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
};

use crate::min_max::min_max_bytes::MinMaxBytesAccumulator;
use datafusion_common::ScalarValue;
use datafusion_expr::{
function::AccumulatorArgs, Accumulator, AggregateUDFImpl, Documentation, Signature,
Expand Down Expand Up @@ -104,7 +107,7 @@ impl Default for Max {
/// the specified [`ArrowPrimitiveType`].
///
/// [`ArrowPrimitiveType`]: arrow::datatypes::ArrowPrimitiveType
macro_rules! instantiate_max_accumulator {
macro_rules! primitive_max_accumulator {
($DATA_TYPE:ident, $NATIVE:ident, $PRIMTYPE:ident) => {{
Ok(Box::new(
PrimitiveGroupsAccumulator::<$PRIMTYPE, _>::new($DATA_TYPE, |cur, new| {
Expand All @@ -123,7 +126,7 @@ macro_rules! instantiate_max_accumulator {
///
///
/// [`ArrowPrimitiveType`]: arrow::datatypes::ArrowPrimitiveType
macro_rules! instantiate_min_accumulator {
macro_rules! primitive_min_accumulator {
($DATA_TYPE:ident, $NATIVE:ident, $PRIMTYPE:ident) => {{
Ok(Box::new(
PrimitiveGroupsAccumulator::<$PRIMTYPE, _>::new(&$DATA_TYPE, |cur, new| {
Expand Down Expand Up @@ -231,6 +234,12 @@ impl AggregateUDFImpl for Max {
| Time32(_)
| Time64(_)
| Timestamp(_, _)
| Utf8
| LargeUtf8
| Utf8View
| Binary
| LargeBinary
| BinaryView
)
}

Expand All @@ -242,58 +251,58 @@ impl AggregateUDFImpl for Max {
use TimeUnit::*;
let data_type = args.return_type;
match data_type {
Int8 => instantiate_max_accumulator!(data_type, i8, Int8Type),
Int16 => instantiate_max_accumulator!(data_type, i16, Int16Type),
Int32 => instantiate_max_accumulator!(data_type, i32, Int32Type),
Int64 => instantiate_max_accumulator!(data_type, i64, Int64Type),
UInt8 => instantiate_max_accumulator!(data_type, u8, UInt8Type),
UInt16 => instantiate_max_accumulator!(data_type, u16, UInt16Type),
UInt32 => instantiate_max_accumulator!(data_type, u32, UInt32Type),
UInt64 => instantiate_max_accumulator!(data_type, u64, UInt64Type),
Int8 => primitive_max_accumulator!(data_type, i8, Int8Type),
Int16 => primitive_max_accumulator!(data_type, i16, Int16Type),
Int32 => primitive_max_accumulator!(data_type, i32, Int32Type),
Int64 => primitive_max_accumulator!(data_type, i64, Int64Type),
UInt8 => primitive_max_accumulator!(data_type, u8, UInt8Type),
UInt16 => primitive_max_accumulator!(data_type, u16, UInt16Type),
UInt32 => primitive_max_accumulator!(data_type, u32, UInt32Type),
UInt64 => primitive_max_accumulator!(data_type, u64, UInt64Type),
Float16 => {
instantiate_max_accumulator!(data_type, f16, Float16Type)
primitive_max_accumulator!(data_type, f16, Float16Type)
}
Float32 => {
instantiate_max_accumulator!(data_type, f32, Float32Type)
primitive_max_accumulator!(data_type, f32, Float32Type)
}
Float64 => {
instantiate_max_accumulator!(data_type, f64, Float64Type)
primitive_max_accumulator!(data_type, f64, Float64Type)
}
Date32 => instantiate_max_accumulator!(data_type, i32, Date32Type),
Date64 => instantiate_max_accumulator!(data_type, i64, Date64Type),
Date32 => primitive_max_accumulator!(data_type, i32, Date32Type),
Date64 => primitive_max_accumulator!(data_type, i64, Date64Type),
Time32(Second) => {
instantiate_max_accumulator!(data_type, i32, Time32SecondType)
primitive_max_accumulator!(data_type, i32, Time32SecondType)
}
Time32(Millisecond) => {
instantiate_max_accumulator!(data_type, i32, Time32MillisecondType)
primitive_max_accumulator!(data_type, i32, Time32MillisecondType)
}
Time64(Microsecond) => {
instantiate_max_accumulator!(data_type, i64, Time64MicrosecondType)
primitive_max_accumulator!(data_type, i64, Time64MicrosecondType)
}
Time64(Nanosecond) => {
instantiate_max_accumulator!(data_type, i64, Time64NanosecondType)
primitive_max_accumulator!(data_type, i64, Time64NanosecondType)
}
Timestamp(Second, _) => {
instantiate_max_accumulator!(data_type, i64, TimestampSecondType)
primitive_max_accumulator!(data_type, i64, TimestampSecondType)
}
Timestamp(Millisecond, _) => {
instantiate_max_accumulator!(data_type, i64, TimestampMillisecondType)
primitive_max_accumulator!(data_type, i64, TimestampMillisecondType)
}
Timestamp(Microsecond, _) => {
instantiate_max_accumulator!(data_type, i64, TimestampMicrosecondType)
primitive_max_accumulator!(data_type, i64, TimestampMicrosecondType)
}
Timestamp(Nanosecond, _) => {
instantiate_max_accumulator!(data_type, i64, TimestampNanosecondType)
primitive_max_accumulator!(data_type, i64, TimestampNanosecondType)
}
Decimal128(_, _) => {
instantiate_max_accumulator!(data_type, i128, Decimal128Type)
primitive_max_accumulator!(data_type, i128, Decimal128Type)
}
Decimal256(_, _) => {
instantiate_max_accumulator!(data_type, i256, Decimal256Type)
primitive_max_accumulator!(data_type, i256, Decimal256Type)
}
Utf8 | LargeUtf8 | Utf8View | Binary | LargeBinary | BinaryView => {
Ok(Box::new(MinMaxBytesAccumulator::new_max(data_type.clone())))
}

// It would be nice to have a fast implementation for Strings as well
// https://github.com/apache/datafusion/issues/6906

// This is only reached if groups_accumulator_supported is out of sync
_ => internal_err!("GroupsAccumulator not supported for max({})", data_type),
Expand Down Expand Up @@ -1057,6 +1066,12 @@ impl AggregateUDFImpl for Min {
| Time32(_)
| Time64(_)
| Timestamp(_, _)
| Utf8
| LargeUtf8
| Utf8View
| Binary
| LargeBinary
| BinaryView
)
}

Expand All @@ -1068,58 +1083,58 @@ impl AggregateUDFImpl for Min {
use TimeUnit::*;
let data_type = args.return_type;
match data_type {
Int8 => instantiate_min_accumulator!(data_type, i8, Int8Type),
Int16 => instantiate_min_accumulator!(data_type, i16, Int16Type),
Int32 => instantiate_min_accumulator!(data_type, i32, Int32Type),
Int64 => instantiate_min_accumulator!(data_type, i64, Int64Type),
UInt8 => instantiate_min_accumulator!(data_type, u8, UInt8Type),
UInt16 => instantiate_min_accumulator!(data_type, u16, UInt16Type),
UInt32 => instantiate_min_accumulator!(data_type, u32, UInt32Type),
UInt64 => instantiate_min_accumulator!(data_type, u64, UInt64Type),
Int8 => primitive_min_accumulator!(data_type, i8, Int8Type),
Int16 => primitive_min_accumulator!(data_type, i16, Int16Type),
Int32 => primitive_min_accumulator!(data_type, i32, Int32Type),
Int64 => primitive_min_accumulator!(data_type, i64, Int64Type),
UInt8 => primitive_min_accumulator!(data_type, u8, UInt8Type),
UInt16 => primitive_min_accumulator!(data_type, u16, UInt16Type),
UInt32 => primitive_min_accumulator!(data_type, u32, UInt32Type),
UInt64 => primitive_min_accumulator!(data_type, u64, UInt64Type),
Float16 => {
instantiate_min_accumulator!(data_type, f16, Float16Type)
primitive_min_accumulator!(data_type, f16, Float16Type)
}
Float32 => {
instantiate_min_accumulator!(data_type, f32, Float32Type)
primitive_min_accumulator!(data_type, f32, Float32Type)
}
Float64 => {
instantiate_min_accumulator!(data_type, f64, Float64Type)
primitive_min_accumulator!(data_type, f64, Float64Type)
}
Date32 => instantiate_min_accumulator!(data_type, i32, Date32Type),
Date64 => instantiate_min_accumulator!(data_type, i64, Date64Type),
Date32 => primitive_min_accumulator!(data_type, i32, Date32Type),
Date64 => primitive_min_accumulator!(data_type, i64, Date64Type),
Time32(Second) => {
instantiate_min_accumulator!(data_type, i32, Time32SecondType)
primitive_min_accumulator!(data_type, i32, Time32SecondType)
}
Time32(Millisecond) => {
instantiate_min_accumulator!(data_type, i32, Time32MillisecondType)
primitive_min_accumulator!(data_type, i32, Time32MillisecondType)
}
Time64(Microsecond) => {
instantiate_min_accumulator!(data_type, i64, Time64MicrosecondType)
primitive_min_accumulator!(data_type, i64, Time64MicrosecondType)
}
Time64(Nanosecond) => {
instantiate_min_accumulator!(data_type, i64, Time64NanosecondType)
primitive_min_accumulator!(data_type, i64, Time64NanosecondType)
}
Timestamp(Second, _) => {
instantiate_min_accumulator!(data_type, i64, TimestampSecondType)
primitive_min_accumulator!(data_type, i64, TimestampSecondType)
}
Timestamp(Millisecond, _) => {
instantiate_min_accumulator!(data_type, i64, TimestampMillisecondType)
primitive_min_accumulator!(data_type, i64, TimestampMillisecondType)
}
Timestamp(Microsecond, _) => {
instantiate_min_accumulator!(data_type, i64, TimestampMicrosecondType)
primitive_min_accumulator!(data_type, i64, TimestampMicrosecondType)
}
Timestamp(Nanosecond, _) => {
instantiate_min_accumulator!(data_type, i64, TimestampNanosecondType)
primitive_min_accumulator!(data_type, i64, TimestampNanosecondType)
}
Decimal128(_, _) => {
instantiate_min_accumulator!(data_type, i128, Decimal128Type)
primitive_min_accumulator!(data_type, i128, Decimal128Type)
}
Decimal256(_, _) => {
instantiate_min_accumulator!(data_type, i256, Decimal256Type)
primitive_min_accumulator!(data_type, i256, Decimal256Type)
}
Utf8 | LargeUtf8 | Utf8View | Binary | LargeBinary | BinaryView => {
Ok(Box::new(MinMaxBytesAccumulator::new_min(data_type.clone())))
}

// It would be nice to have a fast implementation for Strings as well
// https://github.com/apache/datafusion/issues/6906

// This is only reached if groups_accumulator_supported is out of sync
_ => internal_err!("GroupsAccumulator not supported for min({})", data_type),
Expand Down
Loading