One Await Per Batch, Not Per Row: How DataFusion Runs Async UDFs Without Blocking Its Workers

TL;DR: An async UDF is called with a whole RecordBatch of arguments, not one row at a time: one await per batch means one round trip per batch instead of per row. The UDF can declare an ideal_batch_size and DataFusion will slice the batch into chunks of that size, invoking the UDF once per chunk and concatenating the results. The async call happens inside a Stream::poll_next, so while the future is pending the Tokio worker thread is free to run other queries — the runtime is not blocked. Physical planning lifts async functions out of projections, filters and aggregates into a dedicated AsyncFuncExec node, so the async call is always batch-sized. And the sibling trick: a correlated scalar subquery in a filter is rewritten into a LEFT JOIN instead of being evaluated once per row.

The problem: SQL wants to call a remote service

The canonical use case is “call an LLM / a REST API / a geocoder from SQL”. The naive implementation is a scalar UDF evaluated once per row. That is wrong in three separate ways:

1. Round trips. One HTTP request per row. A million-row table is a million requests.
2. Runtime blocking. If the UDF uses a synchronous client (or blocks inside an async fn), it occupies a Tokio worker thread for the entire duration of the request. The worker cannot poll other tasks. One slow API call stalls every query sharing that runtime — the classic “one query killed the whole server” incident.
3. No vectorization. Arrow kernels work on arrays. A per-row UDF throws the array structure away and hands the function a scalar at a time.

DataFusion’s answer is a separate trait with an explicitly async, explicitly batch-shaped signature.

The trait: async, batch-shaped, self-describing

// datafusion/expr/src/async_udf.rs:37-50
#[async_trait]
pub trait AsyncScalarUDFImpl: ScalarUDFImpl {
    /// The ideal batch size for this function.
    ///
    /// This is used to determine what size of data to be evaluated at once.
    /// If None, the whole batch will be evaluated at once.
    fn ideal_batch_size(&self) -> Option<usize> {
        None
    }

    /// Invoke the function asynchronously with the async arguments
    async fn invoke_async_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue>;
}

Three design decisions:

• ScalarFunctionArgs, not scalars. The arguments arrive as already-evaluated ColumnarValues — Arrow arrays with a length. The UDF receives the whole column, not one element of it.

• ideal_batch_size is a hint from the UDF to the engine. Some remote APIs have a natural request size (an embeddings endpoint that takes 100 texts per call). The UDF says “give me 32 rows at a time” and DataFusion honours it.

• The synchronous path is forbidden. AsyncScalarUDF implements ScalarUDFImpl::invoke_with_args by returning an error:

  rust Copy

  // datafusion/expr/src/async_udf.rs:120-123
  fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
      internal_err!("async functions should not be called directly")
  }

  You cannot accidentally invoke an async function on a synchronous path and have it silently block.

The evaluation: chunk, invoke, concatenate

The physical-expression layer is where “per batch, not per row” is actually implemented. AsyncFuncExpr::invoke_with_args (datafusion/physical-expr/src/async_scalar_function.rs:109-196) evaluates the function arguments once per batch and then, if the UDF asked for a chunk size, slices the batch:

// datafusion/physical-expr/src/async_scalar_function.rs:139-166 (abridged)
let mut result_batches = vec![];
if let Some(ideal_batch_size) = self.ideal_batch_size()? {
    let mut remainder = batch.clone();
    while remainder.num_rows() > 0 {
        let size = if ideal_batch_size > remainder.num_rows() {
            remainder.num_rows()
        } else {
            ideal_batch_size
        };

        let current_batch = remainder.slice(0, size); // get next 10 rows
        remainder = remainder.slice(size, remainder.num_rows() - size);
        let args = scalar_function_expr
            .args()
            .iter()
            .map(|e| e.evaluate(&current_batch))
            .collect::<Result<Vec<_>>>()?;
        result_batches.push(
            async_udf
                .invoke_async_with_args(ScalarFunctionArgs {
                    args,
                    arg_fields: arg_fields.clone(),
                    number_rows: current_batch.num_rows(),
                    return_field: Arc::clone(&self.return_field),
                    config_options: Arc::clone(&config_options),
                })
                .await?
                .into_array_of_size(current_batch.num_rows())?,
        );
    }
} else {
    // whole batch at once
}

Note the .await? inside the loop: the chunked case performs one await per chunk, sequentially. The chunk boundaries come from the UDF, not from the engine, because only the UDF knows the remote API’s natural batch size.

The per-chunk results are then stitched back together:

// datafusion/physical-expr/src/async_scalar_function.rs:190-196
let dyn_arrays = result_batches
    .iter()
    .map(|arr| arr as &dyn arrow::array::Array)
    .collect::<Vec<_>>();
let result_array = concat(&dyn_arrays)?;
Ok(ColumnarValue::Array(result_array))

concat is zero-copy for the common case: Arrow arrays are buffers plus offsets, so concatenating N chunk results is one memcpy per column, not a re-materialisation of the values.

The operator: AsyncFuncExec

The evaluation above is wrapped in a dedicated physical operator, datafusion/physical-plan/src/async_func.rs. Its execute (:213-270) does two things: coalesce the input to a sensible batch size, then evaluate the async expressions per batch.

// datafusion/physical-plan/src/async_func.rs:237-243
let coalesced_input_stream = CoalesceInputStream {
    input_stream,
    batch_coalescer: LimitedBatchCoalescer::new(
        Arc::clone(&self.input.schema()),
        config_options_ref.execution.batch_size.get(),
        None,
    ),
};

This is the link to the previous article: the async function is fed datafusion.execution.batch_size (default 8192) row batches, so a remote call carries 8192 rows’ worth of arguments unless the UDF asks for smaller chunks.

Then the per-batch evaluation:

// datafusion/physical-plan/src/async_func.rs:245-268 (abridged)
let stream_with_async_functions = coalesced_input_stream.then(move |batch| {
    async move {
        let batch = batch?;
        // append the result of evaluating the async expressions to the output
        let mut output_arrays = batch.columns().to_vec();
        for async_expr in async_exprs_captured.iter() {
            let output = async_expr
                .invoke_with_args(&batch, Arc::clone(&config_options))
                .await?;
            output_arrays.push(output.to_array(batch.num_rows())?);
        }
        let batch = RecordBatch::try_new(schema_captured, output_arrays)?;
        Ok(batch.record_output(&baseline_metrics_captured))
    }
});

futures::stream::StreamExt::then gives each batch its own future. The .await on the UDF happens inside poll_next, and when the UDF’s future returns Pending, control returns to the Tokio scheduler. That is the whole “doesn’t block the runtime” mechanism: it is not that the call is cheap, it is that the worker thread is released while the socket is waiting.

   Tokio worker thread
   ─────────────────────────────────────────────────────────────────────
   poll_next() ──► coalescer ──► batch of 8192 rows
                                │
                                ▼
                       invoke_with_args(batch).await
                                │
              ┌─────────────────┴──────────────────┐
              │  UDF future returns Pending        │
              │  (HTTP request in flight)          │
              └─────────────────┬──────────────────┘
                                ▼
                   worker thread released ──► runs OTHER queries
                                │
                   socket readable ──► waker fires
                                ▼
                   poll_next() resumes ──► batch with the extra column

Planning: lifting async calls out of arbitrary expressions

The hard part is that an async function can appear anywhere in an expression: SELECT my_async(col) FROM t, WHERE my_async(col) > 0, SUM(my_async(col)). You cannot await inside a PhysicalExpr::evaluate, which is a synchronous function returning ColumnarValue.

DataFusion solves this by lifting the async calls into their own operator during physical planning. try_plan_async_exprs (datafusion/core/src/physical_planner.rs:3566) walks the expression tree, finds async functions, and replaces them with plain column references:

// datafusion/physical-plan/src/async_func.rs:424-470 (abridged)
pub struct AsyncMapper {
    /// the number of columns in the input plan
    /// used to generate the output column names.
    /// the first async expr is `__async_fn_0`, the second is `__async_fn_1`, etc
    num_input_columns: usize,
    /// the expressions to map
    pub async_exprs: Vec<Arc<AsyncFuncExpr>>,
}

The generated names (__async_fn_0, __async_fn_1, …) are how the rest of the plan refers to the results: the planner rewrites my_async(col) into __async_fn_0, and the AsyncFuncExec below it produces that column. The operator then appends the async columns to the end of the schema, and the consuming operator projects them back into place.

This happens for all three expression sites:

• Projection (datafusion/core/src/physical_planner.rs:3550): the async call is hoisted below the projection.

• Filter (:1205): AsyncFuncExec is placed below the filter, and the filter is given a projection that hides the extra columns:

  rust Copy

  // datafusion/core/src/physical_planner.rs:1204-1222 (abridged)
  let async_exec = AsyncFuncExec::try_new(async_map.async_exprs, physical_input)?;
  FilterExecBuilder::new(
      Arc::clone(&runtime_expr[0]),
      Arc::new(async_exec),
  )
  // project the output columns excluding the async functions
  // The async functions are always appended to the end of the schema.
  .apply_projection(Some((0..input.schema().fields().len()).collect::<Vec<_>>()))?

• Aggregate (:2155): the async call is evaluated before the aggregation, once per batch.

The ordering is the interesting consequence. Because the async call sits below the filter, WHERE my_async(col) = 'x' evaluates my_async on every input row, not just the ones that pass some cheaper predicate. Pushdown rules that reorder predicates will move the async call up or down the tree accordingly — and a predicate that could have filtered 99% of rows before the remote call will now, by construction, run after it. This is a real and under-appreciated cost of the “always batch-sized” design.

The sibling trick: scalar subquery → join

The same “don’t do per-row work in a filter” principle shows up in a different place: scalar subqueries.

SELECT * FROM t WHERE x > (SELECT max(y) FROM u) — the subquery is uncorrelated, so it can be evaluated once. But SELECT * FROM t WHERE x > (SELECT max(y) FROM u WHERE u.id = t.id) is correlated: it depends on the outer row. The naive execution evaluates it once per row of t.

DataFusion rewrites the correlated case into a join:

// datafusion/optimizer/src/scalar_subquery_to_join.rs:39-48
/// Optimizer rule that rewrites scalar subquery filters to joins and places an
/// additional projection on top of the filter to preserve the original schema.
///
/// When [`datafusion_common::config::OptimizerOptions::enable_physical_uncorrelated_scalar_subquery`] is
/// true (the default), only *correlated* scalar subqueries are rewritten here;
/// uncorrelated ones are left for physical execution via `ScalarSubqueryExec`.
/// When the option is false, all scalar subqueries — correlated and
/// uncorrelated — are rewritten to left joins by this rule.

The correlated subquery becomes a LEFT JOIN with the aggregate on the outer side, and the comparison moves into the join condition. The rule then adds a projection on top to restore the original schema, and compensates for the NULL semantics that the LEFT JOIN introduces (UN_MATCHED_ROW_INDICATOR, datafusion/optimizer/src/scalar_subquery_to_join.rs:24).

The uncorrelated case is handled physically instead, and it is a nice piece of engineering. ScalarSubqueryExec (datafusion/physical-plan/src/scalar_subquery.rs:64-90) is a pass-through node whose only job is to run the subquery once and publish the scalar:

Before any batches from the main input are yielded, all subquery plans are executed and their scalar results are stored in a shared [ScalarSubqueryResults] container owned by this node.

and the extraction enforces SQL’s cardinality rules:

// datafusion/physical-plan/src/scalar_subquery.rs:396-430 (abridged)
async fn execute_scalar_subquery(
    plan: Arc<dyn ExecutionPlan>,
    context: Arc<TaskContext>,
) -> Result<ScalarValue> {
    ...
    while let Some(batch) = stream.next().await.transpose()? {
        if batch.num_rows() == 0 { continue; }
        if result.is_some() || batch.num_rows() > 1 {
            return exec_err!("Scalar subquery returned more than one row");
        }
        result = Some(ScalarValue::try_from_array(batch.column(0), 0)?);
    }
    // 0 rows → typed NULL per SQL semantics
    match result {
        Some(v) => Ok(v),
        None => ScalarValue::try_from(schema.field(0).data_type()),
    }
}

Zero rows → a typed NULL (not a generic NULL — the type comes from the subquery’s schema), one row → the value, more than one → an error. The result is stored in a shared container indexed by SubqueryIndex, and the ScalarSubqueryExpr nodes embedded in the main input read from it by index. The whole computation is guarded by a OnceAsync (datafusion/physical-plan/src/joins/utils.rs:371) so that concurrent partitions of the main input do not each execute the subquery.

The config datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery (default true, datafusion/common/src/config.rs:1676) decides which of these two strategies is used for uncorrelated subqueries.

A third trick: the one-slot async generator

While we are in the runtime, one more piece of DataFusion’s async machinery deserves a mention, because it is used by nearly every operator that needs “push” semantics: datafusion/execution/src/async_stream.rs.

// datafusion/execution/src/async_stream.rs:29-33 (abridged)
pub fn async_stream<T, F: Future<Output = ()>>(
    generator: impl FnOnce(Emitter<T>) -> F,
) -> impl FusedStream<Item = T> {
    let (emitter, receiver) = tx_rx();
    AsyncStream::new(receiver, generator(emitter))
}

The implementation is a hand-rolled rendezvous channel with exactly one slot:

// datafusion/execution/src/async_stream.rs:173-186
pub fn emit(&mut self, value: T) -> impl FusedFuture<Output = ()> {
    self.set(value);
    Emit { done: false }
}

fn set(&mut self, value: T) {
    let mut guard = self.slot.lock();
    match &mut *guard {
        Some(_) => panic!("Misuse: await was not called after calling emit"),
        slot => *slot = Some(value),
    }
}

and the Emit future returns Poll::Pending on its first poll, which is how control is transferred back to the consumer:

// datafusion/execution/src/async_stream.rs:221-236
fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
    if !self.done {
        self.done = true;
        // Poll::Pending causes the generator to yield, returning control back to the
        // calling Stream
        Poll::Pending
    } else {
        Poll::Ready(())
    }
}

The consumer’s poll_next polls the generator, then takes the value out of the slot (datafusion/execution/src/async_stream.rs:262-290). One Arc<Mutex<Option<T>>>, no channel, no waker registration, and — crucially — the generator cannot run ahead, so backpressure is exact: the producer is suspended until the consumer takes the item. The “Misuse” panic catches the one way to misuse it (emitting twice without awaiting).

Every operator that needs to push batches into a stream uses this: hash_stream.rs, cross_join.rs, merge.rs, the clustered aggregate streams, and more.

Trade-offs and failure modes

• Blocking inside an async fn still blocks the worker. The await only helps if the underlying operation is genuinely async. A reqwest::blocking call inside invoke_async_with_args will park the worker thread for the duration of the request, and no amount of DataFusion machinery will save you. Use a genuinely async client, or spawn_blocking.
• AsyncFuncExec does not measure elapsed_compute. There is an explicit TODO at datafusion/physical-plan/src/async_func.rs:228. Its BaselineMetrics only records output rows, bytes and batches, so an EXPLAIN ANALYZE of a query with a remote UDF will show you how many rows went through the async call but not how long it took.
• The coalescer delays the first call. Because the input is coalesced to batch_size rows before the UDF is invoked, a LIMIT 1 query over a slow source waits for the batch to fill (or for the stream to end). The LIMIT is not pushed into the async path.
• The async call sits below the filter. As described above, a selective predicate that could avoid the remote call for most rows cannot, because the async evaluation is materialised as a column below the operator that consumes it. Predicate pushdown will move the async call around the tree, but it cannot move it above its own consumer.
• Errors abort the plan. An Err from the UDF propagates through the stream, the JoinSet in RecordBatchReceiverStreamBuilder sees it, and the whole query stops rather than emitting partial results.
• Async is scalar-only. There is no async aggregate UDF in this tree: datafusion/functions-aggregate/src/ contains no async aggregate implementation, and the only async trait is AsyncScalarUDFImpl. An async function used inside an aggregate is evaluated per input batch before the aggregation, which is the physical_planner.rs:2155 path.
• Chunking is sequential. ideal_batch_size produces N sequential awaits per batch, not N concurrent ones. For a UDF whose remote API has high per-request latency and low per-row cost, issuing the chunk requests concurrently would be much faster — but then the ordering of the results has to be reassembled, and the current implementation does not do that.

The takeaway

The trick is not “make the UDF async”. It is: make the unit of work a batch, make the unit of waiting a Pending, and lift the async call to a place in the plan where the batch is already materialised. Once the arguments arrive as Arrow arrays and the engine guarantees the call happens once per batch, everything else — the round-trip count, the vectorization, the runtime not being blocked — falls out for free. The scalar-subquery rewrite is the same idea in a different costume: the only way to avoid evaluating a subquery per row is to stop it from being a per-row operation at all.