andygrove commented on code in PR #5567:
URL: https://github.com/apache/datafusion-comet/pull/5567#discussion_r3895434859
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala:
##########
@@ -672,6 +680,194 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle on struct hash partitioning key") {
+ 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") {
+ 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") {
+ // struct<array<int>, string> and array<struct<int, string>>: 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") {
+ // Four levels of nesting, mixing all three recursive branches:
+ // struct< array< struct< m: map<string, array<int>>, 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<STRUCT<m: MAP<STRING, ARRAY<INT>>, 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") {
Review Comment:
Every map in the new tests has a single entry, so `mapsort` is a no-op and
the property this case is gated on never actually gets exercised. The reason we
only admit maps on 4.0+ is that two equal maps with different physical entry
order have to hash alike, and nothing here would notice if that stopped being
true.
Could you add multi-entry maps written in different key orders, something
like `map('a', 1, 'b', 2)` and `map('b', 2, 'a', 1)`, and assert the two rows
get the same `spark_partition_id()`?
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala:
##########
@@ -672,6 +680,194 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle on struct hash partitioning key") {
+ 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") {
+ 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") {
+ // struct<array<int>, string> and array<struct<int, string>>: 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") {
+ // Four levels of nesting, mixing all three recursive branches:
+ // struct< array< struct< m: map<string, array<int>>, 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<STRUCT<m: MAP<STRING, ARRAY<INT>>, 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") {
+ // 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<a: STRING COLLATE UTF8_LCASE, b:
INT>) 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") {
+ 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") {
Review Comment:
This is the only test that checks routing rather than the answer, so I would
like it to be hard to fool, and right now it never asserts that Comet ran the
shuffle. If the gate ever stops admitting these keys this quietly becomes Spark
compared against Spark and still passes. Can you add a `checkCometExchange(df,
1, true)` alongside the `sparkRows.nonEmpty` check?
The other gap is nulls. None of the 200 shallow rows or the 40 deep rows
have a null key at any level, and that is the case I would most want covered.
The struct branch in `hash_funcs/utils.rs` is the one nested branch that does
not look at its own null mask. `List` and `Map` both guard on
`is_null(row_idx)`, but `DataType::Struct(_)` just takes
`struct_array.columns()` and recurses, while Spark is `case null => seed`. So
the two only agree when the children happen to be null under a null parent.
Parquet gives us that, which is why `CometHashExpressionSuite` is green, but a
struct built in the plan may not, and Spark's own rewrite for a nested map key
is `If(IsNull(k), null, named_struct(...))`, which is exactly that shape.
If a null struct ever hashes off leftover child values then two rows with
the same key can land in different partitions, and that breaks grouping and
joins whatever Spark would have done. Could you add null keys at each level
here?
##########
spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala:
##########
@@ -171,14 +171,9 @@ class CometFuzzTestSuite extends CometFuzzTestBase {
// check for Comet shuffle
val plan =
df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan
val cometShuffleExchanges = collectCometShuffleExchanges(plan)
- val expectedNumCometShuffles = CometConf.COMET_SHUFFLE_MODE.get() match {
- case "jvm" =>
- 1
- case "native" =>
- // native shuffle does not support complex types as partitioning keys
- 0
- }
- assert(cometShuffleExchanges.length == expectedNumCometShuffles)
+ // Both shuffle modes support the struct and array partitioning keys in
this file (it is
+ // generated without maps and without nested complex types), so one
Comet shuffle either way.
Review Comment:
This file does have nested complex types. `FuzzDataGenerator.generateSchema`
adds `StructType(arraysOfPrimitives)` and
`createArrayType(StructType(primitives))` when both `generateArray` and
`generateStruct` are set, so `struct<array<...>>` and `array<struct<...>>` are
both in there. Maps are the only thing missing. The comment in
`CometFuzzTestBase` says the same thing and is equally wrong, so this was
inherited, but could you fix both while you are here?
The knock-on is that this now asserts a native shuffle for
`array<decimal(36,18)>` and `struct<... decimal(36,18) ...>`, where we hash 16
LE `i128` bytes and Spark hashes the minimal BE `BigDecimal` bytes. I closed
#3079 on the grounds that we do not need to allocate the same partitions as
Spark, so I am not asking you to change the gate. But the description says this
matches the `hash` expression's coverage and it does not, since
`HashUtils.unsupportedReasonFor` rejects `decimal(precision > 18)` and
`TimeType` at any depth. Worth saying so explicitly.
Last thing, this test only asserts a shuffle count and then calls
`collect()`. Now that these keys run natively, is it worth comparing
`spark_partition_id()` against Spark here too? This is the widest nested type
coverage we have, and it is the one place a leaf type difference would show up
across all of them at once.
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala:
##########
@@ -672,6 +680,194 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle on struct hash partitioning key") {
+ 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") {
+ 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") {
+ // struct<array<int>, string> and array<struct<int, string>>: 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") {
+ // Four levels of nesting, mixing all three recursive branches:
+ // struct< array< struct< m: map<string, array<int>>, 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<STRUCT<m: MAP<STRING, ARRAY<INT>>, 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") {
+ // 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<a: STRING COLLATE UTF8_LCASE, b:
INT>) 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") {
+ 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") {
+ // 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.
+ 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")
+ 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<a: ARRAY<STRUCT<m: MAP<STRING, ARRAY<INT>>, 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")
+ 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
Review Comment:
I do not think this is testing `CalendarIntervalType`. `INTERVAL '1' MONTH`
goes through `constructMultiUnitsIntervalLiteral` and comes out as
`YearMonthIntervalType(MONTH, MONTH)`. The gate rejects both so the assertion
still holds, but #5059 and the `CalendarIntervalType` case in
`supportedSerializableDataType` are about the calendar type, and that one stays
uncovered.
Could you use `make_interval(...)`, or set
`spark.sql.legacy.interval.enabled=true`, so the test matches its comment?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -405,6 +414,22 @@ object CometShuffleExchangeExec
true
case dt if isTimeType(dt) =>
true
+ case StructType(fields) if nestedHashPartitioningEnabled =>
Review Comment:
Do you have any numbers for this? What I am worried about is the shapes that
miss the vectorized paths. `array<struct<...>>` and a map nested inside a
struct both fall through `hash_list_with_primitive_elements!` into
`hash_list_array!`, which slices a one element array and re-enters
`create_murmur3_hashes` for every element, plus a `.columns().to_vec()` per
struct element on top of that. The four level key in the new tests pays that at
every level.
Since the config defaults to true we are opting everyone into this, and if
it turns out slower than just letting Spark do the shuffle then the default is
wrong. `struct<int, string>` is not the interesting case here.
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala:
##########
@@ -672,6 +680,194 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle on struct hash partitioning key") {
+ 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") {
+ 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") {
+ // struct<array<int>, string> and array<struct<int, string>>: 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") {
+ // Four levels of nesting, mixing all three recursive branches:
+ // struct< array< struct< m: map<string, array<int>>, 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<STRUCT<m: MAP<STRING, ARRAY<INT>>, 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") {
+ // 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<a: STRING COLLATE UTF8_LCASE, b:
INT>) 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") {
+ 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") {
+ // 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.
Review Comment:
`withSQLConf` returns `T`, not Unit. `SQLTestUtils` declares `protected
override def withSQLConf[T](pairs: (String, String)*)(f: => T): T` and
`CometTestBase` mixes it in, so both of these can be `val sparkRows =
withSQLConf(...) { ... }`. Also missing a blank line before this test and
before the interval one below it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
