Skip to content

Navigation Menu

Sign in
Sign up

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

carloea2 started this conversation in Ideas
Discussion options

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.

You must be logged in to vote

Replies: 0 comments

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Category
Ideas
Labels
None yet
1 participant

AltStyle によって変換されたページ (->オリジナル) /