diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 0a4ad0b804f0c..7a7cb25317c1d 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 dictionary_filter; mod fixed_size_binary_filter; mod primitive_filter; @@ -3648,6 +3649,23 @@ mod tests { )? ); + // Utf8View in_array, Utf8View and Dict(Utf8View) needles + let utf8view_in = + Arc::new(StringViewArray::from(vec!["a", "b", "c", "d", "e"])) as ArrayRef; + let utf8view_needle = + Arc::new(StringViewArray::from(vec!["a", "missing", "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/byte_view_filter.rs b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs new file mode 100644 index 0000000000000..9d4ea1fdfbb42 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs @@ -0,0 +1,227 @@ +// 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. + +//! Filters for `Utf8View` and `BinaryView` `IN` lists whose non-null values are +//! at most 12 bytes. +//! +//! Arrow stores such values directly in a `u128` view: the low 32 bits contain +//! the length and the remaining bits contain zero-padded bytes. Comparing two +//! inline views compares their complete values. Long views contain a prefix and +//! buffer location instead, so a list containing one uses the generic filter. +//! +//! Lists with up to four non-null values use the existing 128-bit branchless +//! filter. Larger lists use the primitive hash-set filter over the same 128 +//! bits. `Decimal128Type` is only a carrier; no decimal operations are used. +//! A long input value cannot match an inline list value because its length is +//! part of the view. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, MAX_INLINE_VIEW_LEN, PrimitiveArray, +}; +use arrow::buffer::{NullBuffer, ScalarBuffer}; +use arrow::datatypes::{DataType, Decimal128Type}; +use arrow::util::bit_iterator::BitIndexIterator; +use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; + +use super::primitive_filter::instantiate_primitive_filter; +use super::static_filter::{StaticFilter, StaticFilterRef}; + +fn is_inline(view: u128) -> bool { + (view as u32) <= MAX_INLINE_VIEW_LEN +} + +fn all_inline(views: &ScalarBuffer, nulls: Option<&NullBuffer>) -> bool { + match nulls { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .all(|idx| is_inline(views[idx])) + } + None => views.iter().copied().all(is_inline), + } +} + +fn as_decimal128( + views: &ScalarBuffer, + nulls: Option<&NullBuffer>, +) -> PrimitiveArray { + // `views.inner()` is already sliced to the array's offset. + let values = ScalarBuffer::::new(views.inner().clone(), 0, views.len()); + PrimitiveArray::::new(values, nulls.cloned()) +} + +/// Adapts the selected primitive filter to the original byte-view type. +/// +/// Arrow requires unused inline bytes to be zero, so equal values have equal +/// `u128` views. +struct ByteViewFilter { + data_type: DataType, + inner: StaticFilterRef, +} + +impl StaticFilter for ByteViewFilter { + fn null_count(&self) -> usize { + self.inner.null_count() + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + let (views, nulls) = match &self.data_type { + DataType::Utf8View => v.as_string_view_opt().map(|a| (a.views(), a.nulls())), + DataType::BinaryView => { + v.as_binary_view_opt().map(|a| (a.views(), a.nulls())) + } + _ => unreachable!(), + } + .ok_or_else(|| { + exec_datafusion_err!( + "Expected concrete {} array, got {}", + self.data_type, + v.data_type() + ) + })?; + // List values are all inline (len <= 12). Long input views cannot match + // because their encoded length (> 12) is part of the 128-bit key, so + // input lengths do not need to be checked. + let decimal = as_decimal128(views, nulls); + self.inner.contains(&decimal, negated) + } +} + +/// Returns a filter when every non-null byte view is inline. +pub(super) fn instantiate_byte_view_filter( + in_array: &ArrayRef, +) -> Result> { + let views = match in_array.data_type() { + DataType::Utf8View => in_array.as_string_view().views(), + DataType::BinaryView => in_array.as_binary_view().views(), + _ => return Ok(None), + }; + + if !all_inline(views, in_array.nulls()) { + return Ok(None); + } + + let primitive: ArrayRef = Arc::new(as_decimal128(views, in_array.nulls())); + let inner = instantiate_primitive_filter(&primitive)?.ok_or_else(|| { + internal_datafusion_err!( + "Byte view filter: no primitive filter for {}", + primitive.data_type() + ) + })?; + Ok(Some(Arc::new(ByteViewFilter { + data_type: in_array.data_type().clone(), + inner, + }))) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{BinaryViewArray, StringViewArray}; + + fn assert_contains( + filter: &dyn StaticFilter, + needles: &dyn Array, + expected: Vec>, + ) -> Result<()> { + assert_eq!( + filter.contains(needles, false)?, + BooleanArray::from(expected) + ); + Ok(()) + } + + #[test] + fn large_list_nulls_and_slices() -> 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 = instantiate_byte_view_filter(&haystack)?.unwrap(); + let needles = StringViewArray::from(vec![ + Some("outside"), + Some("b"), + Some("missing"), + None, + Some("e"), + Some("tail"), + ]) + .slice(1, 4); + + 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 routing_and_boundaries() -> Result<()> { + let inline = "abcdefghijkl"; + let long = "abcdefghijklm"; + let needles = StringViewArray::from(vec![inline, long, "missing"]); + + // Covers both sides of the shared primitive-filter cutoff. + for values in [ + vec![inline, "one", "two", "three"], + vec![inline, "one", "two", "three", "four"], + ] { + let haystack = Arc::new(StringViewArray::from(values)) as ArrayRef; + let filter = instantiate_byte_view_filter(&haystack)?.unwrap(); + assert_contains( + &*filter, + &needles, + vec![Some(true), Some(false), Some(false)], + )?; + } + + let mixed: ArrayRef = Arc::new(StringViewArray::from(vec!["short", long])); + assert!(instantiate_byte_view_filter(&mixed)?.is_none()); + + let binary: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + b"a".as_slice(), + b"b".as_slice(), + b"c".as_slice(), + b"d".as_slice(), + b"e".as_slice(), + ])); + let filter = instantiate_byte_view_filter(&binary)?.unwrap(); + let needles = BinaryViewArray::from(vec![b"a".as_slice(), b"z".as_slice()]); + assert_contains(&*filter, &needles, vec![Some(true), Some(false)])?; + + let mismatch_err = filter + .contains(&StringViewArray::from(vec!["a"]), false) + .unwrap_err(); + assert!( + mismatch_err + .to_string() + .contains("Expected concrete BinaryView array, got Utf8View") + ); + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index 093752cce04f2..5cbdf8cd7486a 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -23,6 +23,7 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use super::array_static_filter::ArrayStaticFilter; +use super::byte_view_filter::instantiate_byte_view_filter; use super::dictionary_filter::DictionaryFilter; use super::fixed_size_binary_filter::instantiate_fixed_size_binary_filter; use super::primitive_filter::instantiate_primitive_filter; @@ -34,7 +35,11 @@ pub(super) fn instantiate_static_filter( ) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; - let filter = if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { + let filter = if view_types_match(needle_data_type, in_array.data_type()) + && let Some(filter) = instantiate_byte_view_filter(&in_array)? + { + filter + } else if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { filter } else if let Some(filter) = instantiate_primitive_filter(&in_array)? { filter @@ -51,6 +56,20 @@ pub(super) fn instantiate_static_filter( } } +/// Raw view access requires the expression and list to use the same view type. +/// Dictionary wrappers do not change the expression's value type. +fn view_types_match(needle_type: &DataType, list_type: &DataType) -> bool { + matches!(list_type, DataType::Utf8View | DataType::BinaryView) + && dictionary_value_type(needle_type) == list_type +} + +fn dictionary_value_type(mut data_type: &DataType) -> &DataType { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type +} + fn flatten_dictionary_haystack(mut in_array: ArrayRef) -> Result { // Flatten every dictionary layer so the final value type can use a // specialized filter. @@ -109,4 +128,18 @@ mod tests { Ok(()) } + + #[test] + fn byte_view_routing_requires_same_physical_type() { + let dict_view = + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)); + + assert!(view_types_match(&DataType::Utf8View, &DataType::Utf8View)); + assert!(view_types_match(&dict_view, &DataType::Utf8View)); + assert!(!view_types_match( + &DataType::Utf8View, + &DataType::BinaryView + )); + assert!(!view_types_match(&DataType::Utf8, &DataType::Utf8View)); + } }