perf: opt substr into make_scalar_function dict fast path

datafusion/functions/src/unicode/substr.rs

@@ -20,19 +20,21 @@ use std::sync::Arc;
 use crate::strings::sub_view;
 use crate::utils::make_scalar_function;
 use arrow::array::{
-    Array, ArrayRef, AsArray, GenericStringArray, Int64Array, OffsetSizeTrait,
-    StringArrayType, StringViewArray,
+    Array, ArrayRef, AsArray, DictionaryArray, GenericStringArray, Int64Array,
+    OffsetSizeTrait, PrimitiveArray, StringArrayType, StringViewArray,
 };
 use arrow::buffer::{NullBuffer, ScalarBuffer};
-use arrow::datatypes::DataType;
+use arrow::compute::take;
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
+use arrow::downcast_dictionary_array;
 use datafusion_common::cast::as_int64_array;
 use datafusion_common::types::{
     NativeType, logical_int32, logical_int64, logical_string,
 };
 use datafusion_common::{Result, exec_err};
 use datafusion_expr::{
-    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
-    TypeSignature, TypeSignatureClass, Volatility,
+    Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
+    ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility,
 };
 use datafusion_macros::user_doc;
 
@@ -65,6 +67,14 @@ pub struct SubstrFunc {
     aliases: Vec<String>,
 }
 
+/// Whether `key_type` can index one dictionary value per row of any batch.
+fn is_wide_key(key_type: &DataType) -> bool {
+    matches!(
+        key_type,
+        DataType::Int32 | DataType::UInt32 | DataType::Int64 | DataType::UInt64
+    )
+}
+
 impl Default for SubstrFunc {
     fn default() -> Self {
         Self::new()
@@ -73,7 +83,8 @@ impl Default for SubstrFunc {
 
 impl SubstrFunc {
     pub fn new() -> Self {
-        let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string()));
+        let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
+            .with_encoding_preservation(EncodingPreservation::dictionary());
         let int64 = Coercion::new_implicit(
             TypeSignatureClass::Native(logical_int64()),
             vec![TypeSignatureClass::Native(logical_int32())],
@@ -112,10 +123,43 @@ impl ScalarUDFImpl for SubstrFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        Ok(arg_types[0].clone())
+        Ok(match &arg_types[0] {
+            DataType::Dictionary(key_type, value_type) if !is_wide_key(key_type) => {
+                value_type.as_ref().clone()
+            }
+            other => other.clone(),
+        })
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        if let DataType::Dictionary(key_type, value_type) = args.args[0].data_type()
+            && !is_wide_key(&key_type)
+        {
+            // Per-row start/length can produce more distinct values than a
+            // narrow key type can index, so decode instead.
+            let mut args = args.args;
+            args[0] = args[0].cast_to(&value_type, None)?;
+            return make_scalar_function(substr, vec![])(&args);
+        }
+
+        if let ColumnarValue::Array(array) = &args.args[0]
+            && matches!(array.data_type(), DataType::Dictionary(_, _))
+            && array.as_any_dictionary().values().len() <= array.len()
+            && args.args[1..]
+                .iter()
+                .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..]);
+            let new_values = make_scalar_function(substr, vec![])(&inner)?
+                .into_array(dict.values().len())?;
+            return Ok(ColumnarValue::Array(dict.with_values(new_values)));
+        }
+
         make_scalar_function(substr, vec![])(&args.args)
     }
 
@@ -143,13 +187,36 @@ fn substr(args: &[ArrayRef]) -> Result<ArrayRef> {
             let string_array = args[0].as_string_view();
             string_view_substr(string_array, &args[1..])
         }
+        DataType::Dictionary(_, _) => {
+            let array = &args[0];
+            downcast_dictionary_array! {
+                array => dictionary_substr(array, &args[1..]),
+                _ => unreachable!("substr matched a dictionary array"),
+            }
+        }
         other => exec_err!(
             "Unsupported data type {other:?} for function substr,\
             expected Utf8View, Utf8 or LargeUtf8."
         ),
     }
 }
 
+fn dictionary_substr<K: ArrowDictionaryKeyType>(
+    dictionary: &DictionaryArray<K>,
+    args: &[ArrayRef],
+) -> Result<ArrayRef> {
+    // Only wide key types reach here, so one key per row always fits.
+    let decoded = take(dictionary.values().as_ref(), dictionary.keys(), None)?;
+    let mut row_args = vec![decoded];
+    row_args.extend(args.iter().cloned());
+    let values = substr(&row_args)?;
+    let keys = (0..values.len())
+        .map(|i| K::Native::from_usize(i).unwrap())
+        .collect::<Vec<_>>();
+    let keys = PrimitiveArray::<K>::new(keys.into(), values.nulls().cloned());
+    Ok(Arc::new(DictionaryArray::<K>::try_new(keys, values)?))
+}
+
 /// Convert the given `start` and `count` to valid byte indices within `input` string.
 ///
 /// Input `start` and `count` are equivalent to PostgreSQL's `substr(s, start, count)`.
@@ -354,13 +421,15 @@ mod tests {
     use std::sync::Arc;
 
     use arrow::array::{
-        Array, ArrayRef, AsArray, Int64Array, LargeStringArray, StringArray,
-        StringViewArray,
+        Array, ArrayRef, AsArray, DictionaryArray, Int64Array, LargeStringArray,
+        StringArray, StringViewArray,
     };
     use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View};
+    use arrow::datatypes::{DataType, Field, Int32Type};
 
+    use datafusion_common::config::ConfigOptions;
     use datafusion_common::{Result, ScalarValue, exec_err};
-    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
 
     use crate::unicode::substr::SubstrFunc;
     use crate::utils::test::test_function;
@@ -766,4 +835,40 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_substr_dict_fast_path_preserves_encoding() -> Result<()> {
+        let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(Utf8));
+        let keys = arrow::array::Int32Array::from(vec![0, 1, 0, 2, 1]);
+        let values = StringArray::from(vec!["silver", "copper", "golden"]);
+        let dict: ArrayRef = Arc::new(
+            DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values)).unwrap(),
+        );
+
+        let result = SubstrFunc::new()
+            .invoke_with_args(ScalarFunctionArgs {
+                args: vec![
+                    ColumnarValue::Array(dict),
+                    ColumnarValue::Scalar(ScalarValue::from(1i64)),
+                    ColumnarValue::Scalar(ScalarValue::from(3i64)),
+                ],
+                arg_fields: vec![
+                    Arc::new(Field::new("s", dict_type.clone(), true)),
+                    Arc::new(Field::new("start", DataType::Int64, false)),
+                    Arc::new(Field::new("len", DataType::Int64, false)),
+                ],
+                number_rows: 5,
+                return_field: Arc::new(Field::new("r", dict_type, true)),
+                config_options: Arc::new(ConfigOptions::default()),
+            })?
+            .into_array(5)?;
+
+        assert!(matches!(result.data_type(), DataType::Dictionary(_, v) if **v == Utf8));
+        let out = result.as_dictionary::<Int32Type>();
+        let v = out.values().as_string::<i32>();
+        assert_eq!(v.len(), 3);
+        assert_eq!((v.value(0), v.value(1), v.value(2)), ("sil", "cop", "gol"));
+
+        Ok(())
+    }
 }

datafusion/sqllogictest/test_files/string/dictionary_utf8.slt

@@ -57,6 +57,76 @@ SELECT arrow_cast('', 'Dictionary(Int32, Utf8)');
 #
 include ./string_query.slt.part
 
+# Scalar start/length preserves dict encoding end-to-end.
+query T
+SELECT DISTINCT arrow_typeof(substr(c1, 1, 3)) FROM test_substr;
+----
+Dictionary(Int32, Utf8)
+
+# Runtime scalars that are not literals must agree with the planned return type.
+query TT
+SELECT substr(c1, (SELECT 2)), arrow_typeof(substr(c1, (SELECT 2)))
+FROM test_substr
+WHERE c1 = 'foo';
+----
+oo Dictionary(Int32, Utf8)
+
+# Scalar dictionary input must use the same dictionary-aware dispatch as arrays.
+query TT
+SELECT
+  substr(arrow_cast('hello', 'Dictionary(Int32, Utf8)'), 1, 3),
+  arrow_typeof(substr(arrow_cast('hello', 'Dictionary(Int32, Utf8)'), 1, 3));
+----
+hel Dictionary(Int32, Utf8)
+
+statement ok
+create table test_substr_arguments (
+  c VARCHAR NOT NULL,
+  start_pos BIGINT,
+  length BIGINT
+);
+
+statement ok
+insert into test_substr_arguments values
+  ('hello', 1, 2),
+  ('hello', 2, 1),
+  ('world', NULL, 3),
+  ('arrow', 2, NULL);
+
+# Nullable per-row arguments propagate nulls without changing dictionary encoding.
+query TTT
+SELECT
+  substr(c, start_pos),
+  substr(arrow_cast(c, 'Dictionary(Int32, Utf8)'), start_pos),
+  substr(arrow_cast(c, 'Dictionary(Int32, Utf8)'), 1, length)
+FROM test_substr_arguments
+ORDER BY c, start_pos;
+----
+rrow rrow NULL
+hello hello he
+ello ello h
+NULL NULL wor
+
+query TT
+SELECT DISTINCT
+  arrow_typeof(substr(arrow_cast(c, 'Dictionary(Int32, Utf8)'), start_pos)),
+  arrow_typeof(substr(arrow_cast(c, 'Dictionary(Int32, Utf8)'), 1, length))
+FROM test_substr_arguments;
+----
+Dictionary(Int32, Utf8) Dictionary(Int32, Utf8)
+
+statement ok
+create table t8 as
+select
+  arrow_cast(arrow_cast('hello world', 'Utf8'), 'Dictionary(Int8, Utf8)') c,
+  value s
+from generate_series(1, 200);
+
+query II
+select count(*), count(distinct substr(c, s)) from t8;
+----
+200 12
+
 #
 # Clean up
 #
@@ -68,3 +138,9 @@ drop table test_substr_base;
 
 statement ok
 drop table test_datetime_base;
+
+statement ok
+drop table test_substr_arguments;
+
+statement ok
+drop table t8;