Skip to content

Commit 7055e43

Browse files
Handle dictionary IN LIST inputs in one place
Wrap the selected filter when the input is dictionary encoded, instead of teaching every filter implementation how to unpack dictionary keys. Plain inputs continue to use their filter directly. Flatten nested dictionary values before choosing a filter, and cover dictionary key types, nested dictionaries, nulls, timestamps, decimals, and scalar inputs.
1 parent c396e6a commit 7055e43

7 files changed

Lines changed: 267 additions & 79 deletions

File tree

‎datafusion/physical-expr/src/expressions/in_list.rs‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,7 @@ impl InListExpr {
216216
expr,
217217
list,
218218
negated,
219-
Some(instantiate_static_filter(array)?),
219+
Some(instantiate_static_filter(array, &expr_data_type)?),
220220
))
221221
}
222222

@@ -243,7 +243,7 @@ impl InListExpr {
243243

244244
// Try to create a static filter if all list expressions are constants
245245
let static_filter = match try_evaluate_constant_list(&list, schema)? {
246-
Some(in_array) => Some(instantiate_static_filter(in_array)?),
246+
Some(in_array) => Some(instantiate_static_filter(in_array, &expr_data_type)?),
247247
None => None, // Non-constant expressions, fall back to dynamic evaluation
248248
};
249249

@@ -1265,6 +1265,31 @@ mod tests {
12651265
Ok(())
12661266
}
12671267

1268+
#[test]
1269+
fn in_list_nested_dictionary_scalar() -> Result<()> {
1270+
let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1271+
let batch = RecordBatch::try_new(
1272+
Arc::new(schema.clone()),
1273+
vec![Arc::new(Int32Array::from(vec![0, 0, 0]))],
1274+
)?;
1275+
let needle = lit(ScalarValue::Dictionary(
1276+
Box::new(DataType::Int16),
1277+
Box::new(ScalarValue::Dictionary(
1278+
Box::new(DataType::Int8),
1279+
Box::new(ScalarValue::Int32(Some(2))),
1280+
)),
1281+
));
1282+
1283+
let expr = in_list(needle, vec![lit(1_i32), lit(2_i32)], &false, &schema)?;
1284+
let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?;
1285+
assert_eq!(
1286+
as_boolean_array(&result),
1287+
&BooleanArray::from(vec![true, true, true])
1288+
);
1289+
1290+
Ok(())
1291+
}
1292+
12681293
#[test]
12691294
fn in_list_utf8_with_dict_types() -> Result<()> {
12701295
fn dict_lit(key_type: DataType, value: &str) -> Arc<dyn PhysicalExpr> {

‎datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs‎

Lines changed: 2 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use arrow::array::{
19-
Array, ArrayRef, BooleanArray, downcast_array, downcast_dictionary_array,
20-
make_comparator,
21-
};
18+
use arrow::array::{Array, ArrayRef, BooleanArray, make_comparator};
2219
use arrow::buffer::{BooleanBuffer, NullBuffer};
23-
use arrow::compute::{SortOptions, take};
20+
use arrow::compute::SortOptions;
2421
use arrow::datatypes::DataType;
2522
use arrow::util::bit_iterator::BitIndexIterator;
2623
use datafusion_common::Result;
@@ -141,22 +138,6 @@ impl StaticFilter for ArrayStaticFilter {
141138
));
142139
}
143140

144-
// Unwrap dictionary-encoded needles when the value type matches
145-
// in_array, evaluating against the dictionary values and mapping
146-
// back via keys.
147-
downcast_dictionary_array! {
148-
v => {
149-
// Only unwrap when the haystack (in_array) type matches
150-
// the dictionary value type
151-
if v.values().data_type() == self.in_array.data_type() {
152-
let values_contains = self.contains(v.values().as_ref(), negated)?;
153-
let result = take(&values_contains, v.keys(), None)?;
154-
return Ok(downcast_array(result.as_ref()));
155-
}
156-
}
157-
_ => {}
158-
}
159-
160141
self.find_needles_in_haystack(v, negated)
161142
}
162143
}

‎datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ use arrow::util::bit_iterator::BitIndexIterator;
7777
use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err};
7878

7979
use super::result::build_result_from_contains;
80-
use super::static_filter::{StaticFilter, handle_dictionary};
80+
use super::static_filter::StaticFilter;
8181

8282
pub(super) type BranchlessNative<T> =
8383
<<T as BranchlessFilterType>::CompareType as ArrowPrimitiveType>::Native;
@@ -256,8 +256,6 @@ where
256256
}
257257

258258
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
259-
handle_dictionary!(self, v, negated);
260-
261259
// Arrow compatibility ignores timestamp timezone and decimal precision/scale
262260
// while still requiring the same primitive representation.
263261
if !PrimitiveArray::<T>::is_compatible(v.data_type()) {

‎datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ use arrow::datatypes::{
5353
use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err};
5454

5555
use super::primitive_filter::instantiate_primitive_filter;
56-
use super::static_filter::{StaticFilter, StaticFilterRef, handle_dictionary};
56+
use super::static_filter::{StaticFilter, StaticFilterRef};
5757

5858
/// Reinterpret fixed-size binary values as same-width primitive values.
5959
///
@@ -99,8 +99,6 @@ where
9999
}
100100

101101
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
102-
handle_dictionary!(self, v, negated);
103-
104102
if v.data_type() != &self.data_type {
105103
return Err(exec_datafusion_err!(
106104
"FixedSizeBinary filter: expected {} array, got {}",
@@ -165,6 +163,7 @@ mod tests {
165163
use arrow::buffer::{Buffer, MutableBuffer, NullBuffer};
166164
use arrow::datatypes::Int8Type;
167165

166+
use super::super::static_filter::DictionaryFilter;
168167
use super::*;
169168

170169
fn value(width: i32, index: usize, miss: bool) -> Vec<u8> {
@@ -268,7 +267,7 @@ mod tests {
268267

269268
#[test]
270269
fn handles_dictionary_needles() -> Result<()> {
271-
let filter = make_filter(4, &[Some(value(4, 7, false))])?;
270+
let filter = DictionaryFilter::new(make_filter(4, &[Some(value(4, 7, false))])?);
272271
let dictionary_values: ArrayRef = Arc::new(array(
273272
4,
274273
&[Some(value(4, 7, false)), Some(value(4, 8, false))],

‎datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,18 @@
1919
//!
2020
//! This module provides membership tests for Arrow primitive types.
2121
22+
use std::hash::{Hash, Hasher};
23+
use std::marker::PhantomData;
24+
use std::sync::Arc;
25+
2226
use arrow::array::{Array, ArrayRef, AsArray, BooleanArray};
2327
use arrow::datatypes::*;
2428
use arrow::util::bit_iterator::BitIndexIterator;
2529
use datafusion_common::{HashSet, Result, exec_datafusion_err};
26-
use std::hash::{Hash, Hasher};
27-
use std::marker::PhantomData;
28-
use std::sync::Arc;
2930

3031
use super::branchless_filter::{BranchlessFilter, BranchlessFilterType};
3132
use super::result::build_in_list_result;
32-
use super::static_filter::{StaticFilter, StaticFilterRef, handle_dictionary};
33+
use super::static_filter::{StaticFilter, StaticFilterRef};
3334

3435
/// Selects an optimized filter for a primitive representation.
3536
///
@@ -311,7 +312,6 @@ where
311312
}
312313

313314
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
314-
handle_dictionary!(self, v, negated);
315315
let v = v.as_primitive_opt::<T>().ok_or_else(|| {
316316
exec_datafusion_err!("BitmapFilter: expected {} array", T::DATA_TYPE)
317317
})?;
@@ -432,8 +432,6 @@ where
432432
}
433433

434434
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
435-
handle_dictionary!(self, v, negated);
436-
437435
let v = v.as_primitive_opt::<T>().ok_or_else(|| {
438436
exec_datafusion_err!(
439437
"PrimitiveHashSetFilter: expected {} array",
@@ -465,6 +463,8 @@ mod tests {
465463
};
466464
use half::f16;
467465

466+
use super::super::static_filter::DictionaryFilter;
467+
468468
fn uint32_array(values: Vec<Option<u32>>) -> ArrayRef {
469469
Arc::new(UInt32Array::from(values))
470470
}
@@ -553,13 +553,19 @@ mod tests {
553553
#[test]
554554
fn bitmap_filter_u8_handles_dictionary_needles() -> Result<()> {
555555
let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)]));
556-
let filter = BitmapFilter::<UInt8Type>::try_new(&haystack)?;
556+
let inner: StaticFilterRef =
557+
Arc::new(BitmapFilter::<UInt8Type>::try_new(&haystack)?);
558+
let filter = DictionaryFilter::new(inner);
557559

558560
let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]);
559561
let values = Arc::new(UInt8Array::from(vec![Some(1), Some(2), Some(3)]));
560562
let needles = DictionaryArray::try_new(keys, values)?;
561563

562-
assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])
564+
assert_eq!(
565+
filter.contains(&needles, false)?,
566+
BooleanArray::from(vec![Some(true), None, None, Some(true)])
567+
);
568+
Ok(())
563569
}
564570

565571
#[test]

‎datafusion/physical-expr/src/expressions/in_list/static_filter.rs‎

Lines changed: 144 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -17,42 +17,162 @@
1717

1818
use std::sync::Arc;
1919

20-
use arrow::array::{Array, BooleanArray};
20+
use arrow::array::{Array, AsArray, BooleanArray};
21+
use arrow::compute::take;
22+
use arrow::datatypes::DataType;
2123
use datafusion_common::Result;
2224

23-
pub(super) type StaticFilterRef = Arc<dyn StaticFilter + Send + Sync>;
25+
pub(super) type StaticFilterRef = Arc<dyn StaticFilter>;
26+
27+
/// Adds dictionary-encoded needle support to a filter that expects plain arrays.
28+
/// This wrapper is only used when the input expression returns dictionaries.
29+
pub(super) struct DictionaryFilter {
30+
inner: StaticFilterRef,
31+
}
32+
33+
impl DictionaryFilter {
34+
pub(super) fn new(inner: StaticFilterRef) -> Self {
35+
Self { inner }
36+
}
37+
}
38+
39+
impl StaticFilter for DictionaryFilter {
40+
fn null_count(&self) -> usize {
41+
self.inner.null_count()
42+
}
43+
44+
fn contains(&self, needles: &dyn Array, negated: bool) -> Result<BooleanArray> {
45+
let Some(dictionary) = needles.as_any_dictionary_opt() else {
46+
return self.inner.contains(needles, negated);
47+
};
48+
let values = dictionary.values();
49+
let values_contains = if matches!(values.data_type(), DataType::Dictionary(_, _))
50+
{
51+
self.contains(values.as_ref(), negated)?
52+
} else {
53+
self.inner.contains(values.as_ref(), negated)?
54+
};
55+
let result = take(&values_contains, dictionary.keys(), None)?;
56+
Ok(result.as_boolean().clone())
57+
}
58+
}
2459

2560
/// Trait for InList static filters.
2661
///
2762
/// Static filters store a pre-computed set of values (the haystack) and check
2863
/// whether needle values are contained in that set. The haystack is always
2964
/// represented in its non-dictionary (value) type. Dictionary haystacks are
30-
/// flattened via `cast()` before construction.
65+
/// flattened before construction.
3166
///
32-
/// Dictionary-encoded needles are unwrapped inside `contains()` and
33-
/// evaluated against the dictionary's values.
34-
pub(super) trait StaticFilter {
67+
/// Dictionary-encoded needles are handled by [`DictionaryFilter`] before the
68+
/// concrete filter is called.
69+
pub(super) trait StaticFilter: Send + Sync {
3570
fn null_count(&self) -> usize;
3671

37-
/// Checks if values in `v` (needle) are contained in this filter's
38-
/// haystack. `v` may be dictionary-encoded, in which case the
39-
/// implementation unwraps the dictionary and operates on its values.
40-
fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray>;
72+
/// Checks if non-dictionary needles are contained in this filter's haystack.
73+
fn contains(&self, needles: &dyn Array, negated: bool) -> Result<BooleanArray>;
4174
}
4275

43-
/// Evaluate dictionary-encoded needles by applying a filter to dictionary
44-
/// values and remapping the result through the keys.
45-
macro_rules! handle_dictionary {
46-
($self:ident, $v:ident, $negated:ident) => {
47-
arrow::array::downcast_dictionary_array! {
48-
$v => {
49-
let values_contains = $self.contains($v.values().as_ref(), $negated)?;
50-
let result = arrow::compute::take(&values_contains, $v.keys(), None)?;
51-
return Ok(arrow::array::downcast_array(result.as_ref()))
52-
}
53-
_ => {}
54-
}
76+
#[cfg(test)]
77+
mod tests {
78+
use super::*;
79+
80+
use arrow::array::{
81+
ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int16Array, Int32Array,
82+
Int64Array, TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array,
83+
UInt64Array,
5584
};
56-
}
5785

58-
pub(super) use handle_dictionary;
86+
use super::super::strategy::instantiate_static_filter;
87+
88+
#[test]
89+
fn dictionary_needles_support_all_key_types() -> Result<()> {
90+
let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
91+
let expected = BooleanArray::from(vec![Some(true), Some(false), None]);
92+
93+
macro_rules! check_keys {
94+
($keys:expr) => {{
95+
let needles = DictionaryArray::try_new($keys, Arc::clone(&values))?;
96+
let filter = instantiate_static_filter(
97+
Arc::new(Int32Array::from(vec![1, 3])),
98+
needles.data_type(),
99+
)?;
100+
assert_eq!(filter.contains(&needles, false)?, expected);
101+
}};
102+
}
103+
104+
check_keys!(Int8Array::from(vec![Some(0), Some(1), None]));
105+
check_keys!(Int16Array::from(vec![Some(0), Some(1), None]));
106+
check_keys!(Int32Array::from(vec![Some(0), Some(1), None]));
107+
check_keys!(Int64Array::from(vec![Some(0), Some(1), None]));
108+
check_keys!(UInt8Array::from(vec![Some(0), Some(1), None]));
109+
check_keys!(UInt16Array::from(vec![Some(0), Some(1), None]));
110+
check_keys!(UInt32Array::from(vec![Some(0), Some(1), None]));
111+
check_keys!(UInt64Array::from(vec![Some(0), Some(1), None]));
112+
113+
Ok(())
114+
}
115+
116+
#[test]
117+
fn nested_dictionary_needles_preserve_nulls() -> Result<()> {
118+
let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None]));
119+
let inner: ArrayRef = Arc::new(DictionaryArray::try_new(
120+
Int8Array::from(vec![0, 1, 2]),
121+
values,
122+
)?);
123+
let needles = DictionaryArray::try_new(
124+
Int16Array::from(vec![Some(0), Some(1), Some(2), None]),
125+
inner,
126+
)?;
127+
let filter = instantiate_static_filter(
128+
Arc::new(Int32Array::from(vec![1, 3])),
129+
needles.data_type(),
130+
)?;
131+
132+
assert_eq!(
133+
filter.contains(&needles, false)?,
134+
BooleanArray::from(vec![Some(true), Some(false), None, None])
135+
);
136+
assert_eq!(
137+
filter.contains(&needles, true)?,
138+
BooleanArray::from(vec![Some(false), Some(true), None, None])
139+
);
140+
141+
Ok(())
142+
}
143+
144+
#[test]
145+
fn dictionary_timestamp_needles_keep_timezone_compatibility() -> Result<()> {
146+
let timestamps: ArrayRef =
147+
Arc::new(TimestampNanosecondArray::from(vec![1, 3]).with_timezone("UTC"));
148+
let values: ArrayRef = Arc::new(
149+
TimestampNanosecondArray::from(vec![1, 2]).with_timezone("Europe/Paris"),
150+
);
151+
let needles = DictionaryArray::try_new(Int8Array::from(vec![0, 1]), values)?;
152+
let filter = instantiate_static_filter(timestamps, needles.data_type())?;
153+
assert_eq!(
154+
filter.contains(&needles, false)?,
155+
BooleanArray::from(vec![true, false])
156+
);
157+
158+
Ok(())
159+
}
160+
161+
#[test]
162+
fn dictionary_decimal_needles_keep_precision_scale_compatibility() -> Result<()> {
163+
// Five list values use the hash-set filter.
164+
let decimals: ArrayRef = Arc::new(
165+
Decimal128Array::from(vec![1, 3, 5, 7, 9]).with_precision_and_scale(10, 2)?,
166+
);
167+
let values: ArrayRef =
168+
Arc::new(Decimal128Array::from(vec![1, 2]).with_precision_and_scale(11, 3)?);
169+
let needles = DictionaryArray::try_new(Int8Array::from(vec![0, 1]), values)?;
170+
let filter = instantiate_static_filter(decimals, needles.data_type())?;
171+
assert_eq!(
172+
filter.contains(&needles, false)?,
173+
BooleanArray::from(vec![true, false])
174+
);
175+
176+
Ok(())
177+
}
178+
}

0 commit comments

Comments
 (0)