From 734600e2c9c64f23bd6a94c6047ed0ea2ebad3a1 Mon Sep 17 00:00:00 2001 From: Hung Date: Fri, 11 Sep 2026 00:35:52 +0800 Subject: [PATCH 1/7] feat: route wide-decimal hash and non-literal sha2 through the codegen dispatcher (#5581) Keep enclosing projections native when hash/xxhash64 see decimal precision > 18, or sha2 gets a non-foldable numBits, by running Spark's own doGenCode instead of falling the operator back to Spark. --- .../expression-audits/hash_funcs.md | 8 +- docs/source/user-guide/latest/expressions.md | 10 +- .../scala/org/apache/comet/serde/hash.scala | 11 +- .../comet/CometHashExpressionSuite.scala | 102 +++++++++++++++--- 4 files changed, 102 insertions(+), 29 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index 35337b8fe86..b78c7eca2a6 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -34,7 +34,7 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `Murmur3Hash(children, seed) extends HashExpression[Int]`; produces a Murmur3 hash with a configurable Int seed and `IntegerType` result. Comet routes via `CometMurmur3Hash` to the native `murmur3_hash` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; some inner helper refactors only. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Known limitation: `DecimalType` children with precision > 18 fall back because Spark hashes them through Java `BigDecimal`; `TimeType` (Spark 4.0+) is also unsupported. The same limitations apply to `xxhash64`, `sha1`, `sha2` through the shared `HashUtils`. +- Known limitation: the native kernel does not hash `DecimalType` children with precision > 18 (Spark hashes them through Java `BigDecimal`), including when nested in array, struct, or map. With the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `HashExpression.doGenCode` inside the Comet pipeline so the enclosing operator stays native. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. `TimeType` (Spark 4.0+) is still unsupported by the native kernel and is not routed through the dispatcher here. The same `HashUtils` type walk applies to `xxhash64`, `sha1`, and `sha2`. ## md5 @@ -53,21 +53,21 @@ ## sha1 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha1(child) extends UnaryExpression with NullIntolerant`; `inputTypes = Seq(BinaryType) -> StringType`. Comet routes via `CometSha1` to the native `sha1` UDF. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha1(child) extends UnaryExpression with NullIntolerant`; `inputTypes = Seq(BinaryType) -> StringType`. Comet routes via `CometSha1` to the native `sha1` UDF. Mixes in `CodegenDispatchFallback` for the shared `HashUtils` walk; typical calls coerce to binary before that walk is reached. - Spark 4.0.1 (audited 2026-05-27): trait set gains `DefaultStringProducingExpression` and `NullIntolerant` is replaced by `nullIntolerant: Boolean`. Runtime unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. ## sha2 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha2(left, right) extends BinaryExpression`; `inputTypes = Seq(BinaryType, IntegerType) -> StringType`. The `numBits` argument selects SHA-224/256/384/512 (0 is treated as 256); other values return NULL. Comet routes via `CometSha2` to the native `sha2` UDF; non-foldable `numBits` falls back to Spark. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha2(left, right) extends BinaryExpression`; `inputTypes = Seq(BinaryType, IntegerType) -> StringType`. The `numBits` argument selects SHA-224/256/384/512 (0 is treated as 256); other values return NULL. Comet routes via `CometSha2` to the native `sha2` UDF when `numBits` is foldable. A non-foldable `numBits` has no native path; with the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `Sha2.doGenCode` inside the Comet pipeline. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. - Spark 4.0.1 (audited 2026-05-27): trait set gains `DefaultStringProducingExpression` and the `nullIntolerant: Boolean` refactor. Runtime unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. ## xxhash64 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. +- Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. Wide-decimal routing matches `hash` via the shared `HashUtils` and `CodegenDispatchFallback`. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 208a5f3f124..c4561e4d6b0 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -349,12 +349,12 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `crc32` | ✅ | Native | | -| `hash` | ✅ | Native | | +| `hash` | ✅ | Hybrid | Decimal precision > 18 (including nested) routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#hash)) | | `md5` | ✅ | Native | | -| `sha` | ✅ | Native | | -| `sha1` | ✅ | Native | | -| `sha2` | ✅ | Native | | -| `xxhash64` | ✅ | Native | | +| `sha` | ✅ | Hybrid | Alias of `sha1` | +| `sha1` | ✅ | Hybrid | | +| `sha2` | ✅ | Hybrid | Non-foldable `numBits` routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#sha2)) | +| `xxhash64` | ✅ | Hybrid | Decimal precision > 18 (including nested) routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#hash)) | --- diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index ee3e80059d5..215f00ddd01 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -24,7 +24,10 @@ import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} -object CometXxHash64 extends CometExpressionSerde[XxHash64] { +// Native-unsupported, Spark-codegen-compatible cases (`DecimalType` precision > 18, including +// nested, and `sha2` with a non-foldable `numBits`) stay in the Comet pipeline via +// `CodegenDispatchFallback` on the four hash serdes below. +object CometXxHash64 extends CometExpressionSerde[XxHash64] with CodegenDispatchFallback { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons @@ -46,7 +49,7 @@ object CometXxHash64 extends CometExpressionSerde[XxHash64] { } } -object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] { +object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] with CodegenDispatchFallback { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons @@ -72,7 +75,7 @@ object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] { } } -object CometSha2 extends CometExpressionSerde[Sha2] { +object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback { private val nonFoldableNumBitsReason = "The `numBits` argument must be a foldable literal value" @@ -98,7 +101,7 @@ object CometSha2 extends CometExpressionSerde[Sha2] { } } -object CometSha1 extends CometExpressionSerde[Sha1] { +object CometSha1 extends CometExpressionSerde[Sha1] with CodegenDispatchFallback { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index 68c4471e05d..433e4d2efa6 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -26,14 +26,19 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** - * Test suite for Spark murmur3 hash function compatibility between Spark and Comet. + * Test suite for Spark hash function compatibility between Spark and Comet. * - * These tests verify that Comet's native implementation of murmur3 hash produces identical - * results to Spark's implementation for all supported data types. + * Native kernels are asserted for supported input shapes. Cases the native path declines + * (`DecimalType` precision > 18, including nested, and `sha2` with a non-foldable `numBits`) must + * stay in the Comet pipeline via the JVM codegen dispatcher. */ -class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { +class CometHashExpressionSuite + extends CometTestBase + with AdaptiveSparkPlanHelper + with CometCodegenAssertions { test("hash - boolean") { withTable("t") { @@ -134,7 +139,7 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe withTable("t") { sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c) FROM t ORDER BY c") + assertNativeHash("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") } } } @@ -144,37 +149,94 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe withTable("t") { sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - // Large decimals may fall back to Spark, so just check the answer - checkSparkAnswer("SELECT c, hash(c) FROM t ORDER BY c") + assertCodegenRan { + checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") + } } } } - test("hash - array of decimal (precision > 18) falls back to Spark") { + test("hash - array of decimal (precision > 18) routes through the codegen dispatcher") { withTable("t") { sql("CREATE TABLE t(c ARRAY) USING parquet") sql("INSERT INTO t VALUES (array(1.23, 2.34)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c) FROM t", "precision > 18") + assertCodegenRan { + checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t") + } } } - test("hash - struct with decimal (precision > 18) falls back to Spark") { + test("hash - struct with decimal (precision > 18) routes through the codegen dispatcher") { withTable("t") { sql("CREATE TABLE t(c STRUCT) USING parquet") sql("INSERT INTO t VALUES (named_struct('a', 1, 'b', 1.23)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c) FROM t", "precision > 18") + assertCodegenRan { + checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t") + } } } - test("hash - map with decimal (precision > 18) value falls back to Spark") { + test("hash - map with decimal (precision > 18) value routes through the codegen dispatcher") { withSQLConf("spark.sql.legacy.allowHashOnMapType" -> "true") { withTable("t") { sql("CREATE TABLE t(c MAP) USING parquet") sql("INSERT INTO t VALUES (map('a', 1.23)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c) FROM t", "precision > 18") + assertCodegenRan { + checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t") + } + } + } + } + + test("hash - wide decimal falls back when codegen dispatcher is disabled") { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + withTable("t") { + sql("CREATE TABLE t(c DECIMAL(20, 2)) USING parquet") + sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") + checkSparkAnswerAndFallbackReasons( + "SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c", + Set( + s"hash: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false", + s"xxhash64: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false")) + } + } + } + + test("sha2 - non-foldable numBits routes through the codegen dispatcher") { + withTable("t") { + sql("CREATE TABLE t(payload STRING, num_bits INT) USING parquet") + sql("""INSERT INTO t VALUES + ('hello', 0), + ('hello', 224), + ('hello', 256), + ('hello', 384), + ('hello', 512), + ('hello', 128), + ('hello', -1), + (NULL, 256), + ('hello', NULL)""") + assertCodegenRan { + checkSparkAnswerAndOperator("SELECT sha2(payload, num_bits) FROM t") + } + } + } + + test("sha2 - literal numBits stays native") { + withTable("t") { + sql("CREATE TABLE t(payload STRING) USING parquet") + sql("INSERT INTO t VALUES ('hello'), (''), (NULL)") + assertNativeHash("SELECT sha2(payload, 256) FROM t") + } + } + + test("sha2 - non-foldable numBits falls back when codegen dispatcher is disabled") { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + withTable("t") { + sql("CREATE TABLE t(payload STRING, num_bits INT) USING parquet") + sql("INSERT INTO t VALUES ('hello', 256), (NULL, 256), ('hello', NULL)") + checkSparkAnswerAndFallbackReason( + "SELECT sha2(payload, num_bits) FROM t", + s"sha2: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") } } } @@ -587,4 +649,12 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } } + + private def assertNativeHash(query: String): Unit = { + CometScalaUDFCodegen.resetStats() + checkSparkAnswerAndOperator(query) + assert( + CometScalaUDFCodegen.stats().totalLookups == 0, + s"expected native hash execution for $query, got ${CometScalaUDFCodegen.stats()}") + } } From 21cdab84027b6cd43ee4b87541823cc054f1e0c2 Mon Sep 17 00:00:00 2001 From: Hung Date: Fri, 11 Sep 2026 22:10:02 +0800 Subject: [PATCH 2/7] fix: keep TimeType off hash dispatch and add a focused microbenchmark (#5581) TimeType stays on Spark fallback instead of CodegenDispatchFallback. Widen decimal fixtures past 64-bit unscaled values and add a three-arm hash/sha2 dispatcher microbenchmark with mixed projection and first-use. --- .../expression-audits/hash_funcs.md | 6 +- .../scala/org/apache/comet/serde/hash.scala | 111 +++-- .../comet/CometHashExpressionSuite.scala | 90 ++++- .../CometHashCodegenDispatchBenchmark.scala | 380 ++++++++++++++++++ 4 files changed, 527 insertions(+), 60 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index b78c7eca2a6..6c1040fb113 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -34,7 +34,7 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `Murmur3Hash(children, seed) extends HashExpression[Int]`; produces a Murmur3 hash with a configurable Int seed and `IntegerType` result. Comet routes via `CometMurmur3Hash` to the native `murmur3_hash` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; some inner helper refactors only. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Known limitation: the native kernel does not hash `DecimalType` children with precision > 18 (Spark hashes them through Java `BigDecimal`), including when nested in array, struct, or map. With the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `HashExpression.doGenCode` inside the Comet pipeline so the enclosing operator stays native. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. `TimeType` (Spark 4.0+) is still unsupported by the native kernel and is not routed through the dispatcher here. The same `HashUtils` type walk applies to `xxhash64`, `sha1`, and `sha2`. +- Known limitation: the native kernel does not hash `DecimalType` children with precision > 18 (Spark hashes them through Java `BigDecimal`), including when nested in array, struct, or map. With the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `HashExpression.doGenCode` inside the Comet pipeline so the enclosing operator stays native. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. `TimeType` is out of scope for that dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not intercept it, and `convert` declines the native path so the projection falls back to Spark. The same `HashUtils` type walk applies to `xxhash64`, `sha1`, and `sha2`. ## md5 @@ -53,7 +53,7 @@ ## sha1 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha1(child) extends UnaryExpression with NullIntolerant`; `inputTypes = Seq(BinaryType) -> StringType`. Comet routes via `CometSha1` to the native `sha1` UDF. Mixes in `CodegenDispatchFallback` for the shared `HashUtils` walk; typical calls coerce to binary before that walk is reached. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha1(child) extends UnaryExpression with NullIntolerant`; `inputTypes = Seq(BinaryType) -> StringType`. Comet routes via `CometSha1` to the native `sha1` UDF. Mixes in `CodegenDispatchFallback` for the shared `HashUtils` walk (wide decimal); `TimeType` is excluded from that enrollment. Typical calls coerce to binary before that walk is reached. - Spark 4.0.1 (audited 2026-05-27): trait set gains `DefaultStringProducingExpression` and `NullIntolerant` is replaced by `nullIntolerant: Boolean`. Runtime unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. @@ -67,7 +67,7 @@ ## xxhash64 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. Wide-decimal routing matches `hash` via the shared `HashUtils` and `CodegenDispatchFallback`. +- Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. Wide-decimal routing matches `hash` via the shared `HashUtils` and `CodegenDispatchFallback`; `TimeType` is likewise excluded from dispatcher enrollment. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index 215f00ddd01..a0e15d29b02 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -22,11 +22,14 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64} import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructType} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} // Native-unsupported, Spark-codegen-compatible cases (`DecimalType` precision > 18, including // nested, and `sha2` with a non-foldable `numBits`) stay in the Comet pipeline via -// `CodegenDispatchFallback` on the four hash serdes below. +// `CodegenDispatchFallback` on the four hash serdes below. `TimeType` is out of scope for that +// dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not intercept, +// and `convert` declines the native path so the projection falls back to Spark. object CometXxHash64 extends CometExpressionSerde[XxHash64] with CodegenDispatchFallback { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons @@ -38,14 +41,16 @@ object CometXxHash64 extends CometExpressionSerde[XxHash64] with CodegenDispatch expr: XxHash64, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) - val seedBuilder = LiteralOuterClass.Literal - .newBuilder() - .setDatatype(serializeDataType(LongType).get) - .setLongVal(expr.seed) - val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) - // the seed is put at the end of the arguments - scalarFunctionExprToProtoWithReturnType("xxhash64", LongType, false, exprs :+ seedExpr: _*) + HashUtils.convertNativeOrSparkFallback(expr) { + val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) + val seedBuilder = LiteralOuterClass.Literal + .newBuilder() + .setDatatype(serializeDataType(LongType).get) + .setLongVal(expr.seed) + val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) + // the seed is put at the end of the arguments + scalarFunctionExprToProtoWithReturnType("xxhash64", LongType, false, exprs :+ seedExpr: _*) + } } } @@ -60,18 +65,20 @@ object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] with CodegenDi expr: Murmur3Hash, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) - val seedBuilder = LiteralOuterClass.Literal - .newBuilder() - .setDatatype(serializeDataType(IntegerType).get) - .setIntVal(expr.seed) - val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) - // the seed is put at the end of the arguments - scalarFunctionExprToProtoWithReturnType( - "murmur3_hash", - IntegerType, - false, - exprs :+ seedExpr: _*) + HashUtils.convertNativeOrSparkFallback(expr) { + val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) + val seedBuilder = LiteralOuterClass.Literal + .newBuilder() + .setDatatype(serializeDataType(IntegerType).get) + .setIntVal(expr.seed) + val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) + // the seed is put at the end of the arguments + scalarFunctionExprToProtoWithReturnType( + "murmur3_hash", + IntegerType, + false, + exprs :+ seedExpr: _*) + } } } @@ -84,7 +91,11 @@ object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback HashUtils.unsupportedReasons :+ nonFoldableNumBitsReason override def getSupportLevel(expr: Sha2): SupportLevel = { - if (!expr.right.foldable) { + // TimeType is not enrolled in the dispatcher; check it before the non-foldable `numBits` + // `Unsupported` so a mixed tree does not take the mixin path. + if (HashUtils.containsTimeTypeInChildren(expr)) { + Compatible() + } else if (!expr.right.foldable) { Unsupported(Some(nonFoldableNumBitsReason)) } else { HashUtils.supportLevelForChildren(expr) @@ -95,9 +106,11 @@ object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback expr: Sha2, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val leftExpr = exprToProtoInternal(expr.left, inputs, binding) - val numBitsExpr = exprToProtoInternal(expr.right, inputs, binding) - scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr) + HashUtils.convertNativeOrSparkFallback(expr) { + val leftExpr = exprToProtoInternal(expr.left, inputs, binding) + val numBitsExpr = exprToProtoInternal(expr.right, inputs, binding) + scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr) + } } } @@ -112,8 +125,10 @@ object CometSha1 extends CometExpressionSerde[Sha1] with CodegenDispatchFallback expr: Sha1, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - val childExpr = exprToProtoInternal(expr.child, inputs, binding) - scalarFunctionExprToProtoWithReturnType("sha1", StringType, false, childExpr) + HashUtils.convertNativeOrSparkFallback(expr) { + val childExpr = exprToProtoInternal(expr.child, inputs, binding) + scalarFunctionExprToProtoWithReturnType("sha1", StringType, false, childExpr) + } } } @@ -123,19 +138,46 @@ private object HashUtils { "`DecimalType` with precision > 18 is not supported (Spark hashes via Java `BigDecimal`)" private val unsupportedTimeTypeReason = "`TimeType` is not supported" + // `TimeType` is omitted: `CodegenDispatchFallback` documents this list as JVM-dispatch cases. val unsupportedReasons: Seq[String] = - Seq(unsupportedDecimalReason, unsupportedTimeTypeReason, "Unsupported child data type") + Seq(unsupportedDecimalReason, "Unsupported child data type") + + def containsTimeTypeInChildren(expr: Expression): Boolean = + expr.children.exists(c => containsTimeType(c.dataType)) def supportLevelForChildren(expr: Expression): SupportLevel = { - expr.children.iterator - .flatMap(c => unsupportedReasonFor(c.dataType).iterator) - .toSeq - .headOption match { - case Some(reason) => Unsupported(Some(reason)) - case None => Compatible() + // Compatible (not Unsupported) so `CodegenDispatchFallback` does not enroll TimeType. + if (containsTimeTypeInChildren(expr)) { + Compatible() + } else { + expr.children.iterator + .flatMap(c => unsupportedReasonFor(c.dataType).iterator) + .toSeq + .headOption match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible() + } } } + def convertNativeOrSparkFallback(expr: Expression)( + native: => Option[ExprOuterClass.Expr]): Option[ExprOuterClass.Expr] = { + if (containsTimeTypeInChildren(expr)) { + withFallbackReason(expr, unsupportedTimeTypeReason) + None + } else { + native + } + } + + private def containsTimeType(dt: DataType): Boolean = dt match { + case t if isTimeType(t) => true + case s: StructType => s.fields.exists(f => containsTimeType(f.dataType)) + case a: ArrayType => containsTimeType(a.elementType) + case m: MapType => containsTimeType(m.keyType) || containsTimeType(m.valueType) + case _ => false + } + private def unsupportedReasonFor(dt: DataType): Option[String] = dt match { case d: DecimalType if d.precision > 18 => Some(unsupportedDecimalReason) case s: StructType => @@ -143,7 +185,6 @@ private object HashUtils { case a: ArrayType => unsupportedReasonFor(a.elementType) case m: MapType => unsupportedReasonFor(m.keyType).orElse(unsupportedReasonFor(m.valueType)) - case dt if isTimeType(dt) => Some(unsupportedTimeTypeReason) case _ if !supportedDataType(dt, allowComplex = true) => Some(s"Unsupported child data type: $dt") case _ => None diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index 433e4d2efa6..1343d344d41 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.types.{IntegerType, StructField, StructType} +import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} import org.apache.comet.udf.codegen.CometScalaUDFCodegen @@ -33,7 +34,8 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * * Native kernels are asserted for supported input shapes. Cases the native path declines * (`DecimalType` precision > 18, including nested, and `sha2` with a non-foldable `numBits`) must - * stay in the Comet pipeline via the JVM codegen dispatcher. + * stay in the Comet pipeline via the JVM codegen dispatcher. `TimeType` is out of scope for that + * enrollment and falls the projection back to Spark. */ class CometHashExpressionSuite extends CometTestBase @@ -144,35 +146,55 @@ class CometHashExpressionSuite } } - test("hash - decimal (precision > 18)") { - Seq((20, 2), (38, 10)).foreach { case (precision, scale) => - withTable("t") { - sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") - sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - assertCodegenRan { - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") - } - } + test("hash - decimal (precision 20, unscaled > 64-bit) routes through the codegen dispatcher") { + withTable("t") { + sql("CREATE TABLE t(c DECIMAL(20, 2)) USING parquet") + sql("""INSERT INTO t VALUES + (CAST('999999999999999999.99' AS DECIMAL(20, 2))), + (CAST('-999999999999999999.99' AS DECIMAL(20, 2))), + (0.0), + (null)""") + assertDispatchedHash("SELECT hash(c) FROM t") + assertDispatchedHash("SELECT xxhash64(c) FROM t") + } + } + + test("hash - decimal (precision 38, unscaled > 64-bit) routes through the codegen dispatcher") { + withTable("t") { + sql("CREATE TABLE t(c DECIMAL(38, 10)) USING parquet") + sql("""INSERT INTO t VALUES + (CAST('9999999999999999999999999999.9999999999' AS DECIMAL(38, 10))), + (CAST('-9999999999999999999999999999.9999999999' AS DECIMAL(38, 10))), + (0.0), + (null)""") + assertDispatchedHash("SELECT hash(c) FROM t") + assertDispatchedHash("SELECT xxhash64(c) FROM t") } } test("hash - array of decimal (precision > 18) routes through the codegen dispatcher") { withTable("t") { sql("CREATE TABLE t(c ARRAY) USING parquet") - sql("INSERT INTO t VALUES (array(1.23, 2.34)), (null)") - assertCodegenRan { - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t") - } + sql("""INSERT INTO t VALUES + (array( + CAST('999999999999999999.99' AS DECIMAL(20, 2)), + CAST(NULL AS DECIMAL(20, 2)), + CAST('-999999999999999999.99' AS DECIMAL(20, 2)))), + (null)""") + assertDispatchedHash("SELECT hash(c) FROM t") + assertDispatchedHash("SELECT xxhash64(c) FROM t") } } test("hash - struct with decimal (precision > 18) routes through the codegen dispatcher") { withTable("t") { sql("CREATE TABLE t(c STRUCT) USING parquet") - sql("INSERT INTO t VALUES (named_struct('a', 1, 'b', 1.23)), (null)") - assertCodegenRan { - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t") - } + sql("""INSERT INTO t VALUES + (named_struct('a', 1, 'b', CAST('999999999999999999.99' AS DECIMAL(20, 2)))), + (named_struct('a', 1, 'b', CAST(NULL AS DECIMAL(20, 2)))), + (null)""") + assertDispatchedHash("SELECT hash(c) FROM t") + assertDispatchedHash("SELECT xxhash64(c) FROM t") } } @@ -180,10 +202,28 @@ class CometHashExpressionSuite withSQLConf("spark.sql.legacy.allowHashOnMapType" -> "true") { withTable("t") { sql("CREATE TABLE t(c MAP) USING parquet") - sql("INSERT INTO t VALUES (map('a', 1.23)), (null)") - assertCodegenRan { - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t") - } + sql("""INSERT INTO t VALUES + (map('a', CAST('999999999999999999.99' AS DECIMAL(20, 2)))), + (map('a', CAST(NULL AS DECIMAL(20, 2)))), + (null)""") + assertDispatchedHash("SELECT hash(c) FROM t") + assertDispatchedHash("SELECT xxhash64(c) FROM t") + } + } + } + + test("hash - TimeType falls back to Spark") { + assume(isSpark41Plus, "TimeType requires Spark 4.1+") + withSQLConf("spark.sql.timeType.enabled" -> "true") { + withTable("t") { + sql("CREATE TABLE t(c STRING) USING parquet") + sql("INSERT INTO t VALUES ('12:34:56'), ('00:00:00'), (null)") + checkSparkAnswerAndFallbackReasons( + "SELECT hash(to_time(c)) FROM t", + Set("`TimeType` is not supported")) + checkSparkAnswerAndFallbackReasons( + "SELECT xxhash64(to_time(c)) FROM t", + Set("`TimeType` is not supported")) } } } @@ -657,4 +697,10 @@ class CometHashExpressionSuite CometScalaUDFCodegen.stats().totalLookups == 0, s"expected native hash execution for $query, got ${CometScalaUDFCodegen.stats()}") } + + private def assertDispatchedHash(query: String): Unit = { + assertCodegenRan { + checkSparkAnswerAndOperator(query) + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala new file mode 100644 index 00000000000..dd66a287bbe --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala @@ -0,0 +1,380 @@ +/* + * 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.benchmark + +import java.nio.charset.StandardCharsets + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.optimizer.ConstantFolding +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + +/** + * Focused benchmark of hash expressions that `CodegenDispatchFallback` keeps in the Comet + * pipeline: wide-decimal `hash` / `xxhash64` and column-valued `sha2(payload, numBits)`. + * + * The enrollment exists so that one unhandled hash does not drop bucketing, partitioning or dedup + * -- the surrounding aggregate, exchange and (for a mixed projection) native neighbors -- back to + * Spark. Both Comet arms run Spark's `doGenCode` for the hash itself, so an isolated `SELECT + * hash(wide_decimal)` is expected to be close; the grouped cases are the comparison the issue is + * about. + * + * The interesting comparison is therefore + * + * - `codegen dispatch` -- Spark's `doGenCode` runs as a Janino kernel inside the Comet + * pipeline, so the enclosing operator stays native, and + * - `dispatch off` -- `spark.comet.exec.scalaUDF.codegen.enabled=false`, so the enclosing + * projection falls back to Spark and every operator above it leaves the Comet pipeline with + * it. + * + * A pure Spark case is included as a third reference point. A `dispatch off (repeat)` case + * repeats the baseline at the end of every table so the spread between the two is this machine's + * noise floor. + * + * Compilation is a one-time cost, and the steady-state tables warm up before timing, so it does + * not appear there. The `first use` table at the top measures it directly: the same query run + * twice back to back on a cold kernel, after the dispatcher itself has been warmed on an + * unrelated expression. + * + * Before timing anything, every case is run through all three arms and the rows are compared, so + * a timing cannot come from an arm that computed something else. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometHashCodegenDispatchBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometHashCodegenDispatchBenchmark-**results.txt". + */ +object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { + + /** Fewer rows than one Comet batch, so the query is a single batch of real work. */ + private val SmallRows = 1024 + + /** ~128 batches at the default `spark.comet.batchSize`. */ + private val LargeRows = 1024 * 1024 + + /** + * @param name + * Case name, as it appears in the results table. + * @param query + * The query to time. Every argument is a column so that the expression is evaluated per row. + * @param extraConfigs + * Applied to all three arms, so they never account for a difference between them. + */ + private case class DispatchCase( + name: String, + query: String, + extraConfigs: Seq[(String, String)] = Nil) + + /** + * Grouped cases shuffle. AQE is off and the shuffle is one partition so that both arms plan the + * same shape every iteration, and so that the plan check is not looking at `AQEShuffleRead`. + */ + private val groupByConfigs: Seq[(String, String)] = + Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "1") + + private def cases: Seq[DispatchCase] = Seq( + // Isolated projections. Both Comet arms run Spark's `doGenCode` for the hash itself, so they + // are expected to be close; they exist so a results file can tell kernel cost from + // enclosing-operator cost. + DispatchCase("hash(wide decimal)", "select hash(c_wide_dec) from parquetV1Table"), + DispatchCase("xxhash64(wide decimal)", "select xxhash64(c_wide_dec) from parquetV1Table"), + DispatchCase("sha2(payload, numBits)", "select sha2(c_str, c_bits) from parquetV1Table"), + // One unhandled hash used to cost the whole projection, including the native expressions + // next to it. `sha2` is not in this mix: it is so much slower than murmur3 that it hides + // whether the native neighbors stayed native. + DispatchCase( + "mixed projection", + "select length(c_str), c_int + 1, hash(c_int), hash(c_wide_dec) from parquetV1Table"), + // The case the enrollment is really about: hash as a grouping key in bucketing / dedup. + // With the projection gone the aggregate above it has a row-based child, so the partial + // aggregate, the exchange and the final aggregate all leave the Comet pipeline with it. + // `pmod(..., 1024)` keeps the grouping cheap so the per-row hash still dominates. + DispatchCase( + "group by hash", + "select pmod(hash(c_wide_dec), 1024) as k, count(*) from parquetV1Table group by 1", + extraConfigs = groupByConfigs), + DispatchCase( + "group by hash + aggs", + "select pmod(hash(c_wide_dec), 1024) as k, sum(c_int), count(*) " + + "from parquetV1Table group by 1", + extraConfigs = groupByConfigs)) + + private def noConstantFolding: (String, String) = + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName) + + private def sparkConfigs(c: DispatchCase): Seq[(String, String)] = + Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false") ++ c.extraConfigs + + /** Comet on, dispatcher on: the behaviour this benchmark is validating. */ + private def dispatchConfigs(c: DispatchCase): Seq[(String, String)] = + cometConfigs(c, dispatch = true) + + /** Comet on, dispatcher off: the enclosing projection falls back to Spark. */ + private def fallbackConfigs(c: DispatchCase): Seq[(String, String)] = + cometConfigs(c, dispatch = false) + + private def cometConfigs(c: DispatchCase, dispatch: Boolean): Seq[(String, String)] = Seq( + noConstantFolding, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) ++ c.extraConfigs + + private val DispatchCaseName = "Comet, codegen dispatch" + private val FallbackCaseName = "Comet, dispatch off (Spark fallback)" + private val SparkCaseName = "Spark (Comet disabled)" + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + val selected = cases + runBenchmark("Hash codegen dispatch: environment") { + emitEnvironment(selected) + } + + // The kernels must be uncompiled when `runFirstUse` starts, so it runs before anything else + // executes these queries: `CodeGenerator.compile` caches on the generated source JVM-wide, + // and a single earlier execution would make every "first use" number a cache hit. + withCorpus(SmallRows) { + runBenchmark(s"Hash codegen dispatch: first use, $SmallRows rows") { + runFirstUse(selected) + } + selected.foreach(verifyArmsAgree(_, SmallRows)) + selected.foreach(runSteadyState(_, SmallRows)) + } + withCorpus(LargeRows) { + selected.foreach(verifyArmsAgree(_, LargeRows)) + selected.foreach(runSteadyState(_, LargeRows)) + } + } + + private def emitEnvironment(selected: Seq[DispatchCase]): Unit = { + emit(s"Spark version: ${spark.version}") + emit( + s"Java version: ${System.getProperty("java.version")} " + + s"(${System.getProperty("java.vm.name")})") + emit(s"Scala version: ${scala.util.Properties.versionNumberString}") + emit(s"spark.master: ${spark.conf.get("spark.master", "")}") + emit( + s"${CometConf.COMET_BATCH_SIZE.key}: " + + CometConf.COMET_BATCH_SIZE.get(spark.sessionState.conf)) + emit( + s"${SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key}: " + + spark.conf.get(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key)) + emit(s"Dispatcher conf: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}") + emit("Steady-state tables: Spark's Benchmark defaults -- 2s of untimed warmup per case, then") + emit( + " at least 2 iterations and at least 2s of timed iterations; the table reports best and") + emit(" average of the timed iterations.") + emit(s"Row counts: $SmallRows (sub-batch) and $LargeRows (multi-batch).") + emit( + "Wide decimal unscaled values exceed 64 bits " + + "(id + 10^19, DECIMAL(38,10)), so hashing goes through Java BigDecimal.") + emit( + "Isolated hash / xxhash64 / sha2 projections are kernel-cost controls " + + "(both Comet arms run Spark's doGenCode; expected near 1.0X).") + emit( + "mixed projection and the group-by cases keep surrounding native work in the Comet " + + "pipeline only on the dispatch arm -- those are the comparison the enrollment is for.") + emit(s"Cases: ${selected.map(_.name).mkString(", ")}") + } + + /** + * Cost of the first execution of a query whose kernel has never been compiled, against the + * second execution of the same query. + * + * The dispatcher is warmed on an expression that is not in the case list first, so that Janino, + * `CodeGenerator`, the Arrow bridge and the FFI boundary are all loaded and JIT-warm before the + * first case runs and the difference is dominated by compiling that one kernel. It is still an + * upper bound on the compile: the first case in the list absorbs whatever class loading the + * warmup missed, and any case that reaches machinery no earlier case did -- the grouped cases + * are the first to shuffle -- pays for that here too. + */ + private def runFirstUse(selected: Seq[DispatchCase]): Unit = { + // `rlike` routes through the same dispatcher and is not one of the cases below. + val warmup = DispatchCase("warmup", "select c_str rlike '[0-9]+' from parquetV1Table") + (0 until 5).foreach(_ => runQuery(warmup.query, dispatchConfigs(warmup))) + + emit( + f"${"case"}%-24s ${"1st run (ms)"}%14s ${"2nd run (ms)"}%14s " + + f"${"one-time (ms)"}%14s ${"compiles"}%9s") + emit("-" * 84) + selected.foreach { c => + CometScalaUDFCodegen.resetStats() + val cold = timeMillis(runQuery(c.query, dispatchConfigs(c))) + val compiles = CometScalaUDFCodegen.stats().compileCount + val warm = timeMillis(runQuery(c.query, dispatchConfigs(c))) + emit(f"${c.name}%-24s $cold%14.1f $warm%14.1f ${cold - warm}%14.1f $compiles%9d") + } + emit("") + emit("`compiles` counts dispatcher cache misses, which is one per task. The Janino work") + emit("itself is deduplicated JVM-wide by Spark's CodeGenerator source cache, so only the") + emit("first task to reach a given kernel source pays for it.") + } + + /** + * Fails if the three arms disagree on `query`. Rows are compared as a sorted multiset rather + * than positionally, because the grouped cases shuffle and their output order is a property of + * the plan, which is the one thing that differs between the arms. + */ + private def verifyArmsAgree(c: DispatchCase, rows: Int): Unit = { + def collect(configs: Seq[(String, String)]): Array[String] = { + // Assigned to a local rather than returned from the block: Spark 3.4 and 3.5 declare + // `SQLHelper.withSQLConf` as returning `Unit`; only Spark 4 has the result-returning form. + var collected: Array[Row] = Array.empty + withSQLConf(configs: _*) { + collected = spark.sql(c.query).collect() + } + collected + .map( + _.toSeq + .map { + case bytes: Array[Byte] => bytes.mkString("[", ",", "]") + case other => String.valueOf(other) + } + .mkString("|")) + .sorted + } + + val expected = collect(sparkConfigs(c)) + Seq(DispatchCaseName -> dispatchConfigs(c), FallbackCaseName -> fallbackConfigs(c)).foreach { + case (armName, configs) => + val actual = collect(configs) + assert( + expected.length == actual.length, + s"${c.name} @ $rows rows: Spark produced ${expected.length} rows, " + + s"$armName ${actual.length}") + expected.indices.find(i => expected(i) != actual(i)).foreach { i => + throw new AssertionError( + s"${c.name} @ $rows rows: row $i differs -- Spark ${expected(i)}, " + + s"$armName ${actual(i)}") + } + } + } + + private def runSteadyState(c: DispatchCase, rows: Int): Unit = { + runBenchmark(s"${c.name} -- $rows rows") { + val benchmark = new Benchmark(s"${c.name} -- $rows rows", rows, output = output) + checkPlans(benchmark, c) + // The dispatch-off arm goes first so the `Relative` column reads as the speedup this + // change buys over falling the enclosing operator back to Spark. + benchmark.addCase(FallbackCaseName)(_ => runQuery(c.query, fallbackConfigs(c))) + benchmark.addCase(DispatchCaseName)(_ => runQuery(c.query, dispatchConfigs(c))) + benchmark.addCase(SparkCaseName)(_ => runQuery(c.query, sparkConfigs(c))) + benchmark.addCase(s"$FallbackCaseName (repeat)")(_ => runQuery(c.query, fallbackConfigs(c))) + benchmark.run() + } + } + + /** + * Warns rather than fails, so a routing surprise degrades the table to a note instead of + * aborting the run. A case where the dispatch arm is not fully native, or where the + * dispatch-off arm is, is not measuring what its name says. + */ + private def checkPlans(benchmark: Benchmark, c: DispatchCase): Unit = { + var dispatchNonComet: Option[String] = None + var fallbackIsFullyComet = false + var dispatcherRan = false + withSQLConf(dispatchConfigs(c): _*) { + CometScalaUDFCodegen.resetStats() + val df = spark.sql(c.query) + df.noop() + val stats = CometScalaUDFCodegen.stats() + dispatcherRan = stats.compileCount + stats.cacheHitCount > 0 + dispatchNonComet = + findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)).map(_.nodeName) + } + withSQLConf(fallbackConfigs(c): _*) { + val df = spark.sql(c.query) + df.noop() + fallbackIsFullyComet = + findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)).isEmpty + } + dispatchNonComet.foreach(op => + warn( + benchmark, + "WARNING: the codegen-dispatch plan is not fully Comet native (first " + + s"non-Comet operator: $op), so that case is partly measuring Spark.")) + if (!dispatcherRan) { + warn( + benchmark, + "WARNING: the codegen dispatcher did not run for this query, so the two " + + "Comet cases below are measuring the same plan.") + } + if (fallbackIsFullyComet) { + warn( + benchmark, + "WARNING: the dispatch-off plan is fully Comet native, so this case is " + + "not exercising the operator fallback it is meant to be compared against.") + } + } + + private def runQuery(query: String, configs: Seq[(String, String)]): Unit = + withSQLConf(configs: _*) { + spark.sql(query).noop() + } + + private def timeMillis(f: => Unit): Double = { + val start = System.nanoTime() + f + (System.nanoTime() - start) / 1e6 + } + + /** Builds `parquetV1Table` with `rows` rows of the corpus and drops it afterwards. */ + private def withCorpus(rows: Int)(f: => Unit): Unit = { + withTempPath { dir => + withTempTable(tbl, "parquetV1Table") { + spark.range(rows).createOrReplaceTempView(tbl) + prepareTable(dir, spark.sql(corpusQuery)) + f + } + } + } + + /** + * Every column varies per row, so no argument to a benchmarked expression is loop-invariant and + * neither engine can hoist the call out of the row loop. `c_wide_dec` is `id + 10^19` as + * `DECIMAL(38,10)` so the unscaled integer exceeds 64 bits (Spark hashes via `BigDecimal`). + */ + private def corpusQuery: String = { + val columns = Seq( + "c_str" -> "REPEAT(CAST(id AS STRING), 4)", + "c_int" -> "CAST(id AS INT)", + "c_bits" -> "element_at(array(0, 224, 256, 384, 512), CAST(pmod(id, 5) AS INT) + 1)", + "c_wide_dec" -> ("CAST(id AS DECIMAL(38, 10)) + " + + "CAST('10000000000000000000.0000000000' AS DECIMAL(38, 10))")) + s"SELECT ${columns.map { case (name, expr) => s"$expr AS $name" }.mkString(", ")} FROM $tbl" + } + + /** Writes a warning to the results file as well as the console, ordered against the table. */ + private def warn(benchmark: Benchmark, message: String): Unit = { + val border = "=" * 80 + benchmark.out.println(s"\n$border\n$message\n$border") + } + + /** [[Benchmark]] tees console and results file; this benchmark's own tables need the same. */ + private def emit(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes(StandardCharsets.UTF_8))) + } +} From 3b1547f5d61c422351ebec21c17547fabbd29dd0 Mon Sep 17 00:00:00 2001 From: Hung Date: Mon, 14 Sep 2026 09:22:17 +0800 Subject: [PATCH 3/7] fix: keep sha1 native and cover split hash dispatcher codegen Analyzer already casts sha1 inputs to BinaryType, so CodegenDispatchFallback was misclassifying sha/sha1 as Hybrid. Add multi-arg wide-decimal tests that force HashExpression method splitting. --- .../expression-audits/hash_funcs.md | 4 +- docs/source/user-guide/latest/expressions.md | 4 +- .../scala/org/apache/comet/serde/hash.scala | 16 +++---- .../comet/CometHashExpressionSuite.scala | 44 +++++++++++++++++++ 4 files changed, 54 insertions(+), 14 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index 6c1040fb113..ae615808eaa 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -34,7 +34,7 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `Murmur3Hash(children, seed) extends HashExpression[Int]`; produces a Murmur3 hash with a configurable Int seed and `IntegerType` result. Comet routes via `CometMurmur3Hash` to the native `murmur3_hash` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; some inner helper refactors only. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Known limitation: the native kernel does not hash `DecimalType` children with precision > 18 (Spark hashes them through Java `BigDecimal`), including when nested in array, struct, or map. With the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `HashExpression.doGenCode` inside the Comet pipeline so the enclosing operator stays native. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. `TimeType` is out of scope for that dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not intercept it, and `convert` declines the native path so the projection falls back to Spark. The same `HashUtils` type walk applies to `xxhash64`, `sha1`, and `sha2`. +- Known limitation: the native kernel does not hash `DecimalType` children with precision > 18 (Spark hashes them through Java `BigDecimal`), including when nested in array, struct, or map. With the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `HashExpression.doGenCode` inside the Comet pipeline so the enclosing operator stays native. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. `TimeType` is out of scope for that dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not intercept it, and `convert` declines the native path so the projection falls back to Spark. The same wide-decimal routing applies to `xxhash64`. ## md5 @@ -53,7 +53,7 @@ ## sha1 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha1(child) extends UnaryExpression with NullIntolerant`; `inputTypes = Seq(BinaryType) -> StringType`. Comet routes via `CometSha1` to the native `sha1` UDF. Mixes in `CodegenDispatchFallback` for the shared `HashUtils` walk (wide decimal); `TimeType` is excluded from that enrollment. Typical calls coerce to binary before that walk is reached. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha1(child) extends UnaryExpression with NullIntolerant`; `inputTypes = Seq(BinaryType) -> StringType`. The Analyzer casts accepted inputs to `BinaryType` before Comet serde, and `CometSha1` routes the resulting binary child to the native `sha1` UDF. - Spark 4.0.1 (audited 2026-05-27): trait set gains `DefaultStringProducingExpression` and `NullIntolerant` is replaced by `nullIntolerant: Boolean`. Runtime unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index c4561e4d6b0..a4503720aab 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -351,8 +351,8 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `crc32` | ✅ | Native | | | `hash` | ✅ | Hybrid | Decimal precision > 18 (including nested) routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#hash)) | | `md5` | ✅ | Native | | -| `sha` | ✅ | Hybrid | Alias of `sha1` | -| `sha1` | ✅ | Hybrid | | +| `sha` | ✅ | Native | Alias of `sha1` | +| `sha1` | ✅ | Native | | | `sha2` | ✅ | Hybrid | Non-foldable `numBits` routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#sha2)) | | `xxhash64` | ✅ | Hybrid | Decimal precision > 18 (including nested) routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#hash)) | diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index a0e15d29b02..6de0968db0d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -27,9 +27,9 @@ import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, s // Native-unsupported, Spark-codegen-compatible cases (`DecimalType` precision > 18, including // nested, and `sha2` with a non-foldable `numBits`) stay in the Comet pipeline via -// `CodegenDispatchFallback` on the four hash serdes below. `TimeType` is out of scope for that -// dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not intercept, -// and `convert` declines the native path so the projection falls back to Spark. +// `CodegenDispatchFallback` on the applicable hash serdes below. `TimeType` is out of scope for +// that dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not +// intercept, and `convert` declines the native path so the projection falls back to Spark. object CometXxHash64 extends CometExpressionSerde[XxHash64] with CodegenDispatchFallback { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons @@ -114,9 +114,7 @@ object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback } } -object CometSha1 extends CometExpressionSerde[Sha1] with CodegenDispatchFallback { - - override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons +object CometSha1 extends CometExpressionSerde[Sha1] { override def getSupportLevel(expr: Sha1): SupportLevel = HashUtils.supportLevelForChildren(expr) @@ -125,10 +123,8 @@ object CometSha1 extends CometExpressionSerde[Sha1] with CodegenDispatchFallback expr: Sha1, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - HashUtils.convertNativeOrSparkFallback(expr) { - val childExpr = exprToProtoInternal(expr.child, inputs, binding) - scalarFunctionExprToProtoWithReturnType("sha1", StringType, false, childExpr) - } + val childExpr = exprToProtoInternal(expr.child, inputs, binding) + scalarFunctionExprToProtoWithReturnType("sha1", StringType, false, childExpr) } } diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index 1343d344d41..f67ec5fffdc 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -172,6 +172,30 @@ class CometHashExpressionSuite } } + test("hash - multiple columns with a wide decimal uses split dispatcher codegen") { + withMultiArgumentWideDecimalTable { + // Spark checks the accumulated source length before adding each child to a code block. + // With threshold 1, the first materialized child already exceeds the threshold, so four + // children force HashExpression.doGenCode to emit helper methods that take InternalRow row. + // Scope the setting to this SELECT so fixture creation uses normal Spark codegen settings. + withSQLConf("spark.sql.codegen.methodSplitThreshold" -> "1") { + assertDispatchedHash( + "SELECT hash(c_str, c_int, c_ts, c_wide_dec) FROM multi_arg_wide_decimal_t") + } + } + } + + test("xxhash64 - multiple columns with a wide decimal uses split dispatcher codegen") { + withMultiArgumentWideDecimalTable { + // See the hash test above. Successful execution also proves Janino compiled the split + // helpers with the dispatcher's `row` alias in both Spark 3.x and Spark 4.x profiles. + withSQLConf("spark.sql.codegen.methodSplitThreshold" -> "1") { + assertDispatchedHash( + "SELECT xxhash64(c_str, c_int, c_ts, c_wide_dec) FROM multi_arg_wide_decimal_t") + } + } + } + test("hash - array of decimal (precision > 18) routes through the codegen dispatcher") { withTable("t") { sql("CREATE TABLE t(c ARRAY) USING parquet") @@ -703,4 +727,24 @@ class CometHashExpressionSuite checkSparkAnswerAndOperator(query) } } + + private def withMultiArgumentWideDecimalTable(f: => Unit): Unit = { + withTable("multi_arg_wide_decimal_t") { + sql("""CREATE TABLE multi_arg_wide_decimal_t( + c_str STRING, + c_int INT, + c_ts TIMESTAMP, + c_wide_dec DECIMAL(38, 10)) + USING parquet""") + sql("""INSERT INTO multi_arg_wide_decimal_t VALUES + ('alpha', 1, TIMESTAMP '2023-01-01 12:00:00', + CAST('9999999999999999999999999999.9999999999' AS DECIMAL(38, 10))), + ('beta', -7, TIMESTAMP '1970-01-01 00:00:01', + CAST('-9999999999999999999999999999.9999999999' AS DECIMAL(38, 10))), + ('gamma', 42, TIMESTAMP '2000-12-31 23:59:59', + CAST('12345678901234567890.1234567890' AS DECIMAL(38, 10))), + (NULL, NULL, NULL, NULL)""") + f + } + } } From 5c109457d66f6835f1e00933afa4c8acaccd4908 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Fri, 23 Oct 2026 14:37:47 +0800 Subject: [PATCH 4/7] fix: narrow codegen dispatch to non-foldable sha2 --- .../expression-audits/hash_funcs.md | 2 +- .../scala/org/apache/comet/serde/hash.scala | 116 ++++------- .../comet/CometHashExpressionSuite.scala | 28 ++- .../CometHashCodegenDispatchBenchmark.scala | 185 +++++------------- .../CometHashExpressionBenchmark.scala | 2 +- 5 files changed, 116 insertions(+), 217 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index a8a432a9897..5a6dbc9ccaa 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -70,7 +70,7 @@ - Performance (tuned 2026-09-28, issue #5994): the shared typed wide-decimal list path was 2.1–15.8x faster in local Criterion measurements across 8,192-row `decimal(38,0)` List, LargeList and FixedSizeList inputs with 2/32 elements and no/sparse/dense nulls; paired scalar remeasurement found no significant slowdown. Benchmark: `benches/decimal_hash.rs`. - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. Wide decimals, including nested decimals, hash natively as for `hash`; `TimeType` is excluded from dispatcher enrollment and falls back to Spark. +- Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. Wide decimals, including nested decimals, hash natively as for `hash`; `TimeType` falls back to Spark. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. - Native routing: Comet's native `xxhash64` UDF delegates compatible arguments at Spark's default seed (`42`) to `datafusion-spark::SparkXxhash64`. Differential tests in `native/spark-expr/src/hash_funcs/xxhash64_diff.rs` compare Comet's kernel against `SparkXxhash64` for primitives, narrow decimals, dictionaries, lists, maps, and nested combinations, and record where the two differ on a NaN whose bits are not canonical. Wide-decimal tests verify that the expression stays on Comet's kernel: the pinned DataFusion implementation uses fixed-width little-endian bytes rather than Spark's minimal signed big-endian bytes. The Comet kernel is retained for wide decimals (also within nested types), a non-default seed (`SparkXxhash64` hardcodes 42), `Struct` (upstream does not push a parent null mask into children; see #5753), a `Dictionary` nested in a list/map (upstream restarts those hashes from 42), `Time64`, and `Float32`/`Float64`, including values nested in lists and maps (upstream hashes the raw bits of a NaN, where Spark hashes every NaN as the canonical NaN; see [apache/datafusion#25913](https://github.com/apache/datafusion/issues/25913)). `create_xxhash64_hashes` remains for `approx_count_distinct` and the non-delegated native path. diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index bd2e20e2f64..c438305e945 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -22,14 +22,9 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64} import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructType} -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} -// Native-unsupported, Spark-codegen-compatible `sha2` with a non-foldable `numBits` stays -// in the Comet pipeline via `CodegenDispatchFallback`. Wide decimals, including nested -// decimals, hash natively. `TimeType` is out of scope for that dispatcher enrollment: `getSupportLevel` reports `Compatible` so the mixin does not -// intercept, and `convert` declines the native path so the projection falls back to Spark. -object CometXxHash64 extends CometExpressionSerde[XxHash64] with CodegenDispatchFallback { +object CometXxHash64 extends CometExpressionSerde[XxHash64] { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons @@ -40,20 +35,18 @@ object CometXxHash64 extends CometExpressionSerde[XxHash64] with CodegenDispatch expr: XxHash64, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - HashUtils.convertNativeOrSparkFallback(expr) { - val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) - val seedBuilder = LiteralOuterClass.Literal - .newBuilder() - .setDatatype(serializeDataType(LongType).get) - .setLongVal(expr.seed) - val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) - // the seed is put at the end of the arguments - scalarFunctionExprToProtoWithReturnType("xxhash64", LongType, false, exprs :+ seedExpr: _*) - } + val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) + val seedBuilder = LiteralOuterClass.Literal + .newBuilder() + .setDatatype(serializeDataType(LongType).get) + .setLongVal(expr.seed) + val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) + // the seed is put at the end of the arguments + scalarFunctionExprToProtoWithReturnType("xxhash64", LongType, false, exprs :+ seedExpr: _*) } } -object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] with CodegenDispatchFallback { +object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] { override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons @@ -64,37 +57,34 @@ object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] with CodegenDi expr: Murmur3Hash, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - HashUtils.convertNativeOrSparkFallback(expr) { - val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) - val seedBuilder = LiteralOuterClass.Literal - .newBuilder() - .setDatatype(serializeDataType(IntegerType).get) - .setIntVal(expr.seed) - val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) - // the seed is put at the end of the arguments - scalarFunctionExprToProtoWithReturnType( - "murmur3_hash", - IntegerType, - false, - exprs :+ seedExpr: _*) - } + val exprs = expr.children.map(exprToProtoInternal(_, inputs, binding)) + val seedBuilder = LiteralOuterClass.Literal + .newBuilder() + .setDatatype(serializeDataType(IntegerType).get) + .setIntVal(expr.seed) + val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build()) + // the seed is put at the end of the arguments + scalarFunctionExprToProtoWithReturnType( + "murmur3_hash", + IntegerType, + false, + exprs :+ seedExpr: _*) } } +// Non-foldable numBits uses Spark codegen inside the Comet pipeline; foldable values stay native. object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback { + override def hasConditionalNativeDefault: Boolean = true + private val nonFoldableNumBitsReason = "The `numBits` argument must be a foldable literal value" override def getUnsupportedReasons(): Seq[String] = - HashUtils.unsupportedReasons :+ nonFoldableNumBitsReason + Seq(nonFoldableNumBitsReason) override def getSupportLevel(expr: Sha2): SupportLevel = { - // TimeType is not enrolled in the dispatcher; check it before the non-foldable `numBits` - // `Unsupported` so a mixed tree does not take the mixin path. - if (HashUtils.containsTimeTypeInChildren(expr)) { - Compatible() - } else if (!expr.right.foldable) { + if (!expr.right.foldable) { Unsupported(Some(nonFoldableNumBitsReason)) } else { HashUtils.supportLevelForChildren(expr) @@ -105,16 +95,16 @@ object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback expr: Sha2, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { - HashUtils.convertNativeOrSparkFallback(expr) { - val leftExpr = exprToProtoInternal(expr.left, inputs, binding) - val numBitsExpr = exprToProtoInternal(expr.right, inputs, binding) - scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr) - } + val leftExpr = exprToProtoInternal(expr.left, inputs, binding) + val numBitsExpr = exprToProtoInternal(expr.right, inputs, binding) + scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr) } } object CometSha1 extends CometExpressionSerde[Sha1] { + override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons + override def getSupportLevel(expr: Sha1): SupportLevel = HashUtils.supportLevelForChildren(expr) @@ -131,52 +121,26 @@ private object HashUtils { private val unsupportedTimeTypeReason = "`TimeType` is not supported" - // `TimeType` is omitted: `CodegenDispatchFallback` documents this list as JVM-dispatch cases. val unsupportedReasons: Seq[String] = - Seq("Unsupported child data type") - - def containsTimeTypeInChildren(expr: Expression): Boolean = - expr.children.exists(c => containsTimeType(c.dataType)) + Seq(unsupportedTimeTypeReason, "Unsupported child data type") def supportLevelForChildren(expr: Expression): SupportLevel = { - // Compatible (not Unsupported) so `CodegenDispatchFallback` does not enroll TimeType. - if (containsTimeTypeInChildren(expr)) { - Compatible() - } else { - expr.children.iterator - .flatMap(c => unsupportedReasonFor(c.dataType).iterator) - .toSeq - .headOption match { - case Some(reason) => Unsupported(Some(reason)) - case None => Compatible() - } + expr.children.iterator + .flatMap(c => unsupportedReasonFor(c.dataType).iterator) + .toSeq + .headOption match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible() } } - def convertNativeOrSparkFallback(expr: Expression)( - native: => Option[ExprOuterClass.Expr]): Option[ExprOuterClass.Expr] = { - if (containsTimeTypeInChildren(expr)) { - withFallbackReason(expr, unsupportedTimeTypeReason) - None - } else { - native - } - } - - private def containsTimeType(dt: DataType): Boolean = dt match { - case t if isTimeType(t) => true - case s: StructType => s.fields.exists(f => containsTimeType(f.dataType)) - case a: ArrayType => containsTimeType(a.elementType) - case m: MapType => containsTimeType(m.keyType) || containsTimeType(m.valueType) - case _ => false - } - private def unsupportedReasonFor(dt: DataType): Option[String] = dt match { case s: StructType => s.fields.iterator.flatMap(f => unsupportedReasonFor(f.dataType).iterator).toSeq.headOption case a: ArrayType => unsupportedReasonFor(a.elementType) case m: MapType => unsupportedReasonFor(m.keyType).orElse(unsupportedReasonFor(m.valueType)) + case dt if isTimeType(dt) => Some(unsupportedTimeTypeReason) case _ if !supportedDataType(dt, allowComplex = true) => Some(s"Unsupported child data type: $dt") case _ => None diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index 09ed567a384..c0ae85368d9 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -32,8 +32,8 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * Test suite for Spark `hash` (murmur3) and `xxhash64` compatibility between Spark and Comet. * * Native kernels are asserted for supported input shapes, including wide and nested decimals. - * `sha2` with a non-foldable `numBits` stays in the Comet pipeline via the JVM codegen dispatcher. - * `TimeType` is out of scope for that enrollment and falls the projection back to Spark. + * `sha2` with a non-foldable `numBits` stays in the Comet pipeline via the JVM codegen + * dispatcher. `TimeType` hash inputs fall back to Spark. */ class CometHashExpressionSuite extends CometTestBase @@ -145,10 +145,19 @@ class CometHashExpressionSuite ('hello', 512), ('hello', 128), ('hello', -1), + ('', 256), + ('', 512), (NULL, 256), ('hello', NULL)""") - assertCodegenRan { - checkSparkAnswerAndOperator("SELECT sha2(payload, num_bits) FROM t") + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + Seq( + "SELECT sha2(payload, num_bits) FROM t", + "SELECT sha2(payload, num_bits + 0) FROM t", + "SELECT sha2('hello', num_bits) FROM t").foreach { query => + assertCodegenRan { + checkSparkAnswerAndOperator(query) + } + } } } } @@ -157,7 +166,13 @@ class CometHashExpressionSuite withTable("t") { sql("CREATE TABLE t(payload STRING) USING parquet") sql("INSERT INTO t VALUES ('hello'), (''), (NULL)") - assertNativeHash("SELECT sha2(payload, 256) FROM t") + Seq("true", "false").foreach { enabled => + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> enabled) { + Seq("0", "224", "256", "384", "512", "128", "-1", "128 + 128").foreach { bits => + assertNativeHash(s"SELECT sha2(payload, $bits) FROM t") + } + } + } } } @@ -462,4 +477,5 @@ class CometHashExpressionSuite assert( CometScalaUDFCodegen.stats().totalLookups == 0, s"expected native hash execution for $query, got ${CometScalaUDFCodegen.stats()}") - }} + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala index dd66a287bbe..97c1c2d1cad 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala @@ -30,22 +30,11 @@ import org.apache.comet.CometConf import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** - * Focused benchmark of hash expressions that `CodegenDispatchFallback` keeps in the Comet - * pipeline: wide-decimal `hash` / `xxhash64` and column-valued `sha2(payload, numBits)`. + * Benchmark of non-foldable `sha2(payload, numBits)` through `CodegenDispatchFallback`. * - * The enrollment exists so that one unhandled hash does not drop bucketing, partitioning or dedup - * -- the surrounding aggregate, exchange and (for a mixed projection) native neighbors -- back to - * Spark. Both Comet arms run Spark's `doGenCode` for the hash itself, so an isolated `SELECT - * hash(wide_decimal)` is expected to be close; the grouped cases are the comparison the issue is - * about. - * - * The interesting comparison is therefore - * - * - `codegen dispatch` -- Spark's `doGenCode` runs as a Janino kernel inside the Comet - * pipeline, so the enclosing operator stays native, and - * - `dispatch off` -- `spark.comet.exec.scalaUDF.codegen.enabled=false`, so the enclosing - * projection falls back to Spark and every operator above it leaves the Comet pipeline with - * it. + * The dispatcher runs Spark's `doGenCode` inside the Comet pipeline. With dispatch disabled, the + * projection falls back to Spark. The mixed projection also measures native expressions alongside + * `sha2`; it does not imply that SHA-2 itself executes natively. * * A pure Spark case is included as a third reference point. A `dispatch off (repeat)` case * repeats the baseline at the end of every table so the spread between the two is this machine's @@ -78,67 +67,34 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { * Case name, as it appears in the results table. * @param query * The query to time. Every argument is a column so that the expression is evaluated per row. - * @param extraConfigs - * Applied to all three arms, so they never account for a difference between them. - */ - private case class DispatchCase( - name: String, - query: String, - extraConfigs: Seq[(String, String)] = Nil) - - /** - * Grouped cases shuffle. AQE is off and the shuffle is one partition so that both arms plan the - * same shape every iteration, and so that the plan check is not looking at `AQEShuffleRead`. */ - private val groupByConfigs: Seq[(String, String)] = - Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "1") + private case class DispatchCase(name: String, query: String) private def cases: Seq[DispatchCase] = Seq( - // Isolated projections. Both Comet arms run Spark's `doGenCode` for the hash itself, so they - // are expected to be close; they exist so a results file can tell kernel cost from - // enclosing-operator cost. - DispatchCase("hash(wide decimal)", "select hash(c_wide_dec) from parquetV1Table"), - DispatchCase("xxhash64(wide decimal)", "select xxhash64(c_wide_dec) from parquetV1Table"), DispatchCase("sha2(payload, numBits)", "select sha2(c_str, c_bits) from parquetV1Table"), - // One unhandled hash used to cost the whole projection, including the native expressions - // next to it. `sha2` is not in this mix: it is so much slower than murmur3 that it hides - // whether the native neighbors stayed native. DispatchCase( - "mixed projection", - "select length(c_str), c_int + 1, hash(c_int), hash(c_wide_dec) from parquetV1Table"), - // The case the enrollment is really about: hash as a grouping key in bucketing / dedup. - // With the projection gone the aggregate above it has a row-based child, so the partial - // aggregate, the exchange and the final aggregate all leave the Comet pipeline with it. - // `pmod(..., 1024)` keeps the grouping cheap so the per-row hash still dominates. - DispatchCase( - "group by hash", - "select pmod(hash(c_wide_dec), 1024) as k, count(*) from parquetV1Table group by 1", - extraConfigs = groupByConfigs), - DispatchCase( - "group by hash + aggs", - "select pmod(hash(c_wide_dec), 1024) as k, sum(c_int), count(*) " + - "from parquetV1Table group by 1", - extraConfigs = groupByConfigs)) + "sha2 mixed projection", + "select length(c_str), c_int + 1, hash(c_int), sha2(c_str, c_bits) from parquetV1Table")) private def noConstantFolding: (String, String) = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> excludedRulesWith(ConstantFolding.ruleName) - private def sparkConfigs(c: DispatchCase): Seq[(String, String)] = - Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false") ++ c.extraConfigs + private def sparkConfigs: Seq[(String, String)] = + Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false") /** Comet on, dispatcher on: the behaviour this benchmark is validating. */ - private def dispatchConfigs(c: DispatchCase): Seq[(String, String)] = - cometConfigs(c, dispatch = true) + private def dispatchConfigs: Seq[(String, String)] = + cometConfigs(dispatch = true) /** Comet on, dispatcher off: the enclosing projection falls back to Spark. */ - private def fallbackConfigs(c: DispatchCase): Seq[(String, String)] = - cometConfigs(c, dispatch = false) + private def fallbackConfigs: Seq[(String, String)] = + cometConfigs(dispatch = false) - private def cometConfigs(c: DispatchCase, dispatch: Boolean): Seq[(String, String)] = Seq( + private def cometConfigs(dispatch: Boolean): Seq[(String, String)] = Seq( noConstantFolding, CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", - CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) ++ c.extraConfigs + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) private val DispatchCaseName = "Comet, codegen dispatch" private val FallbackCaseName = "Comet, dispatch off (Spark fallback)" @@ -146,16 +102,17 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { override def runCometBenchmark(mainArgs: Array[String]): Unit = { val selected = cases - runBenchmark("Hash codegen dispatch: environment") { + runBenchmark("SHA-2 codegen dispatch: environment") { emitEnvironment(selected) } // The kernels must be uncompiled when `runFirstUse` starts, so it runs before anything else // executes these queries: `CodeGenerator.compile` caches on the generated source JVM-wide, - // and a single earlier execution would make every "first use" number a cache hit. + // and a single earlier execution would make the "first use" number a cache hit. Only the + // isolated case is measured here: the mixed projection reuses the same SHA-2 kernel. withCorpus(SmallRows) { - runBenchmark(s"Hash codegen dispatch: first use, $SmallRows rows") { - runFirstUse(selected) + runBenchmark(s"SHA-2 codegen dispatch: first use, $SmallRows rows") { + runFirstUse(selected.take(1)) } selected.foreach(verifyArmsAgree(_, SmallRows)) selected.foreach(runSteadyState(_, SmallRows)) @@ -185,15 +142,8 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { " at least 2 iterations and at least 2s of timed iterations; the table reports best and") emit(" average of the timed iterations.") emit(s"Row counts: $SmallRows (sub-batch) and $LargeRows (multi-batch).") - emit( - "Wide decimal unscaled values exceed 64 bits " + - "(id + 10^19, DECIMAL(38,10)), so hashing goes through Java BigDecimal.") - emit( - "Isolated hash / xxhash64 / sha2 projections are kernel-cost controls " + - "(both Comet arms run Spark's doGenCode; expected near 1.0X).") - emit( - "mixed projection and the group-by cases keep surrounding native work in the Comet " + - "pipeline only on the dispatch arm -- those are the comparison the enrollment is for.") + emit("Payloads include nulls, empty strings and varying lengths; numBits varies per row.") + emit("Both Comet arms use Spark codegen for SHA-2; dispatch keeps the projection in Comet.") emit(s"Cases: ${selected.map(_.name).mkString(", ")}") } @@ -205,13 +155,12 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { * `CodeGenerator`, the Arrow bridge and the FFI boundary are all loaded and JIT-warm before the * first case runs and the difference is dominated by compiling that one kernel. It is still an * upper bound on the compile: the first case in the list absorbs whatever class loading the - * warmup missed, and any case that reaches machinery no earlier case did -- the grouped cases - * are the first to shuffle -- pays for that here too. + * warmup missed. It is a query-level measurement, not an isolated compiler measurement. */ private def runFirstUse(selected: Seq[DispatchCase]): Unit = { // `rlike` routes through the same dispatcher and is not one of the cases below. val warmup = DispatchCase("warmup", "select c_str rlike '[0-9]+' from parquetV1Table") - (0 until 5).foreach(_ => runQuery(warmup.query, dispatchConfigs(warmup))) + (0 until 5).foreach(_ => runQuery(warmup.query, dispatchConfigs)) emit( f"${"case"}%-24s ${"1st run (ms)"}%14s ${"2nd run (ms)"}%14s " + @@ -219,9 +168,9 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { emit("-" * 84) selected.foreach { c => CometScalaUDFCodegen.resetStats() - val cold = timeMillis(runQuery(c.query, dispatchConfigs(c))) + val cold = timeMillis(runQuery(c.query, dispatchConfigs)) val compiles = CometScalaUDFCodegen.stats().compileCount - val warm = timeMillis(runQuery(c.query, dispatchConfigs(c))) + val warm = timeMillis(runQuery(c.query, dispatchConfigs)) emit(f"${c.name}%-24s $cold%14.1f $warm%14.1f ${cold - warm}%14.1f $compiles%9d") } emit("") @@ -231,9 +180,8 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { } /** - * Fails if the three arms disagree on `query`. Rows are compared as a sorted multiset rather - * than positionally, because the grouped cases shuffle and their output order is a property of - * the plan, which is the one thing that differs between the arms. + * Fails if the three arms disagree on `query`. Rows are compared as a sorted multiset because + * output order is not guaranteed across the different plans. */ private def verifyArmsAgree(c: DispatchCase, rows: Int): Unit = { def collect(configs: Seq[(String, String)]): Array[String] = { @@ -254,8 +202,8 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { .sorted } - val expected = collect(sparkConfigs(c)) - Seq(DispatchCaseName -> dispatchConfigs(c), FallbackCaseName -> fallbackConfigs(c)).foreach { + val expected = collect(sparkConfigs) + Seq(DispatchCaseName -> dispatchConfigs, FallbackCaseName -> fallbackConfigs).foreach { case (armName, configs) => val actual = collect(configs) assert( @@ -273,57 +221,36 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { private def runSteadyState(c: DispatchCase, rows: Int): Unit = { runBenchmark(s"${c.name} -- $rows rows") { val benchmark = new Benchmark(s"${c.name} -- $rows rows", rows, output = output) - checkPlans(benchmark, c) + checkPlans(c) // The dispatch-off arm goes first so the `Relative` column reads as the speedup this // change buys over falling the enclosing operator back to Spark. - benchmark.addCase(FallbackCaseName)(_ => runQuery(c.query, fallbackConfigs(c))) - benchmark.addCase(DispatchCaseName)(_ => runQuery(c.query, dispatchConfigs(c))) - benchmark.addCase(SparkCaseName)(_ => runQuery(c.query, sparkConfigs(c))) - benchmark.addCase(s"$FallbackCaseName (repeat)")(_ => runQuery(c.query, fallbackConfigs(c))) + benchmark.addCase(FallbackCaseName)(_ => runQuery(c.query, fallbackConfigs)) + benchmark.addCase(DispatchCaseName)(_ => runQuery(c.query, dispatchConfigs)) + benchmark.addCase(SparkCaseName)(_ => runQuery(c.query, sparkConfigs)) + benchmark.addCase(s"$FallbackCaseName (repeat)")(_ => runQuery(c.query, fallbackConfigs)) benchmark.run() } } - /** - * Warns rather than fails, so a routing surprise degrades the table to a note instead of - * aborting the run. A case where the dispatch arm is not fully native, or where the - * dispatch-off arm is, is not measuring what its name says. - */ - private def checkPlans(benchmark: Benchmark, c: DispatchCase): Unit = { - var dispatchNonComet: Option[String] = None - var fallbackIsFullyComet = false - var dispatcherRan = false - withSQLConf(dispatchConfigs(c): _*) { + /** Refuses to time a case whose actual routing does not match its label. */ + private def checkPlans(c: DispatchCase): Unit = { + withSQLConf(dispatchConfigs: _*) { CometScalaUDFCodegen.resetStats() val df = spark.sql(c.query) df.noop() val stats = CometScalaUDFCodegen.stats() - dispatcherRan = stats.compileCount + stats.cacheHitCount > 0 - dispatchNonComet = - findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)).map(_.nodeName) + assert(stats.compileCount + stats.cacheHitCount > 0, s"dispatcher did not run: $stats") + val nonComet = findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)) + assert(nonComet.isEmpty, s"dispatch projection fell back to Spark: $nonComet") } - withSQLConf(fallbackConfigs(c): _*) { + withSQLConf(fallbackConfigs: _*) { + CometScalaUDFCodegen.resetStats() val df = spark.sql(c.query) df.noop() - fallbackIsFullyComet = - findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)).isEmpty - } - dispatchNonComet.foreach(op => - warn( - benchmark, - "WARNING: the codegen-dispatch plan is not fully Comet native (first " + - s"non-Comet operator: $op), so that case is partly measuring Spark.")) - if (!dispatcherRan) { - warn( - benchmark, - "WARNING: the codegen dispatcher did not run for this query, so the two " + - "Comet cases below are measuring the same plan.") - } - if (fallbackIsFullyComet) { - warn( - benchmark, - "WARNING: the dispatch-off plan is fully Comet native, so this case is " + - "not exercising the operator fallback it is meant to be compared against.") + assert( + findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)).nonEmpty, + "dispatch-off case did not fall back to Spark") + assert(CometScalaUDFCodegen.stats().totalLookups == 0, "dispatcher ran while disabled") } } @@ -350,26 +277,18 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { } /** - * Every column varies per row, so no argument to a benchmarked expression is loop-invariant and - * neither engine can hoist the call out of the row loop. `c_wide_dec` is `id + 10^19` as - * `DECIMAL(38,10)` so the unscaled integer exceeds 64 bits (Spark hashes via `BigDecimal`). + * Inputs vary per row, including nulls and empty payloads, so neither engine can hoist the + * expression out of the row loop. */ private def corpusQuery: String = { val columns = Seq( - "c_str" -> "REPEAT(CAST(id AS STRING), 4)", + "c_str" -> ("CASE WHEN pmod(id, 31) = 0 THEN NULL WHEN pmod(id, 31) = 1 THEN '' " + + "ELSE REPEAT(CAST(id AS STRING), CAST(pmod(id, 16) AS INT) + 1) END"), "c_int" -> "CAST(id AS INT)", - "c_bits" -> "element_at(array(0, 224, 256, 384, 512), CAST(pmod(id, 5) AS INT) + 1)", - "c_wide_dec" -> ("CAST(id AS DECIMAL(38, 10)) + " + - "CAST('10000000000000000000.0000000000' AS DECIMAL(38, 10))")) + "c_bits" -> "element_at(array(0, 224, 256, 384, 512, 128, NULL), CAST(pmod(id, 7) AS INT) + 1)") s"SELECT ${columns.map { case (name, expr) => s"$expr AS $name" }.mkString(", ")} FROM $tbl" } - /** Writes a warning to the results file as well as the console, ordered against the table. */ - private def warn(benchmark: Benchmark, message: String): Unit = { - val border = "=" * 80 - benchmark.out.println(s"\n$border\n$message\n$border") - } - /** [[Benchmark]] tees console and results file; this benchmark's own tables need the same. */ private def emit(line: String): Unit = { // scalastyle:off println diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala index 6450b822fe2..e07e62f11c9 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala @@ -84,7 +84,7 @@ object CometHashExpressionBenchmark extends CometBenchmarkBase { } // Run only these cases with --wide-decimal. The fallback arm is the path used by - // wide-decimal hash on main before native support; codegen dispatch is tracked in #5835. + // wide-decimal hash on main before native support. private def runWideDecimalBenchmarks(values: Int): Unit = { runBenchmarkWithTable("Wide decimal hash", values) { rows => withTempPath { dir => From 057884204064848b7247c81ef1ae6abd757b208c Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Fri, 23 Oct 2026 16:52:37 +0800 Subject: [PATCH 5/7] fix: explicitly widen benchmark row counts to Long --- .../sql/benchmark/CometHashCodegenDispatchBenchmark.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala index 97c1c2d1cad..2dff9d0d044 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala @@ -220,7 +220,7 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { private def runSteadyState(c: DispatchCase, rows: Int): Unit = { runBenchmark(s"${c.name} -- $rows rows") { - val benchmark = new Benchmark(s"${c.name} -- $rows rows", rows, output = output) + val benchmark = new Benchmark(s"${c.name} -- $rows rows", rows.toLong, output = output) checkPlans(c) // The dispatch-off arm goes first so the `Relative` column reads as the speedup this // change buys over falling the enclosing operator back to Spark. @@ -269,7 +269,7 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { private def withCorpus(rows: Int)(f: => Unit): Unit = { withTempPath { dir => withTempTable(tbl, "parquetV1Table") { - spark.range(rows).createOrReplaceTempView(tbl) + spark.range(rows.toLong).createOrReplaceTempView(tbl) prepareTable(dir, spark.sql(corpusQuery)) f } From b7d27cc13a4f8a544608293594e73859a8ce5f8c Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Sat, 24 Oct 2026 05:33:43 +0800 Subject: [PATCH 6/7] Support non-foldable SHA-2 arguments natively --- .../expression-audits/hash_funcs.md | 2 +- docs/source/user-guide/latest/expressions.md | 2 +- .../scala/org/apache/comet/serde/hash.scala | 20 +--- .../comet/CometHashExpressionSuite.scala | 59 +++------- .../CometHashCodegenDispatchBenchmark.scala | 111 ++++-------------- 5 files changed, 48 insertions(+), 146 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index 5a6dbc9ccaa..d6eddffeee5 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -61,7 +61,7 @@ ## sha2 - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha2(left, right) extends BinaryExpression`; `inputTypes = Seq(BinaryType, IntegerType) -> StringType`. The `numBits` argument selects SHA-224/256/384/512 (0 is treated as 256); other values return NULL. Comet routes via `CometSha2` to the native `sha2` UDF when `numBits` is foldable. A non-foldable `numBits` has no native path; with the JVM codegen dispatcher enabled (the default), `CodegenDispatchFallback` runs Spark's `Sha2.doGenCode` inside the Comet pipeline. The projection falls back to Spark only when the dispatcher is disabled or refuses the tree. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Sha2(left, right) extends BinaryExpression`; `inputTypes = Seq(BinaryType, IntegerType) -> StringType`. The `numBits` argument selects SHA-224/256/384/512 (0 is treated as 256); other values return NULL. Comet routes via `CometSha2` to the native `SparkSha2` UDF for both foldable and column `numBits`. Tests cover STRING and BINARY payload columns with row-varying bit lengths, including invalid values and NULL, with the JVM codegen dispatcher enabled and disabled. - Spark 4.0.1 (audited 2026-05-27): trait set gains `DefaultStringProducingExpression` and the `nullIntolerant: Boolean` refactor. Runtime unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f5c2ba9c558..36d3ddafe39 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -362,7 +362,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | `md5` | ✅ | Native | | | `sha` | ✅ | Native | Alias of `sha1` | | `sha1` | ✅ | Native | | -| `sha2` | ✅ | Hybrid | Non-foldable `numBits` routes through the JVM codegen dispatcher ([audit](../../contributor-guide/expression-audits/hash_funcs.md#sha2)) | +| `sha2` | ✅ | Native | Supports literal and column `numBits` ([audit](../../contributor-guide/expression-audits/hash_funcs.md#sha2)) | | `xxhash64` | ✅ | Native | Includes decimals with precision 19–38 and nested decimals | --- diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index c438305e945..f0a0209e08d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -72,24 +72,12 @@ object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] { } } -// Non-foldable numBits uses Spark codegen inside the Comet pipeline; foldable values stay native. -object CometSha2 extends CometExpressionSerde[Sha2] with CodegenDispatchFallback { +object CometSha2 extends CometExpressionSerde[Sha2] { - override def hasConditionalNativeDefault: Boolean = true - - private val nonFoldableNumBitsReason = - "The `numBits` argument must be a foldable literal value" - - override def getUnsupportedReasons(): Seq[String] = - Seq(nonFoldableNumBitsReason) + override def getUnsupportedReasons(): Seq[String] = HashUtils.unsupportedReasons - override def getSupportLevel(expr: Sha2): SupportLevel = { - if (!expr.right.foldable) { - Unsupported(Some(nonFoldableNumBitsReason)) - } else { - HashUtils.supportLevelForChildren(expr) - } - } + override def getSupportLevel(expr: Sha2): SupportLevel = + HashUtils.supportLevelForChildren(expr) override def convert( expr: Sha2, diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index c0ae85368d9..7592dae12c9 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -32,13 +32,10 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen * Test suite for Spark `hash` (murmur3) and `xxhash64` compatibility between Spark and Comet. * * Native kernels are asserted for supported input shapes, including wide and nested decimals. - * `sha2` with a non-foldable `numBits` stays in the Comet pipeline via the JVM codegen - * dispatcher. `TimeType` hash inputs fall back to Spark. + * `sha2` supports both literal and column `numBits` natively. `TimeType` hash inputs fall back to + * Spark. */ -class CometHashExpressionSuite - extends CometTestBase - with AdaptiveSparkPlanHelper - with CometCodegenAssertions { +class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { private def checkNativeHashes(query: String): Unit = { withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { @@ -134,28 +131,22 @@ class CometHashExpressionSuite } } - test("sha2 - non-foldable numBits routes through the codegen dispatcher") { - withTable("t") { - sql("CREATE TABLE t(payload STRING, num_bits INT) USING parquet") - sql("""INSERT INTO t VALUES - ('hello', 0), - ('hello', 224), - ('hello', 256), - ('hello', 384), - ('hello', 512), - ('hello', 128), - ('hello', -1), - ('', 256), - ('', 512), - (NULL, 256), - ('hello', NULL)""") - withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { - Seq( - "SELECT sha2(payload, num_bits) FROM t", - "SELECT sha2(payload, num_bits + 0) FROM t", - "SELECT sha2('hello', num_bits) FROM t").foreach { query => - assertCodegenRan { - checkSparkAnswerAndOperator(query) + Seq("STRING", "BINARY").foreach { payloadType => + test(s"sha2 - non-foldable numBits stays native with $payloadType payload") { + withTable("t") { + sql(s"CREATE TABLE t(payload $payloadType, num_bits INT) USING parquet") + val payload = if (payloadType == "BINARY") "X'0080FF68656C6C6F'" else "'hello'" + val empty = if (payloadType == "BINARY") "X''" else "''" + val values = Seq("0", "224", "256", "384", "512", "128", "-1", "NULL") + .map(bits => s"($payload, $bits)") ++ + Seq(s"($empty, 256)", s"($empty, 512)", "(NULL, 256)") + sql(s"INSERT INTO t VALUES ${values.mkString(",")}") + Seq("true", "false").foreach { enabled => + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> enabled) { + Seq( + "SELECT sha2(payload, num_bits) FROM t", + "SELECT sha2(payload, num_bits + 0) FROM t", + s"SELECT sha2($payload, num_bits) FROM t").foreach(assertNativeHash) } } } @@ -176,18 +167,6 @@ class CometHashExpressionSuite } } - test("sha2 - non-foldable numBits falls back when codegen dispatcher is disabled") { - withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { - withTable("t") { - sql("CREATE TABLE t(payload STRING, num_bits INT) USING parquet") - sql("INSERT INTO t VALUES ('hello', 256), (NULL, 256), ('hello', NULL)") - checkSparkAnswerAndFallbackReason( - "SELECT sha2(payload, num_bits) FROM t", - s"sha2: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") - } - } - } - test("hash - array of integers") { withTable("t") { sql("CREATE TABLE t(c ARRAY) USING parquet") diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala index 2dff9d0d044..97a9d43cc24 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashCodegenDispatchBenchmark.scala @@ -30,20 +30,10 @@ import org.apache.comet.CometConf import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** - * Benchmark of non-foldable `sha2(payload, numBits)` through `CodegenDispatchFallback`. + * Benchmark of native non-foldable `sha2(payload, numBits)` against Spark. * - * The dispatcher runs Spark's `doGenCode` inside the Comet pipeline. With dispatch disabled, the - * projection falls back to Spark. The mixed projection also measures native expressions alongside - * `sha2`; it does not imply that SHA-2 itself executes natively. - * - * A pure Spark case is included as a third reference point. A `dispatch off (repeat)` case - * repeats the baseline at the end of every table so the spread between the two is this machine's - * noise floor. - * - * Compilation is a one-time cost, and the steady-state tables warm up before timing, so it does - * not appear there. The `first use` table at the top measures it directly: the same query run - * twice back to back on a cold kernel, after the dispatcher itself has been warmed on an - * unrelated expression. + * Both Comet arms execute SHA-2 natively, with the codegen dispatcher enabled and disabled. The + * mixed projection also measures native expressions alongside `sha2`. * * Before timing anything, every case is run through all three arms and the rows are compared, so * a timing cannot come from an arm that computed something else. @@ -82,12 +72,12 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { private def sparkConfigs: Seq[(String, String)] = Seq(noConstantFolding, CometConf.COMET_ENABLED.key -> "false") - /** Comet on, dispatcher on: the behaviour this benchmark is validating. */ + /** Native execution with the dispatcher enabled. */ private def dispatchConfigs: Seq[(String, String)] = cometConfigs(dispatch = true) - /** Comet on, dispatcher off: the enclosing projection falls back to Spark. */ - private def fallbackConfigs: Seq[(String, String)] = + /** Native execution with the dispatcher disabled. */ + private def nativeConfigs: Seq[(String, String)] = cometConfigs(dispatch = false) private def cometConfigs(dispatch: Boolean): Seq[(String, String)] = Seq( @@ -96,24 +86,17 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) - private val DispatchCaseName = "Comet, codegen dispatch" - private val FallbackCaseName = "Comet, dispatch off (Spark fallback)" + private val DispatchCaseName = "Comet native, dispatch on" + private val NativeCaseName = "Comet native, dispatch off" private val SparkCaseName = "Spark (Comet disabled)" override def runCometBenchmark(mainArgs: Array[String]): Unit = { val selected = cases - runBenchmark("SHA-2 codegen dispatch: environment") { + runBenchmark("SHA-2 native: environment") { emitEnvironment(selected) } - // The kernels must be uncompiled when `runFirstUse` starts, so it runs before anything else - // executes these queries: `CodeGenerator.compile` caches on the generated source JVM-wide, - // and a single earlier execution would make the "first use" number a cache hit. Only the - // isolated case is measured here: the mixed projection reuses the same SHA-2 kernel. withCorpus(SmallRows) { - runBenchmark(s"SHA-2 codegen dispatch: first use, $SmallRows rows") { - runFirstUse(selected.take(1)) - } selected.foreach(verifyArmsAgree(_, SmallRows)) selected.foreach(runSteadyState(_, SmallRows)) } @@ -143,42 +126,10 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { emit(" average of the timed iterations.") emit(s"Row counts: $SmallRows (sub-batch) and $LargeRows (multi-batch).") emit("Payloads include nulls, empty strings and varying lengths; numBits varies per row.") - emit("Both Comet arms use Spark codegen for SHA-2; dispatch keeps the projection in Comet.") + emit("Both Comet arms execute SHA-2 natively without using the dispatcher.") emit(s"Cases: ${selected.map(_.name).mkString(", ")}") } - /** - * Cost of the first execution of a query whose kernel has never been compiled, against the - * second execution of the same query. - * - * The dispatcher is warmed on an expression that is not in the case list first, so that Janino, - * `CodeGenerator`, the Arrow bridge and the FFI boundary are all loaded and JIT-warm before the - * first case runs and the difference is dominated by compiling that one kernel. It is still an - * upper bound on the compile: the first case in the list absorbs whatever class loading the - * warmup missed. It is a query-level measurement, not an isolated compiler measurement. - */ - private def runFirstUse(selected: Seq[DispatchCase]): Unit = { - // `rlike` routes through the same dispatcher and is not one of the cases below. - val warmup = DispatchCase("warmup", "select c_str rlike '[0-9]+' from parquetV1Table") - (0 until 5).foreach(_ => runQuery(warmup.query, dispatchConfigs)) - - emit( - f"${"case"}%-24s ${"1st run (ms)"}%14s ${"2nd run (ms)"}%14s " + - f"${"one-time (ms)"}%14s ${"compiles"}%9s") - emit("-" * 84) - selected.foreach { c => - CometScalaUDFCodegen.resetStats() - val cold = timeMillis(runQuery(c.query, dispatchConfigs)) - val compiles = CometScalaUDFCodegen.stats().compileCount - val warm = timeMillis(runQuery(c.query, dispatchConfigs)) - emit(f"${c.name}%-24s $cold%14.1f $warm%14.1f ${cold - warm}%14.1f $compiles%9d") - } - emit("") - emit("`compiles` counts dispatcher cache misses, which is one per task. The Janino work") - emit("itself is deduplicated JVM-wide by Spark's CodeGenerator source cache, so only the") - emit("first task to reach a given kernel source pays for it.") - } - /** * Fails if the three arms disagree on `query`. Rows are compared as a sorted multiset because * output order is not guaranteed across the different plans. @@ -203,7 +154,7 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { } val expected = collect(sparkConfigs) - Seq(DispatchCaseName -> dispatchConfigs, FallbackCaseName -> fallbackConfigs).foreach { + Seq(DispatchCaseName -> dispatchConfigs, NativeCaseName -> nativeConfigs).foreach { case (armName, configs) => val actual = collect(configs) assert( @@ -222,35 +173,25 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { runBenchmark(s"${c.name} -- $rows rows") { val benchmark = new Benchmark(s"${c.name} -- $rows rows", rows.toLong, output = output) checkPlans(c) - // The dispatch-off arm goes first so the `Relative` column reads as the speedup this - // change buys over falling the enclosing operator back to Spark. - benchmark.addCase(FallbackCaseName)(_ => runQuery(c.query, fallbackConfigs)) - benchmark.addCase(DispatchCaseName)(_ => runQuery(c.query, dispatchConfigs)) + // Spark goes first so `Relative` reports the native speedup. benchmark.addCase(SparkCaseName)(_ => runQuery(c.query, sparkConfigs)) - benchmark.addCase(s"$FallbackCaseName (repeat)")(_ => runQuery(c.query, fallbackConfigs)) + benchmark.addCase(NativeCaseName)(_ => runQuery(c.query, nativeConfigs)) + benchmark.addCase(DispatchCaseName)(_ => runQuery(c.query, dispatchConfigs)) benchmark.run() } } /** Refuses to time a case whose actual routing does not match its label. */ private def checkPlans(c: DispatchCase): Unit = { - withSQLConf(dispatchConfigs: _*) { - CometScalaUDFCodegen.resetStats() - val df = spark.sql(c.query) - df.noop() - val stats = CometScalaUDFCodegen.stats() - assert(stats.compileCount + stats.cacheHitCount > 0, s"dispatcher did not run: $stats") - val nonComet = findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)) - assert(nonComet.isEmpty, s"dispatch projection fell back to Spark: $nonComet") - } - withSQLConf(fallbackConfigs: _*) { - CometScalaUDFCodegen.resetStats() - val df = spark.sql(c.query) - df.noop() - assert( - findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)).nonEmpty, - "dispatch-off case did not fall back to Spark") - assert(CometScalaUDFCodegen.stats().totalLookups == 0, "dispatcher ran while disabled") + Seq(dispatchConfigs, nativeConfigs).foreach { configs => + withSQLConf(configs: _*) { + CometScalaUDFCodegen.resetStats() + val df = spark.sql(c.query) + df.noop() + val nonComet = findFirstNonCometOperator(stripAQEPlan(df.queryExecution.executedPlan)) + assert(nonComet.isEmpty, s"native projection fell back to Spark: $nonComet") + assert(CometScalaUDFCodegen.stats().totalLookups == 0, "native SHA-2 used the dispatcher") + } } } @@ -259,12 +200,6 @@ object CometHashCodegenDispatchBenchmark extends CometBenchmarkBase { spark.sql(query).noop() } - private def timeMillis(f: => Unit): Double = { - val start = System.nanoTime() - f - (System.nanoTime() - start) / 1e6 - } - /** Builds `parquetV1Table` with `rows` rows of the corpus and drops it afterwards. */ private def withCorpus(rows: Int)(f: => Unit): Unit = { withTempPath { dir => From e13159a3652fbfa8e01db337d636199a734a1dd6 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Sat, 24 Oct 2026 05:57:34 +0800 Subject: [PATCH 7/7] fix: preserve SHA-2 null short-circuit semantics --- .../scala/org/apache/comet/serde/hash.scala | 6 ++++-- .../comet/CometHashExpressionSuite.scala | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index f0a0209e08d..a0227a8a241 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64} import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructType} -import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} +import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType, withNullShortCircuit} object CometXxHash64 extends CometExpressionSerde[XxHash64] { @@ -85,7 +85,9 @@ object CometSha2 extends CometExpressionSerde[Sha2] { binding: Boolean): Option[ExprOuterClass.Expr] = { val leftExpr = exprToProtoInternal(expr.left, inputs, binding) val numBitsExpr = exprToProtoInternal(expr.right, inputs, binding) - scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr) + withNullShortCircuit( + expr, + scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr)) } } diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index 7592dae12c9..fc4d868d8d2 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -153,6 +153,25 @@ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelpe } } + Seq("STRING", "BINARY").foreach { payloadType => + test(s"sha2 - NULL $payloadType payload short-circuits ANSI numBits cast") { + withSQLConf("spark.sql.ansi.enabled" -> "true") { + withTable("t") { + sql(s"CREATE TABLE t(payload $payloadType, bits STRING) USING parquet") + val payload = if (payloadType == "BINARY") "X'0080FF'" else "'hello'" + sql( + s"INSERT INTO t VALUES (NULL, 'invalid'), ($payload, '256'), " + + s"($payload, NULL), (NULL, '256')") + Seq("true", "false").foreach { enabled => + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> enabled) { + assertNativeHash("SELECT sha2(payload, CAST(bits AS INT)) FROM t") + } + } + } + } + } + } + test("sha2 - literal numBits stays native") { withTable("t") { sql("CREATE TABLE t(payload STRING) USING parquet")