diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b58328..4a74f18be8665 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,6 +38,7 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; mod branchless_filter; +mod byte_view_filter; mod primitive_filter; mod result; mod static_filter; @@ -215,7 +216,7 @@ impl InListExpr { expr, list, negated, - Some(instantiate_static_filter(array)?), + Some(instantiate_static_filter(array, &expr_data_type)?), )) } @@ -242,7 +243,7 @@ impl InListExpr { // Try to create a static filter if all list expressions are constants let static_filter = match try_evaluate_constant_list(&list, schema)? { - Some(in_array) => Some(instantiate_static_filter(in_array)?), + Some(in_array) => Some(instantiate_static_filter(in_array, &expr_data_type)?), None => None, // Non-constant expressions, fall back to dynamic evaluation }; @@ -3576,6 +3577,23 @@ mod tests { )? ); + // Utf8View in_array, Utf8View and Dict(Utf8View) needles + let utf8view_in = + Arc::new(StringViewArray::from(vec!["a", "b", "c"])) as ArrayRef; + let utf8view_needle = + Arc::new(StringViewArray::from(vec!["a", "d", "b"])) as ArrayRef; + assert_eq!( + expected, + eval_in_list_from_array( + Arc::clone(&utf8view_needle), + Arc::clone(&utf8view_in), + )? + ); + assert_eq!( + expected, + eval_in_list_from_array(wrap_in_dict(utf8view_needle), utf8view_in)? + ); + // Struct in_array, Struct needle: multi-column join let struct_fields = Fields::from(vec![ Field::new("c0", DataType::Utf8, true), diff --git a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs index cd0cbd0de59a8..ff010e4603a38 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -71,7 +71,7 @@ use std::mem::size_of; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow::buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; @@ -244,6 +244,17 @@ where check_values, }) } + + #[inline] + pub(super) fn contains_slice( + &self, + input_values: &[BranchlessNative], + nulls: Option<&NullBuffer>, + negated: bool, + ) -> BooleanArray { + let matches = (self.check_values)(self.in_list_values.as_ref(), input_values); + build_result_from_contains(nulls, self.null_count > 0, negated, matches) + } } impl StaticFilter for BranchlessFilter @@ -272,14 +283,7 @@ where exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) })?; let input_values = branchless_values::(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, - )) + Ok(self.contains_slice(input_values.as_ref(), v.nulls(), negated)) } } diff --git a/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs new file mode 100644 index 0000000000000..8d076d695607a --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs @@ -0,0 +1,452 @@ +// 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. + +//! Optimized filters for Utf8View and BinaryView IN lists. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, GenericByteViewArray, MAX_INLINE_VIEW_LEN, + PrimitiveArray, +}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::{ + BinaryViewType, ByteViewType, DataType, Decimal128Type, StringViewType, +}; +use arrow::util::bit_iterator::BitIndexIterator; +use datafusion_common::{HashSet, Result, exec_datafusion_err}; + +use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; +use super::result::build_in_list_result; +use super::static_filter::{StaticFilter, handle_dictionary}; + +#[inline(always)] +fn view_len(view: u128) -> u32 { + view as u32 +} + +fn downcast_byte_view( + array: &dyn Array, +) -> Result<&GenericByteViewArray> { + array + .as_byte_view_opt::() + .ok_or_else(|| exec_datafusion_err!("Expected concrete {} array", T::DATA_TYPE)) +} + +fn all_views_inline(array: &ArrayRef) -> Result { + let array = downcast_byte_view::(array.as_ref())?; + let is_inline = |idx: usize| view_len(array.views()[idx]) <= MAX_INLINE_VIEW_LEN; + Ok(match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .all(is_inline) + } + None => (0..array.len()).all(is_inline), + }) +} + +fn reinterpret_byte_view_as_decimal128( + array: &dyn Array, +) -> Result { + let array = downcast_byte_view::(array)?; + let views = array.views(); + let values = ScalarBuffer::::new(views.inner().clone(), 0, views.len()); + Ok(Arc::new(PrimitiveArray::::new( + values, + array.nulls().cloned(), + ))) +} + +fn make_byte_view_branchless_filter( + in_array: &ArrayRef, +) -> Result> { + let values = reinterpret_byte_view_as_decimal128::(in_array.as_ref())?; + + Ok(Arc::new(ByteViewBranchless:: { + inner: BranchlessFilter::::try_new(&values)?, + _marker: PhantomData, + })) +} + +struct ByteViewBranchless { + inner: BranchlessFilter, + _marker: PhantomData, +} + +/// Exact set membership for inline views. +/// +/// Arrow validates that unused inline bytes are zero, so equal values have the +/// same canonical `u128` representation. +struct InlineByteViewFilter { + set: HashSet, + null_count: usize, + _marker: PhantomData, +} + +impl InlineByteViewFilter { + fn try_new(in_array: &ArrayRef) -> Result { + if in_array.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "InlineByteViewFilter: expected {} array, got {}", + T::DATA_TYPE, + in_array.data_type() + )); + } + + let array = downcast_byte_view::(in_array.as_ref())?; + let mut set = HashSet::with_capacity(array.len() - array.null_count()); + match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .for_each(|idx| { + set.insert(array.views()[idx]); + }); + } + None => set.extend(array.views().iter().copied()), + } + + Ok(Self { + set, + null_count: array.null_count(), + _marker: PhantomData, + }) + } +} + +impl StaticFilter for InlineByteViewFilter { + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "InlineByteViewFilter: expected {} array, got {}", + T::DATA_TYPE, + v.data_type() + )); + } + + let array = downcast_byte_view::(v)?; + let views = array.views(); + Ok(build_in_list_result( + array.len(), + array.nulls(), + self.null_count > 0, + negated, + |idx| { + // SAFETY: `build_in_list_result` visits indices in `0..array.len()`. + self.set.contains(unsafe { views.get_unchecked(idx) }) + }, + )) + } +} + +impl StaticFilter for ByteViewBranchless { + fn null_count(&self) -> usize { + self.inner.null_count() + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "ByteViewBranchless: expected {} array, got {}", + T::DATA_TYPE, + v.data_type() + )); + } + + let array = downcast_byte_view::(v)?; + let values: &[i128] = array.views().inner().typed_data(); + Ok(self.inner.contains_slice(values, array.nulls(), negated)) + } +} + +fn make_inline_byte_view_filter( + in_array: &ArrayRef, +) -> Result> { + Ok(Arc::new(InlineByteViewFilter::::try_new(in_array)?)) +} + +fn instantiate_typed_byte_view_filter( + in_array: &ArrayRef, +) -> Result>> { + if !all_views_inline::(in_array)? { + return Ok(None); + } + + let non_null_count = in_array.len() - in_array.null_count(); + if non_null_count <= ::MAX_LIST_LEN { + make_byte_view_branchless_filter::(in_array).map(Some) + } else { + make_inline_byte_view_filter::(in_array).map(Some) + } +} + +/// Creates an optimized byte-view filter for supported list compositions. +/// +/// All-inline lists use direct view comparisons or a hash set. Lists containing +/// long values stay on the generic filter because long views are not exact keys. +pub(super) fn instantiate_byte_view_filter( + in_array: &ArrayRef, +) -> Result>> { + match in_array.data_type() { + DataType::Utf8View => { + instantiate_typed_byte_view_filter::(in_array) + } + DataType::BinaryView => { + instantiate_typed_byte_view_filter::(in_array) + } + _ => Ok(None), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{BinaryViewArray, DictionaryArray, Int8Array, StringViewArray}; + use arrow::datatypes::{BinaryViewType, StringViewType}; + + fn assert_contains( + filter: &dyn StaticFilter, + needles: &dyn Array, + expected: Vec>, + ) -> Result<()> { + assert_eq!( + filter.contains(needles, false)?, + BooleanArray::from(expected) + ); + Ok(()) + } + + #[test] + fn inline_filter_handles_slices_nulls_and_not_in() -> Result<()> { + let haystack: ArrayRef = Arc::new( + StringViewArray::from(vec![ + Some("outside"), + Some("a"), + Some("b"), + None, + Some("c"), + Some("d"), + Some("e"), + Some("tail"), + ]) + .slice(1, 6), + ); + let filter = InlineByteViewFilter::::try_new(&haystack)?; + let needles = + StringViewArray::from(vec![Some("b"), Some("missing"), None, Some("e")]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + Ok(()) + } + + #[test] + fn hash_set_filter_handles_binary_views() -> Result<()> { + let haystack: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some([0xff].as_slice()), + Some([0x00].as_slice()), + Some([0x01, 0x02].as_slice()), + Some([0x03].as_slice()), + Some([0x04].as_slice()), + None, + ])); + let filter = instantiate_byte_view_filter(&haystack)? + .expect("five inline values select the hash-set filter"); + let needles = BinaryViewArray::from(vec![ + Some([0xff].as_slice()), + Some([0x05].as_slice()), + None, + ]); + + assert_contains(&*filter, &needles, vec![Some(true), None, None])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None]) + ); + Ok(()) + } + + #[test] + fn hash_set_filter_handles_dictionary_needles() -> Result<()> { + let haystack: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("one"), + Some("two"), + Some("three"), + Some("four"), + Some("five"), + None, + ])); + let filter = instantiate_byte_view_filter(&haystack)? + .expect("five inline values select the hash-set filter"); + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); + let values = Arc::new(StringViewArray::from(vec![ + Some("one"), + Some("five"), + Some("missing"), + ])); + let needles = DictionaryArray::try_new(keys, values)?; + + assert_contains(&*filter, &needles, vec![Some(true), Some(true), None, None]) + } + + #[test] + fn inline_filters_handle_twelve_and_thirteen_byte_values() -> Result<()> { + let inline_value = "abcdefghijkl"; + let long_value = "abcdefghijklm"; + let needles = StringViewArray::from(vec![ + Some(inline_value), + Some(long_value), + Some("missing"), + ]); + let expected = vec![Some(true), Some(false), Some(false)]; + + let branchless_haystack: ArrayRef = Arc::new(StringViewArray::from(vec![ + inline_value, + "one", + "two", + "three", + ])); + let branchless = instantiate_byte_view_filter(&branchless_haystack)? + .expect("four inline values select the branchless filter"); + assert_contains(&*branchless, &needles, expected.clone())?; + + let hash_set_haystack: ArrayRef = Arc::new(StringViewArray::from(vec![ + inline_value, + "one", + "two", + "three", + "four", + ])); + let hash_set = instantiate_byte_view_filter(&hash_set_haystack)? + .expect("five inline values select the hash-set filter"); + assert_contains(&*hash_set, &needles, expected)?; + + let long_haystack: ArrayRef = Arc::new(StringViewArray::from(vec![long_value])); + assert!(instantiate_byte_view_filter(&long_haystack)?.is_none()); + Ok(()) + } + + #[test] + fn branchless_filter_handles_utf8_and_binary_views() -> Result<()> { + let utf8_haystack: ArrayRef = + Arc::new(StringViewArray::from(vec![Some("one"), None, Some("two")])); + let utf8_filter = + make_byte_view_branchless_filter::(&utf8_haystack)?; + let utf8_needles = StringViewArray::from(vec![Some("two"), Some("three"), None]); + assert_contains(&*utf8_filter, &utf8_needles, vec![Some(true), None, None])?; + + let binary_haystack: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some([0xff].as_slice()), + Some([0x00].as_slice()), + ])); + let binary_filter = + make_byte_view_branchless_filter::(&binary_haystack)?; + let binary_needles = + BinaryViewArray::from(vec![Some([0x00].as_slice()), Some([0x01].as_slice())]); + assert_contains( + &*binary_filter, + &binary_needles, + vec![Some(true), Some(false)], + ) + } + + #[test] + fn byte_view_filters_reject_other_types() -> Result<()> { + let haystack: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("one"), + Some("two"), + Some("three"), + Some("four"), + Some("five"), + ])); + let hash_set = InlineByteViewFilter::::try_new(&haystack)?; + let branchless_haystack: ArrayRef = + Arc::new(StringViewArray::from(vec![Some("short")])); + let branchless = + make_byte_view_branchless_filter::(&branchless_haystack)?; + let needles = BinaryViewArray::from(vec![Some(b"short".as_slice())]); + + assert!(hash_set.contains(&needles, false).is_err()); + assert!(branchless.contains(&needles, false).is_err()); + Ok(()) + } + + #[test] + fn byte_view_routing_selects_only_inline_lists() -> Result<()> { + let inline_four: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("a"), + None, + Some("b"), + Some("c"), + Some("d"), + ])); + assert!(instantiate_byte_view_filter(&inline_four)?.is_some()); + + let inline_five: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + assert!(instantiate_byte_view_filter(&inline_five)?.is_some()); + + let all_null: ArrayRef = + Arc::new(StringViewArray::from(vec![None::<&str>, None::<&str>])); + assert!(instantiate_byte_view_filter(&all_null)?.is_some()); + + let all_long: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("abcdefghijklmn1"), + Some("abcdefghijklmn2"), + ])); + assert!(instantiate_byte_view_filter(&all_long)?.is_none()); + + let mixed: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("short"), + Some("abcdefghijklmn1"), + ])); + assert!(instantiate_byte_view_filter(&mixed)?.is_none()); + + let inline_binary: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some(b"a".as_slice()), + Some(b"b".as_slice()), + Some(b"c".as_slice()), + Some(b"d".as_slice()), + Some(b"e".as_slice()), + ])); + assert!(instantiate_byte_view_filter(&inline_binary)?.is_some()); + + let mixed_binary: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some([0xff].as_slice()), + Some(b"abcdefghijklmn1".as_slice()), + ])); + assert!(instantiate_byte_view_filter(&mixed_binary)?.is_none()); + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d5ca8154a92f6..50df83b7b573e 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -34,14 +34,25 @@ use super::array_static_filter::ArrayStaticFilter; use super::branchless_filter::{ BranchlessFilter, BranchlessFilterType, BranchlessNative, }; +use super::byte_view_filter::instantiate_byte_view_filter; use super::primitive_filter::*; use super::static_filter::StaticFilter; type StaticFilterRef = Arc; -pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result { +pub(super) fn instantiate_static_filter( + in_array: ArrayRef, + expr_data_type: &DataType, +) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; + // Byte-view filters inspect the physical view representation directly. + if dictionary_value_type(expr_data_type) == in_array.data_type() + && let Some(filter) = instantiate_byte_view_filter(&in_array)? + { + return Ok(filter); + } + if let Some(filter) = instantiate_branchless_filter(&in_array)? { return Ok(filter); } @@ -49,6 +60,13 @@ pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result &DataType { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type +} + fn flatten_dictionary_haystack(in_array: ArrayRef) -> Result { // Flatten dictionary-encoded haystacks to their value type so that // specialized filters (e.g. Int32StaticFilter) are used instead of