jayzhan211 commented on code in PR #21737:
URL: https://github.com/apache/datafusion/pull/21737#discussion_r4184652052
##########
datafusion/catalog/src/information_schema.rs:
##########
@@ -454,118 +455,98 @@ impl InformationSchemaConfig {
}
}
-/// get the arguments and return types of a UDF
-/// returns a tuple of (arg_types, return_type)
+/// Origins used to enumerate the physical types a native type can take
+const RESOLVE_CAST_SOURCES: [DataType; 2] = [DataType::Null,
DataType::LargeUtf8];
+
+/// Build argument fields for `information_schema` to provide possible return
types
+fn resolve_informational_fields(idx: usize, t: &NativeType) ->
Result<Vec<FieldRef>> {
+ // Since native types map to several physical types, resolve it against
+ // ambiguous types to get canonical `DataType`s for the native type
+ let data_types = RESOLVE_CAST_SOURCES
Review Comment:
`default_cast_for(&LargeUtf8)` has no arm for `Struct`/`Map`/`Union` (falls
through to `_internal_err!`), and the `?` here propagates it out of
`make_routines`/`make_parameters`. So registering one UDF with such an argument
(also nested, e.g. `List<Struct>`) makes `information_schema.routines`,
`information_schema.parameters` and `SHOW FUNCTIONS` fail for every function.
Before this PR the path was infallible. On head this test fails with
`Internal("Unavailable default cast for native type Struct(\"a\": Int32) from
physical type LargeUtf8")`:
```rs
#[test]
fn test_get_udf_args_and_return_types_nested() -> Result<()> {
use arrow::datatypes::Fields;
let struct_type =
DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32,
true)]));
let signature = Signature::exact(vec![struct_type], Volatility::Stable);
let udf = Arc::new(ScalarUDF::from(TestScalarUDF { signature }));
let result = get_udf_args_and_return_types(&udf)?;
assert_eq!(result.len(), 1);
Ok(())
}
```
Fix: skip origins the type has no cast from:
```diff
-fn resolve_informational_fields(idx: usize, t: &NativeType) ->
Result<Vec<FieldRef>> {
+fn resolve_informational_fields(idx: usize, t: &NativeType) ->
Vec<FieldRef> {
// Since native types map to several physical types, resolve it against
- // ambiguous types to get canonical `DataType`s for the native type
- let data_types = RESOLVE_CAST_SOURCES
+ // ambiguous types to get canonical `DataType`s for the native type.
+ // Skip origins the type has no cast from (e.g. `Struct` from
`LargeUtf8`)
+ RESOLVE_CAST_SOURCES
.iter()
- .map(|source| t.default_cast_for(source))
- .collect::<Result<Vec<DataType>, _>>()?;
- Ok(data_types
- .into_iter()
+ .filter_map(|source| t.default_cast_for(source).ok())
.unique()
.map(|dt| Arc::new(Field::new(format!("arg_{idx}"), dt, true)))
- .collect())
+ .collect()
}
```
```diff
.map(|(i, t)| resolve_informational_fields(i, t))
- .collect::<Result<Vec<_>>>()?;
+ .collect::<Vec<_>>();
```
##########
datafusion/catalog/src/information_schema.rs:
##########
@@ -454,118 +455,98 @@ impl InformationSchemaConfig {
}
}
-/// get the arguments and return types of a UDF
-/// returns a tuple of (arg_types, return_type)
+/// Origins used to enumerate the physical types a native type can take
+const RESOLVE_CAST_SOURCES: [DataType; 2] = [DataType::Null,
DataType::LargeUtf8];
+
+/// Build argument fields for `information_schema` to provide possible return
types
+fn resolve_informational_fields(idx: usize, t: &NativeType) ->
Result<Vec<FieldRef>> {
+ // Since native types map to several physical types, resolve it against
+ // ambiguous types to get canonical `DataType`s for the native type
+ let data_types = RESOLVE_CAST_SOURCES
+ .iter()
+ .map(|source| t.default_cast_for(source))
+ .collect::<Result<Vec<DataType>, _>>()?;
+ Ok(data_types
+ .into_iter()
+ .unique()
+ .map(|dt| Arc::new(Field::new(format!("arg_{idx}"), dt, true)))
+ .collect())
+}
+
+/// Function information schema is a set of tuples - argument types and an
optional return type
+type FunctionInformationSchema = BTreeSet<(Vec<String>, Option<String>)>;
+
+/// Get the arguments and return types of a function from its signature
+fn get_args_and_return_types(
+ signature: &Signature,
+ return_field: impl Fn(&[FieldRef]) -> Result<FieldRef>,
+) -> Result<FunctionInformationSchema> {
+ let arg_types = signature.type_signature.get_representative_types();
+ if arg_types.is_empty() {
+ // Edge case if function doesn't have arguments
+ return Ok(BTreeSet::from([(vec![], None)]));
+ }
+ arg_types
+ .into_iter()
+ .map(|arg_types| {
+ // Get possible types for each input arg
+ let arg_fields = arg_types
+ .iter()
+ .enumerate()
+ .map(|(i, t)| resolve_informational_fields(i, t))
+ .collect::<Result<Vec<_>>>()?;
+ // Build combinations of arg types with the return type
+ let return_types = arg_fields
Review Comment:
`generate_series`/`range` now report only `(String, Date, Interval) →
List(String)`, but the query returns `List(Date32)`. Base also listed
`List(Date)`, but only because its `Date64` example happened to hit `(_,
Some(Date64))` in `Range::return_type`. That match is lost now that `Date`
resolves only to `Date32`. The underlying cause predates this PR: the return
type is computed on arguments that haven't been coerced, which the
`get_representative_types` docs say callers must do. Fine to handle in a
follow-up.
```sql
select arrow_typeof(generate_series('2020-01-01', DATE '2020-01-03',
INTERVAL '1 day'));
-- List(Date32)
select data_type from information_schema.parameters
where specific_name = 'generate_series' and parameter_mode = 'OUT';
-- has List(String) for (String, Date, Interval), no List(Date)
```
Fix: coerce the arguments the same way the planner does before asking for
the return type. I tried this locally: it fixes this row and also corrects
`sum(Int32)` (`NULL` → `Int64`), `trunc(Int64)` (`NULL` → `Float64`),
`median(Int64)` (`Int64` → `Float64`) and `date_trunc(String, Date)` (`Date` →
`Timestamp(ns)`), all checked against `arrow_typeof`. It does change the
existing `date_trunc` expectations in `information_schema.slt`.
```diff
use datafusion_expr::function::WindowUDFFieldArgs;
+use datafusion_expr::type_coercion::functions::fields_with_udf;
```
```diff
get_args_and_return_types(udf.signature(), |arg_fields| {
+ let arg_fields = &fields_with_udf(arg_fields, udf.as_ref())?;
let scalar_arguments = &vec![None; arg_fields.len()];
```
```diff
get_args_and_return_types(udaf.signature(), |arg_fields| {
- udaf.return_field(arg_fields)
+ udaf.return_field(&fields_with_udf(arg_fields, udaf.as_ref())?)
})
```
```diff
get_args_and_return_types(udwf.signature(), |arg_fields| {
+ let arg_fields = &fields_with_udf(arg_fields, udwf.as_ref())?;
udwf.field(WindowUDFFieldArgs::new(arg_fields, udwf.name()))
```
--
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]
