srielau commented on code in PR #58549:
URL: https://github.com/apache/spark/pull/58549#discussion_r4125448602
##########
python/pyspark/sql/tests/connect/test_parity_udf.py:
##########
@@ -46,6 +46,10 @@ def test_udf_with_input_file_name_for_hadooprdd(self):
def test_same_accumulator_in_udfs(self):
super().test_same_accumulator_in_udfs()
+ @unittest.skip("Spark Connect resolves both UDFs when the plan is
submitted.")
Review Comment:
Kept the skip with a more precise reason. Encoding the captured CHAR/VARCHAR
check policy needs a new optional field on the Connect `PythonUDF` proto plus
stub regen, which I am not bundling into this landing. Classic still captures
the policy per expression; Connect continues to rebuild from submission-time
SQLConf.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -189,12 +216,64 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+ val passThruClosed = new AtomicBoolean(false)
+
+ def closeCols(cols: Array[ColumnVector]): Unit = {
+ var i = 0
+ while (i < cols.length) {
+ cols(i).close()
+ i += 1
+ }
+ }
+
+ def drainPassThru(): Unit = {
+ passThruClosed.set(true)
+ var entry = passThruQueue.poll()
+ while (entry != null) {
+ closeCols(entry._1)
+ entry = passThruQueue.poll()
+ }
+ }
+
+ context.addTaskCompletionListener[Unit] { _ =>
Review Comment:
Added an Arrow-backed multi-batch case
(`ARROW_EXECUTION_MAX_RECORDS_PER_BATCH=3`, 10 rows) that asserts pass-through
ids stay aligned, plus `take(1)` so consumption can stop with later batches
still outstanding. Task completion now closes the live transferred vectors as
well as anything still queued.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -189,12 +216,64 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+ val passThruClosed = new AtomicBoolean(false)
+
+ def closeCols(cols: Array[ColumnVector]): Unit = {
+ var i = 0
+ while (i < cols.length) {
+ cols(i).close()
+ i += 1
+ }
+ }
+
+ def drainPassThru(): Unit = {
+ passThruClosed.set(true)
+ var entry = passThruQueue.poll()
+ while (entry != null) {
+ closeCols(entry._1)
+ entry = passThruQueue.poll()
+ }
+ }
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ drainPassThru()
+ }
val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
+ // The input reader owns and may close its vectors as soon as the
Python runner consumes
+ // the input iterator. Create independent vector views whose buffers
remain valid until
+ // the corresponding output batch is closed.
+ val cols = Array.ofDim[ColumnVector](childOutput.length)
+ var allocated = 0
+ try {
+ var i = 0
+ while (i < childOutput.length) {
+ val vector =
batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector
+ val transferPair = vector.getTransferPair(ArrowUtils.rootAllocator)
+ transferPair.splitAndTransfer(0, batch.numRows())
Review Comment:
`splitAndTransfer`/`getTo` now take ownership immediately: if wrapping the
transferred vector throws, that target is closed before the outer catch runs.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -86,6 +87,9 @@ case class UserDefinedPythonFunction(
val optionInputTypes: List[List[String]] =
transpiledInputTypes.asScala.map(_.asScala.toList).toList
+ if (CharVarcharUtils.hasCharVarcharInUDT(dataType)) {
Review Comment:
The JVM builder now rejects logical CHAR/VARCHAR results for the four
grouped-aggregate eval types with
`PYTHON_UDF_CHAR_VARCHAR_AGGREGATE_RETURN_TYPE`. Scalar batched/pandas/Arrow
builders still accept and check CHAR/VARCHAR output.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala:
##########
@@ -267,19 +314,27 @@ object EvaluatePython {
}
case MapType(keyType, valueType, _) =>
- val keyFromJava = makeFromJava(keyType)
- val valueFromJava = makeFromJava(valueType)
+ val keyFromJava = makeFromJava(keyType, applyCharVarcharChecks)
+ val valueFromJava = makeFromJava(valueType, applyCharVarcharChecks)
(obj: Any) => nullSafeConvert(obj) {
case javaMap: java.util.Map[_, _] =>
- ArrayBasedMapData(
- javaMap,
- (key: Any) => keyFromJava(key),
- (value: Any) => valueFromJava(value))
+ if (applyCharVarcharChecks &&
CharVarcharUtils.hasCharVarchar(keyType)) {
Review Comment:
The map-key CHAR/VARCHAR branch condition is computed once when the
converter is built (`checkCharVarcharMapKeys`), so row conversion no longer
rescans `keyType`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -106,9 +110,32 @@ case class UserDefinedPythonFunction(
throw
QueryCompilationErrors.invalidIncrementalPythonAggregatorBufferError(
name, bufferType)
}
+ if (CharVarcharUtils.physicalTypeHasCharVarchar(bufferStruct)) {
+ throw
QueryCompilationErrors.invalidPythonAggregatorBufferType(bufferStruct)
+ }
PythonAggregate(name, func, dataType, e, udfDeterministic, bufferStruct)
} else {
- PythonUDF(name, func, dataType, e, pythonEvalType, udfDeterministic)
+ val conf = SQLConf.get
+ val resolvedDataType = if (conf.charVarcharFirstClassTypes) {
Review Comment:
Added the preserve + legacy-as-string + standard-off case for scalar UDFs
and for Arrow createDataFrame/toArrow, asserting first-class schema, padding,
overflow rejection, and physical STRING Arrow routing.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -189,12 +216,64 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+ val passThruClosed = new AtomicBoolean(false)
+
+ def closeCols(cols: Array[ColumnVector]): Unit = {
+ var i = 0
+ while (i < cols.length) {
+ cols(i).close()
+ i += 1
+ }
+ }
+
+ def drainPassThru(): Unit = {
+ passThruClosed.set(true)
+ var entry = passThruQueue.poll()
+ while (entry != null) {
+ closeCols(entry._1)
+ entry = passThruQueue.poll()
+ }
+ }
+
+ context.addTaskCompletionListener[Unit] { _ =>
+ drainPassThru()
+ }
val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
+ // The input reader owns and may close its vectors as soon as the
Python runner consumes
+ // the input iterator. Create independent vector views whose buffers
remain valid until
+ // the corresponding output batch is closed.
+ val cols = Array.ofDim[ColumnVector](childOutput.length)
+ var allocated = 0
+ try {
+ var i = 0
+ while (i < childOutput.length) {
+ val vector =
batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector
+ val transferPair = vector.getTransferPair(ArrowUtils.rootAllocator)
Review Comment:
The output iterator now closes the previous batch's transferred pass-through
columns when the next batch is produced, and the task-completion listener
closes the live set plus anything still queued. That returns those
`rootAllocator` views even when downstream never calls `ColumnarBatch.close()`.
##########
python/pyspark/sql/udf.py:
##########
@@ -467,7 +498,10 @@ def _check_return_type(returnType: DataType, evalType:
int) -> None:
"return_type": str(returnType),
},
)
- elif evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF:
+ elif (
+ evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF
+ or evalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF
Review Comment:
Unintended. Those iterator eval types are out of the CHAR/VARCHAR support
fence; they now fall through to the CHAR/VARCHAR-only rejection and no longer
run the grouped-agg `StructType` / `to_arrow_type` checks for other types.
--
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]
