grorge123 commented on code in PR #5526:
URL: https://github.com/apache/datafusion-comet/pull/5526#discussion_r4195997002
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -73,6 +82,167 @@ object CometExecRule {
scan.relation.cacheBuilder.serializer.isInstanceOf[ArrowCachedBatchSerializer]
&&
ArrowCachedBatchSerializer.supportsSchema(scan.relation.output)
+ val nondeterministicBesideNullTypeReason: String =
+ "a non-deterministic expression that does not run entirely in the codegen
dispatcher, in " +
+ "an operator whose native plan also runs a dispatcher kernel with a
NullType result: " +
+ "native evaluation computes every argument over the whole batch and can
evaluate an " +
+ "argument twice, while Spark skips arguments that other arguments
decide, so a stateful " +
+ "expression advances on rows Spark does not evaluate"
+
+ /**
+ * True when `op` evaluates a non-deterministic expression outside the
codegen dispatcher, and
+ * `op`'s native plan runs a dispatcher kernel whose result type contains
NullType.
+ *
+ * Spark's generated code skips an argument when another one decides the
result (a NULL operand,
+ * a NULL divisor), while native evaluation computes every argument over the
whole batch, and
+ * some serdes serialize an argument twice (`ln`'s domain check), so a
stateful expression
+ * (`monotonically_increasing_id`, `rand`) advances on rows Spark does not
evaluate. That
+ * happens on main for any native operator. But the dispatcher used to
refuse NullType results,
+ * and an operator with one fell back to Spark together with the native
operators above it up to
+ * a shuffle. Every operator this check refuses therefore ran in Spark on
main, so refusing all
+ * of them costs nothing against main and does not depend on modelling each
expression's
+ * evaluation order. A top-level expression the dispatcher runs as a whole
is exempt: its kernel
+ * evaluates Spark's own generated code, with its own state per occurrence.
The exemption reads
+ * the dispatcher's tag on `op`'s expressions, which a rebuilt tree (a
decimal expression
+ * wrapped in `CheckOverflow`) does not carry; that only makes the operator
fall back.
+ *
+ * The kernels are read from the serialized plans rather than from `op`'s
expressions, since
+ * serialization can dispatch a tree it built itself (a folded map literal
expanded into
+ * `CreateMap`, a decimal addition wrapped in `CheckOverflow`).
+ */
+ private[rules] def nondeterministicBesideNullTypeKernel(
+ op: SparkPlan,
+ nativeOp: OperatorOuterClass.Operator): Boolean =
+ op.expressions.exists(evaluatesNondeterministicNatively) &&
+ (runsNullTypeKernel(nativeOp) || projectionRunsNullTypeKernel(op) ||
+ op.children.exists(nullTypeKernelBelow))
+
+ private def evaluatesNondeterministicNatively(expr: Expression): Boolean =
expr match {
+ case e if !e.exists(isStateful) => false
+ case e if e.getTagValue(CometExplainInfo.DISPATCHED_SELF).isDefined =>
false
+ case alias: Alias => evaluatesNondeterministicNatively(alias.child)
+ case _ => true
+ }
+
+ /**
+ * A node whose value depends on how many rows were evaluated before it: a
Catalyst
+ * `Nondeterministic` node or a user function that keeps state, whatever its
arguments. The
+ * partition id and the input file only depend on the partition, however
many rows Spark skips.
+ */
+ private def isStateful(e: Expression): Boolean = e match {
+ case _: SparkPartitionID | _: InputFileName | _: InputFileBlockStart |
+ _: InputFileBlockLength =>
+ false
+ case _: Nondeterministic => true
+ case other => DispatchOccurrence.holdsSharedState(other)
+ }
+
+ /**
+ * `TakeOrderedAndProjectExec` serializes as a sink, a leaf scan, and builds
its native sort and
+ * projection only when it executes, so those expressions are serialized
again here.
+ */
+ private def projectionRunsNullTypeKernel(op: SparkPlan): Boolean = op match {
+ case top: TakeOrderedAndProjectExec =>
+ deferredRunsNullTypeKernel(top.sortOrder ++ top.projectList,
top.child.output)
+ case _ => false
+ }
+
+ private def deferredRunsNullTypeKernel(
+ exprs: Seq[Expression],
+ inputs: Seq[Attribute]): Boolean =
+ exprs.exists(e => QueryPlanSerde.exprToProto(e,
inputs).exists(runsNullTypeKernel))
+
+ /**
+ * The native plans below `op` up to a shuffle. A sink is a leaf scan of the
serialized plan
+ * above it, and a broadcast, union or coalesce sink exists only when the
plan below it is
+ * native, so on main a NullType kernel there kept the operators above it in
Spark as well; a
+ * shuffle took a Spark child, so only its own partitioning expressions
count.
+ */
+ private def nullTypeKernelBelow(plan: SparkPlan): Boolean = plan match {
+ // A Comet shuffle's partitioning expressions are serialized when it runs.
Main refused a
+ // shuffle keyed on a NullType kernel, and then the operators above it ran
in Spark too.
+ case shuffle: CometShuffleExchangeExec =>
+ shuffle.outputPartitioning match {
+ case partitioning: Expression =>
+ deferredRunsNullTypeKernel(partitioning.children,
shuffle.child.output)
+ case _ => false
+ }
+ case stage: ShuffleQueryStageExec => nullTypeKernelBelow(stage.plan)
+ case read: AQEShuffleReadExec => nullTypeKernelBelow(read.child)
+ case _: ShuffleExchangeLike => false
+ case native: CometNativeExec =>
+ runsNullTypeKernel(native.nativeOp) ||
native.children.exists(nullTypeKernelBelow)
+ case stage: BroadcastQueryStageExec => nullTypeKernelBelow(stage.plan)
+ case reused: ReusedExchangeExec => nullTypeKernelBelow(reused.child)
+ case top: CometTakeOrderedAndProjectExec =>
+ deferredRunsNullTypeKernel(top.sortOrder ++ top.projectList,
top.child.output) ||
+ nullTypeKernelBelow(top.child)
+ case comet: CometPlan => comet.children.exists(nullTypeKernelBelow)
+ case _ => false
+ }
+
+ private val dispatcherClassName = classOf[CometScalaUDFCodegen].getName
+
+ /**
+ * Walks a serialized plan or expression for a dispatcher kernel that
computes a NullType value:
+ * as its result, or anywhere in the tree it runs (a whole `coalesce`
dispatched with an INT
+ * result can hold a `transform` with an `array<void>` result, which main
refused).
+ */
+ private def runsNullTypeKernel(message: Message): Boolean = message match {
+ case udf: ExprOuterClass.JvmScalarUdf
+ if udf.getClassName == dispatcherClassName &&
+ (protoContainsNullType(udf.getReturnType) ||
kernelComputesNullType(udf)) =>
+ true
+ case _ =>
+ message.getAllFields.asScala.exists {
+ case (field, value) if field.getJavaType ==
FieldDescriptor.JavaType.MESSAGE =>
+ value match {
+ case values: java.util.List[_] =>
+ values.asScala.exists(v =>
runsNullTypeKernel(v.asInstanceOf[Message]))
+ case child: Message => runsNullTypeKernel(child)
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+
+ /**
+ * The kernel's first argument is its closure-serialized expression (see
+ * `CometScalaUDF.emitJvmCodegenDispatch`). One that cannot be read counts
as computing
+ * NullType.
+ */
+ private def kernelComputesNullType(udf: ExprOuterClass.JvmScalarUdf):
Boolean =
+ udf.getArgsCount == 0 || !udf.getArgs(0).hasLiteral || {
+ val bytes = udf.getArgs(0).getLiteral.getBytesVal.toByteArray
+ try {
+ val loader = Option(Thread.currentThread().getContextClassLoader)
+ .getOrElse(classOf[Expression].getClassLoader)
+ SparkEnv.get.closureSerializer
+ .newInstance()
+ .deserialize[Expression](ByteBuffer.wrap(bytes), loader)
+ .exists(e => containsNullType(e.dataType))
+ } catch {
+ case NonFatal(_) => true
+ }
+ }
+
+ private def containsNullType(dt: DataType): Boolean = dt match {
Review Comment:
Good idea. I'll record whether a kernel computes a `NullType` value while
`emitJvmCodegenDispatch` serializes it, and carry that up to the operator, so
the rule no longer deserializes closures or walks every message field. It still
has to see expressions the serializer creates, TakeOrderedAndProject's deferred
sort keys and projection, and kernels below broadcast, union and coalesce sinks
and in shuffle keys, so I'll keep those covered rather than narrow the gate.
On the issue: it's related to #6006 but not the same bug. #6006 tracks
expressions that raise an error on rows Spark skips. Here a stateful operand
advances on those rows and the answer is wrong without any error, e.g. `IF(id =
0, CAST(NULL AS BIGINT), id) + monotonically_increasing_id()`. I'll confirm
that on `main`, file it as its own issue linked to #6006, and reference it from
the scaladoc. The rule only shields the shapes this PR newly admits; it doesn't
fix the typed case.
##########
spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala:
##########
@@ -0,0 +1,919 @@
+/*
+ * 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.comet
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+import scala.util.{Failure, Success, Try}
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayExcept,
ArrayIntersect, ArrayRepeat, ArraysZip, ArrayUnion, Coalesce, CreateArray,
Expression, JsonToStructs, MapFromArrays, Murmur3Hash, RuntimeReplaceable,
Sequence, Size, StringToMap, XxHash64}
+import org.apache.spark.sql.catalyst.expressions.aggregate._
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.comet.CometProjectExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
AdaptiveSparkPlanHelper}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.NullType
+
+import org.apache.comet.serde.{CodegenDispatchFallback,
CometAggregateExpressionSerde, CometExpressionSerde, Compatible,
QueryPlanSerde, SupportLevel, Unsupported}
+
+/**
+ * Cross-product sweep of the `NullType` shapes the JVM codegen dispatcher
admits against the
+ * expressions that consume them.
+ *
+ * Three sweeps share one driver, `sweep`: producers under consumers,
producers under operators,
+ * and producers nested in a container and then put under a serializing
operator.
+ */
+class CometNullTypeCompositionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
Review Comment:
Makes sense. The suite is out of this PR as of 267b5ce. In the follow-up
I'll move the sweep into the fuzz suites, reuse `CometFuzzTestBase`'s matrix
and add only the two-row batches, the spilling sort-based writer and AQE.
Dispatch decisions will stay in the `expect_dispatch` fixtures, and the checks
that each serde publishes its reasons will move to focused serde tests, since
the fixtures don't cover those.
--
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]
