apache / apache/datafusion

Async UDF rewrite input to Aggregate breaks UDAFs

Open
#18,149 2 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.