comphead commented on code in PR #7253:
URL: https://github.com/apache/datafusion-comet/pull/7253#discussion_r4414766943
##########
spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala:
##########
@@ -722,4 +737,98 @@ class CometGenerateExecSuite extends CometTestBase {
}
}
+ /** Run `f` over a single-file Parquet table `t` of `(id, s, js)` rows. */
+ private def withGeneratorTable(f: => Unit): Unit = {
+ val rows = Seq(
+ (1, "a", """{"a": 1, "b": "x"}"""),
+ (2, null, """{"a": null}"""),
+ (3, "c", "not json"),
+ (4, "d", null))
+ withTempPath { path =>
+ rows.toDF("id", "s",
"js").repartition(1).write.parquet(path.getAbsolutePath)
+ withParquetTable(path.getAbsolutePath, "t")(f)
+ }
+ }
+
+ /** Check `df` against Spark, and that its generator ran in Comet's native
operator. */
+ private def checkSparkAnswerWithNativeGenerate(df: DataFrame): Unit = {
+ val (_, cometPlan) = checkSparkAnswer(df)
+ assert(
+ collectFirst(cometPlan) { case generate: CometExplodeExec => generate
}.isDefined,
+ s"Expected a native generator in\n$cometPlan")
+ }
+
+ test("Hive UDTF runs natively and emits its terminate rows") {
Review Comment:
The description says the state a Hive UDTF keeps carries across batches, but
`withGeneratorTable` writes four rows into one partition and the batch size is
not set, so every new test runs a single batch. A change that gave each batch a
fresh generator would still pass. Could this test also run with
`CometConf.COMET_BATCH_SIZE.key -> "2"`, as the earlier tests in this file do,
so the running count and the `close` total cross a batch boundary?
##########
spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala:
##########
@@ -722,4 +737,98 @@ class CometGenerateExecSuite extends CometTestBase {
}
}
+ /** Run `f` over a single-file Parquet table `t` of `(id, s, js)` rows. */
+ private def withGeneratorTable(f: => Unit): Unit = {
+ val rows = Seq(
+ (1, "a", """{"a": 1, "b": "x"}"""),
+ (2, null, """{"a": null}"""),
+ (3, "c", "not json"),
+ (4, "d", null))
+ withTempPath { path =>
+ rows.toDF("id", "s",
"js").repartition(1).write.parquet(path.getAbsolutePath)
+ withParquetTable(path.getAbsolutePath, "t")(f)
+ }
+ }
+
+ /** Check `df` against Spark, and that its generator ran in Comet's native
operator. */
+ private def checkSparkAnswerWithNativeGenerate(df: DataFrame): Unit = {
+ val (_, cometPlan) = checkSparkAnswer(df)
+ assert(
+ collectFirst(cometPlan) { case generate: CometExplodeExec => generate
}.isDefined,
+ s"Expected a native generator in\n$cometPlan")
+ }
+
+ test("Hive UDTF runs natively and emits its terminate rows") {
+ CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF])
+ withGeneratorTable {
+ // The row `close` emits carries the last input row's `id`, as in Spark.
+ checkSparkAnswerAndOperator(sql("SELECT id, counting_udtf(s) FROM t"))
+ checkSparkAnswerAndOperator(
+ sql("SELECT id, value, n FROM t LATERAL VIEW OUTER counting_udtf(s) u
AS value, n"))
+ // With no input rows only the `close` row remains.
+ checkSparkAnswerAndOperator(sql("SELECT counting_udtf(s) FROM t WHERE id
> 10"))
+ // Two identical UDTFs in one plan each keep their own count.
+ checkSparkAnswerAndOperator(sql("""
+ |SELECT id, a.value, a.n, b.value, b.n FROM t
+ |LATERAL VIEW counting_udtf(s) a AS value, n
+ |LATERAL VIEW counting_udtf(s) b AS value, n""".stripMargin))
+ }
+ }
+
+ test("Hive UDTF keeps separate state per partition") {
+ CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF])
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "3") {
+ withGeneratorTable {
+ // Each partition ends with one `close` row holding its own count, so
the counts add up
+ // to the input rows whichever partition each row lands in.
+ checkSparkAnswerWithNativeGenerate(sql("""
+ |SELECT count(*), sum(n) FROM (
+ | SELECT counting_udtf(s) FROM (SELECT * FROM t DISTRIBUTE BY id))
+ |WHERE value = 'total'""".stripMargin))
+ }
+ }
+ }
+
+ test("json_tuple runs natively") {
Review Comment:
This and the two `EXCEPT ALL` tests fit the SQL file harness. I could not
find a `json_tuple` fixture or an `EXCEPT ALL`/`INTERSECT ALL` fixture under
`spark/src/test/resources/sql-tests`, so a `json/json_tuple.sql` and an
`except_all.sql` with plain `query` blocks would also check native placement.
`ReplicateRows` carries every input column through the new `ArrowWriter` output
path, but the test uses only a `STRING` column. Could the fixture cover
`DECIMAL`, `BINARY`, `FLOAT` and `DOUBLE` with NaN, `-0.0` and infinity,
`DATE`, `TIMESTAMP`, `TIMESTAMP_NTZ`, and array and struct columns, each with
NULLs? The Hive UDTF cases that need a registered class can stay in Scala.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -1920,59 +1918,102 @@ object CometExplodeExec extends
CometOperatorSerde[GenerateExec] {
if (!op.generator.deterministic) {
return Unsupported(Some("Only deterministic generators are supported"))
}
- if (op.generator.children.length != 1) {
- return Unsupported(Some("generators with multiple inputs are not
supported"))
- }
- val nodeName = op.generator.nodeName.toLowerCase(Locale.ROOT)
- if (nodeName != "explode" && nodeName != "posexplode") {
- return Unsupported(Some(s"Unsupported generator:
${op.generator.nodeName}"))
- }
- op.generator.children.head.dataType match {
- case _: ArrayType =>
+ op.generator match {
+ // Spark only accepts an array of structs for `inline`, and `stack` is
planned as `inline`;
+ // see `stackAsStructArray`.
+ case _: Inline | _: Stack =>
Compatible()
- case _: MapType =>
- // TODO add support for map types
- // https://github.com/apache/datafusion-comet/issues/2837
- Unsupported(Some("Comet only supports explode/explode_outer for
arrays, not maps"))
+ case explode: ExplodeBase =>
+ explode.child.dataType match {
+ case _: ArrayType =>
+ Compatible()
+ case _: MapType =>
+ // TODO add support for map types
+ // https://github.com/apache/datafusion-comet/issues/2837
+ Unsupported(Some("Comet only supports explode/explode_outer for
arrays, not maps"))
+ case other =>
+ Unsupported(Some(s"Unsupported data type: $other"))
+ }
case other =>
- Unsupported(Some(s"Unsupported data type: $other"))
+ CometGeneratorUDF.declineReason(other) match {
+ case Some(reason) => Unsupported(Some(reason))
+ case None => Compatible()
+ }
}
}
override def convert(
op: GenerateExec,
builder: Operator.Builder,
childOp: OperatorOuterClass.Operator*):
Option[OperatorOuterClass.Operator] = {
- val childExpr = op.generator.children.head
- val childExprProto = exprToProto(childExpr, op.child.output)
+ val explodeBuilder =
OperatorOuterClass.Explode.newBuilder().setOuter(op.outer)
+ val childExprProto = op.generator match {
+ case inline: Inline =>
+ explodeBuilder.setInline(true)
+ exprToProto(inline.child, op.child.output)
+ case stack: Stack =>
+ explodeBuilder.setInline(true)
+ val array = stackAsStructArray(stack)
+ exprToProto(array, op.child.output).orElse {
+ // The array of structs is built here rather than taken from the
plan, so carry any
+ // reason recorded on the built nodes to the generator, where the
operator collects it.
+ liftFallbackReasons(array, op.generator)
+ None
+ }
+ case explode: ExplodeBase =>
+ explodeBuilder.setPosition(explode.position)
+ exprToProto(explode.child, op.child.output)
+ case generator =>
Review Comment:
`ReplicateRows.eval` builds `n` copies of its value row, which is
`inline(array_repeat(named_struct(...), n))`, and #7251 already lowers `stack`
to `inline` over a built array in `stackAsStructArray`. Would lowering
`ReplicateRows` the same way keep `EXCEPT ALL` and `INTERSECT ALL` fully
native, without the per-row JVM round trip? I have not checked that native
`array_repeat` matches Spark for struct elements, NULL and negative counts, and
the `Long` count.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala:
##########
@@ -0,0 +1,255 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet
+
+import java.nio.ByteBuffer
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+import org.apache.arrow.vector.{BitVector, ValueVector, VarBinaryVector}
+import org.apache.arrow.vector.types.pojo.Field
+import org.apache.spark.SparkEnv
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, AttributeSeq, BindReferences, Expression, Generator,
GenericInternalRow, Literal, Nondeterministic, Unevaluable}
+import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.comet.execution.arrow.ArrowWriter
+import org.apache.spark.sql.execution.ExecSubqueryExpression
+import org.apache.spark.sql.types.ArrayType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+import org.apache.comet.CometConf
+import org.apache.comet.codegen.CometBatchKernelCodegen
+import org.apache.comet.serde.{CometInvokeTargets, ExprOuterClass}
+import org.apache.comet.serde.CometScalaUDF.closureSerialize
+import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType}
+import org.apache.comet.udf.{CometUDF, NativeUdfCall}
+import org.apache.comet.udf.codegen.CometScalaUDFCodegen
+import org.apache.comet.vector.CometVector
+
+/**
+ * Runs a Spark [[Generator]] for Comet's native generate operator. For each
input row it returns
+ * the rows the generator emits as one `array<struct>` element, which the
native operator unnests
+ * and flattens the way it does for `inline`. An input row that emits nothing
gets an empty array,
+ * so `outer` keeps it as one row of nulls.
+ *
+ * Arg 0 is the digest of the closure-serialized bound generator, arg 1 the
serialized bytes, and
+ * arg 2 a boolean scalar that is true for the call the native operator makes
once its input is
+ * exhausted. That call returns one element holding the rows of
`Generator.terminate`. Args 3..N
+ * are the columns the bound generator reads, in ordinal order.
+ *
+ * A generator can keep state across rows and emit rows from `terminate`, as a
Hive UDTF does, so
+ * every native plan gets its own instance: the key is the digest and the
plan, and `releasePlan`
+ * drops the plan's instances. The serde ships the bound generator rather than
evaluating it on
+ * the driver, so this works in cluster mode without driver-side state.
`evaluate` is synchronized
+ * for the reason `CometScalaUDFCodegen.evaluate` is.
+ */
+class CometGeneratorUDF extends CometUDF {
+ import CometGeneratorUDF._
+
+ private val generators = mutable.HashMap.empty[(ByteBuffer, Long),
BoundGenerator]
+
+ override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector
=
+ evaluate(inputs, numRows, partitionIndex = 0, planId = NoPlan)
+
+ override def evaluate(
+ inputs: Array[ValueVector],
+ numRows: Int,
+ partitionIndex: Int,
+ planId: Long): ValueVector = this.synchronized {
+ require(
+ inputs.length >= 3,
+ "CometGeneratorUDF requires the generator digest, its serialized bytes
and the terminate " +
+ s"flag, got ${inputs.length} inputs")
+ val digest = ByteBuffer.wrap(binaryScalar(inputs(0)))
+ val bound = generators.getOrElseUpdate(
+ (digest, planId),
+ BoundGenerator(deserialize(binaryScalar(inputs(1))), partitionIndex))
+
+ val output = CometBatchKernelCodegen.allocateOutput(bound.outputField,
numRows, 0)
+ try {
+ val writer = ArrowWriter.createFieldWriter(output)
+ val outputRow = new GenericInternalRow(1)
+ def writeRows(rows: Any): Unit = {
+ val copies = new ArrayBuffer[Any]
+ iteratorOf(rows).foreach(row => copies += row.copy())
+ outputRow.update(0, new GenericArrayData(copies.toArray))
+ writer.write(outputRow, 0)
+ }
+
+ if (inputs(2).asInstanceOf[BitVector].get(0) == 1) {
+ writeRows(bound.generator.terminate())
+ } else {
+ val columns: Array[ColumnVector] =
+ inputs.drop(3).map(v => CometVector.getVector(v, null): ColumnVector)
+ val batch = new ColumnarBatch(columns, numRows)
+ var i = 0
+ while (i < numRows) {
+ writeRows(bound.generator.eval(batch.getRow(i)))
+ i += 1
+ }
+ }
+ writer.finish()
+ output
+ } catch {
+ case t: Throwable =>
+ output.close()
+ throw t
+ }
+ }
+
+ override def releasePlan(planId: Long): Unit = this.synchronized {
+ generators.keys.filter(_._2 == planId).toList.foreach(generators.remove)
+ }
+
+ private def deserialize(bytes: Array[Byte]): Generator = {
+ val loader = Option(Thread.currentThread().getContextClassLoader)
+ .getOrElse(classOf[Generator].getClassLoader)
+ SparkEnv.get.closureSerializer
+ .newInstance()
+ .deserialize[Generator](ByteBuffer.wrap(bytes), loader)
+ }
+}
+
+object CometGeneratorUDF {
+
+ /** The plan id of calls from outside a native plan, such as unit tests. */
+ private val NoPlan = -1L
+
+ /** Numbers each generator the serde ships; see [[calls]]. */
+ private val nextOccurrence = new AtomicLong()
+
+ /** Why `generator` cannot run through [[CometGeneratorUDF]], or `None` when
it can. */
+ def declineReason(generator: Generator): Option[String] = {
+ val name = generator.nodeName
+ if (!CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.get()) {
+ return Some(
+ s"Unsupported generator: $name. Set
${CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.key}" +
+ "=true to run it on the JVM inside Comet")
+ }
+ // The generator is serialized at plan time, before any subquery in it has
run, and an
+ // unevaluable or native-only node has nothing to run on the JVM.
Attributes are unevaluable
+ // only until they are bound, which `calls` does.
+ generator
+ .find {
+ case _: Attribute => false
+ case _: ExecSubqueryExpression | _: Unevaluable | _: NativeUdfCall =>
true
+ case _ => false
+ }
+ .map(node => s"$name: ${node.nodeName} cannot run inside a JVM
generator")
+ .orElse(CometInvokeTargets.declineReason(generator))
+ .orElse(Some(s"$name: unsupported output type
${outputType(generator).catalogString}")
+ .filter(_ => serializeDataType(outputType(generator)).isEmpty))
+ }
+
+ /**
+ * The `JvmScalarUdf` calls that run `generator` through
[[CometGeneratorUDF]]: one evaluated
+ * for each input batch, and the one the native operator makes once its
input is exhausted.
+ * Fails with the reason when `generator` cannot be shipped.
+ */
+ def calls(
+ generator: Generator,
+ inputs: Seq[Attribute]): Either[String, (ExprOuterClass.Expr,
ExprOuterClass.Expr)] = {
+ val name = generator.nodeName
+ // Bind against only the attributes the generator reads, which become the
data args.
+ val attrs = generator.collect { case a: AttributeReference => a }.distinct
+ val bound = BindReferences.bindReference(generator, AttributeSeq(attrs))
+ val bytes =
+ try {
+ closureSerialize(bound)
+ } catch {
+ case NonFatal(e) =>
+ return Left(
+ s"$name could not be closure-serialized
(${e.getClass.getSimpleName}: " +
+ s"${e.getMessage})")
+ }
+ // The digest also covers a number unique to this generator, so two
identical generators in
+ // one plan get their own instances and state.
+ val occurrence =
ByteBuffer.allocate(8).putLong(nextOccurrence.getAndIncrement()).array()
+ val digest = CometScalaUDFCodegen.digest(bytes ++ occurrence)
+ val returnType = serializeDataType(outputType(generator)).getOrElse {
+ return Left(s"$name: unsupported output type
${outputType(generator).catalogString}")
+ }
+
+ def call(
+ terminate: Boolean,
+ dataArgs: Seq[Expression]): Either[String, ExprOuterClass.Expr] = {
+ val args = (Seq(Literal(digest), Literal(bytes), Literal(terminate)) ++
dataArgs)
+ .map(exprToProto(_, inputs))
+ if (args.exists(_.isEmpty)) {
+ Left(s"$name: an argument could not be passed to the JVM generator")
+ } else {
+ val udf = ExprOuterClass.JvmScalarUdf
+ .newBuilder()
+ .setClassName(classOf[CometGeneratorUDF].getName)
+ .addAllArgs(args.flatten.asJava)
+ .setReturnType(returnType)
+ .setReturnNullable(false)
+ Right(ExprOuterClass.Expr.newBuilder().setJvmScalarUdf(udf).build())
+ }
+ }
+ for {
+ process <- call(terminate = false, attrs)
+ terminate <- call(terminate = true, Nil)
+ } yield (process, terminate)
+ }
+
+ /**
+ * A generator ready to run in one native plan, and the Arrow field its
output is exported as.
+ */
+ private case class BoundGenerator(generator: Generator, outputField: Field)
+
+ private object BoundGenerator {
+ def apply(generator: Generator, partitionIndex: Int): BoundGenerator = {
+ generator.foreach {
+ case n: Nondeterministic => n.initialize(partitionIndex)
Review Comment:
This branch cannot run. `Nondeterministic` defines `deterministic` as false,
so any such node makes `op.generator.deterministic` false, and
`CometExplodeExec.getSupportLevel` returns `Unsupported` for that before a
`CometGeneratorUDF` is built. Could `BoundGenerator` drop the companion
`apply`, so `partitionIndex` is no longer threaded into it?
--
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]
