Async UDF rewrite input to Aggregate breaks UDAFs
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
When using an async UDF as an input to certain UDAFs we can encounter an issue with schema.
### To Reproduce
```rust
use arrow::compute::cast_with_options;
use datafusion::arrow::datatypes::DataType;
use datafusion::error::Result;
use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
use datafusion::logical_expr::{
ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};
use datafusion::prelude::*;
use std::any::Any;
use std::sync::Arc;
use tonic::async_trait;
#[derive(Debug, PartialEq, Eq, Hash)]
struct CustomUDF {
signature: Signature,
}
impl CustomUDF {
fn new() -> Self {
CustomUDF {
signature: Signature::exact(vec![DataType::Int64], Volatility::Immutable),
}
}
}
impl ScalarUDFImpl for CustomUDF {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"custom_udf"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result {
Ok(DataType::Utf8)
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result {
fn_impl(args)
}
}
#[async_trait]
impl AsyncScalarUDFImpl for CustomUDF {
async fn invoke_async_with_args(
&self,
args: ScalarFunctionArgs,
) -> Result {
fn_impl(args)
}
}
fn fn_impl(args: ScalarFunctionArgs) -> Result {
let arg = &args.args[0];
let array = arg.to_array(args.number_rows)?;
let array = cast_with_options(&array, &DataType::Utf8, &Default::default())?;
Ok(ColumnarValue::Array(array))
}
#[tokio::main]
async fn main() -> Result<()> {
let ctx = SessionContext::new_with_config(
SessionConfig::new().set_bool("datafusion.explain.show_schema", true),
);
ctx.sql("create table data(x int) as values (-10), (2)")
.await?
.collect()
.await?;
// sync works fine
ctx.register_udf(ScalarUDF::new_from_impl(CustomUDF::new()));
ctx.sql("select approx_distinct(custom_udf(x)) from data")
.await?
.explain(false, false)?
.show()
.await?;
ctx.sql("select approx_distinct(custom_udf(x)) from data")
.await?
.show()
.await?;
// issue with async
ctx.register_udf(AsyncScalarUDF::new(Arc::new(CustomUDF::new())).into_scalar_udf());
ctx.sql("select approx_distinct(custom_udf(x)) from data")
.await?
.explain(false, false)?
.show()
.await?;
ctx.sql("select approx_distinct(custom_udf(x)) from data")
.await?
.show()
.await?;
Ok(())
}
```
- Here we test a sync version and the async version; we expect the same result for both
Output:
```sh
datafusion (main)$ cargo run --example dataframe
Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.43s
Running `/Users/jeffrey/.cargo_target_cache/debug/examples/dataframe`
+---------------+---------------------------------------------------------------------------------------------------------------------------------------+
| plan_type | plan |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------+
| logical_plan | Aggregate: groupBy=[[]], aggr=[[approx_distinct(custom_udf(CAST(data.x AS Int64)))]] |
| | TableScan: data projection=[x] |
| physical_plan | AggregateExec: mode=Single, gby=[], aggr=[approx_distinct(custom_udf(data.x))], schema=[approx_distinct(custom_udf(data.x)):UInt64;N] |
| | DataSourceExec: partitions=1, partition_sizes=[1], schema=[x:Int32;N] |
| | |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------+
+-------------------------------------+
| approx_distinct(custom_udf(data.x)) |
+-------------------------------------+
| 2 |
+-------------------------------------+
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+
| plan_type | plan |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+
| logical_plan | Aggregate: groupBy=[[]], aggr=[[approx_distinct(custom_udf(CAST(data.x AS Int64)))]] |
| | TableScan: data projection=[x] |
| physical_plan | AggregateExec: mode=Final, gby=[], aggr=[approx_distinct(custom_udf(data.x))], schema=[approx_distinct(custom_udf(data.x)):UInt64;N] |
| | CoalescePartitionsExec, schema=[approx_distinct(custom_udf(data.x))[hll_registers]:Binary] |
| | AggregateExec: mode=Partial, gby=[], aggr=[approx_distinct(custom_udf(data.x))], schema=[approx_distinct(custom_udf(data.x))[hll_registers]:Binary] |
| | RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1, schema=[x:Int32;N, __async_fn_0:Utf8;N] |
| | AsyncFuncExec: async_expr=[async_expr(name=__async_fn_0, expr=custom_udf(CAST(x@0 AS Int64)))], schema=[x:Int32;N, __async_fn_0:Utf8;N] |
| | CoalesceBatchesExec: target_batch_size=8192, schema=[x:Int32;N] |
| | DataSourceExec: partitions=1, partition_sizes=[1], schema=[x:Int32;N] |
| | |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+
Error: Internal("PhysicalExpr Column references column '__async_fn_0' at index 1 (zero-based) but input schema only has 1 columns: [\"x\"]")
```
The error comes from here:
https://github.com/apache/datafusion/blob/e323357b1e245d8651183e42747cb92709cb1998/datafusion/functions-aggregate/src/approx_distinct.rs#L363-L366
- line 364
### Expected behavior
Should show same result as sync version instead of error.
### Additional context
I want to emphasize that this is not an issue with UDAFs (in this case `approx_distinct`). This is because during physical planning, async UDFs rewrite the aggregate with new expressions:
https://github.com/apache/datafusion/blob/e323357b1e245d8651183e42747cb92709cb1998/datafusion/core/src/physical_planner.rs#L745-L785
- Introduced by #17619
`agg_func` has new expressions written in, to replace the async function call. However, the schema it has does not change; therefore later in the pipeline (see `approx_distinct` above) it has a mismatch between schema and input expressions.
Overall I question this API:
https://github.com/apache/datafusion/blob/e323357b1e245d8651183e42747cb92709cb1998/datafusion/physical-expr/src/aggregate.rs#L613-L620
It allows rewriting the expressions but it does not enforce or check that the new expressions conform with the schema it already has.
Contributor guide
Research direction
Reproduce the async and sync cases from the issue, then read physical_planner.rs around lines 745-785 and physical-expr/src/aggregate.rs around lines 613-620. Inspect functions-aggregate/src/approx_distinct.rs around lines 363-366 to trace the schema and rewritten expressions. Done means the async query produces the same result as the sync query without the column-reference schema error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend-api-design, databases
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100