diff --git a/Cargo.lock b/Cargo.lock index be40df978369d..d1aecaccdc133 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1925,6 +1925,7 @@ dependencies = [ "arrow", "arrow-ipc", "arrow-schema", + "arrow-select", "chrono", "criterion", "foldhash 0.2.0", diff --git a/clippy.toml b/clippy.toml index 7b781d9b6605f..63d6906b8a70c 100644 --- a/clippy.toml +++ b/clippy.toml @@ -1,6 +1,8 @@ disallowed-methods = [ { path = "tokio::task::spawn", reason = "To provide cancel-safety, use `SpawnedTask::spawn` instead (https://github.com/apache/datafusion/issues/6513)" }, { path = "tokio::task::spawn_blocking", reason = "To provide cancel-safety, use `SpawnedTask::spawn_blocking` instead (https://github.com/apache/datafusion/issues/6513)" }, + { path = "arrow_array::array::AnyDictionaryArray::with_values", reason = "Running code on every dictionary value can cost far more than the batch has rows; use `datafusion_common::utils::map_dictionary_values`" }, + { path = "arrow_array::array::DictionaryArray::with_values", reason = "Running code on every dictionary value can cost far more than the batch has rows; use `datafusion_common::utils::map_dictionary_values`" }, { path = "std::vec::Vec::reserve", reason = "Use `Vec::try_reserve` so allocation failures can be reported instead of panicking", replacement = "try_reserve" }, ] diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 9ee199fe82f28..4cb3ff570fb4e 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -72,6 +72,7 @@ name = "record_batch_memory" arrow = { workspace = true } arrow-ipc = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } +arrow-select = { workspace = true } chrono = { workspace = true } foldhash = "0.2" half = { workspace = true } diff --git a/datafusion/common/benches/with_hashes.rs b/datafusion/common/benches/with_hashes.rs index fe7ddca2f9ebb..ee08f2af3a0ac 100644 --- a/datafusion/common/benches/with_hashes.rs +++ b/datafusion/common/benches/with_hashes.rs @@ -20,7 +20,7 @@ use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, DictionaryArray, FixedSizeListArray, GenericStringArray, Int32Array, Int64Array, ListArray, ListViewArray, MapArray, - NullArray, NullBufferBuilder, OffsetSizeTrait, PrimitiveArray, RunArray, + NullArray, NullBufferBuilder, OffsetSizeTrait, PrimitiveArray, RunArray, StringArray, StringViewArray, StructArray, UnionArray, make_array, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; @@ -407,6 +407,22 @@ fn sliced_array_benchmark(c: &mut Criterion) { ); } } + + // Sliced DictionaryArray with one value per row of the full array, like the + // group keys an aggregation emits in `batch_size` slices. Hashing reduces + // it to the referenced values only above 16 values per row, so 1/8 keeps + // the full dictionary and 1/32 doesn't. + for ratio in [8, 32] { + let total_rows = BATCH_SIZE * ratio; + let full_array = distinct_dictionary_array(total_rows); + let sliced = full_array.slice(BATCH_SIZE * (ratio / 2), BATCH_SIZE); + c.bench_function( + &format!("dictionary_sliced: 1/{ratio} of {total_rows} values"), + |b| { + do_hash_test_with_len(b, std::slice::from_ref(&sliced), BATCH_SIZE); + }, + ); + } } fn do_hash_test_with_len(b: &mut Bencher, arrays: &[ArrayRef], expected_len: usize) { @@ -420,6 +436,14 @@ fn do_hash_test_with_len(b: &mut Bencher, arrays: &[ArrayRef], expected_len: usi }); } +/// A `Dictionary(Int32, Utf8)` with a distinct value for every row +fn distinct_dictionary_array(num_rows: usize) -> ArrayRef { + let values = + StringArray::from_iter_values((0..num_rows).map(|i| format!("group_{i}"))); + let keys = Int32Array::from_iter_values(0..num_rows as i32); + Arc::new(DictionaryArray::::try_new(keys, Arc::new(values)).unwrap()) +} + fn list_array(num_rows: usize) -> ArrayRef { let mut rng = make_rng(); let elements_per_row = 5; diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index ef29cc3ef8eb0..be17fd112e4ff 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -17,6 +17,8 @@ //! Functionality used both on logical and physical plans +#[cfg(not(feature = "force_hash_collisions"))] +use crate::utils::compact_oversized_dictionary; use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; use arrow::array::*; #[cfg(not(feature = "force_hash_collisions"))] @@ -671,7 +673,11 @@ fn hash_dictionary( ) -> Result<()> { // Hash each dictionary value once, and then use that computed // hash for each key value to avoid a potentially expensive - // redundant hashing for large dictionary elements (e.g. strings) + // redundant hashing for large dictionary elements (e.g. strings). + // A dictionary with many more values than rows (e.g. a slice) is first + // reduced to the values its rows reference. + let compacted = compact_oversized_dictionary(array)?; + let array = compacted.as_ref().unwrap_or(array); let dict_values = array.values(); let mut dict_hashes = vec![0; dict_values.len()]; create_hashes([dict_values], random_state, &mut dict_hashes)?; @@ -686,6 +692,8 @@ fn hash_dictionary_with_child_hashing( hashes_buffer: &mut [u64], multi_col: bool, ) -> Result<()> { + let compacted = compact_oversized_dictionary(array)?; + let array = compacted.as_ref().unwrap_or(array); let dict_values = array.values(); let mut dict_hashes = vec![0; dict_values.len()]; child_hashing.create_hashes([dict_values], &mut dict_hashes)?; @@ -1851,6 +1859,42 @@ mod tests { assert_ne!(dict_hashes[0], dict_hashes[2]); } + #[test] + // Tests actual values of hashes, which are different if forcing collisions + #[cfg(not(feature = "force_hash_collisions"))] + fn create_hashes_for_dict_slices_with_more_values_than_rows() { + use crate::utils::MAX_DICTIONARY_VALUES_PER_ROW; + + let values: ArrayRef = Arc::new(StringArray::from_iter( + (0..40).map(|i| (i != 2).then(|| format!("v{i}"))), + )); + let keys = Int8Array::from_iter((0..48).map(|i| (i % 7 != 1).then_some(i % 40))); + let dict: ArrayRef = Arc::new(DictionaryArray::try_new(keys, values).unwrap()); + let ints: ArrayRef = Arc::new(Int32Array::from_iter_values(0..48)); + let random_state = RandomState::with_seed(0); + + let mut single = vec![0; dict.len()]; + create_hashes([&dict], &random_state, &mut single).unwrap(); + let mut multi = vec![0; dict.len()]; + create_hashes([&ints, &dict], &random_state, &mut multi).unwrap(); + + for (offset, len) in [(0, 2), (1, 2), (8, 1), (46, 2)] { + let slice = dict.slice(offset, len); + assert!( + slice.as_any_dictionary().values().len() + > MAX_DICTIONARY_VALUES_PER_ROW * len + ); + let mut hashes = vec![0; len]; + create_hashes([&slice], &random_state, &mut hashes).unwrap(); + assert_eq!(hashes, single[offset..offset + len]); + + let ints = ints.slice(offset, len); + let mut hashes = vec![0; len]; + create_hashes([&ints, &slice], &random_state, &mut hashes).unwrap(); + assert_eq!(hashes, multi[offset..offset + len]); + } + } + #[test] // Tests actual values of hashes, which are different if forcing collisions #[cfg(not(feature = "force_hash_collisions"))] diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 7811f522ea121..65fc8478cbbdd 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -26,7 +26,8 @@ pub mod string_utils; use crate::assert_or_internal_err; use crate::error::{ - _exec_datafusion_err, _exec_err, _internal_datafusion_err, _plan_datafusion_err, + _exec_datafusion_err, _exec_err, _internal_datafusion_err, _internal_err, + _plan_datafusion_err, }; use crate::{Result, ScalarValue}; use arrow::array::{ @@ -34,17 +35,20 @@ use arrow::array::{ cast::AsArray, }; use arrow::array::{ - ArrowPrimitiveType, BooleanArray, Datum, GenericListArray, Int32Array, Int64Array, - MutableArrayData, PrimitiveArray, layout, make_array, + ArrowPrimitiveType, BooleanArray, Datum, DictionaryArray, GenericListArray, + Int32Array, Int64Array, MutableArrayData, PrimitiveArray, layout, make_array, }; use arrow::array::{LargeListViewArray, ListViewArray}; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::kernels::cmp::eq; use arrow::compute::kernels::length::length; -use arrow::compute::{SortColumn, SortOptions, nullif, partition}; +use arrow::compute::{SortColumn, SortOptions, nullif, partition, take}; use arrow::datatypes::{ - ArrowNativeType, DataType, Field, Int32Type, Int64Type, Metadata, SchemaRef, + ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int32Type, Int64Type, + Metadata, SchemaRef, }; +use arrow::downcast_dictionary_array; +use arrow_select::dictionary::garbage_collect_any_dictionary; #[cfg(feature = "sql")] use sqlparser::{ast::Ident, dialect::GenericDialect, parser::Parser}; use std::borrow::{Borrow, Cow}; @@ -416,6 +420,70 @@ pub fn split_vec_min_alloc(vec: &mut Vec, n: usize) -> Vec { } } +/// Applies `f` to the values of the dictionary `array` and keeps its keys. +/// +/// Slices, filters and repartitioned batches keep the whole values array of +/// the dictionary they came from, so running `f` on every value can cost far +/// more than the batch has rows. `f` never runs on more than +/// [`MAX_DICTIONARY_VALUES_PER_ROW`] values per row: past that, it runs on the +/// values the rows reference instead. +/// +/// `f` can also see values that no row references. If it fails, it runs once +/// more on only the referenced values, so it fails only when running it on +/// each row would. +pub fn map_dictionary_values(array: &ArrayRef, mut f: F) -> Result +where + F: FnMut(&ArrayRef) -> Result, +{ + let compacted = downcast_dictionary_array! { + array => compact_oversized_dictionary(array)?.map(|d| Arc::new(d) as ArrayRef), + dt => return _internal_err!("expected a dictionary array, got {dt}"), + }; + let mut dictionary = compacted.unwrap_or_else(|| Arc::clone(array)); + let values = match f(dictionary.as_any_dictionary().values()) { + Ok(values) => values, + Err(e) => { + let referenced = + garbage_collect_any_dictionary(dictionary.as_any_dictionary())?; + if referenced.as_any_dictionary().values().len() + == dictionary.as_any_dictionary().values().len() + { + return Err(e); + } + dictionary = referenced; + f(dictionary.as_any_dictionary().values())? + } + }; + #[expect(clippy::disallowed_methods)] // the values are bounded above + Ok(dictionary.as_any_dictionary().with_values(values)) +} + +/// The most dictionary values per row that code processing every value of a +/// dictionary should handle. Past this, taking the values the rows reference +/// is cheaper than processing all of them, even for the cheapest per-value +/// work: the break-even is about 4 values per row for hashing and about 13 for +/// ASCII `lower`/`upper`, which convert the whole values buffer at once. Any +/// bound removes the cost that grows with the dictionary size; the exact value +/// only trades a constant factor. +pub const MAX_DICTIONARY_VALUES_PER_ROW: usize = 16; + +/// Returns `None` when `array` has at most [`MAX_DICTIONARY_VALUES_PER_ROW`] +/// values per row, and otherwise a copy whose values are `take(values, keys)` +/// and whose keys are `0..len`, so that work over its values costs O(rows). +pub(crate) fn compact_oversized_dictionary( + array: &DictionaryArray, +) -> Result>> { + if array.values().len() <= MAX_DICTIONARY_VALUES_PER_ROW * array.len() { + return Ok(None); + } + let values = take(array.values().as_ref(), array.keys(), None)?; + // `len < values.len()`, so every row index fits in the key type + let keys: ScalarBuffer = + (0..array.len()).map(K::Native::usize_as).collect(); + let keys = PrimitiveArray::::new(keys, array.keys().nulls().cloned()); + Ok(Some(DictionaryArray::try_new(keys, values)?)) +} + #[cfg(test)] mod split_vec_min_alloc_tests { use super::split_vec_min_alloc; @@ -1501,14 +1569,15 @@ pub fn apply_parent_nulls( #[cfg(test)] mod tests { + use std::cell::Cell; use std::collections::HashMap; use std::sync::Arc; use super::*; use crate::ScalarValue::Null; use arrow::{ - array::{Float64Array, Int32Array, NullArray}, - datatypes::Int32Type, + array::{Float64Array, Int32Array, NullArray, StringArray}, + datatypes::{Int8Type, Int32Type}, }; #[cfg(feature = "sql")] use sqlparser::ast::Ident; @@ -2106,4 +2175,71 @@ mod tests { fsl_values_row_number(-1, 2).unwrap_err(); fsl_values_row_number(-1, 0).unwrap_err(); } + + #[test] + fn map_dictionary_values_bounds_values_per_row() -> Result<()> { + let strings = (0..24) + .map(|i| (i % 5 != 1).then(|| "x".repeat(i))) + .collect::>(); + let dict: DictionaryArray = + strings.iter().map(|s| s.as_deref()).collect(); + let dict: ArrayRef = Arc::new(dict); + for (offset, len) in [(0, 24), (1, 1), (6, 1)] { + let slice = dict.slice(offset, len); + let mut seen = 0; + let mapped = map_dictionary_values(&slice, |values| { + seen = values.len(); + Ok(length(values)?) + })?; + assert!( + seen <= MAX_DICTIONARY_VALUES_PER_ROW * len, + "{seen} values for {len} rows" + ); + let expected = length(&arrow::compute::cast(&slice, &DataType::Utf8)?)?; + assert_eq!(&arrow::compute::cast(&mapped, &DataType::Int32)?, &expected); + } + Ok(()) + } + + #[test] + fn map_dictionary_values_fails_only_on_referenced_values() -> Result<()> { + let calls = Cell::new(0); + let fail_on_bad = |values: &ArrayRef| { + calls.set(calls.get() + 1); + if values.as_string::().iter().any(|v| v == Some("bad")) { + return _exec_err!("bad value"); + } + Ok(length(values)?) + }; + let values: ArrayRef = + Arc::new(StringArray::from(vec!["ok", "bad", "fine", "bad"])); + + // "bad" at 1 is behind null keys only, and at 3 has no key. There are + // more rows than Int8 keys can index, so the retry can't give each row + // its own key. + let keys = (0..200) + .map(|i| [1, 0, 2][i % 3]) + .collect::>(); + let nulls = NullBuffer::from_iter((0..200).map(|i| i % 3 != 0)); + let keys = PrimitiveArray::::new(keys, Some(nulls)); + let dict: ArrayRef = + Arc::new(DictionaryArray::try_new(keys, Arc::clone(&values))?); + let mapped = map_dictionary_values(&dict, fail_on_bad)?; + assert_eq!( + mapped.data_type(), + &DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)) + ); + let expected = length(&arrow::compute::cast(&dict, &DataType::Utf8)?)?; + assert_eq!(&arrow::compute::cast(&mapped, &DataType::Int32)?, &expected); + assert_eq!(calls.get(), 2); + + // A row references "bad", so `f` must still fail, without a retry when + // every value is referenced + calls.set(0); + let keys = PrimitiveArray::::from(vec![0, 1, 2, 3]); + let dict: ArrayRef = Arc::new(DictionaryArray::try_new(keys, values)?); + map_dictionary_values(&dict, fail_on_bad).unwrap_err(); + assert_eq!(calls.get(), 1); + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 9c27b8b249f33..8ce9718aca7d3 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -36,6 +36,7 @@ use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion_common::cast::as_list_array; use datafusion_common::hash_utils::{RandomState, create_hashes}; use datafusion_common::scalar::copy_array_data; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::utils::memory::estimate_memory_size; use datafusion_common::utils::proxy::HashTableAllocExt; use datafusion_common::utils::{SingleRowListArrayBuilder, offset_span}; @@ -2061,11 +2062,11 @@ fn compact_payload(array: ArrayRef) -> Result { DataType::BinaryView => Ok(Arc::new(array.as_binary_view().gc())), DataType::Dictionary(_, _) => { let array = garbage_collect_any_dictionary(array.as_any_dictionary())?; - let dictionary = array.as_any_dictionary(); // GC can return the original values if every entry is referenced. // Detach offset slices as well as any nested view/dictionary buffers. - let values = make_array(copy_array_data(&dictionary.values().to_data())); - Ok(dictionary.with_values(compact_payload(values)?)) + map_dictionary_values(&array, |values| { + compact_payload(make_array(copy_array_data(&values.to_data()))) + }) } data_type if data_type.is_nested() => { let data = array.to_data(); diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index 07ae4b1f2e57e..137e36d06381d 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -121,6 +121,8 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result { diff --git a/datafusion/functions/src/regex/regexpreplace.rs b/datafusion/functions/src/regex/regexpreplace.rs index 26683a89cb091..ad1c2efab68ff 100644 --- a/datafusion/functions/src/regex/regexpreplace.rs +++ b/datafusion/functions/src/regex/regexpreplace.rs @@ -34,7 +34,7 @@ use datafusion_common::cast::{ }; use datafusion_common::exec_err; use datafusion_common::plan_err; -use datafusion_common::utils::offset_span_len; +use datafusion_common::utils::{map_dictionary_values, offset_span_len}; use datafusion_common::{Result, cast::as_generic_string_array, internal_err}; use datafusion_expr::ColumnarValue; use datafusion_expr::function::Hint; @@ -188,18 +188,21 @@ impl ScalarUDFImpl for RegexpReplaceFunc { { if preserve_dictionary { let compacted; - let mut dict = array.as_any_dictionary(); - if dict.values().len() > dict.keys().len() { + let mut dict = array; + if array.as_any_dictionary().values().len() > array.len() { // Filtered batches and shared Parquet dictionaries can have more // values than rows, so drop unreferenced values first. - compacted = garbage_collect_any_dictionary(dict)?; - dict = compacted.as_any_dictionary(); + compacted = + garbage_collect_any_dictionary(array.as_any_dictionary())?; + dict = &compacted; } - let mut inner: Vec = Vec::with_capacity(args.len()); - inner.push(ColumnarValue::Array(Arc::clone(dict.values()))); - inner.extend_from_slice(&args[1..]); - let values_out = regexp_replace_func(&inner)?; - return Ok(ColumnarValue::Array(dict.with_values(values_out))); + let result = map_dictionary_values(dict, |values| { + let mut inner: Vec = Vec::with_capacity(args.len()); + inner.push(ColumnarValue::Array(Arc::clone(values))); + inner.extend_from_slice(&args[1..]); + regexp_replace_func(&inner) + })?; + return Ok(ColumnarValue::Array(result)); } args[0] = ColumnarValue::Array(arrow::compute::cast(array, &value_type)?); } diff --git a/datafusion/functions/src/string/ascii.rs b/datafusion/functions/src/string/ascii.rs index db539a4d11719..adf3363721f83 100644 --- a/datafusion/functions/src/string/ascii.rs +++ b/datafusion/functions/src/string/ascii.rs @@ -20,6 +20,7 @@ use arrow::array::{ArrayRef, AsArray, Int32Array, StringArrayType}; use arrow::datatypes::DataType; use arrow::error::ArrowError; use datafusion_common::types::logical_string; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::{ @@ -190,9 +191,7 @@ pub fn ascii(args: &[ArrayRef]) -> Result { Ok(calculate_ascii(&string_array)?) } DataType::Dictionary(_, _) => { - let dictionary = args[0].as_any_dictionary(); - let converted = ascii(&[Arc::clone(dictionary.values())])?; - Ok(dictionary.with_values(converted)) + map_dictionary_values(&args[0], |values| ascii(&[Arc::clone(values)])) } _ => internal_err!("Unsupported data type"), } diff --git a/datafusion/functions/src/string/btrim.rs b/datafusion/functions/src/string/btrim.rs index 82e1f7d2c778d..7ef588b2db8a0 100644 --- a/datafusion/functions/src/string/btrim.rs +++ b/datafusion/functions/src/string/btrim.rs @@ -17,9 +17,10 @@ use crate::string::common::*; use crate::utils::make_scalar_function; -use arrow::array::{ArrayRef, AsArray}; +use arrow::array::ArrayRef; use arrow::datatypes::DataType; use datafusion_common::types::logical_string; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ @@ -43,9 +44,7 @@ fn btrim(args: &[ArrayRef]) -> Result { DataType::LargeUtf8 => general_trim::(&args, false), DataType::Utf8View => general_trim::(&args, true), DataType::Dictionary(_, _) => { - let dictionary = args[0].as_any_dictionary(); - let trimmed = btrim(&[Arc::clone(dictionary.values())])?; - Ok(dictionary.with_values(trimmed)) + map_dictionary_values(&args[0], |values| btrim(&[Arc::clone(values)])) } other => exec_err!( "Unsupported data type {other:?} for function btrim, expected Utf8, LargeUtf8 or Utf8View." diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 73d66045a0ecc..bf72d82933037 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -24,13 +24,14 @@ use crate::strings::{ StringViewArrayBuilder, substr_view, }; use arrow::array::{ - Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, StringViewArray, new_null_array, }; use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::utils::{offset_span, offset_span_len}; use datafusion_common::{ScalarValue, exec_err}; use datafusion_expr::ColumnarValue; @@ -541,9 +542,9 @@ fn case_conversion_dictionary( lower: bool, name: &str, ) -> Result { - let dictionary = array.as_any_dictionary(); - let converted = case_conversion_columnar_array(dictionary.values(), lower, name)?; - Ok(dictionary.with_values(converted)) + map_dictionary_values(array, |values| { + case_conversion_columnar_array(values, lower, name) + }) } fn case_conversion_array( diff --git a/datafusion/functions/src/string/ltrim.rs b/datafusion/functions/src/string/ltrim.rs index 04e33253ed6df..2fabf83821ae3 100644 --- a/datafusion/functions/src/string/ltrim.rs +++ b/datafusion/functions/src/string/ltrim.rs @@ -15,13 +15,14 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, AsArray}; +use arrow::array::ArrayRef; use arrow::datatypes::DataType; use std::sync::Arc; use crate::string::common::*; use crate::utils::make_scalar_function; use datafusion_common::types::logical_string; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ @@ -44,9 +45,7 @@ fn ltrim(args: &[ArrayRef]) -> Result { DataType::LargeUtf8 => general_trim::(&args, false), DataType::Utf8View => general_trim::(&args, true), DataType::Dictionary(_, _) => { - let dictionary = args[0].as_any_dictionary(); - let trimmed = ltrim(&[Arc::clone(dictionary.values())])?; - Ok(dictionary.with_values(trimmed)) + map_dictionary_values(&args[0], |values| ltrim(&[Arc::clone(values)])) } other => exec_err!( "Unsupported data type {other:?} for function ltrim, expected Utf8, LargeUtf8 or Utf8View." diff --git a/datafusion/functions/src/string/rtrim.rs b/datafusion/functions/src/string/rtrim.rs index d1126cb2418ae..0166dddf3a9df 100644 --- a/datafusion/functions/src/string/rtrim.rs +++ b/datafusion/functions/src/string/rtrim.rs @@ -15,13 +15,14 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, AsArray}; +use arrow::array::ArrayRef; use arrow::datatypes::DataType; use std::sync::Arc; use crate::string::common::*; use crate::utils::make_scalar_function; use datafusion_common::types::logical_string; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ @@ -44,9 +45,7 @@ fn rtrim(args: &[ArrayRef]) -> Result { DataType::LargeUtf8 => general_trim::(&args, false), DataType::Utf8View => general_trim::(&args, true), DataType::Dictionary(_, _) => { - let dictionary = args[0].as_any_dictionary(); - let trimmed = rtrim(&[Arc::clone(dictionary.values())])?; - Ok(dictionary.with_values(trimmed)) + map_dictionary_values(&args[0], |values| rtrim(&[Arc::clone(values)])) } other => exec_err!( "Unsupported data type {other:?} for function rtrim, expected Utf8, LargeUtf8 or Utf8View." diff --git a/datafusion/functions/src/unicode/character_length.rs b/datafusion/functions/src/unicode/character_length.rs index 9f0d952a02636..8b70cd9242ca3 100644 --- a/datafusion/functions/src/unicode/character_length.rs +++ b/datafusion/functions/src/unicode/character_length.rs @@ -25,6 +25,7 @@ use arrow::array::{ use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type}; use datafusion_common::Result; use datafusion_common::types::{NativeType, logical_string}; +use datafusion_common::utils::map_dictionary_values; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, @@ -124,11 +125,9 @@ fn character_length(args: &[ArrayRef]) -> Result { let string_array = args[0].as_string_view(); character_length_general::(&string_array) } - DataType::Dictionary(_, _) => { - let dictionary = args[0].as_any_dictionary(); - let converted = character_length(&[Arc::clone(dictionary.values())])?; - Ok(dictionary.with_values(converted)) - } + DataType::Dictionary(_, _) => map_dictionary_values(&args[0], |values| { + character_length(&[Arc::clone(values)]) + }), _ => unreachable!("CharacterLengthFunc"), } } diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs index d4f4f15323875..e8b0b656db29a 100644 --- a/datafusion/functions/src/unicode/initcap.rs +++ b/datafusion/functions/src/unicode/initcap.rs @@ -17,13 +17,14 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, AsArray, GenericStringArray, OffsetSizeTrait}; +use arrow::array::{Array, ArrayRef, GenericStringArray, OffsetSizeTrait}; use arrow::buffer::Buffer; use arrow::datatypes::DataType; use crate::strings::{GenericStringArrayBuilder, StringViewArrayBuilder}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; use datafusion_common::types::logical_string; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::utils::offset_span_len; use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ @@ -141,11 +142,7 @@ fn initcap_array(array: &ArrayRef) -> Result { DataType::Utf8 => initcap::(&[Arc::clone(array)]), DataType::LargeUtf8 => initcap::(&[Arc::clone(array)]), DataType::Utf8View => initcap_utf8view(&[Arc::clone(array)]), - DataType::Dictionary(_, _) => { - let dictionary = array.as_any_dictionary(); - let converted = initcap_array(dictionary.values())?; - Ok(dictionary.with_values(converted)) - } + DataType::Dictionary(_, _) => map_dictionary_values(array, initcap_array), other => { exec_err!("Unsupported data type {other:?} for function `initcap`") } diff --git a/datafusion/functions/src/unicode/reverse.rs b/datafusion/functions/src/unicode/reverse.rs index b05d8f847daa9..7cb7a2810dd8a 100644 --- a/datafusion/functions/src/unicode/reverse.rs +++ b/datafusion/functions/src/unicode/reverse.rs @@ -26,6 +26,7 @@ use arrow::array::{Array, ArrayRef, AsArray, StringArrayType}; use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::types::logical_string; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::utils::offset_span_len; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, @@ -128,9 +129,7 @@ fn reverse(args: &[ArrayRef]) -> Result { StringViewArrayBuilder::with_capacity(len), ), DataType::Dictionary(_, _) => { - let dictionary = args[0].as_any_dictionary(); - let converted = reverse(&[Arc::clone(dictionary.values())])?; - Ok(dictionary.with_values(converted)) + map_dictionary_values(&args[0], |values| reverse(&[Arc::clone(values)])) } _ => unreachable!( "Reverse can only be applied to Utf8View, Utf8 and LargeUtf8 types" diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs index 15c6324b8f3af..31555dc7a4414 100644 --- a/datafusion/functions/src/unicode/substr.rs +++ b/datafusion/functions/src/unicode/substr.rs @@ -31,6 +31,7 @@ use datafusion_common::cast::as_int64_array; use datafusion_common::types::{ NativeType, logical_int32, logical_int64, logical_string, }; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::{Result, exec_err}; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, @@ -150,17 +151,13 @@ impl ScalarUDFImpl for SubstrFunc { .all(|v| matches!(v, ColumnarValue::Scalar(_))) { // Fast path: substr the N_distinct values, reuse the keys. - let dict = array.as_any_dictionary(); - let values = Arc::clone(dict.values()); - let mut inner = Vec::with_capacity(args.args.len()); - inner.push(ColumnarValue::Array(values)); - inner.extend_from_slice(&args.args[1..]); - // Values that no valid row references can still fail (e.g. a - // negative length), so fall back to the row-wise path on error. - if let Ok(new_values) = make_scalar_function(substr, vec![])(&inner) { - let new_values = new_values.into_array(dict.values().len())?; - return Ok(ColumnarValue::Array(dict.with_values(new_values))); - } + let result = map_dictionary_values(array, |values| { + let mut inner = Vec::with_capacity(args.args.len()); + inner.push(ColumnarValue::Array(Arc::clone(values))); + inner.extend_from_slice(&args.args[1..]); + make_scalar_function(substr, vec![])(&inner)?.into_array(values.len()) + })?; + return Ok(ColumnarValue::Array(result)); } make_scalar_function(substr, vec![])(&args.args) diff --git a/datafusion/spark/src/function/math/hex.rs b/datafusion/spark/src/function/math/hex.rs index c3f1cce2abd6d..fafc47a41aa45 100644 --- a/datafusion/spark/src/function/math/hex.rs +++ b/datafusion/spark/src/function/math/hex.rs @@ -19,16 +19,14 @@ use std::str::from_utf8_unchecked; use std::sync::Arc; use arrow::array::{Array, ArrayAccessor, ArrayRef, StringArray, StringBuilder}; +use arrow::array::{as_largestring_array, as_string_array}; use arrow::buffer::{Buffer, OffsetBuffer}; use arrow::datatypes::DataType; -use arrow::{ - array::{as_dictionary_array, as_largestring_array, as_string_array}, - datatypes::Int32Type, -}; use datafusion_common::cast::as_large_binary_array; use datafusion_common::cast::as_string_view_array; use datafusion_common::types::{NativeType, logical_int64, logical_string}; use datafusion_common::utils::hex::{HexCase, ToHex, encode_bytes_into}; +use datafusion_common::utils::map_dictionary_values; use datafusion_common::utils::take_function_args; use datafusion_common::{ DataFusionError, @@ -271,48 +269,45 @@ pub fn compute_hex( ); } - let dict = as_dictionary_array::(&array); - let dict_values = dict.values(); - - let encoded_values = match dict_values.data_type() { - DataType::Int64 => { - let arr = as_int64_array(dict_values)?; - hex_encode_int64(arr.iter(), arr.len())? - } - DataType::Utf8 => { - let arr = as_string_array(dict_values); - hex_encode_bytes(&arr, lowercase)? - } - DataType::LargeUtf8 => { - let arr = as_largestring_array(dict_values); - hex_encode_bytes(&arr, lowercase)? - } - DataType::Utf8View => { - let arr = as_string_view_array(dict_values)?; - hex_encode_bytes(&arr, lowercase)? - } - DataType::Binary => { - let arr = as_binary_array(dict_values)?; - hex_encode_bytes(&arr, lowercase)? - } - DataType::LargeBinary => { - let arr = as_large_binary_array(dict_values)?; - hex_encode_bytes(&arr, lowercase)? - } - DataType::FixedSizeBinary(_) => { - let arr = as_fixed_size_binary_array(dict_values)?; - hex_encode_bytes(&arr, lowercase)? - } - _ => { - return exec_err!( - "hex got an unexpected argument type: {}", - dict_values.data_type() - ); - } - }; - - let new_dict = dict.with_values(encoded_values); - Ok(ColumnarValue::Array(Arc::new(new_dict))) + let encoded = map_dictionary_values(array, |dict_values| { + Ok(match dict_values.data_type() { + DataType::Int64 => { + let arr = as_int64_array(dict_values)?; + hex_encode_int64(arr.iter(), arr.len())? + } + DataType::Utf8 => { + let arr = as_string_array(dict_values); + hex_encode_bytes(&arr, lowercase)? + } + DataType::LargeUtf8 => { + let arr = as_largestring_array(dict_values); + hex_encode_bytes(&arr, lowercase)? + } + DataType::Utf8View => { + let arr = as_string_view_array(dict_values)?; + hex_encode_bytes(&arr, lowercase)? + } + DataType::Binary => { + let arr = as_binary_array(dict_values)?; + hex_encode_bytes(&arr, lowercase)? + } + DataType::LargeBinary => { + let arr = as_large_binary_array(dict_values)?; + hex_encode_bytes(&arr, lowercase)? + } + DataType::FixedSizeBinary(_) => { + let arr = as_fixed_size_binary_array(dict_values)?; + hex_encode_bytes(&arr, lowercase)? + } + _ => { + return exec_err!( + "hex got an unexpected argument type: {}", + dict_values.data_type() + ); + } + }) + })?; + Ok(ColumnarValue::Array(encoded)) } _ => exec_err!("hex got an unexpected argument type: {}", array.data_type()), }, diff --git a/datafusion/sqllogictest/test_files/dictionary.slt b/datafusion/sqllogictest/test_files/dictionary.slt index daa27dba860c8..22ed47bb1c7af 100644 --- a/datafusion/sqllogictest/test_files/dictionary.slt +++ b/datafusion/sqllogictest/test_files/dictionary.slt @@ -674,3 +674,62 @@ SELECT min(dict_of_ree), max(dict_of_ree), arrow_typeof(max(dict_of_ree)) FROM d statement ok DROP TABLE dict_encoded_values; + +# Aggregate output is emitted in `batch_size` slices that keep every group's +# dictionary value, so these batches have many more dictionary values than +# rows. Hashing them and running string functions on them must give the same +# results as on the decoded values. +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +set datafusion.execution.target_partitions = 2; + +statement ok +create table padded_groups as +select arrow_cast( + case when value % 13 = 0 then null else ' Grp' || (value % 100) || ' ' end, + 'Dictionary(Int32, Utf8)' +) as p +from generate_series(1, 1000); + +query I +select count(*) from (select p from padded_groups group by p); +---- +101 + +# Each function on the grouped dictionary must match the same function on the +# decoded value +query I +select count(*) +from ( + select p, arrow_cast(p, 'Utf8') as u + from (select p from padded_groups group by p) +) +where lower(p) is distinct from lower(u) + or upper(p) is distinct from upper(u) + or reverse(p) is distinct from reverse(u) + or character_length(p) is distinct from character_length(u) + or ascii(p) is distinct from ascii(u) + or initcap(p) is distinct from initcap(u) + or btrim(p) is distinct from btrim(u) + or ltrim(p) is distinct from ltrim(u) + or rtrim(p) is distinct from rtrim(u); +---- +0 + +query T +select distinct arrow_typeof(lower(p)) from (select p from padded_groups group by p); +---- +Dictionary(Int32, Utf8) + +statement ok +drop table padded_groups; + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# reset it explicitly. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.execution.batch_size;