dongjoon-hyun commented on PR #57804:
URL: https://github.com/apache/spark/pull/57804#issuecomment-5235039620
I went through the whole diff. Overall this is a well-designed and
thoroughly tested PR: the shared predicate
(`PythonUDF.isElementwiseRewritableUDF` / `canRewritePythonUDFInLambda`) keeps
"what analysis accepts" aligned with "what the rule rewrites", the marker
traits make the rewrite generic for future HOFs, and the SPARK-48706 failure
mode (a raw `PythonUDF` left inside a generated lambda) is explicitly guarded
and regression-tested. A few points below.
### Verified while reviewing
- Dedup semantics: deterministic duplicate calls collapse to one lifted UDF
while nondeterministic calls stay distinct via `liftKey` (`resultId`-bearing
node), and the key-form comparator's `udf(a)`/`udf(b)` correctly merge after
lifting — all pinned by tests.
- Worker re-nesting is consistent with `pa.ListArray.flatten()` (null
entries consume no offsets), and fused UDFs re-nest by their own first
argument's shape; `verify_result_row_count` defends against misalignment.
- Ragged `zip_with`/`map_zip_with`: the common `arrays_zip` alignment and
the `ArrayUnion` + `element_at(..., failOnError=false)` desugar match the
native null-padding / key-union semantics.
- The conf's `.version("4.4.0")` matches the next feature release branch.
### Suggestions
1. **Document the eager-evaluation semantics.** After the rewrite, the UDF
runs over *every* element regardless of `exists` short-circuiting, `when`
branches, or boolean short-circuit around it. If the UDF raises on inputs a
user expects to be skipped conditionally, the query now fails. Not a regression
(these queries used to fail analysis), but the conf doc only partially covers
it — a note in the SQL migration guide and/or the PySpark UDF docs would help
(this PR has no doc changes).
2. **Consider a size guard for the pairwise `array_sort` comparator.** `(a,
b) -> udf(a, b)` costs n² Python calls plus an n²-cell array per row; n = 10k
means 100M calls per row, which will effectively hang or OOM an executor. Also,
the "O(n²) carry via shared references" claim holds for `GenericArrayData`, but
if the carrier is ever serialized to the Unsafe format the per-element copies
degrade to O(n³). A cap with a clear error (or at least a warning in the conf
doc) seems worthwhile.
3. **Defense-in-depth in the rule itself.** `liftableHof` doesn't check for
free lambda variables — the nested-lambda guard lives only in `CheckAnalysis`.
If the rule ever sees a plan that didn't go through analysis (internally
constructed, or a shape reintroduced by a future rule), it would mis-rewrite
`transform(arr, i -> transform(i, x -> f(x)))` and resurrect SPARK-48706.
Reusing `hasFreeLambdaVariable` in `liftableHof` is cheap.
4. **Test gaps worth covering:**
- The SQL string syntax path, e.g. `spark.sql("SELECT transform(v, x ->
pyUdf(x)) ...")`.
- An end-to-end test that a kwargs call (`udf(x, k=1)`) still fails
analysis (only the predicate unit test exists).
- A rewritable HOF sitting *inside another HOF's lambda* while iterating
a real column, e.g. `transform(arr2, i -> array_max(transform(arr, x -> f(x)))
+ i)` — analysis accepts it, the lifted element-wise UDF stays inside the outer
lambda, and `ExtractPythonUDFs` then extracts it per-row. Deterministic UDFs
give the same result, but the path is unusual enough to deserve a test.
- Decimal/timestamp/struct element types (Arrow conversion edge cases);
current e2e coverage is int/double/string/array.
5. Minor: `transform_keys` rebuilds through `MapFromArrays`, so null-key /
duplicate-key handling matches native dedup-policy behavior, but the error
class on a null key may differ slightly from `TransformKeys`' own — fine for a
new feature, just noting it.
None of these block the PR in my view — #1 and #2 are the substantive ones,
the rest are nice-to-haves.
--
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]
