From 525815935c7638b925de1f2ec359adff2a0bbbfe Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Tue, 29 Sep 2026 13:28:14 +0200 Subject: [PATCH] perf: add dict fast path to make_scalar_function When exactly one argument is dictionary-encoded and NDV < batch size, evaluate `inner` on the values array (O(NDV)) and scatter the result back with `take` (O(N)) instead of running `inner` on every row. When NDV >= batch size the optimisation pays no dividend, so we cast the dict to its value type before calling `inner`, keeping the same O(N) complexity as the pre-encoding path. Callers opt in by setting EncodingPreservation::dictionary() on the string argument coercion and stripping the dict wrapper in return_type so the declared output type matches the plain-array result of `take`. --- datafusion/functions/src/utils.rs | 88 ++++++++++++++++++++++++++++--- 1 file changed, 81 insertions(+), 7 deletions(-) diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index b6d9aacaadf66..3523df0db7c5e 100644 --- a/datafusion/functions/src/utils.rs +++ b/datafusion/functions/src/utils.rs @@ -21,7 +21,7 @@ use arrow::array::{ }; use arrow::buffer::NullBuffer; use arrow::compute::kernels::cmp::eq; -use arrow::compute::{SortOptions, cast, try_binary}; +use arrow::compute::{SortOptions, cast, take, try_binary}; use arrow::datatypes::{DataType, DecimalType}; use arrow::error::ArrowError; use datafusion_common::{DataFusionError, Result, ScalarValue, exec_err, internal_err}; @@ -120,6 +120,78 @@ where move |args: &[ColumnarValue]| { // first, identify if any of the arguments is an Array. If yes, store its `len`, // as any scalar will need to be converted to an array of len `len`. + // + // When exactly one array arg is dict-encoded and no plain-array args + // exist alongside it, and NDV < batch size, evaluate `inner` on the + // values array (O(NDV)) and scatter the result back with `take` (O(N)). + // Callers opt in by setting EncodingPreservation::dictionary() on the + // coercion; their return_type must also strip the dict wrapper so the + // declared output type matches the plain-array result of `take`. + // When NDV >= batch size the optimization pays no dividend, so we fall + // through and cast the dict to its value type before calling `inner`. + let mut n_dict = 0usize; + let mut n_plain = 0usize; + for arg in args { + match arg { + ColumnarValue::Array(arr) => { + if arr.as_any_dictionary_opt().is_some() { + n_dict += 1; + } else { + n_plain += 1; + } + } + ColumnarValue::Scalar(_) => {} + } + } + + if n_dict == 1 && n_plain == 0 { + let dict_arr = args + .iter() + .find_map(|arg| match arg { + ColumnarValue::Array(arr) + if arr.as_any_dictionary_opt().is_some() => + { + Some(arr) + } + _ => None, + }) + .expect("n_dict == 1"); + let dict = dict_arr.as_any_dictionary_opt().unwrap(); + let ndv = dict.values().len(); + let batch_n = dict.keys().len(); + + if ndv < batch_n { + let ndv_args: Vec = args + .iter() + .map(|arg| match arg { + ColumnarValue::Array(arr) + if arr.as_any_dictionary_opt().is_some() => + { + ColumnarValue::Array(Arc::clone(dict.values())) + } + other => other.clone(), + }) + .collect(); + + let ndv_arrays = ndv_args + .iter() + .zip(hints.iter().chain(std::iter::repeat(&Hint::Pad))) + .map(|(arg, hint)| { + let expansion_len = match hint { + Hint::AcceptsSingular => 1, + Hint::Pad => ndv, + }; + arg.to_array(expansion_len) + }) + .collect::>>()?; + + let ndv_result = (inner)(&ndv_arrays)?; + let result = take(&ndv_result, dict.keys(), None)?; + return Ok(ColumnarValue::Array(result)); + } + // NDV >= batch size: fall through; cast dict to its value type below. + } + let len = args .iter() .fold(Option::::None, |acc, arg| match arg { @@ -134,19 +206,21 @@ where .iter() .zip(hints.iter().chain(std::iter::repeat(&Hint::Pad))) .map(|(arg, hint)| { - // Decide on the length to expand this scalar to depending - // on the given hints. let expansion_len = match hint { Hint::AcceptsSingular => 1, Hint::Pad => inferred_length, }; - arg.to_array(expansion_len) + let arr = arg.to_array(expansion_len)?; + if let Some(d) = arr.as_any_dictionary_opt() { + cast(&arr, d.values().data_type()).map_err(Into::into) + } else { + Ok(arr) + } }) .collect::>>()?; let result = (inner)(&args); if is_scalar { - // If all inputs are scalar, keeps output as scalar let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0)); result.map(ColumnarValue::Scalar) } else { @@ -838,6 +912,7 @@ mod map_lookup_tests { StringArray, StructArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::{Field, Int32Type}; /// A map whose rows have the given lengths, drawing keys and values in @@ -964,8 +1039,7 @@ mod map_lookup_tests { } }) .collect(); - let repeated = - arrow::compute::take(key, &UInt32Array::from(vec![0; map.len()]), None)?; + let repeated = take(key, &UInt32Array::from(vec![0; map.len()]), None)?; for keys in [key, repeated.as_ref()] { let result = map_lookup(map, keys)?; assert_eq!(result.len(), matches.len());