[SPARK-55411][SQL] SPJ may throw ArrayIndexOutOfBoundsException when join keys are less than cluster keys

[pan3793]

What changes were proposed in this pull request?

Fix a java.lang.ArrayIndexOutOfBoundsException when spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled=true, by correcting the expression(should pass the full partition expression instead of the projected one) passed to KeyGroupedPartitioning#project.

Also, fix a test code issue, change the calculation result of BucketTransform defined at InMemoryBaseTable.scala to match BucketFunctions defined at transformFunctions.scala (thanks @peter-toth for pointing out this!)

Why are the changes needed?

It's a bug fix.

Does this PR introduce any user-facing change?

Some queries that failed when spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled=true now run normally.

How was this patch tested?

New UT is added, previously it failed with ArrayIndexOutOfBoundsException, now passed.

$ build/sbt "sql/testOnly *KeyGroupedPartitioningSuite -- -z SPARK=55411"
...
[info] - bug *** FAILED *** (1 second, 884 milliseconds)
[info]   java.lang.ArrayIndexOutOfBoundsException: Index 1 out of bounds for length 1
[info]   at scala.collection.immutable.ArraySeq$ofRef.apply(ArraySeq.scala:331)
[info]   at org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning$.$anonfun$project$1(partitioning.scala:471)
[info]   at org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning$.$anonfun$project$1$adapted(partitioning.scala:471)
[info]   at scala.collection.immutable.ArraySeq.map(ArraySeq.scala:75)
[info]   at scala.collection.immutable.ArraySeq.map(ArraySeq.scala:35)
[info]   at org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning$.project(partitioning.scala:471)
[info]   at org.apache.spark.sql.execution.KeyGroupedPartitionedScan.$anonfun$getOutputKeyGroupedPartitioning$5(KeyGroupedPartitionedScan.scala:58)
...

UTs affected by bucket() calculate logic change are tuned.

Was this patch authored or co-authored using generative AI tooling?

No.

[github-actions]

JIRA Issue Information

=== Bug SPARK-55411 ===
Summary: SPJ may throw ArrayIndexOutOfBoundsException when join keys are less than cluster keys
Assignee: None
Status: Open
Affected: ["4.0.2","4.1.1"]

---

This comment was automatically generated by GitHub Actions

[szehon-ho]

thanks for the repo, ill try to take a look.

[pan3793]

I found originalPartitionValues is not always populated. is it intentional?

[peter-toth]

Actually the wrong projected excepression is the root cause of the ArrayIndexOutOfBoundsException you hit and passing in basePartitioning.expressions looks good.

But the test you added will unlikely pass as there is an issue with the test framework.
I left a note here:

spark/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala

Lines 2801 to 2802 in 3405255

	// Do not use `bucket()` in "one side partition" tests as its implementation in
	// `InMemoryBaseTable` conflicts with `BucketFunction`

, but forgot to open a fix for the problem with using bucket() in these one side shuffle tests.

The problem is that the bucket() implementation here:

spark/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala

Lines 93 to 95 in 3405255

	override def produceResult(input: InternalRow): Int = {
	(input.getLong(1) % input.getInt(0)).toInt
	}

and in InMemoryBaseTable:

spark/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala

Lines 240 to 247 in 3405255

	val valueTypePairs = cols.map(col => extractor(col.fieldNames, cleanedSchema, row))
	var valueHashCode = 0
	valueTypePairs.foreach( pair =>
	if ( pair._1 != null) valueHashCode += pair._1.hashCode()
	)
	var dataTypeHashCode = 0
	valueTypePairs.foreach(dataTypeHashCode += _._2.hashCode())
	((valueHashCode + 31 * dataTypeHashCode) & Integer.MAX_VALUE) % numBuckets

mismatch.
So technically the partition keys that the datasource reports and the calculated key of the partition where the partitioner puts the shuffled records don't match.

@pan3793, could you please keep your fix in KeyGroupedPartitionedScan.scala‎ and fix the BucketTransform key calculation in InMemoryBaseTable?
You don't need need the other changes. originalPartitionValues seems unrelated as it is used only when partially clustered distribution is enabled.

BTW, I'm working on refactoring SPJ based on this idea: #53859 (comment) and it looks prosmising so far, but I need some more days to wrap it up.

[pan3793]

// Do not use bucket() in "one side partition" tests as its implementation in
// InMemoryBaseTable conflicts with BucketFunction

Oh, god, @peter-toth, thanks a lot for pointing this out, I wasn't aware of it and have spent a few hours trying to figure out why SMJ partition key value mismatch and produce wrong result after fixing the ArrayIndexOutOfBoundsException ...

Actually, the current code changes are just a draft; the test cases have not yet passed. I will try to fix it following your guidance. Thank you again, @peter-toth!

[pan3793]

scala> Long.MaxValue + 1L
res0: Long = -9223372036854775808

scala> (Long.MaxValue + 1L) & 0xFFFFFFFFFFFFL
res1: Long = 0

scala> (Long.MaxValue + 2L) & 0xFFFFFFFFFFFFL
res2: Long = 1

[peter-toth]

Ah, this is needed because % N can return negative results, isn't it? That seems like problem at both places as bucket N should return max N different values.

[peter-toth]

Should we use Math.floorMod()?

[pan3793]

the bucket num should be >=1 (seems we don't have such a check though), then (non_negative_long % positive_int) should always be positive?

[peter-toth]

Yeah, that's correct, but

spark/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala

Lines 93 to 95 in 3405255

	override def produceResult(input: InternalRow): Int = {
	(input.getLong(1) % input.getInt(0)).toInt
	}

seems also wrong as it can return values between -N+1 and N-1 so we should probably fix both places. If we used Math.floorMod() then we wouldn't need that & 0xFFFFFFFFFFFFL non-negative conversion.

[peter-toth]

Looks good to me, let's wait for CI.