andygrove opened a new pull request, #5744:
URL: https://github.com/apache/datafusion-comet/pull/5744
## Which issue does this PR close?
Closes #5386.
## Rationale for this change
Comet's opt-in Arrow UDF path (`spark.comet.exec.pyarrowUDF.enabled`)
covered only `mapInArrow` and `mapInPandas`. Scalar Python UDFs fell back
entirely, so a Comet scan feeding one paid a `ColumnarToRow` transition and
Spark then re-encoded those rows back into Arrow to reach the Python worker.
One operator, `ArrowEvalPythonExec`, covers three user-facing UDF families,
because `ExtractPythonUDFs` routes all of them to it:
- a plain `udf()` when `spark.sql.execution.pythonUDF.arrow.enabled=true`
- `@pandas_udf` (scalar)
- Spark 4.1's `@arrow_udf`
The first is what makes ordinary PySpark UDF code benefit from Comet, so
this is the operator that turns "Comet accelerates PyArrow UDFs" into "Comet
accelerates Python UDFs".
The win is larger than it was for `mapInArrow`, because
`EvalPythonEvaluatorFactory` is considerably more row-bound than
`MapInBatchExec` was. Per row the vanilla path does a
`HybridRowQueue.add(UnsafeRow)` (a full row copy that can spill to disk), a
`MutableProjection` to materialize the UDF arguments, then on the way back a
`queue.remove()`, a `JoinedRow`, and an `UnsafeProjection` over the full output
schema.
## What changes are included in this PR?
`CometArrowEvalPythonExec` sends the UDF argument columns straight from the
input batch to a `CometArrowEvalPythonRunner`, and appends the columns the
worker returns to the input batch's own columns to form the output batch. No
row is materialized, and the row queue and its spill path are gone.
One design point worth flagging, since it differs from what the issue
proposed. The issue suggested retaining the input `ColumnarBatch` and appending
to it. That is not safe: `CometExecIterator.hasNext` closes the previous batch
before fetching the next one, precisely because "the buffer memory [is] shared
across batches in the native side", and the Python runner writes from its own
thread and pipelines ahead. So the input batch is deep-copied on the way in
(`CometBatchDeepCopy`) — a bulk `memcpy` per Arrow buffer, with dictionary
columns decoded first. That copy replaces Spark's per-row `UnsafeRow` copy into
a spillable queue, so it stays well ahead of the unoptimized path, but it is a
copy rather than the zero-copy the issue envisaged.
Scope of this first version. Each restriction is enforced in the shim's
matcher, so an unsupported shape falls back to vanilla Spark rather than
failing:
- **Non-iterator eval types only** (`SQL_ARROW_BATCHED_UDF`,
`SQL_SCALAR_PANDAS_UDF`, and `SQL_SCALAR_ARROW_UDF` on 4.1+). The iterator
variants guarantee only equal total row counts, not equal batching, which the
batch pairing relies on.
- **Every UDF argument must be a plain attribute of the child.** Spark does
not project UDF arguments below the operator, so `udf(col("a"))` qualifies but
`udf(col("a") + 1)` does not. This restriction also excludes chained UDFs
(`f(g(x))`, folded into one operator whose argument is a `PythonUDF`) and
keyword arguments (`NamedArgumentExpression`), neither of which is an
`Attribute`.
- **Pickled `BatchEvalPythonExec` stays out of scope**, as the issue
specified: per-row boxing, pickle, and interpreter cost dominate, and the
unpickled Java-object results would have to be converted into Arrow vectors to
keep the batch all-Comet.
Comet does not flip `spark.sql.execution.pythonUDF.arrow.enabled` on the
user's behalf, since that conf changes Spark's own type coercion and error
semantics rather than just the transport. The user guide now explains what it
does and why enabling it is a prerequisite for accelerating a plain `udf()`.
Supporting changes:
- `CometArrowPythonRunnerBase` gains a `wrapInputInStruct` hook. Scalar eval
types exchange one top-level column per argument, matching the flat `_0`, `_1`,
... schema Spark builds, rather than the single struct `mapInArrow` uses. The
runner writes the input schema JSON for `SQL_ARROW_BATCHED_UDF` and uses the
`ArgumentMetadata` form of `writeUDFs`, both of which the worker's protocol
requires for these eval types.
- The shared driver-side runner-input resolution moves out of
`Spark4xMapInBatchSupport` into `ShimPythonRunnerInputs`, so a class can mix in
both Python operator shims without a conflicting `RunnerInputs`.
- `EliminateRedundantTransitions` rewrites the operator behind the existing
feature flag, annotates it with the same `[COMET-INFO]` opt-in hint when the
flag is off, and honours the same `useLargeVarTypes` fallback.
`extractColumnarChild` also accepts a `CometArrowEvalPythonExec` child, so
stacked operators stay columnar end to end.
- The shared `spark` / `accelerated` pytest fixtures move into `conftest.py`
so both UDF test modules run against one Spark session.
## How are these changes tested?
`CometArrowEvalPythonSuite` (11 tests, new) covers the plan rule without
spinning up Python: the rewrite and its output attributes, transition
stripping, each supported eval type, and a negative case for every fallback
condition (iterator eval type, non-attribute argument, chained UDF,
`useLargeVarTypes`, feature flag off, and the opt-in hint), plus stacked
operators.
`test_scalar_python_udf.py` (29 tests, new) runs a real Python worker end to
end, each test in both accelerated and fallback modes: all three UDF families,
pass-through of a wide mixed-type child, nulls, multiple arguments, a repeated
argument that must deduplicate to one exchanged column, several UDFs in one
operator, stacked operators of different eval types (where the outer UDF's
argument is the column the inner worker produced), empty input, and a
multi-batch run with `maxRecordsPerBatch=100`. The fallback cases assert the
operator is left to Spark and still produces correct results.
`CometArrowPythonRunnerSuite` gains a test that the flat path emits the
columns as top-level fields with no struct field node or validity buffer
prepended, and allocates nothing of its own.
Locally: all 162 pytest cases across both modules pass against Spark 4.1.3;
the 28 JVM tests in the three suites pass; `CometExecSuite` +
`CometSparkSessionExtensionsSuite` (155 tests) pass as a regression check on
the modified plan rule; and `test-compile` is clean on all five Spark profiles
(3.4, 3.5, 4.0, 4.1, 4.2). CI runs the pytest modules against a real worker on
4.0, 4.1, and 4.2, whose runner constructor and command framing differ.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
