Skip to content

feat: support nested types as native shuffle hash partitioning keys #5567

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
viirya wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: support nested types as native shuffle hash partitioning keys #5567

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions spark/src/main/scala/org/apache/comet/CometConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<struct<...>>`, 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")
Expand Down
31 changes: 28 additions & 3 deletions ...rc/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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,
Expand All @@ -424,6 +433,22 @@ object CometShuffleExchangeExec
true
case dt if isTimeType(dt) =>
true
case StructType(fields) if nestedHashPartitioningEnabled =>

@andygrove andygrove Aug 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

// `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)
Comment on lines +443 to +451

@sunchao sunchao Sep 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Rebase sliced map offsets before enabling map shuffle keys

Could we fix spark_map_sort for sliced maps before admitting this case? With the nested-hash config enabled on Spark 4.0, a native OFFSET below a map repartition can pass a MapArray whose first entry offset is nonzero. DataFusion's limit uses batch.slice(skip, ...), and Arrow preserves the map's original entry offsets. The unified native shuffle evaluates mapsort(m) on that slice without a JVM round trip.

map_sort.rs builds its take indices only for the visible maps, but passes the original offsets to MapArray::try_new. For two two-entry maps, skipping the first leaves offsets [2,4]. Sorting creates two entries, so Arrow rejects the result with Max offset of 4 exceeds length of entries 2. The query fails instead of repartitioning. The map-sort code predates this PR, but the base gate rejected map hash keys, so the newly admitted path exposes it.

Please rebase the output offsets, or retain fallback for this case, and cover OFFSET followed by map repartition while asserting both operators run natively. This finding is source-derived. I have not executed the query.

case _ =>
false
}
Expand Down
5 changes: 3 additions & 2 deletions spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<array<primitive>>` and `array<struct<primitive>>` 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",
Expand Down
13 changes: 12 additions & 1 deletion spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}

Expand Down
Loading
Loading