From e819c2f98c1983ec3180d8027791cc23ec420b63 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 2 Sep 2026 11:31:36 -0700 Subject: [PATCH] feat: support nested types as native shuffle hash partitioning keys `CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct, array and map partitioning keys, so any query repartitioning on a nested column fell back to Spark for the whole shuffle. The comment said "Native code does not support hashing complex types, see hash_funcs/utils.rs", but that file hashes nested types recursively (struct fields, list elements, map keys and values), and shuffle partitioning shares that kernel, and Spark's seed, with the `hash` expression. Adds the recursive struct/array/map cases to the gate, behind `spark.comet.shuffle.native.partitioning.hash.nested.enabled` (default true). Nesting is checked recursively through the same predicate, so a leaf type that cannot be hashed natively disqualifies the whole key and the shuffle still falls back: - collated strings, which Comet hashes as raw bytes (see #1947 / #4035, where rows equal under the collation reached different partitions and a downstream collation-aware DISTINCT produced a wrong answer) - CalendarInterval, which the native hasher has no branch for (#5059) Map keys are additionally restricted to Spark 4.0+. Map entry order is not semantically meaningful, so two equal maps must hash alike, and Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`. Earlier versions insert no such normalization, so Comet would hash physical entry order. When the `mapsort` itself is not convertible -- CometMapSort supports scalar map keys only -- the existing expression check fails and the shuffle falls back. The config defaults to false. The native hasher only vectorizes nested shapes whose leaves are primitives; `array>` and a map inside a struct fall through to a per-element path that re-enters `create_murmur3_hashes` for every element, so enabling this by default before measuring could make these shuffles slower than letting Spark do them. `CometFuzzTestSuite`'s "distribute by single column (complex types)" keeps its existing expectation, since the keys still fall back by default, and additionally asserts that they are admitted with the config enabled. Also corrects the comment in `CometFuzzTestBase` claiming that file has no nested complex types -- `generateSchema` does add `struct>` and `array>` when both array and struct generation are on; maps are the only thing missing. Also covers the repartition-on-map route into the sliced-map `mapsort` defect fixed in #5630: a native OFFSET below a map shuffle key, asserting that both the exchange and the offset stay native so the sliced map actually reaches the native mapsort. Co-authored-by: Claude Code --- .../scala/org/apache/comet/CometConf.scala | 16 + .../shuffle/CometShuffleExchangeExec.scala | 31 +- .../org/apache/comet/CometFuzzTestBase.scala | 5 +- .../org/apache/comet/CometFuzzTestSuite.scala | 13 +- .../comet/exec/CometNativeShuffleSuite.scala | 325 +++++++++++++++++- 5 files changed, 382 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a7d7b80db1f..3db0bf4e7ed 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -404,6 +404,22 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.shuffle.native.partitioning.hash.nested.enabled") + .category(CATEGORY_SHUFFLE) + .doc( + "Whether to allow nested types (struct, array, map) as hash partitioning keys in " + + "Comet native shuffle. Comet's native Murmur3 kernel hashes nested types " + + "recursively and shares that kernel, and Spark's seed, with the `hash` expression, " + + "so partition assignment matches Spark. A map key additionally requires the " + + "`mapsort` normalization that Spark 4.0 and later insert, so maps are rejected on " + + "earlier versions. Disabled by default until the performance of the nested hashing " + + "paths has been measured: shapes whose leaves are not primitives, such as " + + "`array>`, fall back to a per-element code path in the native hasher " + + "rather than a vectorized one.") + .booleanConf + .createWithDefault(false) + val COMET_SHUFFLE_NATIVE_RANGE_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.range.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.range.enabled") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 46fe621e481..5d5010f3f6f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -51,7 +51,7 @@ import com.google.common.base.Objects import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometConf.{COMET_SHUFFLE_ENABLED, COMET_SHUFFLE_MODE} -import org.apache.comet.CometSparkSessionExtensions.{cometCelebornShuffleFallbackReason, hasFallbackReason, isCometCelebornShuffleManagerEnabled, isCometShuffleManagerEnabled, withFallbackReasons} +import org.apache.comet.CometSparkSessionExtensions.{cometCelebornShuffleFallbackReason, hasFallbackReason, isCometCelebornShuffleManagerEnabled, isCometShuffleManagerEnabled, isSpark40Plus, withFallbackReasons} import org.apache.comet.serde.{Compatible, OperatorOuterClass, QueryPlanSerde, SupportLevel, Unsupported} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.{CometTypeShim, ShimCometShuffleExchangeExec} @@ -401,12 +401,21 @@ object CometShuffleExchangeExec private def nativeShuffleFailureReasons(s: ShuffleExchangeExec): Seq[String] = { val conf = SQLConf.get + val nestedHashPartitioningEnabled = + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.get(conf) + /** * Determine which data types are supported as partition columns in native shuffle. * * For HashPartitioning this defines the key that determines how data should be collocated for - * operations like `groupByKey`, `reduceByKey`, or `join`. Native code does not support - * hashing complex types, see hash_funcs/utils.rs + * operations like `groupByKey`, `reduceByKey`, or `join`. + * + * Nested types (struct/array/map) are supported when + * `spark.comet.shuffle.native.partitioning.hash.nested.enabled` is enabled: the native + * Murmur3 kernel in hash_funcs/utils.rs hashes them recursively. Nesting is checked + * recursively, so a leaf type that cannot be hashed natively -- a collated string, or an + * interval the hasher has no branch for -- disqualifies the whole key and the shuffle falls + * back to Spark. */ def supportedHashPartitioningDataType(dt: DataType): Boolean = dt match { // Collated strings require collation-aware hashing; Comet only hashes raw bytes, @@ -424,6 +433,22 @@ object CometShuffleExchangeExec true case dt if isTimeType(dt) => true + case StructType(fields) if nestedHashPartitioningEnabled => + // `fields.nonEmpty` mirrors the guard on the data-column gate below. An empty struct is + // not reachable end-to-end anyway: Parquet cannot store an empty group, and an in-memory + // relation with one does not survive scan conversion. + fields.nonEmpty && fields.forall(f => supportedHashPartitioningDataType(f.dataType)) + case ArrayType(elementType, _) if nestedHashPartitioningEnabled => + supportedHashPartitioningDataType(elementType) + case MapType(keyType, valueType, _) if nestedHashPartitioningEnabled => + // Map entry order is not semantically meaningful, so two equal maps must hash alike. + // Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`, which is + // gated separately by CometMapSort (scalar map keys only) and, when unsupported, fails + // the expression check below. Earlier Spark versions insert no such normalization, so + // Comet would hash physical entry order and could route equal maps differently. + isSpark40Plus && + supportedHashPartitioningDataType(keyType) && + supportedHashPartitioningDataType(valueType) case _ => false } diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala index 6b6f02f9fd6..b112a11d548 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala @@ -62,8 +62,9 @@ class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { // override base date due to known issues with experimental scans baseDate = new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime) - // generate Parquet file with primitives, structs, and arrays, but no maps - // and no nested complex types + // generate Parquet file with primitives, structs, and arrays, but no maps. Note that + // `generateSchema` does add `struct>` and `array>` when both + // `generateArray` and `generateStruct` are set, so nested complex types are present. filename = s"$tempDir/CometFuzzTestSuite_${System.currentTimeMillis()}.parquet" withSQLConf( CometConf.COMET_ENABLED.key -> "false", diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala index 1d7ddfdd5db..c84db0396b0 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala @@ -175,10 +175,21 @@ class CometFuzzTestSuite extends CometFuzzTestBase { case "jvm" => 1 case "native" => - // native shuffle does not support complex types as partitioning keys + // Nested hash partitioning keys are off by default, so native shuffle falls back here. 0 } assert(cometShuffleExchanges.length == expectedNumCometShuffles) + + // With the config enabled these keys do run through native shuffle. This is the widest + // nested-type coverage in the repo, so it is worth asserting that they are admitted rather + // than only that they fall back. + withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true") { + val enabledDf = spark.sql(sql) + enabledDf.collect() + val enabledPlan = + enabledDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan + assert(collectCometShuffleExchanges(enabledPlan).length == 1) + } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 10cdbc7932e..d6e2656d04b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -35,7 +35,7 @@ import org.apache.arrow.vector.types.pojo.{Field, Schema} import org.apache.hadoop.fs.Path import org.apache.spark.SparkEnv import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} -import org.apache.spark.sql.comet.{CometExec, CometMetricNode} +import org.apache.spark.sql.comet.{CometExec, CometMetricNode, CometTakeOrderedAndProjectExec} import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -44,6 +44,7 @@ import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, Native} +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass} import org.apache.comet.shuffle.ShufflePartitionPusher @@ -62,6 +63,13 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper import testImplicits._ + /** + * Runs `f` with nested hash partitioning keys enabled. The config is disabled by default, so + * every test that expects a nested key to reach native shuffle has to opt in. + */ + private def withNestedHashPartitioning(f: => Unit): Unit = + withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true")(f) + private def rssShufflePlanBytes: Array[Byte] = { val scan = OperatorOuterClass.Operator .newBuilder() @@ -444,8 +452,14 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper .filter($"_3" > 10) .repartition(numPartitions, $"_2") - // Partitioning on nested array falls back to Spark + // Partitioning on a nested array falls back to Spark by default, and is native + // when nested hash partitioning is enabled. checkShuffleAnswer(df, 0) + withNestedHashPartitioning { + checkShuffleAnswer( + sql("SELECT * FROM tbl").filter($"_3" > 10).repartition(numPartitions, $"_2"), + 1) + } df = sql("SELECT * FROM tbl") .filter($"_3" > 10) @@ -672,6 +686,313 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("native shuffle on struct hash partitioning key") { + withNestedHashPartitioning { + Seq(10, 201).foreach { numPartitions => + withParquetTable((0 until 50).map(i => (i, (i % 7, (i % 5).toString))), "tbl") { + val df = sql("SELECT * FROM tbl") + .repartition(numPartitions, $"_2") + .sortWithinPartitions($"_1") + + checkShuffleAnswer(df, 1) + } + } + } + } + + test("native shuffle on array hash partitioning key") { + withNestedHashPartitioning { + Seq(10, 201).foreach { numPartitions => + withParquetTable((0 until 50).map(i => (i, Seq(i % 7, i % 5))), "tbl") { + val df = sql("SELECT * FROM tbl") + .repartition(numPartitions, $"_2") + .sortWithinPartitions($"_1") + + checkShuffleAnswer(df, 1) + } + } + } + } + + test("native shuffle on two-level nested hash partitioning key") { + withNestedHashPartitioning { + // struct, string> and array>: one level of nesting inside the + // top-level type, covering both recursive branches of the type gate. + Seq(10, 201).foreach { numPartitions => + withParquetTable( + (0 until 50).map(i => (i, (Seq(i % 7, i % 3), (i % 5).toString), Seq((i % 4, "x")))), + "tbl") { + val df = sql("SELECT * FROM tbl") + .repartition(numPartitions, $"_2", $"_3") + .sortWithinPartitions($"_1") + + checkShuffleAnswer(df, 1) + } + } + } + } + + test("native shuffle on deeply nested hash partitioning key") { + withNestedHashPartitioning { + // Four levels of nesting, mixing all three recursive branches: + // struct< array< struct< m: map>, s: string > >, i: int > + // so the gate and the native hasher both have to descend struct -> array -> struct -> map + // -> array -> int. + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withTable("tbl") { + sql("""CREATE TABLE tbl( + id INT, + k STRUCT< + a: ARRAY>, s: STRING>>, + i: INT>) + USING parquet""") + sql("""INSERT INTO tbl VALUES + (1, named_struct('a', array(named_struct('m', map('x', array(1, 2)), 's', 'p')), 'i', 1)), + (2, named_struct('a', array(named_struct('m', map('y', array(3)), 's', 'q')), 'i', 2)), + (3, named_struct('a', array(named_struct('m', map('x', array(1, 2)), 's', 'p')), 'i', 1)), + (4, named_struct('a', array(), 'i', 4)), + (5, null)""") + val df = sql("SELECT * FROM tbl").repartition(10, $"k").sortWithinPartitions($"id") + + checkShuffleAnswer(df, 1) + } + } + } + + test("native shuffle on map hash partitioning key below an offset") { + // A native OFFSET slices the batch and Arrow keeps a sliced MapArray's original entry offsets, + // so the `mapsort` Spark inserts for a map shuffle key sees a map whose first entry offset is + // nonzero. That overran the sorted entries until #5630; this is the repartition-on-map route + // into it, alongside the group-by route covered in CometMapExpressionSuite. + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withNestedHashPartitioning { + withParquetTable( + (0 until 20).map(i => (i, Map(s"b${i % 5}" -> i, s"a${i % 5}" -> (i + 1)))), + "tbl") { + val df = + sql("SELECT * FROM (SELECT * FROM tbl ORDER BY _1 LIMIT 15 OFFSET 5) DISTRIBUTE BY _2") + + // Both operators have to stay native, otherwise the sliced map never reaches the native + // mapsort and this stops covering the fix. + checkCometExchange(df, 1, true) + assert( + collectFirst(stripAQEPlan(df.queryExecution.executedPlan)) { + case t: CometTakeOrderedAndProjectExec => t + }.isDefined, + "expected a native offset below the shuffle") + checkSparkAnswer(df) + } + } + } + test("native shuffle on map hash partitioning key") { + withNestedHashPartitioning { + // Map entry order carries no meaning, so equal maps must hash alike. Spark 4.0+ normalizes a + // map shuffle key with `mapsort(...)`; earlier versions do not, so Comet must not hash a raw + // map there. The gate therefore only admits map keys on Spark 4.0+, and only when the + // `mapsort` itself is convertible (CometMapSort supports scalar map keys only). + withParquetTable((0 until 50).map(i => (i, Map(i % 7 -> (i % 5)))), "tbl") { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + + checkShuffleAnswer(df, if (isSpark40Plus) 1 else 0) + } + } + } + + test("native shuffle on map hash partitioning key with non-scalar map key falls back") { + // A map whose own key is nested cannot be `mapsort`ed by Comet (Arrow's sort_to_indices + // handles scalar keys only), so the normalization Spark 4.0+ requires is unavailable and the + // shuffle must fall back rather than hash an unnormalized map. + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + + checkShuffleAnswer(df, 0) + } + } + + test("native shuffle on struct hash partitioning key with collated string falls back") { + // The top-level gate rejects collated strings because Comet hashes raw bytes, which would + // misroute rows that are equal under the collation. Recursing through the nested cases must + // preserve that: a collated leaf disqualifies the whole key. + assume(isSpark40Plus, "string collation requires Spark 4.0+") + withTable("tbl") { + sql( + "CREATE TABLE tbl(id INT, s STRUCT) USING parquet") + sql("INSERT INTO tbl VALUES (1, named_struct('a', 'x', 'b', 1))") + sql("INSERT INTO tbl VALUES (2, named_struct('a', 'X', 'b', 2))") + val df = sql("SELECT * FROM tbl").repartition(10, $"s") + + checkShuffleAnswer(df, 0) + } + } + + test("native shuffle nested hash partitioning key honors its config") { + withParquetTable((0 until 50).map(i => (i, (i % 7, (i % 5).toString))), "tbl") { + // "false" is the default; "true" opts in. + Seq("true" -> 1, "false" -> 0).foreach { case (enabled, expectedShuffles) => + withSQLConf( + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> enabled) { + val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1") + + checkShuffleAnswer(df, expectedShuffles) + } + } + } + } + test("native shuffle nested hash partitioning key matches Spark's partition assignment") { + withNestedHashPartitioning { + // checkShuffleAnswer only compares the query answer, which is order-insensitive and so would + // pass even if Comet routed rows to different partitions than Spark. Nested keys are only safe + // if partition ASSIGNMENT matches, so compare spark_partition_id() per row against Spark. + withParquetTable( + (0 until 200).map(i => (i, (i % 13, (i % 7).toString), Seq(i % 11, i % 5))), + "tbl") { + Seq("_2", "_3", "_2, _3").foreach { keys => + val query = + "SELECT _1, spark_partition_id() AS pid FROM (" + + s"SELECT /*+ REPARTITION(10, $keys) */ * FROM tbl)" + val cometRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).sorted + // `withSQLConf` returns Unit, so capture the Spark-side rows via a var rather than + // relying on the block's value. + // `SQLHelper.withSQLConf` returns T on Spark 4.x but Unit on Spark 3.x, so capture the + // Spark-side rows via a var to keep this compiling on both. + var sparkRows: Array[(Int, Int)] = Array.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).sorted + } + assert(sparkRows.nonEmpty, "Spark produced no rows; the comparison would be vacuous") + // Without this the test would still pass if the gate stopped admitting these keys, having + // quietly compared Spark against Spark. + checkCometExchange(sql(s"SELECT /*+ REPARTITION(10, $keys) */ * FROM tbl"), 1, true) + assert( + cometRows === sparkRows, + s"partition assignment differs from Spark for keys ($keys)") + } + } + + // Same check for a deeply nested key, where Spark rewrites the partitioning expression into a + // transform(...) containing a nested mapsort(...). + if (isSpark40Plus) { + withTable("deep") { + sql("""CREATE TABLE deep( + id INT, + k STRUCT>, s: STRING>>, i: INT>) + USING parquet""") + (0 until 40).foreach { i => + sql(s"""INSERT INTO deep VALUES + ($i, named_struct('a', array(named_struct( + 'm', map('k${i % 6}', array(${i % 4}, ${i % 3})), 's', 's${i % 5}')), + 'i', ${i % 7}))""") + } + val query = + "SELECT id, spark_partition_id() AS pid FROM (" + + "SELECT /*+ REPARTITION(10, k) */ * FROM deep)" + val cometRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).sorted + var sparkRows: Array[(Int, Int)] = Array.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).sorted + } + assert(sparkRows.nonEmpty, "Spark produced no rows; the comparison would be vacuous") + checkCometExchange(sql("SELECT /*+ REPARTITION(10, k) */ * FROM deep"), 1, true) + assert( + cometRows === sparkRows, + "partition assignment differs from Spark for a deep key") + } + } + } + } + test("native shuffle on nested hash partitioning key with interval leaf falls back") { + // CalendarIntervalType is allowed as a shuffle DATA column but the native hasher has no + // branch for it (https://github.com/apache/datafusion-comet/issues/5059). Because the nested + // cases recurse through this same predicate, an interval leaf disqualifies the whole key and + // the shuffle falls back instead of failing in native code. + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val nested = spark + .sql("SELECT id, named_struct('i', INTERVAL '1' MONTH, 'n', id) AS s " + + "FROM VALUES (1), (2), (3) AS t(id)") + .repartition(2, $"s") + checkShuffleAnswer(nested, 0) + + val topLevel = spark + .sql("SELECT id, INTERVAL '1' MONTH AS i FROM VALUES (1), (2), (3) AS t(id)") + .repartition(2, $"i") + checkShuffleAnswer(topLevel, 0) + } + } + + test("native shuffle map hash partitioning key ignores entry order") { + withNestedHashPartitioning { + // The reason map keys are only admitted on Spark 4.0+ is that two equal maps with different + // physical entry order must hash alike, which relies on Spark's `mapsort` normalization. A + // single-entry map cannot show that, so use multi-entry maps written in opposite key orders + // and assert both rows land in the same partition. + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val query = """ + SELECT id, spark_partition_id() AS pid FROM ( + SELECT /*+ REPARTITION(10, m) */ * FROM VALUES + (1, map('a', 1, 'b', 2)), + (2, map('b', 2, 'a', 1)), + (3, map('a', 1, 'b', 2, 'c', 3)), + (4, map('c', 3, 'b', 2, 'a', 1)) AS t(id, m))""" + val rows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + assert(rows(1) == rows(2), "equal maps in different entry order must share a partition") + assert(rows(3) == rows(4), "equal maps in different entry order must share a partition") + + var sparkRows: Map[Int, Int] = Map.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + } + assert(sparkRows.nonEmpty) + assert(rows == sparkRows, "map key partition assignment differs from Spark") + } + } + } + + test("native shuffle nested hash partitioning key with null keys at each level") { + withNestedHashPartitioning { + // Spark hashes a null struct as the seed. The struct branch in hash_funcs/utils.rs recurses + // into `columns()` without consulting its own null mask (unlike the List and Map branches, + // which guard on `is_null(row_idx)`), so a null struct whose children are not themselves null + // would hash off leftover child values. Cover nulls at the top level, inside a struct, inside + // an array element, and as a map value. + assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+") + withTable("nulls") { + sql("""CREATE TABLE nulls( + id INT, + k STRUCT, s: STRING>>, i: INT>) + USING parquet""") + sql("""INSERT INTO nulls VALUES + (1, NULL), + (2, NULL), + (3, named_struct('a', NULL, 'i', 1)), + (4, named_struct('a', NULL, 'i', 1)), + (5, named_struct('a', array(NULL), 'i', 2)), + (6, named_struct('a', array(NULL), 'i', 2)), + (7, named_struct('a', array(named_struct('m', NULL, 's', NULL)), 'i', 3)), + (8, named_struct('a', array(named_struct('m', NULL, 's', NULL)), 'i', 3)), + (9, named_struct('a', array(named_struct('m', map('k', 1), 's', 'v')), 'i', NULL)), + (10, named_struct('a', array(named_struct('m', map('k', 1), 's', 'v')), 'i', NULL))""") + + val query = + "SELECT id, spark_partition_id() AS pid FROM (" + + "SELECT /*+ REPARTITION(10, k) */ * FROM nulls)" + val rows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + // Each pair shares a key, so each pair must share a partition. + Seq((1, 2), (3, 4), (5, 6), (7, 8), (9, 10)).foreach { case (l, r) => + assert(rows(l) == rows(r), s"rows $l and $r have equal keys but different partitions") + } + + checkCometExchange(sql("SELECT /*+ REPARTITION(10, k) */ * FROM nulls"), 1, true) + var sparkRows: Map[Int, Int] = Map.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = sql(query).collect().map(r => (r.getInt(0), r.getInt(1))).toMap + } + assert(sparkRows.nonEmpty) + assert(rows == sparkRows, "null-key partition assignment differs from Spark") + } + } + } test("fix: Comet native shuffle with binary data") { withParquetTable((0 until 5).map(i => (i, (i + 1).toLong)), "tbl") { val df = sql("SELECT cast(cast(_1 as STRING) as BINARY) as binary, _2 FROM tbl")