sunchao commented on code in PR #5526: URL: https://github.com/apache/datafusion-comet/pull/5526#discussion_r3923980673
########## spark/src/test/scala/org/apache/comet/CometNullTypeCompositionSuite.scala: ########## @@ -0,0 +1,543 @@ +/* + * 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.ArrayBuffer +import scala.util.{Failure, Success, Try} + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{Expression, JsonToStructs, RuntimeReplaceable, Sequence, StringToMap} +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.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.NullType + +import org.apache.comet.serde.{QueryPlanSerde, SupportLevel} + +/** + * 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 { + + /** + * Non-foldable `NullType`-bearing expressions over a live column. Each is a shape that the + * widened gate lets reach native execution; the constant-folded forms are literals and take a + * different path, so every producer here references `id`. + */ + private val arrayOfNull = Seq( + "transform(array(id), x -> NULL)", + "filter(array(CAST(NULL AS int)), x -> id IS NOT NULL)") + + private val mapWithNullValue = Seq("map(id, NULL)") + + private val mapWithNullKey = Seq("transform_values(map(), (k, v) -> id)") + + private val arrayOfStructWithNull = Seq("map_entries(map(id, NULL))") + + private val structWithNull = Seq("named_struct('a', id, 'b', NULL)") + + private val scalarNull = Seq("aggregate(array(id), NULL, (acc, x) -> NULL)") + + /** Consumers valid for a value of any type, written as templates over the producer `%s`. */ + private val anyTypeConsumers = Seq( + "to_json(struct(%s AS c))", + "to_csv(struct(%s AS c))", + "hash(%s)", + "xxhash64(%s)", + "CAST(%s AS string)", + "CASE WHEN id > 2 THEN %s END", + "IF(id > 2, %s, NULL)", + "coalesce(%s, %s)", + "%s IS NULL", + "%s = %s", + "%s <=> %s") + + /** Consumers valid for any `array<T>`. */ + private val arrayConsumers = Seq( + "size(%s)", + "reverse(%s)", + "array_distinct(%s)", + "sort_array(%s)", + "array_sort(%s)", + "element_at(%s, 1)", + "slice(%s, 1, 1)", + "array_repeat(%s, 2)", + "array_union(%s, %s)", + "array_except(%s, %s)", + "array_intersect(%s, %s)", + // Set ops only diverge when a Null-typed side meets a side that actually holds entries, so + // the same-producer templates above cannot reach that branch on their own. + "array_union(%s, array(1))", + "array_union(array(1), %s)", + "array_except(%s, array(1))", + "array_intersect(%s, array(1))", + "arrays_overlap(%s, %s)", + "arrays_zip(%s, %s)", + "concat(%s, %s)", + "flatten(array(%s))", + "array_position(%s, NULL)", + "array_contains(%s, NULL)", + // Spark rejects a NullType needle, so the typed needle is what reaches the serde. + "array_contains(%s, %s[0])", + "array_remove(%s, NULL)", + "array_append(%s, NULL)", + "array_insert(%s, 1, NULL)", + "exists(%s, x -> x IS NULL)", + "forall(%s, x -> x IS NULL)", + "filter(%s, x -> x IS NULL)", + "transform(%s, x -> x)", + "aggregate(%s, 0, (acc, x) -> acc)", + "zip_with(%s, %s, (x, y) -> x)", + "map_from_entries(arrays_zip(%s, %s))", + "%s[0]", + "array_compact(%s)", + "array_max(%s)", + "array_min(%s)", + "array_join(%s, ',')", + // shuffle is non-deterministic, so only the sorted result compares. + "sort_array(shuffle(%s))", + "map_from_arrays(%s, array(id))") ++ anyTypeConsumers + + /** Consumers valid for any `map<K, V>`. */ + private val mapConsumers = Seq( + "size(%s)", + "map_keys(%s)", + "map_values(%s)", + "map_entries(%s)", + "map_concat(%s, %s)", + "map_filter(%s, (k, v) -> k IS NOT NULL)", + "transform_values(%s, (k, v) -> v)", + "transform_keys(%s, (k, v) -> k)", + "map_zip_with(%s, %s, (k, v1, v2) -> v1)", + "map_from_entries(map_entries(%s))", + "element_at(%s, id)", + "%s[id]") ++ anyTypeConsumers + + private val structConsumers = Seq( + "%s.a", + "%s.b", + "struct(%s)", + "array(%s)", + "array(%s).b", + "size(array(%s))") ++ anyTypeConsumers + + private val scalarConsumers = Seq( + "array(%s)", + "map(id, %s)", + "named_struct('a', %s)", + "coalesce(%s, NULL)", + "abs(%s)", + // Spark's analyzer rejects a void scalar in an array position (no implicit cast), so these + // stay skipped; they are here so the sweep notices if that ever changes. + "array_union(%s, array(1))", + "array_distinct(%s)", + "array_contains(%s, 1)", + "CAST(%s AS array<int>)") ++ anyTypeConsumers + + private def cases(wrap: String => String = identity): Seq[(String, String)] = + Seq( + arrayOfNull -> arrayConsumers, + arrayOfStructWithNull -> arrayConsumers, + mapWithNullValue -> mapConsumers, + mapWithNullKey -> mapConsumers, + structWithNull -> structConsumers, + scalarNull -> scalarConsumers).flatMap { case (producers, consumers) => + for (p <- producers; c <- consumers) + yield { + val wrapped = wrap(p) + (wrapped, substitute(c, wrapped, p)) + } + } + + /** + * Fills a consumer template, putting `first` in the leading placeholder and `rest` in any + * others. Wrapping only the first argument keeps a stateful producer to a single occurrence: + * two of them desynchronize inside Spark itself, whereas the first child is one that every + * consumer evaluates, so a divergence there belongs to Comet. + */ + private def substitute(template: String, first: String, rest: String): String = { + val at = template.indexOf("%s") + template.substring(0, at) + first + template.substring(at + 2).replace("%s", rest) + } + + /** + * Makes a producer nullable and non-deterministic, the only input that can tell apart the two + * copies a serde null guard serializes (see `NullGuard.doubleEvaluationReason`). + * + * Applied to the first argument only, so multi-argument consumers stay in the sweep: a serde + * that null-guards every child, as `CometArraysZip` does, needs just one stateful argument to + * diverge, and restricting the sweep to single-placeholder consumers hid exactly that case. + */ + private def nullableNondeterministic(producer: String): String = + s"IF(monotonically_increasing_id() % 2 = 0, $producer, NULL)" + + /** + * Query templates that put a producer under a different physical operator. The consumer sweep + * holds the operator fixed at a projection, so it never sees the paths that serialize and + * re-read a value rather than compute over it. + */ + private val operatorTemplates = Seq( + "project" -> "SELECT %s AS c FROM t", + "filter" -> "SELECT %s AS c FROM t WHERE id > 2", + "sort-by-id" -> "SELECT %s AS c FROM t ORDER BY id", + "sort-by-value" -> "SELECT %s AS c FROM t ORDER BY c", + "limit" -> "SELECT %s AS c FROM t LIMIT 3", + "take-ordered" -> "SELECT %s AS c FROM t ORDER BY id LIMIT 3", + "groupby-key" -> "SELECT %s AS c, count(*) FROM t GROUP BY c", + "groupby-value" -> "SELECT id, first(%s) AS c FROM t GROUP BY id", + "groupby-last" -> "SELECT id, last(%s) AS c FROM t GROUP BY id", + "collect-list" -> "SELECT collect_list(%s) AS c FROM t", + "collect-set" -> "SELECT collect_set(%s) AS c FROM t", + "max" -> "SELECT max(%s) AS c FROM t", + "min" -> "SELECT min(%s) AS c FROM t", + "count-distinct" -> "SELECT count(DISTINCT %s) AS c FROM t", + "distinct" -> "SELECT DISTINCT %s AS c FROM t", + "union-all" -> "SELECT %s AS c FROM t UNION ALL SELECT %s AS c FROM t", + "union-distinct" -> "SELECT %s AS c FROM t UNION SELECT %s AS c FROM t", + "window-order" -> "SELECT %s AS c, row_number() OVER (ORDER BY id) AS r FROM t", + "window-partition" -> "SELECT id, count(*) OVER (PARTITION BY %s) AS n FROM t", + // The producer is computed on the build side, so the value itself crosses the exchange + // (shuffle or broadcast) rather than being projected after the join. + "join-shuffle" -> + "SELECT a.id, b.c FROM t a JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id = b.bid", + "join-broadcast" -> + ("SELECT /*+ BROADCAST(b) */ a.id, b.c FROM t a " + + "JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id = b.bid"), + "join-nested-loop" -> + ("SELECT /*+ BROADCAST(b) */ a.id, b.c FROM t a " + + "JOIN (SELECT id AS bid, %s AS c FROM t) b ON a.id > b.bid"), + "expand-cube" -> "SELECT id, count(%s) AS n FROM t GROUP BY CUBE(id)", + "explode" -> "SELECT explode(%s) AS c FROM t", + "repartition" -> "SELECT /*+ REPARTITION(3) */ %s AS c FROM t", + "coalesce-partitions" -> "SELECT /*+ COALESCE(1) */ %s AS c FROM t", + "subquery" -> "SELECT id FROM t WHERE id IN (SELECT id FROM t WHERE %s IS NOT NULL)", + "nested-project" -> "SELECT c FROM (SELECT %s AS c, id FROM t ORDER BY id) x") + + /** + * The subset of `operatorTemplates` that serializes the value rather than only computing over + * it. + */ + private val serializingOperators = { + val names = Set( + "project", + "sort-by-id", + "groupby-value", + "collect-list", + "collect-set", + "repartition", + "union-all", + "join-shuffle") + require(names.subsetOf(operatorTemplates.map(_._1).toSet), "unknown operator name") + operatorTemplates.filter { case (name, _) => names(name) } + } + + private val arrayOfStructWithNullField = Seq("array(named_struct('a', id, 'b', NULL))") + + private def allProducers: Seq[String] = + arrayOfNull ++ mapWithNullValue ++ mapWithNullKey ++ arrayOfStructWithNull ++ + structWithNull ++ scalarNull ++ arrayOfStructWithNullField + + /** + * Containers to wrap a producer in before putting it under a serializing operator. The nested + * nullability mismatches only show where a container nests the value and an operator re-reads + * the nesting: `collect_list(array(map(k, NULL)))` fails while both `array(map(k, NULL))` and + * `collect_list(map(k, NULL))` pass. + */ + private val nestingWrappers = Seq( + "array(%s)", + "array_repeat(%s, 2)", + "element_at(%s, 1)", + "slice(%s, 1, 1)", + "map(id, %s)", + "named_struct('s', %s)") + + /** + * The one native gap the operator sweeps tolerate: Comet's row shuffle writer has no case for a + * `Null` struct field and panics whenever a struct holding one is shuffled. It reproduces on + * main and belongs to the native writer. Matching on the panic's signature rather than on a + * list of queries keeps every other failure red. + */ + private val knownRowShuffleGap = "Unsupported data type of struct field: Null" + + private def rowsOf( + query: String, + cometEnabled: Boolean, + ansi: Boolean, + nativeColumnarToRow: Boolean): (Seq[String], Boolean) = { + val confs = Seq( + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> nativeColumnarToRow.toString, + SQLConf.ANSI_ENABLED.key -> ansi.toString, + // Keep the producers out of the optimizer's hands so they are evaluated per row by the + // engine under test rather than folded to a literal at plan time. + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") + // withSQLConf's body is typed `=> Unit` on the older supported Spark versions and generic + // only on the newer ones, so the result is captured out of band to stay portable. + var result: (Seq[String], Boolean) = (Seq.empty, false) + withSQLConf(confs: _*) { + val df = spark.sql(query) + val rows = df.collect().map(_.toString()).sorted.toSeq + // The sweep's queries put the NullType expression in a projection, so a native + // CometProjectExec is what distinguishes a case that actually exercised a native kernel + // from one that merely fell back. Without this, a sweep where everything falls back would + // still be green while proving nothing. + val nativeProject = + collectFirst(df.queryExecution.executedPlan) { case _: CometProjectExec => + () + }.isDefined + result = (rows, nativeProject) + } + result + } + + // ANSI is a dimension of the consumer sweeps because it changes which serdes wrap their child + // in a null guard and which kernels raise on out-of-range access; the operators below carry + // no ANSI semantics of their own. + for (ansi <- Seq(false, true)) { + test(s"NullType producers survive every consumer that Spark accepts (ansi=$ansi)") { + sweep( + "consumer", + cases().map { case (producer, expr) => (producer, s"SELECT $expr FROM t") }, + comparedFloor = 120, + nativeFloor = 100, + ansi = ansi) + } + + test(s"nullable non-deterministic NullType producers survive every consumer (ansi=$ansi)") { + sweep( + "non-deterministic", + cases(nullableNondeterministic).map { case (producer, expr) => + (producer, s"SELECT $expr FROM t") + }, + comparedFloor = 140, + nativeFloor = 110, + ansi = ansi) + } + } + + // Native columnar-to-row is off by default, so the operator sweep, whose queries all end in a + // collect, is the one that runs both conversions over every NullType shape. + for (nativeColumnarToRow <- Seq(false, true)) { + test( + "NullType producers survive every operator that Spark accepts " + + s"(nativeColumnarToRow=$nativeColumnarToRow)") { + sweep( + "operator", + for (producer <- allProducers; (op, template) <- operatorTemplates) + yield (op, template.replace("%s", producer)), + comparedFloor = 160, + nativeFloor = 100, + tolerated = Some(knownRowShuffleGap), + nativeColumnarToRow = nativeColumnarToRow) + } + } + + /** + * Registered array, map and struct serdes whose nested-typed argument no `NullType` producer + * can occupy, so no consumer template can reach them: `Sequence` takes integral or temporal + * bounds, and `StringToMap` and `JsonToStructs` take a string. + */ + private val serdesWithoutNestedInput: Set[Class[_ <: Expression]] = + Set(classOf[Sequence], classOf[StringToMap], classOf[JsonToStructs]) + + /** Aggregates whose serde accepts any input type, so a `NullType` shape can reach them. */ + private val anyTypeAggregates: Set[Class[_]] = Set( + classOf[CollectList], + classOf[CollectSet], + classOf[Count], + classOf[First], + classOf[Last], + classOf[Max], + classOf[Min]) + + /** + * Every expression class in the analyzed and optimized plans of `query`, or none if Spark + * rejects it. Both plans, because the optimizer is what inserts some expressions (`MapSort` + * under a map grouping key) and what expands `RuntimeReplaceable` ones. + */ + private def expressionClasses(query: String): Set[Class[_]] = + Try { + val qe = spark.sql(query).queryExecution + Seq(qe.analyzed, qe.optimizedPlan) + }.toOption.toSeq.flatten + .toSet[LogicalPlan] + .flatMap { plan => + plan.flatMap(_.expressions).flatMap { root => + root.collect { case e: Expression => e }.flatMap { + case r: RuntimeReplaceable => r +: r.replacement.collect { case e: Expression => e } + case e => Seq(e) + } + } + } + .map(_.getClass) + + // The consumer and operator lists are written by hand, so this is the check that a serde added + // to the registry later, or one forgotten now, does not silently stay outside the sweep. + test("the sweep reaches every registered array, map, struct and any-type aggregate serde") { + withTempView("t") { + spark.range(0, 8).createOrReplaceTempView("t") + val reached = + (cases().map { case (_, expr) => + s"SELECT $expr FROM t" + } ++ + (for (producer <- allProducers; (_, template) <- operatorTemplates) + yield template.replace("%s", producer))) + .flatMap(expressionClasses) + .toSet + val registered: Set[Class[_]] = + (QueryPlanSerde.arrayExpressions.keySet ++ + QueryPlanSerde.mapExpressions.keySet ++ + QueryPlanSerde.structExpressions.keySet).toSet[Class[_]] -- + serdesWithoutNestedInput ++ anyTypeAggregates + val missing = registered -- reached + assert( + missing.isEmpty, + s"registered serdes no sweep template reaches: ${missing + .map(_.getSimpleName) + .toSeq + .sorted + .mkString(", ")}") + } + } + + test("nested NullType values survive the operators that serialize them") { + // Select producers by the type they actually have: `filter(array(CAST(NULL AS int)), ...)` + // is `array<int>`, and nesting it only exercises a pre-existing nested-container limitation. + val nullBearingProducers = allProducers.filter { p => + Try(spark.range(0, 8).selectExpr(s"$p AS c").schema.head.dataType).toOption + .exists(SupportLevel.containsType(_, classOf[NullType])) + } + assert( + nullBearingProducers.size >= 5, + s"only ${nullBearingProducers.size} producers still carry a NullType; the producer list " + + "has drifted away from what this sweep is meant to cover") + val nested = + for (producer <- nullBearingProducers; wrapper <- nestingWrappers) + yield wrapper.replace("%s", producer) + sweep( + "nesting", + for (value <- nested.distinct; (op, template) <- serializingOperators) + yield (op, template.replace("%s", value)), + comparedFloor = 200, + nativeFloor = 100, + tolerated = Some(knownRowShuffleGap)) + } + + /** + * Runs every `(label, query)` twice, Comet off and on, and reports all divergences together. A + * query Spark itself rejects is skipped; one whose Comet arm throws or disagrees is a failure, + * unless its cause matches `tolerated`. + * + * The floors are floors rather than equalities because which cases Spark accepts varies across + * the supported Spark versions. `comparedFloor` fails a sweep whose templates have gone stale; + * `nativeFloor` fails one where Comet fell back almost everywhere, since falling back is a pass + * and such a sweep would prove nothing about the native kernels. + */ + private def sweep( + name: String, + queries: Seq[(String, String)], + comparedFloor: Int, + nativeFloor: Int, + tolerated: Option[String] = None, + ansi: Boolean = false, + nativeColumnarToRow: Boolean = false): Unit = { + var native = 0 + withTempPath { dir => + // One partition, so every batch holds several rows: a null guard over a non-deterministic + // child only diverges where the CASE sees both matching and non-matching rows in one batch. + spark.range(0, 16, 1, 1).write.parquet(dir.getAbsolutePath) + withTempView("t") { + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t") + + val failures = ArrayBuffer.empty[String] + var compared = 0 + var skipped = 0 + var knownGap = 0 + + for ((label, query) <- queries) { + Try(rowsOf(query, cometEnabled = false, ansi, nativeColumnarToRow)).toOption match { + case None => + skipped += 1 + case Some((sparkRows, _)) => + compared += 1 + Try(rowsOf(query, cometEnabled = true, ansi, nativeColumnarToRow)) match { + case Failure(e) if tolerated.exists(causeText(e).contains) => Review Comment: You're right. Correction to my [P2] example: at `79fc8400`, the array-of-struct path reaches [`append_field`](https://github.com/apache/datafusion-comet/blob/79fc84008a7dca16a4c24f28036649c0af2603b0/native/shuffle/src/spark_unsafe/row.rs#L272-L275), which already has a `DataType::Null` arm. My assertion that this path reaches the unsupported-struct-field panic was incorrect, and I withdraw that specific witness. Thanks for checking it. I am reviewing the two field-major fixes and removal of the sweep waiver separately at `496c3f7f`. This correction is based on the pinned source, not an executed SQL reproduction. -- 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]
