rich7420 commented on code in PR #7253:
URL: https://github.com/apache/datafusion-comet/pull/7253#discussion_r4390955004
##########
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:
Discard empty termination output before creating null placeholders;
nonnullable carried columns currently make empty partitions fail.
##########
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:
Snapshot rows with `UnsafeProjection`; `GenericInternalRow.copy()` retains
reused binary buffers, producing duplicate values instead of the original
outputs.
##########
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:
Keep custom generators on Spark until this bridge preserves laziness; `LIMIT
1` currently evaluates a throwing second row.
--
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]
