sunchao commented on code in PR #59062:
URL: https://github.com/apache/spark/pull/59062#discussion_r4128799985
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala:
##########
@@ -104,3 +105,44 @@ class ArrowPythonUDTFRunner(
"Pandas execution requires more than 4 bytes. Please set higher buffer. " +
s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
}
+
+/**
+ * Similar to [[ArrowPythonRunner]], but for [[PythonUDTF]]s.
+ */
+class ArrowPythonUDTFRunner(
+ udtf: PythonUDTF,
+ evalType: Int,
+ argMetas: Array[ArgumentMetadata],
+ schema: StructType,
+ timeZoneId: String,
+ largeVarTypes: Boolean,
+ pythonRunnerConf: Map[String, String],
+ pythonMetrics: Map[String, SQLMetric],
+ jobArtifactUUID: Option[String],
+ sessionUUID: Option[String])
+ extends BaseArrowPythonUDTFRunner[Iterator[InternalRow]](
+ udtf, evalType, argMetas, schema, timeZoneId, largeVarTypes,
pythonRunnerConf,
+ pythonMetrics, jobArtifactUUID, sessionUUID)
+ with BatchedPythonArrowInput
+
+/**
+ * Similar to [[ArrowPythonUDTFRunner]], but the input is [[ColumnarBatch]]es,
and the UDTF
+ * arguments are the columns at `inputColumnIndices`. Arrow-backed columns are
serialized to the
+ * Python worker directly, see [[ColumnarArrowPythonInput]].
+ */
+private[python] class ColumnarArrowPythonUDTFRunner(
+ udtf: PythonUDTF,
+ evalType: Int,
+ argMetas: Array[ArgumentMetadata],
+ schema: StructType,
+ timeZoneId: String,
+ largeVarTypes: Boolean,
+ pythonRunnerConf: Map[String, String],
+ pythonMetrics: Map[String, SQLMetric],
+ jobArtifactUUID: Option[String],
+ sessionUUID: Option[String],
+ override protected val inputColumnIndices: Array[Int])
+ extends BaseArrowPythonUDTFRunner[ColumnarBatch](
+ udtf, evalType, argMetas, schema, timeZoneId, largeVarTypes,
pythonRunnerConf,
+ pythonMetrics, jobArtifactUUID, sessionUUID)
+ with ColumnarArrowPythonInput
Review Comment:
**[P2] Preserve the existing Arrow input batch limits**
Switching UDTFs from `BatchedPythonArrowInput` to `ColumnarArrowPythonInput`
loses enforcement of `spark.sql.execution.arrow.maxRecordsPerBatch` and
`maxBytesPerBatch`: the reused trait sends each upstream batch whole, including
its non-Arrow fallback. A real IPC probe with ten integers and a record limit
of 3 produces 3/3/3/1 rows through the existing row input, but one 10-row batch
through both new routes. With a byte limit of 16, the row path splits 4/4/2
integers while both columnar routes send all ten (40 data bytes). The Python
UDTF worker eagerly converts whole input columns into Python lists, so
workloads that lower these settings to control worker memory lose that
protection under the default-enabled path. Although the trait already exists
for scalar UDFs, routing UDTFs through it is new here. Please preserve both
limits and cover Arrow-backed and fallback batches larger than the configured
bounds.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonUDTFExec.scala:
##########
@@ -60,6 +62,83 @@ case class ArrowEvalPythonUDTFExec(
}
}
+ override lazy val metrics: Map[String, SQLMetric] = pythonMetrics ++ Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output
rows"),
+ "numOutputBatches" -> SQLMetrics.createMetric(sparkContext, "number of
output batches"))
+
+ /**
+ * The ordinals of the UDTF arguments in the child's output, if all the
arguments are columns of
+ * the child.
+ */
+ @transient private lazy val argumentOrdinals: Option[Array[Int]] = {
+ val ordinals = flattenArguments()._2.map {
+ case a: Attribute => child.output.indexWhere(_.exprId == a.exprId)
+ case _ => -1
+ }
+ if (ordinals.forall(_ >= 0)) Some(ordinals.toArray) else None
+ }
+
+ // When the child supports columnar output (e.g., Arrow-backed DSv2
connectors) and all the UDTF
+ // arguments are columns of the child, accept columnar input to avoid the
ColumnarToRow ->
+ // ArrowWriter round-trip of the arguments. Their Arrow FieldVectors are
serialized to the Python
+ // worker directly. The output rows are still joined row by row with the
input rows, as a UDTF
+ // returns any number of rows per input row.
+ override def supportsColumnar: Boolean =
+ child.supportsColumnar && conf.arrowPySparkUDTFColumnarInputEnabled &&
+ argumentOrdinals.isDefined
+ override def supportsRowBased: Boolean = true
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ if (child.supportsColumnar && argumentOrdinals.isDefined) {
+ executeWithColumnarInput()
+ } else {
+ super.doExecute()
+ }
+ }
+
+ override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ val evaluatorFactory = new RowToColumnarEvaluatorFactory(
+ conf.offHeapColumnVectorEnabled,
+ conf.columnBatchSize,
+ schema,
Review Comment:
**[P2] Gate columnar output on the converter's supported schema**
This newly reachable output path passes the complete UDTF schema to
`RowToColumnConverter`, which rejects `NullType` (and UDTs), although Arrow
UDTFs support these results. With a vectorized Parquet table `p(id)`, let
`nuller(id)` return `x VOID` by yielding `(None,)`, and `echo(x)` return `y
INT` by yielding `(42,)`. Then `SELECT p.id, f.x, g.y FROM p, LATERAL
nuller(p.id) f, LATERAL echo(f.x) g` fails with `[UNSUPPORTED_DATATYPE]
Unsupported data type "VOID"` when the new setting is enabled: the outer UDTF
requests the inner UDTF's columnar output. The same three-row probe succeeds on
base and with the setting disabled on head. Please restrict the advertised
output capability to schemas this conversion supports, or preserve support for
these valid types in the converter, and add a chained-operator regression test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonUDTFExec.scala:
##########
@@ -60,6 +62,83 @@ case class ArrowEvalPythonUDTFExec(
}
}
+ override lazy val metrics: Map[String, SQLMetric] = pythonMetrics ++ Map(
+ "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output
rows"),
+ "numOutputBatches" -> SQLMetrics.createMetric(sparkContext, "number of
output batches"))
+
+ /**
+ * The ordinals of the UDTF arguments in the child's output, if all the
arguments are columns of
+ * the child.
+ */
+ @transient private lazy val argumentOrdinals: Option[Array[Int]] = {
+ val ordinals = flattenArguments()._2.map {
+ case a: Attribute => child.output.indexWhere(_.exprId == a.exprId)
+ case _ => -1
+ }
+ if (ordinals.forall(_ >= 0)) Some(ordinals.toArray) else None
+ }
+
+ // When the child supports columnar output (e.g., Arrow-backed DSv2
connectors) and all the UDTF
+ // arguments are columns of the child, accept columnar input to avoid the
ColumnarToRow ->
+ // ArrowWriter round-trip of the arguments. Their Arrow FieldVectors are
serialized to the Python
+ // worker directly. The output rows are still joined row by row with the
input rows, as a UDTF
+ // returns any number of rows per input row.
+ override def supportsColumnar: Boolean =
+ child.supportsColumnar && conf.arrowPySparkUDTFColumnarInputEnabled &&
+ argumentOrdinals.isDefined
+ override def supportsRowBased: Boolean = true
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ if (child.supportsColumnar && argumentOrdinals.isDefined) {
+ executeWithColumnarInput()
Review Comment:
**[P2] Honor the disable setting in the row execution dispatch**
This condition omits `conf.arrowPySparkUDTFColumnarInputEnabled`. Setting it
to false only changes `supportsColumnar`; a dual-capability child such as
`InMemoryTableScanExec` still supports both row and columnar execution, so
transition insertion leaves it directly beneath this operator and this branch
still selects `executeWithColumnarInput()`. I reproduced this with a
materialized three-row `spark.range(3).coalesce(1).cache()` and an Arrow UDTF
over its `id`: with the flag false and `maxRecordsPerBatch=1`, head still
follows the new route, whereas base uses the bounded row runner. Consequently
the documented opt-out cannot restore the old path for cached data. Please
include the setting in this dispatch and test the disabled case with a child
that supports both execution modes.
--
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]
