andygrove commented on code in PR #7253:
URL: https://github.com/apache/datafusion-comet/pull/7253#discussion_r4456345830
##########
native/operators/src/explode.rs:
##########
@@ -455,13 +493,59 @@ impl ExplodeStream {
if other.is_none() {
let input_schema = self.input.schema();
self.input =
Box::pin(EmptyRecordBatchStream::new(input_schema));
+ if let Some(terminate) = self.terminate.take() {
+ let elapsed_compute =
self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let result = self.terminate_batch(&terminate);
+ timer.done();
+ if let Some(batch) = result.transpose() {
+ if let Ok(batch) = &batch {
+
batch.record_output(&self.baseline_metrics);
+ }
+ return Poll::Ready(Some(batch));
+ }
+ }
}
return Poll::Ready(other);
}
}
}
}
+ /// Unnest the rows `terminate` yields in place of the last input row's
list. With no input
+ /// rows, the other columns are null. Spark's `outer` handling does not
apply to these rows, so
+ /// an empty result adds nothing.
+ fn terminate_batch(
+ &mut self,
+ terminate: &Arc<dyn PhysicalExpr>,
+ ) -> Result<Option<RecordBatch>> {
+ let options = RecordBatchOptions::new().with_row_count(Some(1));
+ let no_columns =
+ RecordBatch::try_new_with_options(Arc::new(Schema::empty()),
vec![], &options)?;
+ let rows = terminate.evaluate(&no_columns)?.into_array(1)?;
+
+ let input_schema = self.input.schema();
+ let mut columns = match self.last_row.take() {
+ Some(row) => row.columns().to_vec(),
+ None => input_schema
Review Comment:
Fixed in 4c335a8a5. `terminate_batch` now returns early when `terminate`
yields a NULL or empty list, before it builds the placeholder row. When it does
yield rows over an empty input, the planner marks the carried columns nullable
(only when `terminate` is set) and the placeholder batch uses a nullable copy
of the input schema, so the null row is valid. Tests: two Rust cases over a
non-nullable `id` (empty `terminate` adds nothing, non-empty emits a null
`id`), and `terminate with no rows over a non-nullable carried column adds
nothing` in `CometGenerateExecSuite`. I confirmed that Scala test fails without
the native change with `Column 'col_0' is declared as non-nullable but contains
null values`.
##########
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())
Review Comment:
Fixed in 4c335a8a5. Each row is now snapshotted with an `UnsafeProjection`
over `elementSchema` and `.copy()`, so the `UnsafeRow` owns its bytes. Added
`ReusedBufferGenerator` (`CometTestGenerators`), which returns one
`GenericInternalRow` over one byte array it overwrites between rows. `generator
rows that share a reused buffer are kept apart` compares against Spark, for
both the plain and `LATERAL VIEW OUTER` forms.
##########
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)
Review Comment:
Not changed in 4c335a8a5. The bridge cannot be lazy per row: the native
operator hands the UDF a whole batch, so every row of the batch is evaluated,
and the rows one input row generates are held together, before anything is
returned. I documented this in the Generators section of `scala_java_udfs.md`,
and `spark.comet.exec.jvmGenerator.enabled=false` keeps such generators on
Spark. A regression test is not possible while the behavior stays. Whether the
default should be off, or custom generators and Hive UDTFs should stay on
Spark, is still open. I would like your view, since it decides whether this PR
keeps its main benefit. `ReplicateRows` no longer goes through this path (see
the thread on `operators.scala`).
##########
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:
Done in 4c335a8a5. `Hive UDTF runs natively and emits its terminate rows`
now runs with `COMET_BATCH_SIZE=2` over its four rows, so the running count and
the `close` total cross a batch boundary. The per-partition test uses a batch
size of 1.
##########
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:
Done in 4c335a8a5. `BoundGenerator` no longer has the companion `apply` and
`partitionIndex` is no longer threaded into it. A nondeterministic generator is
declined by `getSupportLevel`, so `initialize` was never reachable. I left a
comment where `initialize` used to be called.
##########
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:
Done in 4c335a8a5. Added `expressions/json/json_tuple.sql` (valid, missing,
repeated and null fields, invalid and empty documents, `LATERAL VIEW` with and
without `OUTER`, filter and aggregate above) and `operators/except_all.sql`.
The latter runs `EXCEPT ALL` and `INTERSECT ALL` over `DECIMAL`, `BINARY`,
`FLOAT` and `DOUBLE` (NaN, `-0.0`, infinity), `DATE`, `TIMESTAMP`,
`TIMESTAMP_NTZ`, array, struct and string, each with NULLs, one type at a time
and all columns at once. Both use plain `query` blocks and pass with native
placement checked. The Hive UDTF cases stay in Scala.
--
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]
