|
19 | 19 |
|
20 | 20 | use std::sync::Arc; |
21 | 21 |
|
22 | | -use arrow::array::{ArrayRef, StructArray}; |
23 | | -use arrow::datatypes::{Field, FieldRef, Fields}; |
| 22 | +use arrow::array::{ArrayRef, AsArray, StructArray}; |
| 23 | +use arrow::compute::cast; |
| 24 | +use arrow::datatypes::{Field, FieldRef, Fields, Int32Type}; |
24 | 25 | use arrow_schema::DataType; |
25 | 26 | use datafusion_common::Result; |
26 | 27 |
|
@@ -77,6 +78,25 @@ pub(super) fn build_struct_inlist_values( |
77 | 78 | Ok(Some(source_array)) |
78 | 79 | } |
79 | 80 |
|
| 81 | +/// Deduplicates `array`'s values (dropping nulls) via Arrow's dictionary-encoding cast, |
| 82 | +/// cheaper than converting every row to a `ScalarValue` to test uniqueness. |
| 83 | +/// |
| 84 | +/// Returns `None` if the type isn't dictionary-encodable or every value is null - |
| 85 | +/// both safe to treat as "no pruning literals available". |
| 86 | +pub(super) fn dedupe_array_values(array: &ArrayRef) -> Option<ArrayRef> { |
| 87 | + let dict_type = DataType::Dictionary( |
| 88 | + Box::new(DataType::Int32), |
| 89 | + Box::new(array.data_type().clone()), |
| 90 | + ); |
| 91 | + let dict = cast(array.as_ref(), &dict_type).ok()?; |
| 92 | + let values = Arc::clone(dict.as_dictionary::<Int32Type>().values()); |
| 93 | + if values.is_empty() { |
| 94 | + None |
| 95 | + } else { |
| 96 | + Some(values) |
| 97 | + } |
| 98 | +} |
| 99 | + |
80 | 100 | #[cfg(test)] |
81 | 101 | mod tests { |
82 | 102 | use super::*; |
@@ -142,6 +162,27 @@ mod tests { |
142 | 162 | ); |
143 | 163 | } |
144 | 164 |
|
| 165 | + #[test] |
| 166 | + fn test_dedupe_array_values() { |
| 167 | + let array = Arc::new(Int32Array::from(vec![ |
| 168 | + Some(1), |
| 169 | + Some(2), |
| 170 | + Some(1), |
| 171 | + None, |
| 172 | + Some(2), |
| 173 | + ])) as ArrayRef; |
| 174 | + let deduped = dedupe_array_values(&array).expect("dictionary-encodable"); |
| 175 | + assert_eq!(deduped.len(), 2); |
| 176 | + let deduped = deduped.as_ref().as_primitive::<Int32Type>(); |
| 177 | + assert_eq!(deduped.values(), &[1, 2]); |
| 178 | + } |
| 179 | + |
| 180 | + #[test] |
| 181 | + fn test_dedupe_array_values_all_null() { |
| 182 | + let array = Arc::new(Int32Array::from(vec![None, None])) as ArrayRef; |
| 183 | + assert!(dedupe_array_values(&array).is_none()); |
| 184 | + } |
| 185 | + |
145 | 186 | #[test] |
146 | 187 | fn test_build_single_column_dictionary_inlist() { |
147 | 188 | let keys = Int8Array::from(vec![0i8, 0, 0]); |
|
0 commit comments