jayzhan211 commented on code in PR #25874:
URL: https://github.com/apache/datafusion/pull/25874#discussion_r4337236581
##########
datafusion/functions/src/unicode/substr.rs:
##########
@@ -115,8 +117,51 @@ impl ScalarUDFImpl for SubstrFunc {
Ok(arg_types[0].clone())
}
+ fn return_field_from_args(&self, args: ReturnFieldArgs) ->
Result<FieldRef> {
+ // Per-row start/length materializes the dict, so the output type is
the
+ // value type; scalar tail keeps the dict wrapper.
+ let input = &args.arg_fields[0];
+ let output_type = match input.data_type() {
+ DataType::Dictionary(_, value_type)
+ if !args.scalar_arguments.iter().skip(1).all(Option::is_some)
=>
+ {
+ value_type.as_ref().clone()
+ }
+ other => other.clone(),
+ };
+ Ok(Field::new(self.name(), output_type, input.is_nullable()).into())
Review Comment:
Using `input.is_nullable()` here drops the nullability of `start`/`length`,
and since this override applies to every input type, it affects plain `Utf8`
too. The simplifier then folds `IS NULL` to `false`:
```sql
create table tnn (c VARCHAR NOT NULL, s BIGINT);
insert into tnn values ('hello', 1), ('world', NULL);
select substr(c, s) is null from tnn; -- false, false (main: false, true)
```
With a dictionary `c`, `select substr(c, s)` fails with "declared as
non-nullable but contains null values". Removing this override (see my other
comment) fixes it.
##########
datafusion/functions/src/unicode/substr.rs:
##########
@@ -115,8 +117,51 @@ impl ScalarUDFImpl for SubstrFunc {
Ok(arg_types[0].clone())
}
+ fn return_field_from_args(&self, args: ReturnFieldArgs) ->
Result<FieldRef> {
+ // Per-row start/length materializes the dict, so the output type is
the
+ // value type; scalar tail keeps the dict wrapper.
+ let input = &args.arg_fields[0];
+ let output_type = match input.data_type() {
+ DataType::Dictionary(_, value_type)
+ if !args.scalar_arguments.iter().skip(1).all(Option::is_some)
=>
Review Comment:
`scalar_arguments` is only `Some` for literals, but the fast path triggers
on any runtime `Scalar`, so the two can disagree. For example, `select
substr(c, (select 2)) from t` on a dictionary `c` fails with "returned value of
type 'Dictionary(Int32, Utf8)' while ... 'Utf8' was promised at planning time".
It also makes `substr(c, 1, 1+2)` change type when the simplifier folds `1+2`,
which the `return_field_from_args` docs disallow. A scalar dictionary input
like `substr(arrow_cast('hello', 'Dictionary(Int32, Utf8)'), 1, 3)` also errors
now, because it reaches `substr()` without a `Dictionary` arm.
Could we keep `return_type = arg_types[0]`, delete `return_field_from_args`,
and handle `Dictionary` in the dispatch like `rtrim`/`reverse` do? I tried this
locally, and it passes these cases plus the `string/` slt directory:
```rs
DataType::Dictionary(_, value_type) => {
// Per-row start/length: substr the materialized values, then
// re-encode so the output matches `return_type`.
let mut flat_args = args.to_vec();
flat_args[0] = cast(&args[0], value_type)?;
Ok(cast(&substr(&flat_args)?, args[0].data_type())?)
}
```
`invoke_with_args` then keeps your fast path and ends with
`make_scalar_function(substr, vec![])(&args.args)`.
##########
datafusion/sqllogictest/test_files/string/dictionary_utf8.slt:
##########
@@ -57,6 +57,12 @@ 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;
Review Comment:
The new tests only cover literal start/length on an array input. Could you
add a per-row `start`, a NULL start/length, and a scalar dictionary input here
(the repros from my other comments work), so these cases stay covered?
--
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]
