jayzhan211 commented on code in PR #25874:
URL: https://github.com/apache/datafusion/pull/25874#discussion_r4360804755
##########
datafusion/functions/src/unicode/substr.rs:
##########
@@ -143,13 +187,83 @@ 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>(
Review Comment:
Thanks for the Int8 fix, the repro passes now. One more thing on this
row-wise path: per-row `start`/`length` rarely repeat as combinations, so the
`HashMap` dedup mostly adds hashing cost. On 8192 rows (same release build),
`substr(dict, start_array)` is 1.5x slower than `cast` to Utf8 + `substr` (what
main does) with 16 distinct values, and 5.7x slower with 8192.
Only wide key types reach this function, so one key per row always fits.
Decoding and wrapping the result with identity keys runs at 1.01-1.03x of the
decode path, and the existing unit tests and `string/` slt still pass:
```rs
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)?))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
