cloud-fan commented on code in PR #58549:
URL: https://github.com/apache/spark/pull/58549#discussion_r4107605106
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -1316,7 +1316,13 @@ object ConvertToCatalyst extends Rule[LogicalPlan] {
def keepPython: Expression = s.pythonUDFExpr.mapChildren(recurse(_,
parentIsUdf = true))
// We _shouldn't_ have these nodes if ANSI is not enabled or
transpilation is disabled
// but if someone changed it while running we'll want to strip the
nodes out.
- if (!conf.getConf(SQLConf.ANSI_ENABLED)) {
+ val hasCheckedCharVarcharResult = s.pythonUDFExpr match {
+ case udf: PythonUDF => udf.hasCharVarcharResult
+ case _ => false
+ }
+ if (hasCheckedCharVarcharResult) {
+ s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true))
Review Comment:
**Blocking (P1):** This checked-result arm traverses every child through the
public `applyExpr` overload. That overload restarts with `inLambda = false` and
no `preEvaluate` callback, unlike the adjacent `keepPython`/`recurse` path. A
checked outer UDF in a higher-order-function lambda can therefore transpile a
nested child across the protected lambda boundary, and an ordinary nested call
can lose the plan's single-evaluation placement.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -189,11 +215,25 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+ context.addTaskCompletionListener[Unit] { _ =>
+ var entry = passThruQueue.poll()
+ while (entry != null) {
+ entry._1.foreach(_.close())
+ entry = passThruQueue.poll()
+ }
+ }
val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
+ // The input reader owns and may close its vectors as soon as the
Python runner consumes
Review Comment:
**Non-blocking (P2):** These transferred vectors acquire a cleanup owner
only after the whole array is constructed and added to `passThruQueue`. If a
later column transfer throws, earlier vectors never reach any close path.
Separately, cancellation can make the result side finish and drain the queue
before the writer runnable exits, allowing it to publish another transferred
array after cleanup. Both paths leak Arrow off-heap buffers on failed or
cancelled tasks.
**Recommended change:** Introduce one lifecycle owner for transferred
pass-through vectors that acquires each allocation immediately, cleans partial
construction on failure, atomically closes on task completion, and closes
rather than enqueues any late publication.
**Why this works:** Build transfers under a registry/guard that owns every
vector as soon as splitAndTransfer succeeds. Transfer ownership to a queued
batch only after the full array is valid. Make cleanup atomically mark the
registry closed before draining, and make producer publication observe that
state and close rejected late entries. Result consumption continues to remove a
complete entry and hand its vectors to the returned batch.
**Scope:** Make the optimized Arrow pass-through path exception-safe and
cancellation-safe without changing the row-queue fallback or Python worker
protocol.
**Compatibility:** Eligible Arrow-backed scalar UDFs continue using the full
columnar optimization and return the same pass-through and UDF result columns.
**Risks:** Double-close if ownership is not removed atomically when a
complete entry is handed to a result batch. Result/queue row-count pairing
could be disturbed if cleanup semantics alter normal FIFO publication.
**Constraints:** Preserve the optimized Arrow path for eligible unchecked
outputs. Do not close source-owned input vectors or result vectors owned by the
returned ColumnarBatch. Cleanup must remain idempotent across task failure and
normal completion.
**Success:** Every transferred vector is owned from successful allocation
until it is either closed or handed to a returned ColumnarBatch. A later-column
failure closes all earlier transfers from the same input batch. Task completion
prevents leaked late publications even when the writer thread has not yet
observed interruption. Normal FIFO pairing and output values remain unchanged.
##########
python/pyspark/sql/pandas/types.py:
##########
@@ -157,7 +159,7 @@ def to_arrow_type(
arrow_type = pa.float64()
elif isinstance(dt, DecimalType):
arrow_type = pa.decimal128(dt.precision, dt.scale)
- elif isinstance(dt, StringType):
+ elif isinstance(dt, (StringType, CharType, VarcharType)):
Review Comment:
**Blocking (P1):** This mapper widening also admits CHAR/VARCHAR in
incremental `Aggregator.bufferSchema`. That path is outside the supported
scope, yet neither `udaf()` nor the JVM builder rejects the schema, and the
partial worker serializes `reduce()` buffers directly with `pa.array`. An
over-length VARCHAR or unpadded CHAR buffer can therefore cross the shuffle
without assignment checks instead of failing deterministically at construction.
**Recommended change:** Recursively reject logical CHAR/VARCHAR and
CHAR/VARCHAR hidden in UDT storage in incremental bufferSchema at public
construction, with a matching JVM builder preflight before PythonAggregate is
created.
**Why this works:** Reuse the existing Python logical/UDT type predicates
for an early classed udaf() error, and the existing JVM physical/UDT
CHAR/VARCHAR predicate for defense in depth after bufferType is decoded. Leave
the shared transport mapper unchanged for supported consumers.
**Scope:** Restore the scope fence at the incremental aggregate buffer
boundary in both Python and JVM construction paths.
**Compatibility:** Incremental aggregators using supported Arrow buffer
types and all existing supported scalar CHAR/VARCHAR surfaces remain unchanged.
**Risks:** Python and JVM predicates could drift if they do not use the
established shared helpers. Rejecting only top-level fields would leave nested
structs, arrays, maps, or UDT storage silently admitted.
**Constraints:** Use the existing recursive logical and UDT-storage
predicates rather than duplicating schema traversal. Return a classed
unsupported-schema error before launching workers or planning the partial
aggregation. Do not reject ordinary supported scalar UDF return schemas.
**Success:** Public udaf() rejects direct and nested CHAR/VARCHAR in
bufferSchema before UDF creation. The JVM builder rejects the same schemas for
Connect reconstruction and direct/internal construction. No invalid constrained
buffer reaches PARTIAL pa.array serialization.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -456,9 +456,14 @@ case class PythonUDF(
// single lambda, and one more for each enclosing lambda when the UDF is
lifted out of a nested
// lambda (e.g. `transform(arr, i -> transform(i, x -> f(x)))` lifts `f`
to depth 2). Ignored
// for every non-element-wise eval type, where it stays at its default of
1.
- elementwiseNestingDepth: Int = 1)
+ elementwiseNestingDepth: Int = 1,
+ // The original CHAR/VARCHAR result type when write-side checks apply.
This is absent for
+ // unconstrained results, so the CHAR/VARCHAR policy is not part of
PythonUDF equality.
Review Comment:
**Nit (P3):** `charVarcharCheckedResultType` is a case-class constructor
field, so generated equality includes it, and canonicalized copies retain the
field. This sentence says the opposite and can mislead optimizer work that
depends on checked and unchecked UDF expressions remaining distinct.
--
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]
