|
17 | 17 |
|
18 | 18 | use std::sync::Arc; |
19 | 19 |
|
20 | | -use arrow::array::{Array, BooleanArray}; |
| 20 | +use arrow::array::{Array, AsArray, BooleanArray}; |
| 21 | +use arrow::compute::take; |
| 22 | +use arrow::datatypes::DataType; |
21 | 23 | use datafusion_common::Result; |
22 | 24 |
|
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 | +} |
24 | 59 |
|
25 | 60 | /// Trait for InList static filters. |
26 | 61 | /// |
27 | 62 | /// Static filters store a pre-computed set of values (the haystack) and check |
28 | 63 | /// whether needle values are contained in that set. The haystack is always |
29 | 64 | /// represented in its non-dictionary (value) type. Dictionary haystacks are |
30 | | -/// flattened via `cast()` before construction. |
| 65 | +/// flattened before construction. |
31 | 66 | /// |
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 { |
35 | 70 | fn null_count(&self) -> usize; |
36 | 71 |
|
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>; |
41 | 74 | } |
42 | 75 |
|
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, |
55 | 84 | }; |
56 | | -} |
57 | 85 |
|
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