Repository: ArroyoSystems/arroyo (language: rust)

Issue:
Async UDFs on subqueries produces runtime error
A query like this:

```sql
select new_udf(out.seconds) from (
select timestamp_struct as out
from (
    select
        header.ts as timestamp_struct
    from nested2
));
```

with any async UDF, for example

```rust
use arroyo_udf_plugin::udf;

#[udf]
async fn new_udf(x: i64) -> i64 {
    x
}
```

Produces a runtime error:

```
2024-11-01T23:04:54.758632Z  INFO arroyo_controller::states: state transition job_id="job_q5WOEkQvWV" from="Scheduling" to="Running" duration_ms=167
2024-11-01T23:04:56.711565Z  INFO arroyo_api::jobs: Subscribed to output
2024-11-01T23:05:00.698013Z ERROR arroyo_server_common: panicked at crates/arroyo-worker/src/arrow/async_udf.rs:299:22:
called `Result::unwrap()` on an `Err` value: Execution("get indexed field seconds not found in struct") panic.file="crates/arroyo-worker/src/arrow/async_udf.rs" panic.line=299 panic.column=22
```


Discussion / hints:

Fix patch:
diff --git a/crates/arroyo-planner/src/extension/mod.rs b/crates/arroyo-planner/src/extension/mod.rs
index 161a9de21..8213d0bae 100644
--- a/crates/arroyo-planner/src/extension/mod.rs
+++ b/crates/arroyo-planner/src/extension/mod.rs
@@ -15,15 +15,9 @@ use datafusion::logical_expr::{
 };
 use datafusion_proto::physical_plan::to_proto::serialize_physical_expr;
 use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
-use datafusion_proto::protobuf::ProjectionNode;
 use prost::Message;
 use watermark_node::WatermarkNode;
 
-use crate::builder::{NamedNode, Planner};
-use crate::schemas::{add_timestamp_field, has_timestamp_field};
-use crate::{fields_with_qualifiers, schema_from_df_fields, DFField, ASYNC_RESULT_FIELD};
-use join::JoinExtension;
-
 use self::debezium::{DebeziumUnrollingExtension, ToDebeziumExtension};
 use self::updating_aggregate::UpdatingAggregateExtension;
 use self::{
@@ -31,6 +25,10 @@ use self::{
     remote_table::RemoteTableExtension, sink::SinkExtension, table_source::TableSourceExtension,
     window_fn::WindowFunctionExtension,
 };
+use crate::builder::{NamedNode, Planner};
+use crate::schemas::{add_timestamp_field, has_timestamp_field};
+use crate::{fields_with_qualifiers, schema_from_df_fields, DFField, ASYNC_RESULT_FIELD};
+use join::JoinExtension;
 
 pub(crate) mod aggregate;
 pub(crate) mod debezium;
@@ -207,12 +205,6 @@ impl ArroyoExtension for AsyncUDFExtension {
             })
             .collect::<Result<Vec<_>>>()?;
 
-        ProjectionNode {
-            input: None,
-            expr: vec![],
-            optional_alias: None,
-        };
-
         let config = AsyncUdfOperator {
             name: self.name.clone(),
             udf: Some(self.udf.clone().into()),
diff --git a/crates/arroyo-planner/src/rewriters.rs b/crates/arroyo-planner/src/rewriters.rs
index fd3b47cf6..a938292f3 100644
--- a/crates/arroyo-planner/src/rewriters.rs
+++ b/crates/arroyo-planner/src/rewriters.rs
@@ -536,9 +536,25 @@ impl<'a> TreeNodeRewriter for AsyncUdfRewriter<'a> {
 
         let udf = self.provider.dylib_udfs.get(&name).unwrap().clone();
 
+        let input = if matches!(*projection.input, LogicalPlan::Projection(..)) {
+            // if our input is a projection, we need to plan it separately -- this happens
+            // for subqueries
+
+            Arc::new(LogicalPlan::Extension(Extension {
+                node: Arc::new(RemoteTableExtension {
+                    input: (*projection.input).clone(),
+                    name: TableReference::bare("subquery_projection"),
+                    schema: projection.input.schema().clone(),
+                    materialize: false,
+                }),
+            }))
+        } else {
+            projection.input
+        };
+
         Ok(Transformed::yes(LogicalPlan::Extension(Extension {
             node: Arc::new(AsyncUDFExtension {
-                input: projection.input,
+                input,
                 name,
                 udf,
                 arg_exprs: args,


Test patch: