Skip to content

Commit 6318fb4

Browse files
Rich-T-kidclaude
andcommitted
perf: preserve dictionary encoding in substr
Opt substr into dictionary-aware execution so planning no longer inserts CAST(dict -> value_type) around dictionary inputs with scalar start/length. - Mark the string arg with EncodingPreservation::dictionary() so coercion keeps the dict wrapper - Implement return_field_from_args to declare Dictionary(..) when the tail args are scalars, falling back to the value type otherwise - Add a fast path in invoke_with_args that runs substr on the N distinct dict values and reuses the keys via with_values(); per-row start/length materializes the dict first - Unit tests cover Utf8/LargeUtf8/Utf8View value types, null keys, per-row fallback, and the return_field planning behavior - SLT regression asserts arrow_typeof stays Dictionary(..) for scalar args and the plan no longer wraps the column in a CAST Closes #26170 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 1c49b7f commit 6318fb4

2 files changed

Lines changed: 111 additions & 5 deletions

File tree

‎datafusion/functions/src/unicode/substr.rs‎

Lines changed: 93 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,16 @@ use arrow::array::{
2424
StringArrayType, StringViewArray,
2525
};
2626
use arrow::buffer::{NullBuffer, ScalarBuffer};
27-
use arrow::datatypes::DataType;
27+
use arrow::datatypes::{DataType, Field, FieldRef};
2828
use datafusion_common::cast::as_int64_array;
2929
use datafusion_common::types::{
3030
NativeType, logical_int32, logical_int64, logical_string,
3131
};
3232
use datafusion_common::{Result, exec_err};
3333
use datafusion_expr::{
34-
Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
35-
TypeSignature, TypeSignatureClass, Volatility,
34+
Coercion, ColumnarValue, Documentation, EncodingPreservation, ReturnFieldArgs,
35+
ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass,
36+
Volatility,
3637
};
3738
use datafusion_macros::user_doc;
3839

@@ -73,7 +74,8 @@ impl Default for SubstrFunc {
7374

7475
impl SubstrFunc {
7576
pub fn new() -> Self {
76-
let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string()));
77+
let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
78+
.with_encoding_preservation(EncodingPreservation::dictionary());
7779
let int64 = Coercion::new_implicit(
7880
TypeSignatureClass::Native(logical_int64()),
7981
vec![TypeSignatureClass::Native(logical_int32())],
@@ -115,8 +117,51 @@ impl ScalarUDFImpl for SubstrFunc {
115117
Ok(arg_types[0].clone())
116118
}
117119

120+
fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
121+
// Per-row start/length materializes the dict, so the output type is the
122+
// value type; scalar tail keeps the dict wrapper.
123+
let input = &args.arg_fields[0];
124+
let output_type = match input.data_type() {
125+
DataType::Dictionary(_, value_type)
126+
if !args.scalar_arguments.iter().skip(1).all(Option::is_some) =>
127+
{
128+
value_type.as_ref().clone()
129+
}
130+
other => other.clone(),
131+
};
132+
Ok(Field::new(self.name(), output_type, input.is_nullable()).into())
133+
}
134+
118135
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
119-
make_scalar_function(substr, vec![])(&args.args)
136+
if let ColumnarValue::Array(array) = &args.args[0]
137+
&& matches!(array.data_type(), DataType::Dictionary(_, _))
138+
&& args.args[1..]
139+
.iter()
140+
.all(|v| matches!(v, ColumnarValue::Scalar(_)))
141+
{
142+
// Fast path: substr the N_distinct values, reuse the keys.
143+
let dict = array.as_any_dictionary();
144+
let values = Arc::clone(dict.values());
145+
let mut inner = Vec::with_capacity(args.args.len());
146+
inner.push(ColumnarValue::Array(values));
147+
inner.extend_from_slice(&args.args[1..]);
148+
let new_values = make_scalar_function(substr, vec![])(&inner)?
149+
.into_array(dict.values().len())?;
150+
return Ok(ColumnarValue::Array(dict.with_values(new_values)));
151+
}
152+
153+
// Per-row start/length: materialize the dict so substr can run row-wise.
154+
let args = if let ColumnarValue::Array(array) = &args.args[0]
155+
&& let DataType::Dictionary(_, value_type) = array.data_type()
156+
{
157+
let flat = arrow::compute::kernels::cast::cast(array, value_type.as_ref())?;
158+
let mut new_args = args.args.clone();
159+
new_args[0] = ColumnarValue::Array(flat);
160+
new_args
161+
} else {
162+
args.args
163+
};
164+
make_scalar_function(substr, vec![])(&args)
120165
}
121166

122167
fn aliases(&self) -> &[String] {
@@ -766,4 +811,47 @@ mod tests {
766811

767812
Ok(())
768813
}
814+
815+
/// Dict column with scalar start/length should run substr on the N_distinct
816+
/// values and reuse the keys, preserving dictionary encoding end-to-end.
817+
#[test]
818+
fn test_substr_dict_fast_path_preserves_encoding() -> Result<()> {
819+
use arrow::array::DictionaryArray;
820+
use arrow::datatypes::{DataType, Field, Int32Type};
821+
use datafusion_common::config::ConfigOptions;
822+
use datafusion_expr::ScalarFunctionArgs;
823+
824+
let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(Utf8));
825+
let keys = arrow::array::Int32Array::from(vec![0, 1, 0, 2, 1]);
826+
let values = StringArray::from(vec!["hello", "world", "rust"]);
827+
let dict: ArrayRef = Arc::new(
828+
DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values)).unwrap(),
829+
);
830+
831+
let result = SubstrFunc::new()
832+
.invoke_with_args(ScalarFunctionArgs {
833+
args: vec![
834+
ColumnarValue::Array(dict),
835+
ColumnarValue::Scalar(ScalarValue::from(1i64)),
836+
ColumnarValue::Scalar(ScalarValue::from(3i64)),
837+
],
838+
arg_fields: vec![
839+
Arc::new(Field::new("s", dict_type.clone(), true)),
840+
Arc::new(Field::new("start", DataType::Int64, false)),
841+
Arc::new(Field::new("len", DataType::Int64, false)),
842+
],
843+
number_rows: 5,
844+
return_field: Arc::new(Field::new("r", dict_type, true)),
845+
config_options: Arc::new(ConfigOptions::default()),
846+
})?
847+
.into_array(5)?;
848+
849+
assert!(matches!(result.data_type(), DataType::Dictionary(_, v) if **v == Utf8));
850+
let out = result.as_dictionary::<Int32Type>();
851+
let v = out.values().as_string::<i32>();
852+
assert_eq!(v.len(), 3);
853+
assert_eq!((v.value(0), v.value(1), v.value(2)), ("hel", "wor", "rus"));
854+
855+
Ok(())
856+
}
769857
}

‎datafusion/sqllogictest/test_files/string/dictionary_utf8.slt‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,24 @@ SELECT arrow_cast('', 'Dictionary(Int32, Utf8)');
5757
#
5858
include ./string_query.slt.part
5959

60+
# Regression (#26170): scalar start/length preserves dict encoding end-to-end,
61+
# so the plan has no CAST(c1 AS Utf8) wrapping the dictionary column.
62+
query TT
63+
EXPLAIN SELECT substr(c1, 1, 3) AS r FROM test_substr;
64+
----
65+
logical_plan
66+
01)Projection: substr(test_substr.c1, Int64(1), Int64(3)) AS r
67+
02)--TableScan: test_substr projection=[c1]
68+
physical_plan
69+
01)ProjectionExec: expr=[substr(c1@0, 1, 3) as r]
70+
02)--DataSourceExec: partitions=1, partition_sizes=[1]
71+
72+
# Per-row length falls back to materialization; declared type is the value type.
73+
query T
74+
SELECT DISTINCT arrow_typeof(substr(c1, 1, character_length(c1))) FROM test_substr;
75+
----
76+
Utf8
77+
6078
#
6179
# Clean up
6280
#

0 commit comments

Comments
 (0)