Support Arrow batches as an explicit Python UDF input and output mode

[carloea2]

Feature Summary

Allow Python UDFs to explicitly receive and return Arrow batches, avoiding the intermediate Python Tuple conversion while keeping existing UDF APIs unchanged.

PyAmber currently expands incoming Arrow tables into Tuples and converts UDF table output back through Tuples before rebuilding Arrow. This adds substantial overhead for operators that already work on entire batches.

An exploratory local conversion benchmark on upstream commit ec3a9dd, using PyArrow 23.0.1 and 10,000 rows with 10 string columns of 64-character values, measured:

Conversion path	Time
Current Arrow, Tuple, pandas, Tuple, Arrow path	1,566 ms
Direct Arrow, pandas, Arrow path	20.4 ms

This is approximately 77 times faster for the measured conversion path. It indicates potential savings from avoiding row conversion, not a measured 77 times improvement in workflow execution. The proposed engine path has not been implemented. The full Arrow Flight benchmark was blocked locally by a JOOQ schema mismatch, so an end-to-end benchmark is still needed.

Proposed Solution or Design

Introduce an explicit opt-in API, for example:

import pyarrow as pa
import pyarrow.compute as pc
from pyamber import ArrowBatchOperator  # Proposed API


class FilterPrices(ArrowBatchOperator):
    BATCH_SIZE = 4096

    def process_batch(self, batch: pa.Table, port: int):
        yield batch.filter(pc.greater(batch["price"], 100))

The engine would deliver Arrow batches directly and accept Arrow output without expanding every row into a Tuple. Syntax alone would not remove the current input and output conversions.

• Keep existing TupleOperatorV2, BatchOperator, and TableOperator behavior unchanged.
• Respect configured batch sizes, final partial batches, port boundaries, and row order.
• Preserve schema validation, Texera partition hashing, storage writes, and backpressure.
• Define inspection and retry around a batch invocation for the new API, with control handling between invocations and output yields. Do not automatically rerun UDFs to fall back, since they may have side effects.
• Specify null, timestamp, binary, and error behavior explicitly, including what happens to valid output preceding a validation failure.

Start with ArrowBatchOperator. A separate ArrowTableOperator could later support whole-port input at completion. Validate the proposal with reproducible conversion and full engine benchmarks, plus tests for control handling and output equivalence where the APIs share semantics.

Affected Area

Workflow Engine (Amber)

Originally raised in #8476. Continuing the proposal here for discussion.

[sshiv012]

+1 on the direction. The conversion cost is real: DataProcessor._set_output_tuple finalizes every output row individually, OutputManager.tuple_to_batch feeds partitioners one tuple at a time, and tuple_to_frame rebuilds the Arrow table from tuples. For a batch or table UDF that already produced a DataFrame, that round trip is pure overhead.

Please check #8914 (per-port output, I implemented a prototype on a branch). You might want to take into account

1. Output port. A UDF can emit on more than one port, and ports may have different schemas (table on port 0, html-content on port 1). An Arrow output needs a port, e.g. yield PortOutput(batch, port=1), and schema validation has to happen per port, at the batch level (cast the Arrow schema to the port schema once, instead of row by row).
2. Partitioning and completion flush. Hash partitioning hashes field values per row, and emit_ecm flushes each link's buffered tail with that link's port schema. The Arrow path has to keep both, otherwise rows can land on a different worker than today or a tail batch can go out with the wrong schema. Arrow compute can hash the key columns, but the hash must match the current per-row result.

Happy to coordinate so that the two branches don't have conflicts in design, merge conflicts can be resolved :)