mbutrovich commented on code in PR #5526: URL: https://github.com/apache/datafusion-comet/pull/5526#discussion_r4189962358
########## 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: What do you think about moving this sweep in with the fuzz suites? Only its producers go through the codegen dispatcher. Most of what it found was in native kernels, the shuffle writer and broadcast IPC. Its physical profiles also overlap the shuffle mode and native columnar-to-row matrix that [`CometFuzzTestBase`](https://github.com/apache/datafusion-comet/blob/0ac4dadae70838d8675bfd846c2d1c7617c25c78/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala#L113-L126) already runs every test under. Could it reuse that matrix and add only the settings it needs beyond it (two-row batches, the sort-based writer with spills, AQE)? The gate test in this suite also checks the same dispatch decisions as the `expect_dispatch` SQL fixtures this PR adds. Could those checks live only in the fixtures? ########## .github/workflows/pr_build_linux.yml: ########## @@ -619,6 +619,7 @@ jobs: org.apache.comet.serde.CometLiteralSuite org.apache.comet.serde.QueryPlanSerdeSuite org.apache.comet.CometFallbackInvarianceSuite + org.apache.comet.CometNullTypeCompositionSuite Review Comment: Should `CometNullTypeCompositionSuite` run on every PR? You measured about 8 minutes for it locally, and this adds it to the `expressions` job for every Spark profile in the Linux matrix, plus macOS in `pr_build_macos.yml`. The SQL file fixtures already cover each gated shape. Could the composition sweep run in the nightly `ci.yml` schedule instead? ########## native/core/src/execution/planner.rs: ########## @@ -6570,6 +6570,70 @@ mod tests { Ok(()) } + #[test] + fn test_literal_to_list_with_only_empty_children() -> Result<(), DataFusionError> { + // [[], null]: no child contributes values, so the values array is built empty. It must + // still have the child list's element type, as it does when some child is non-empty. + let data = ListLiteral { + list_values: vec![ListLiteral::default(), ListLiteral::default()], + null_mask: vec![true, false], + ..Default::default() + }; + let inner = DataType::List(Arc::new(Field::new("item", DataType::Int64, true))); + let nested_type = DataType::List(Arc::new(Field::new("item", inner.clone(), true))); + + let array = literal_to_array_ref(nested_type, data)?; + assert_eq!(array.data_type(), &inner); + let list = array.as_any().downcast_ref::<ListArray>().unwrap(); + assert_eq!(list.len(), 2); + assert!(list.value(0).is_empty()); + assert!(list.is_null(1)); + + // [[[]], [[[1]]]] with every declared field non-nullable: the first child has no values + // below its own level, the second has. Both must come out as the same type, or + // concatenating them fails. + let list_of = |dt: DataType| DataType::List(Arc::new(Field::new("item", dt, false))); + let one = ListLiteral { + int_values: vec![1], + null_mask: vec![true], + ..Default::default() + }; + let data = ListLiteral { + list_values: vec![ + ListLiteral { + list_values: vec![ListLiteral::default()], + null_mask: vec![true], + ..Default::default() + }, + ListLiteral { + list_values: vec![ListLiteral { + list_values: vec![one], + null_mask: vec![true], + ..Default::default() + }], + null_mask: vec![true], + ..Default::default() + }, + ], + null_mask: vec![true, true], + ..Default::default() + }; + let nested_type = list_of(list_of(list_of(list_of(DataType::Int32)))); + + let array = literal_to_array_ref(nested_type, data)?; + let list = array.as_any().downcast_ref::<ListArray>().unwrap(); + assert_eq!(list.len(), 2); + assert!(list + .value(0) + .as_any() + .downcast_ref::<ListArray>() + .unwrap() + .value(0) + .is_empty()); + assert_eq!(list.value(1).len(), 1); + Ok(()) Review Comment: This second case (`[[[]], [[[1]]]]` with every declared field non-nullable, over `Int32`) has the same input as [`test_literal_to_list_with_empty_nested_children_non_nullable`](https://github.com/apache/datafusion-comet/blob/0ac4dadae70838d8675bfd846c2d1c7617c25c78/native/core/src/execution/planner.rs#L6700-L6768), which #4715 added to `main` along with the fix. Could we drop it and keep only the `[[], null]` case above, which the tests on `main` don't cover? ########## 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: This deserializes each dispatcher kernel's closure bytes during planning to recover its expression, and `runsNullTypeKernel` walks every message field of the serialized operator to find the kernels. `CometScalaUDF.emitJvmCodegenDispatch` already holds the bound expression when it serializes it ([`CometScalaUDF.scala#L141`](https://github.com/apache/datafusion-comet/blob/90d02ebcebb898a4c6afef1a153116d9f7a29e64/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala#L141)). Could it record there whether the tree computes a `NullType` value, so this rule reads that instead of round-tripping through the closure serializer? The rule exists because of a bug on `main`: native evaluation computes the right operand of a binary expression even when the left one is NULL, so a stateful operand advances on rows Spark skips. Is there an issue for that? If it's the same problem as #6006, could the scaladoc reference that issue, so it's clear when this rule can be removed? If not, could you file one? -- 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]
