jayzhan211 commented on code in PR #25874:
URL: https://github.com/apache/datafusion/pull/25874#discussion_r4365790124
##########
datafusion/functions/src/unicode/substr.rs:
##########
@@ -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());
Review Comment:
Returning a dictionary also makes a high-cardinality GROUP BY on the result
slower, with or without the guard. With `enable_rle_to_dictionary=true` on
ClickBench, `SELECT substr("SearchPhrase", 1, 10) s, COUNT(*) c FROM hits WHERE
"SearchPhrase" <> '' GROUP BY s ORDER BY c DESC LIMIT 10` is 1.2x slower than
main (1.3x with 2 partitions), although `substr` itself is faster. The time
goes into repartitioning and re-aggregating dictionary group keys, the same
cost main already pays when grouping on `"SearchPhrase"` directly (#24822).
Compacting the dictionary when the aggregate slices its output (#26696)
fixes the 2-partition case, but the default stays at 1.2x because the
repartition coalescer still rebuilds dictionaries. The next step is to dig into
that further, and until then we might need to hold this PR.
##########
datafusion/functions/src/unicode/substr.rs:
##########
@@ -112,10 +123,42 @@ 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(_, _))
Review Comment:
This fast path runs `substr` over every dictionary value, so its cost
follows the dictionary size rather than the batch size: `filter` keeps the full
dictionary, and Parquet shares one dictionary across batches. That's what the
bot's Q22 regression is. On data where `c_phone` is `Dictionary(Int32, Utf8)`,
Q22 goes from 11.8 ms on main to 15.1 ms here (38.6 → 55.3 ms with one
partition).
Only taking the fast path when the dictionary isn't larger than the batch
brings it back to 12.1 ms (40.4 ms), with the same results:
```diff
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..]
```
--
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]
