alamb commented on code in PR #23014:
URL: https://github.com/apache/datafusion/pull/23014#discussion_r3647891503
##########
datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs:
##########
@@ -222,6 +223,276 @@ where
}
}
+pub(super) type BranchlessNative<T> =
+ <<T as BranchlessFilterType>::CompareType as ArrowPrimitiveType>::Native;
+
+/// Maximum list size for branchless lookup on 1-byte primitives.
+///
+/// Sixteen 1-byte values fit in one 128-bit SIMD vector, so this keeps the
+/// branchless list small enough for a single vectorized membership check.
+const BRANCHLESS_MAX_1B: usize = 16;
+
+/// Maximum list size for branchless lookup on 2-byte primitives.
+///
+/// Eight 2-byte values fit in one 128-bit SIMD vector, so this keeps the
+/// branchless list small enough for a single vectorized membership check.
+const BRANCHLESS_MAX_2B: usize = 8;
+
+/// Maximum list size for branchless lookup on 4-byte primitives.
+///
+/// Thirty-two 4-byte values keep the inline list at 128 bytes. Beyond that,
+/// the comparison chain and filter footprint grow enough that the hash/generic
+/// fallback is a better fit.
+const BRANCHLESS_MAX_4B: usize = 32;
+
+/// Maximum list size for branchless lookup on 8-byte primitives.
+///
+/// Sixteen 8-byte values use the same 128-byte inline-list budget as 4-byte
+/// primitives. Larger lists are left to the hash/generic fallback.
+const BRANCHLESS_MAX_8B: usize = 16;
+
+/// Maximum list size for branchless lookup on 16-byte primitives.
+///
+/// These comparisons are wider, so this path is limited to four values.
+/// Larger lists are left to the generic fallback.
+const BRANCHLESS_MAX_16B: usize = 4;
+
+/// Arrow primitive types supported by [`BranchlessFilter`].
+///
+/// `T` is the logical Arrow type accepted by the filter. `CompareType` is the
+/// same-width type used for the fixed comparison chain. Signed integers,
+/// floats, and temporal values use an unsigned comparison type so they compare
+/// by their raw bit pattern.
+pub(super) trait BranchlessFilterType:
+ ArrowPrimitiveType + Send + Sync + 'static
+{
+ type CompareType: ArrowPrimitiveType + Send + Sync + 'static;
+
+ /// Maximum number of non-null IN-list values to handle with
+ /// [`BranchlessFilter`] for this primitive type.
+ const MAX_LIST_LEN: usize;
+}
+
+macro_rules! branchless_filter_type {
+ ($logical:ty, $compare:ty, $max_len:expr) => {
+ // The branchless filter reads the same Arrow value buffer as the
+ // comparison type. That is only valid when both native types have the
+ // same width, so catch any bad mapping here at compile time.
+ const _: () = assert!(
+ size_of::<<$logical as ArrowPrimitiveType>::Native>()
+ == size_of::<<$compare as ArrowPrimitiveType>::Native>(),
+ "BranchlessFilterType::CompareType must use the same native width"
+ );
+
+ impl BranchlessFilterType for $logical {
+ type CompareType = $compare;
+ const MAX_LIST_LEN: usize = $max_len;
+ }
+ };
+}
+
+branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B);
+branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B);
+branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B);
+branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B);
+branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B);
+
+branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B);
+
+branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampSecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampMillisecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampMicrosecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampNanosecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationSecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationMillisecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationMicrosecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationNanosecondType, UInt64Type, BRANCHLESS_MAX_8B);
+
+branchless_filter_type!(Decimal128Type, Decimal128Type, BRANCHLESS_MAX_16B);
+branchless_filter_type!(
+ IntervalMonthDayNanoType,
+ IntervalMonthDayNanoType,
+ BRANCHLESS_MAX_16B
+);
+
+/// Checks each input value against the `IN`-list values.
+type MembershipCheck<C> = fn(in_list_values: &[C], input_values: &[C]) ->
BooleanBuffer;
+
+/// A branchless filter for fixed-width primitive `IN` lists up to
+/// `T::MAX_LIST_LEN` values.
+///
+/// The filter stores the non-null `IN`-list values in a slice and chooses a
+/// comparison function for that length. Keeping the length out of
+/// `BranchlessFilter` avoids generating a full copy of the filter for every
+/// supported length.
+pub(super) struct BranchlessFilter<T: BranchlessFilterType> {
+ expected_data_type: DataType,
+ null_count: usize,
+ in_list_values: Box<[BranchlessNative<T>]>,
+ check_values: MembershipCheck<BranchlessNative<T>>,
+}
+
+impl<T> BranchlessFilter<T>
+where
+ T: BranchlessFilterType,
+ BranchlessNative<T>: Copy + PartialEq,
+{
+ pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
+ let in_array = in_array.as_primitive_opt::<T>().ok_or_else(|| {
+ exec_datafusion_err!("BranchlessFilter: expected {} array",
T::DATA_TYPE)
+ })?;
+ let non_null_count = in_array.len() - in_array.null_count();
+ // `try_new` can be called on its own, so check the limit here too.
+ if non_null_count > T::MAX_LIST_LEN {
+ return Err(exec_datafusion_err!(
Review Comment:
This is probably more correctly an internal error rather than an execution
error as it isn't really based on user input (it would signal a bug in the code
I think)
##########
datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs:
##########
@@ -222,6 +223,276 @@ where
}
}
+pub(super) type BranchlessNative<T> =
+ <<T as BranchlessFilterType>::CompareType as ArrowPrimitiveType>::Native;
+
+/// Maximum list size for branchless lookup on 1-byte primitives.
+///
+/// Sixteen 1-byte values fit in one 128-bit SIMD vector, so this keeps the
+/// branchless list small enough for a single vectorized membership check.
+const BRANCHLESS_MAX_1B: usize = 16;
+
+/// Maximum list size for branchless lookup on 2-byte primitives.
+///
+/// Eight 2-byte values fit in one 128-bit SIMD vector, so this keeps the
+/// branchless list small enough for a single vectorized membership check.
+const BRANCHLESS_MAX_2B: usize = 8;
+
+/// Maximum list size for branchless lookup on 4-byte primitives.
+///
+/// Thirty-two 4-byte values keep the inline list at 128 bytes. Beyond that,
+/// the comparison chain and filter footprint grow enough that the hash/generic
+/// fallback is a better fit.
+const BRANCHLESS_MAX_4B: usize = 32;
+
+/// Maximum list size for branchless lookup on 8-byte primitives.
+///
+/// Sixteen 8-byte values use the same 128-byte inline-list budget as 4-byte
+/// primitives. Larger lists are left to the hash/generic fallback.
+const BRANCHLESS_MAX_8B: usize = 16;
+
+/// Maximum list size for branchless lookup on 16-byte primitives.
+///
+/// These comparisons are wider, so this path is limited to four values.
+/// Larger lists are left to the generic fallback.
+const BRANCHLESS_MAX_16B: usize = 4;
+
+/// Arrow primitive types supported by [`BranchlessFilter`].
+///
+/// `T` is the logical Arrow type accepted by the filter. `CompareType` is the
+/// same-width type used for the fixed comparison chain. Signed integers,
+/// floats, and temporal values use an unsigned comparison type so they compare
+/// by their raw bit pattern.
+pub(super) trait BranchlessFilterType:
+ ArrowPrimitiveType + Send + Sync + 'static
+{
+ type CompareType: ArrowPrimitiveType + Send + Sync + 'static;
+
+ /// Maximum number of non-null IN-list values to handle with
+ /// [`BranchlessFilter`] for this primitive type.
+ const MAX_LIST_LEN: usize;
+}
+
+macro_rules! branchless_filter_type {
+ ($logical:ty, $compare:ty, $max_len:expr) => {
+ // The branchless filter reads the same Arrow value buffer as the
+ // comparison type. That is only valid when both native types have the
+ // same width, so catch any bad mapping here at compile time.
+ const _: () = assert!(
+ size_of::<<$logical as ArrowPrimitiveType>::Native>()
+ == size_of::<<$compare as ArrowPrimitiveType>::Native>(),
+ "BranchlessFilterType::CompareType must use the same native width"
+ );
+
+ impl BranchlessFilterType for $logical {
+ type CompareType = $compare;
+ const MAX_LIST_LEN: usize = $max_len;
+ }
+ };
+}
+
+branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B);
+branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B);
+branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B);
+branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B);
+branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B);
+
+branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B);
+branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B);
+
+branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampSecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampMillisecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampMicrosecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(TimestampNanosecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationSecondType, UInt64Type, BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationMillisecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationMicrosecondType, UInt64Type,
BRANCHLESS_MAX_8B);
+branchless_filter_type!(DurationNanosecondType, UInt64Type, BRANCHLESS_MAX_8B);
+
+branchless_filter_type!(Decimal128Type, Decimal128Type, BRANCHLESS_MAX_16B);
+branchless_filter_type!(
+ IntervalMonthDayNanoType,
+ IntervalMonthDayNanoType,
+ BRANCHLESS_MAX_16B
+);
+
+/// Checks each input value against the `IN`-list values.
+type MembershipCheck<C> = fn(in_list_values: &[C], input_values: &[C]) ->
BooleanBuffer;
+
+/// A branchless filter for fixed-width primitive `IN` lists up to
+/// `T::MAX_LIST_LEN` values.
+///
+/// The filter stores the non-null `IN`-list values in a slice and chooses a
+/// comparison function for that length. Keeping the length out of
+/// `BranchlessFilter` avoids generating a full copy of the filter for every
+/// supported length.
+pub(super) struct BranchlessFilter<T: BranchlessFilterType> {
+ expected_data_type: DataType,
+ null_count: usize,
+ in_list_values: Box<[BranchlessNative<T>]>,
+ check_values: MembershipCheck<BranchlessNative<T>>,
+}
+
+impl<T> BranchlessFilter<T>
+where
+ T: BranchlessFilterType,
+ BranchlessNative<T>: Copy + PartialEq,
+{
+ pub(super) fn try_new(in_array: &ArrayRef) -> Result<Self> {
+ let in_array = in_array.as_primitive_opt::<T>().ok_or_else(|| {
+ exec_datafusion_err!("BranchlessFilter: expected {} array",
T::DATA_TYPE)
+ })?;
+ let non_null_count = in_array.len() - in_array.null_count();
+ // `try_new` can be called on its own, so check the limit here too.
+ if non_null_count > T::MAX_LIST_LEN {
+ return Err(exec_datafusion_err!(
+ "BranchlessFilter: supports at most {} non-null values, got
{non_null_count}",
+ T::MAX_LIST_LEN
+ ));
+ }
+
+ let all_values = branchless_values::<T>(in_array);
+ let mut in_list_values = Vec::with_capacity(non_null_count);
+
+ match in_array.nulls() {
+ None => {
+ in_list_values.extend(all_values.iter().copied());
+ }
+ Some(nulls) => {
+ for row in
+ BitIndexIterator::new(nulls.validity(), nulls.offset(),
nulls.len())
+ {
+ in_list_values.push(all_values[row]);
+ }
+ }
+ }
+
+ debug_assert_eq!(in_list_values.len(), non_null_count);
+ let in_list_values = in_list_values.into_boxed_slice();
+ let check_values = membership_check_for_len::<T>(in_list_values.len());
+
+ Ok(Self {
+ expected_data_type: in_array.data_type().clone(),
+ null_count: in_array.null_count(),
+ in_list_values,
+ check_values,
+ })
+ }
+}
+
+impl<T> StaticFilter for BranchlessFilter<T>
+where
+ T: BranchlessFilterType,
+ BranchlessNative<T>: Copy + PartialEq + Send + Sync,
+{
+ fn null_count(&self) -> usize {
+ self.null_count
+ }
+
+ fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
+ handle_dictionary!(self, v, negated);
+
+ // Arrow compatibility ignores timestamp timezone and decimal
precision/scale
+ // while still requiring the same primitive representation.
+ if !PrimitiveArray::<T>::is_compatible(v.data_type()) {
+ return Err(exec_datafusion_err!(
+ "BranchlessFilter: expected {} array, got {}",
+ self.expected_data_type,
+ v.data_type()
+ ));
+ }
+
+ let v = v.as_primitive_opt::<T>().ok_or_else(|| {
+ exec_datafusion_err!("BranchlessFilter: expected {} array",
T::DATA_TYPE)
+ })?;
+ let input_values = branchless_values::<T>(v);
+ let matches =
+ (self.check_values)(self.in_list_values.as_ref(),
input_values.as_ref());
+ Ok(build_result_from_contains(
+ v.nulls(),
+ self.null_count > 0,
+ negated,
+ matches,
+ ))
+ }
+}
+
+/// Picks the comparison function for `len` non-null `IN`-list values.
+///
+/// A length of zero is used when the list contains only nulls. The comparisons
+/// return false, and the caller then applies the usual SQL null behavior.
+fn membership_check_for_len<T>(len: usize) ->
MembershipCheck<BranchlessNative<T>>
+where
+ T: BranchlessFilterType,
+ BranchlessNative<T>: Copy + PartialEq,
+{
+ macro_rules! choose {
+ ($($n:literal),* $(,)?) => {
+ match len {
+ $($n => check_values::<BranchlessNative<T>, $n>,)*
+ _ => unreachable!("list length exceeds the configured limit"),
+ }
+ };
+ }
+
+ // Avoid creating checks for lengths a type does not support.
+ match T::MAX_LIST_LEN {
+ 4 => choose!(0, 1, 2, 3, 4),
+ 8 => choose!(0, 1, 2, 3, 4, 5, 6, 7, 8),
+ 16 => choose!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
16),
+ 32 => choose!(
Review Comment:
this is cool -- it makes much more sense to isolate the code repetition to
just the comparison function (which is where it is most helpful)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
