Opensource internals

Once Per Batch: How DataFusion's ColumnarValue API Killed the Per-Row UDF Call

11 min read Jay Hubbert

DataFusion evaluates UDF arguments once per RecordBatch and hands whole Arrow arrays to the function, with ColumnarValue as the currency.

TL;DR: A DataFusion scalar function is not called per row. PhysicalExpr::evaluate returns a ColumnarValue — either an ArrayRef (the whole batch) or a single ScalarValue — and ScalarFunctionExpr::evaluate evaluates each child expression exactly once, then makes one call into the UDF with the whole batch. The UDF body then works with Arrow kernels. The only place per-row work survives is where a new array has to be constructed (list/struct builders), and even there the work is amortised through Arrow’s MutableArrayData.

The API that used to be slow

For years the shape of a UDF in DataFusion (and in Arrow-rs before it) was:

rust 
fn invoke(&self, args: &[ScalarValue]) -> Result<ScalarValue>

One call per row. Every call re-entered the vtable, re-matched on the argument types, allocated a ScalarValue for the result, and then the caller had to un-pack that ScalarValue into a batch. For a UInt8 column, that is 20x the memory traffic of the actual work, plus N virtual calls, plus N allocations.

The modern API is:

rust 
// datafusion/expr/src/udf.rs:702
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue>;

and the difference is not cosmetic — it changes what the function body can even be. The body receives ArrayRefs and is expected to call Arrow’s vectorized compute kernels (arrow::compute::kernels, take, filter, comparison, …). The per-row loop is gone, not moved.

ColumnarValue: the currency

rust 
// datafusion/expr-common/src/columnar_value.rs:102
pub enum ColumnarValue {
    /// Array of values
    Array(ArrayRef),
    /// A single value
    Scalar(ScalarValue),
}

Two variants, because a value in a query can legitimately be either a column of N rows or one constant. The Scalar variant is not a convenience: it is how literals and constant subexpressions avoid being materialised N times.

The leaves of the physical expression tree produce one or the other, and neither copies the data:

rust 
// datafusion/physical-expr/src/expressions/column.rs:121
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    self.bounds_check(batch.schema().as_ref())?;
    Ok(ColumnarValue::Array(Arc::clone(batch.column(self.index))))
}

// datafusion/physical-expr/src/expressions/literal.rs:106
fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
    Ok(ColumnarValue::Scalar(self.value.clone()))
}

A column reference is an Arc clone of the batch’s column — O(1), no data movement. A literal is a ScalarValue clone — usually a small enum, no data movement either. So “evaluating the arguments” of a function is O(number of arguments), not O(rows × arguments).

The per-batch contract, precisely

rust 
// datafusion/physical-expr/src/scalar_function.rs:236
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    let args = self.args.iter().map(|e| e.evaluate(batch)).collect::<Result<Vec<_>>>()?;
    let arg_fields = self.args.iter().map(|e| e.return_field(batch.schema_ref())).collect::<Result<Vec<_>>>()?;

    // evaluate the function
    let output = self.fun.invoke_with_args(ScalarFunctionArgs {
        args,
        arg_fields,
        number_rows: batch.num_rows(),
        return_field: Arc::clone(&self.return_field),
        config_options: Arc::clone(&self.config_options),
    })?;

    if let ColumnarValue::Array(array) = &output
        && array.len() != batch.num_rows()
    {
        // If the arguments are a non-empty slice of scalar values, we can assume that
        // returning a one-element array is equivalent to returning a scalar.
        let preserve_scalar = array.len() == 1 && !input_empty && input_all_scalar;
        return if preserve_scalar {
            ScalarValue::try_from_array(array, 0).map(ColumnarValue::Scalar)
        } else {
            internal_err!("UDF {} returned a different number of rows than expected. ...")
        };
    }
    Ok(output)
}

Four things to notice:

  1. One call. invoke_with_args is called exactly once per batch per function node.
  2. number_rows is passed explicitly. The function is told how many rows the batch has, so it can produce a result of the right length even if all of its inputs were scalars.
  3. The output length is checked. A function that returns the wrong number of rows is an internal error, caught at the boundary rather than producing a corrupt batch downstream.
  4. The scalar round-trip is preserved. If every argument was a scalar and the function returns a one-element array, it is converted back to a ColumnarValue::Scalar. That keeps SELECT 1 + 1 from becoming a one-element array that then has to be broadcast everywhere.

There is also a debug-only assertion in ScalarUDF::invoke_with_args (datafusion/expr/src/udf.rs:267) that the returned data type matches the type promised at planning time — free correctness checking in debug builds, free in release builds.

The cost of the uniform API: broadcasting scalars

Here is the honest trade-off. To write one array kernel and use it for both f(column) and f(1), the scalar has to be expanded into an array of num_rows:

rust 
// datafusion/expr-common/src/columnar_value.rs:192
pub fn to_array(&self, num_rows: usize) -> Result<ArrayRef> {
    Ok(match self {
        ColumnarValue::Array(array) => Arc::clone(array),
        ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(num_rows)?,
    })
}

f(1) over a million rows therefore materialises a million copies of 1 before the kernel runs. That is a real cost, and the type’s own documentation says so:

Because all functions must handle ArrayRef, it is convenient to convert ColumnarValue::Scalars using into_array. However, it is often much more performant to provide a different, implementation that handles scalar values differently — datafusion/expr-common/src/columnar_value.rs:90-100

The standard helper handles the common case for you. make_scalar_function (datafusion/functions/src/utils.rs:113) is the glue between the ColumnarValue world and the ArrayRef world:

rust 
// datafusion/functions/src/utils.rs:120
move |args: &[ColumnarValue]| {
    // first, identify if any of the arguments is an Array. If yes, store its `len`,
    // as any scalar will need to be converted to an array of len `len`.
    let len = args.iter().fold(Option::<usize>::None, |acc, arg| match arg {
        ColumnarValue::Scalar(_) => acc,
        ColumnarValue::Array(a) => Some(a.len()),
    });
    let is_scalar = len.is_none();
    let inferred_length = len.unwrap_or(1);

    let args = args.iter()
        .zip(hints.iter().chain(std::iter::repeat(&Hint::Pad)))
        .map(|(arg, hint)| {
            let expansion_len = match hint {
                Hint::AcceptsSingular => 1,
                Hint::Pad => inferred_length,
            };
            arg.to_array(expansion_len)
        })
        .collect::<Result<Vec<_>>>()?;

    let result = (inner)(&args);
    if is_scalar {
        // If all inputs are scalar, keeps output as scalar
        let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
        result.map(ColumnarValue::Scalar)
    } else {
        result.map(ColumnarValue::Array)
    }
}

Read it as three rules:

  • The batch length is inferred from the first array argument.
  • Scalars are padded to that length (or kept at length 1 if the Hint says the argument is singular, e.g. a “delta” in date_bin).
  • If every argument was a scalar, the result is collapsed back to a scalar.

That last rule is why SELECT 1 + 1 stays a scalar all the way down, and why SELECT f(1) FROM t does not allocate a million-element array just to read element 0.

The performance-minded escape hatch is to look at args.args directly and match on ColumnarValue::Scalar vs ColumnarValue::Array yourself, writing two code paths. datafusion/functions-nested/src/utils.rs:230 has a simpler variant of the same helper for the nested-function package, and datafusion/functions/src/utils.rs:164 (calculate_binary_math) shows the “handle the scalar operand specially” pattern for binary math.

Nested values: where per-row work is unavoidable

Now the interesting case. make_array(1, 2, 3) does not just compute something per row — it builds a new array whose shape depends on the row. Each output row is a list of N elements, one from each input column:

text 
┌─────────┐   ┌─────────┐           ┌──────────────┐
│ ┌─────┐ │   │ ┌─────┐ │           │ ┌──────────┐ │
│ │  A  │ │   │ │  X  │ │           │ │  [A, X]  │ │
│ ├─────┤ │   │ ├─────┤ │           │ ├──────────┤ │
│ │NULL │ │   │ │  Y  │ │──────────▶│ │[NULL, Y] │ │
│ ├─────┤ │   │ ├─────┤ │           │ ├──────────┤ │
│ │  C  │ │   │ │  Z  │ │           │ │  [C, Z]  │ │
│ └─────┘ │   │ └─────┘ │           │ └──────────┘ │
└─────────┘   └─────────┘           └──────────────┘
  col1           col2                    output

The implementation (datafusion/functions-nested/src/make_array.rs:195) builds a MutableArrayData over all the input arrays and then walks the rows appending one element per column:

rust 
// datafusion/functions-nested/src/make_array.rs:227
let num_rows = args[0].len();
for row_idx in 0..num_rows {
    for (arr_idx, arg) in args.iter().enumerate() {
        if !arg.as_any().is::<NullArray>() && !arg.is_null(row_idx) && arg.is_valid(row_idx) {
            mutable.try_extend(arr_idx, row_idx, row_idx + 1)?;
        } else {
            mutable.try_extend_nulls(1)?;
        }
    }
    offsets.push(O::usize_as(mutable.len()));
}
let data = mutable.freeze();
Ok(Arc::new(GenericListArray::<O>::try_new(
    Arc::new(Field::new(field_name, data_type, true)),
    OffsetBuffer::new(offsets.into()),
    arrow::array::make_array(data),
    None,
)?))

This is a per-row loop — it has to be, the output shape is per row. But notice what it is not: it is not a virtual call per row, it does not allocate per row (the MutableArrayData buffer grows amortised and is frozen once), and the null handling is inline. The offsets are built in the same pass, so there is no second traversal. The alternative — per-row ScalarValue construction followed by a bulk ListArray build — would allocate a million ScalarValues.

named_struct (datafusion/functions/src/core/named_struct.rs:155) takes the zero-copy route where it can:

rust 
let values: Vec<ColumnarValue> = args.args.as_chunks::<2>().0.iter()
    .map(|[_name, value]| value.clone()).collect();
let arrays = ColumnarValue::values_to_arrays(&values)?;
Ok(ColumnarValue::Array(Arc::new(StructArray::new(fields.clone(), arrays, None))))

StructArray::new takes the child arrays as they are — for column arguments that is zero copy. Only scalar arguments get broadcast (via values_to_arrays, columnar_value.rs:249), and only because a StructArray requires all children to have the same length.

Note that named_struct also has to reject non-constant field names at planning time (return_field_from_args, named_struct.rs:98-137): the schema of a struct depends on the values of its first argument, so this function’s return type can only be computed from literals. This is a general pattern — when a function’s output type depends on its arguments’ values, the type has to be resolved at plan time, and the argument evaluation has to happen then, not at runtime.

“Per batch” does not mean “per group of rows”

A common misreading: if the function is called once per batch, then random() would return the same value for every row. It does not — because the function is told the row count:

rust 
// datafusion/functions/src/math/random.rs:82
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    let mut rng = rng();
    let mut values = vec![0.0; args.number_rows];
    // Equivalent to set each element with rng.random_range(0.0..1.0), but more efficient
    rng.fill(&mut values[..]);
    Ok(ColumnarValue::Array(Arc::new(Float64Array::from(values))))
}

One RNG, one Vec<f64>, one Arrow array — but N distinct values. The batch-at-a-time API does not force constant semantics; it just moves the decision about how to produce per-row variation into the function, where it can be done efficiently. What it does forbid is the optimizer from assuming the call is repeatable, which is exactly what Volatility::Volatile encodes (see the signature article).

From Expr to ColumnarValue

The logical plan is a tree of Expr (datafusion/expr/src/expr.rs:327), and a scalar call is one node:

rust 
// datafusion/expr/src/expr.rs:950
pub struct ScalarFunction {
    /// The function
    pub func: Arc<crate::ScalarUDF>,
    /// List of expressions to feed to the functions as arguments
    pub args: Vec<Expr>,
}

The Expr tree is a description of what to compute; it is manipulated by the optimizer, hashed, compared, and rewritten. It never runs. At physical planning time it is lowered into a tree of PhysicalExpr — where the scalar call becomes ScalarFunctionExpr with a Vec<Arc<dyn PhysicalExpr>> of children — and that tree is what runs, once per batch, producing ColumnarValues.

The two trees have different jobs:

  • Expr is a description with types, spans, and metadata. It is cheap to clone, expensive to hash, and it is what the optimizer reasons about.
  • PhysicalExpr is an executable plan fragment. It holds an Arc<ScalarUDF> (so the signature and the implementation are shared, not copied), the argument expressions, and the resolved return field. It is built once per query and reused for every batch.

The Expr tree is where a “constant” stays a constant: Expr::Literal is folded by the simplifier (if the function is not volatile) into a ScalarValue, and the physical planner turns it into a Literal physical expression whose evaluate returns ColumnarValue::Scalar. That is the mechanism by which “evaluate the arguments once per batch” also means “don’t evaluate constant arguments at all”.

Trade-offs and rough edges

  • Broadcasting is a real copy. f(1) over N rows materialises N values. The Hint::AcceptsSingular escape hatch and hand-written scalar fast paths exist precisely to avoid it, but they are opt-in and easy to forget.
  • The output length check is a runtime cost. One comparison per function node per batch. Cheap, but it is why the “return a one-element array for scalar inputs” case has to be special-cased rather than just allowed.
  • arg_fields is recomputed per batch. ScalarFunctionExpr::evaluate calls return_field on every child on every batch (scalar_function.rs:243-248). For a Column that is an Arc clone of a schema field; for a Literal it is a field clone. It is not free, and it is pure overhead that could be hoisted into the constructor. It is, however, small compared to the work the function is about to do.
  • Nested builders are per row by necessity. make_array and friends walk rows. The mitigation is amortisation (MutableArrayData, a single freeze(), offsets built in the same pass) and the avoidance of per-row allocation.
  • The scalar/array duality leaks into every UDF. Every author has to decide: do I handle ColumnarValue::Scalar specially, or do I let make_scalar_function broadcast for me? Both are legitimate; the second is simpler and the first is faster. The API does not force a choice, which is exactly why performance varies between functions that look identical from the outside.
  • Bigger batches amortise everything. The per-batch overhead (vtable call, length check, field lookups) is amortised over batch_size rows, so a function that is cheap per row can still be a measurable cost on 8 KiB batches. The vectorised API makes the per-row cost small; it does not make the per-batch cost zero.

DatafusionFunctionsPerformanceInternals