You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "stefankandic (via GitHub)" <gi...@apache.org> on 2026/08/18 16:04:48 UTC
[PR] [WIP] Collation-aware hash aggregation [spark]
stefankandic opened a new pull request, #58083:
URL: https://github.com/apache/spark/pull/58083
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'common/utils/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
Implementing support for binary unstable types (eg collated strings) in hash aggregators (both `HashAggregate` and `ObjectHashAggregate`). Binary unstable types are types for whom we can't derive equality / ordering by their bytes representations alone (eg `a` and `A` are equal in any case-insensitive collation).
### Why are the changes needed?
Hash aggregation is the default aggregation mode in spark, and we've seen regressions in queries after moving from `UTF8_BINARY` to other collations as they would have to fall back to `SortAggregate`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Unit and query-level tests.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: gpt-5.6-sol
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "stefankandic (via GitHub)" <gi...@apache.org>.
stefankandic commented on PR #58083:
URL: https://github.com/apache/spark/pull/58083#issuecomment-5414062950
@cloud-fan Please take a look when you get the chance. Or forward to someone familiar with hash aggregation code. Thanks!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on PR #58083:
URL: https://github.com/apache/spark/pull/58083#issuecomment-5490951011
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/c7190edd2617a128f03050eed3ed0c6e0227421b
- merged into branch-4.x https://github.com/apache/spark/commit/87627ffa2058efccbb15507f7d9c78eee7368250
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847] Collation-aware hash aggregation [spark]
Posted by "stefankandic (via GitHub)" <gi...@apache.org>.
stefankandic commented on PR #58083:
URL: https://github.com/apache/spark/pull/58083#issuecomment-5372102195
@uros-b and @cloud-fan Please take a look when you get the chance.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on PR #58083:
URL: https://github.com/apache/spark/pull/58083#issuecomment-5386089478
Adding @cloud-fan to also take a look please
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on PR #58083:
URL: https://github.com/apache/spark/pull/58083#issuecomment-5386087598
I added the [SQL] component tag to the PR title
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58083:
URL: https://github.com/apache/spark/pull/58083#discussion_r3862857569
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -516,12 +574,36 @@ public Location lookup(Object keyBase, long keyOffset, int keyLength, int hash)
}
/**
- * Looks up a key, and saves the result in provided `loc`.
+ * Looks up a key and saves the result in the provided `loc`.
+ *
+ * This is a thread-safe version of `lookup`, provided that each thread supplies its own
Review Comment:
**Non-blocking:**
Both lookup paths increment the shared `numKeyLookups` and `numProbes` fields with non-atomic read-modify-write operations. Concurrent calls allowed by this contract can therefore lose updates and report an inaccurate average probe count. Please either make these counters concurrency-safe or explicitly exclude probe statistics from the concurrent-use guarantee.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -726,7 +861,7 @@ public int getValueLength() {
* The return value indicates whether the put succeeded or whether it failed because additional
* memory could not be acquired.
* <p>
- * It is only valid to call this method immediately after calling `lookup()` using the same key.
+ * It is only valid to call this method immediately after looking up the same key.
Review Comment:
**Nit:**
While updating this contract, please also broaden the false-return wording above. `append` returns `false` at `MAX_CAPACITY - 1` and when `!canGrowArray && numKeys >= growthThreshold`, before trying to allocate, so failure does not always mean that memory could not be acquired.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -557,31 +639,84 @@ public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location l
}
}
- /**
- * Handle returned by {@link BytesToBytesMap#lookup(Object, long, int)} function.
- */
+ private void safeLookupWithKeyOperations(
+ Object keyBase, long keyOffset, int keyLength, Location loc) {
+ assert(longArray != null);
+ assert(keyOperationsFactory != null);
+
+ final KeyOperations keyOperations = loc.getKeyOperations();
+ assert(keyOperations != null);
+ final int hash = keyOperations.hash(keyBase, keyOffset, keyLength);
+
+ numKeyLookups++;
+
+ int pos = hash & mask;
+ int step = 1;
+ while (true) {
+ numProbes++;
+ if (longArray.get(pos * 2) == 0) {
+ // This is a new key.
+ loc.with(pos, hash, false);
+ return;
+ } else {
+ long stored = longArray.get(pos * 2 + 1);
+ if ((int) (stored) == hash) {
+ // Full hash code matches. Let's compare the keys for equality.
+ loc.with(pos, hash, true);
+ if (loc.getKeyLength() == keyLength &&
+ ByteArrayMethods.arrayEquals(
+ keyBase,
+ keyOffset,
+ loc.getKeyBase(),
+ loc.getKeyOffset(),
+ keyLength)) {
+ return;
+ }
+ if (keyOperations.equals(
+ keyBase,
+ keyOffset,
+ keyLength,
+ loc.getKeyBase(),
+ loc.getKeyOffset(),
+ loc.getKeyLength())) {
+ return;
+ }
+ }
+ }
+ pos = (pos + step) & mask;
+ step++;
+ }
+ }
+
+ /** Handle returned by this map's lookup methods. */
public final class Location {
/** An index into the hash map's Long array */
private int pos;
/** True if this location points to a position where a key is defined, false otherwise */
private boolean isDefined;
- /**
- * The hashcode of the most recent key passed to
- * {@link BytesToBytesMap#lookup(Object, long, int, int)}. Caching this hashcode here allows us
- * to avoid re-hashing the key when storing a value for that key.
- */
+ /** The hash code computed by the most recent lookup. */
Review Comment:
**Nit:**
The precomputed-hash overload can store a caller-supplied value here, so `computed` is too narrow.
```suggestion
/** The hash code used by the most recent lookup. */
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/UnsafeRowKeyOperations.scala:
##########
@@ -0,0 +1,328 @@
+/*
+ * 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.catalyst.util
+
+import com.ibm.icu.text.RawCollationKey
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{SpecializedGetters, UnsafeRow}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.unsafe.array.ByteArrayMethods
+import org.apache.spark.unsafe.hash.Murmur3_x86_32
+import org.apache.spark.unsafe.map.BytesToBytesMap
+import org.apache.spark.unsafe.types.UTF8String
+
+/** Hashes and compares unsafe grouping keys containing non-binary collated strings. */
+final class UnsafeRowKeyOperations(val schema: StructType)
+ extends BytesToBytesMap.KeyOperationsFactory {
+ import UnsafeRowKeyOperations._
+
+ require(
+ schema.forall(field => supportsDataType(field.dataType)),
+ s"Unsupported grouping key schema: $schema")
+
+ private val operations = new RowOperations(schema, new HashScratch)
+
+ def hash(row: UnsafeRow): Int = operations.hash(row, HASH_SEED)
+
+ def areEqual(left: UnsafeRow, right: UnsafeRow): Boolean = {
+ operations.areEqual(left, right)
+ }
+
+ override def create(): BytesToBytesMap.KeyOperations = new BytesToBytesMap.KeyOperations {
+ private val localOperations = new RowOperations(schema, new HashScratch)
+ private val left = new UnsafeRow(schema.length)
+ private val right = new UnsafeRow(schema.length)
+
+ override def hash(base: AnyRef, offset: Long, length: Int): Int = {
+ left.pointTo(base, offset, length)
+ localOperations.hash(left, HASH_SEED)
+ }
+
+ override def equals(
+ leftBase: AnyRef,
+ leftOffset: Long,
+ leftLength: Int,
+ rightBase: AnyRef,
+ rightOffset: Long,
+ rightLength: Int): Boolean = {
+ left.pointTo(leftBase, leftOffset, leftLength)
+ right.pointTo(rightBase, rightOffset, rightLength)
+ localOperations.areEqual(left, right)
+ }
+ }
+}
+
+object UnsafeRowKeyOperations {
+ private val HASH_SEED = 42
+
+ /** Returns whether this type can be hashed and compared as part of an unsafe row key. */
+ def supportsDataType(dataType: DataType): Boolean = {
+ if (UnsafeRowUtils.isBinaryStable(dataType)) {
+ true
+ } else {
+ dataType match {
+ case st: StringType => !st.supportsBinaryEquality
+ case ArrayType(elementType, _) => supportsDataType(elementType)
+ case StructType(fields) => fields.forall(field => supportsDataType(field.dataType))
+ case _ => false
+ }
+ }
+ }
+
+ private final class BinaryRegion {
+ var base: AnyRef = _
+ var offset: Long = 0L
+ var length: Int = 0
+ }
+
+ private final class HashScratch {
+ private var reusableRawCollationKey: RawCollationKey = _
+
+ def rawCollationKey: RawCollationKey = {
+ if (reusableRawCollationKey == null) {
+ reusableRawCollationKey = new RawCollationKey()
+ }
+ reusableRawCollationKey
+ }
+ }
+
+ private sealed trait FieldOperations {
+ final def hash(input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ if (input.isNullAt(ordinal)) seed else hashNonNull(input, ordinal, seed)
+ }
+
+ final def areEqual(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ val leftIsNull = left.isNullAt(leftOrdinal)
+ val rightIsNull = right.isNullAt(rightOrdinal)
+ if (leftIsNull || rightIsNull) {
+ leftIsNull == rightIsNull
+ } else {
+ areEqualNonNull(left, leftOrdinal, right, rightOrdinal)
+ }
+ }
+
+ protected def hashNonNull(input: SpecializedGetters, ordinal: Int, seed: Int): Int
+
+ protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean
+ }
+
+ private final class BinaryStableOperations(dataType: DataType) extends FieldOperations {
+ private val leftRegion = new BinaryRegion
+ private val rightRegion = new BinaryRegion
+
+ override protected def hashNonNull(
+ input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ setRegion(input, ordinal, leftRegion)
+ Murmur3_x86_32.hashUnsafeBytes(
+ leftRegion.base, leftRegion.offset, leftRegion.length, seed)
+ }
+
+ override protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ setRegion(left, leftOrdinal, leftRegion)
+ setRegion(right, rightOrdinal, rightRegion)
+ leftRegion.length == rightRegion.length && ByteArrayMethods.arrayEquals(
+ leftRegion.base,
+ leftRegion.offset,
+ rightRegion.base,
+ rightRegion.offset,
+ leftRegion.length)
+ }
+
+ private def setRegion(
+ input: SpecializedGetters, ordinal: Int, region: BinaryRegion): Unit = input match {
+ case row: UnsafeRow =>
+ region.base = row.getBaseObject
+ if (UnsafeRow.isFixedLength(dataType)) {
+ region.offset = row.getBaseOffset +
+ UnsafeRow.calculateBitSetWidthInBytes(row.numFields()) + ordinal * 8L
+ region.length = 8
+ } else {
+ setVariableLengthRegion(row.getBaseOffset, row.getLong(ordinal), region)
+ }
+
+ case other =>
+ throw SparkException.internalError(
+ s"Expected unsafe grouping-key storage, found ${other.getClass.getName}")
+ }
+ }
+
+ private final class CollatedStringOperations(
+ dataType: StringType,
+ hashScratch: HashScratch) extends FieldOperations {
+ private val collation = CollationFactory.fetchCollation(dataType.collationId)
+ private val useRawCollationKey = collation.provider == CollationFactory.PROVIDER_ICU
+
+ override protected def hashNonNull(
+ input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ val value = input.getUTF8String(ordinal)
+ if (useRawCollationKey) {
+ hashICUCollationKey(value, seed)
+ } else {
+ val key = collation.sortKeyFunction.apply(value)
+ Murmur3_x86_32.hashUnsafeBytes(key, Platform.BYTE_ARRAY_OFFSET, key.length, seed)
+ }
+ }
+
+ /**
+ * Hashes an ICU sort key from a reusable buffer after applying configured space trimming.
+ * `toValidString` applies Spark's replacement policy for malformed UTF-8 before calling ICU.
+ */
+ private def hashICUCollationKey(value: UTF8String, seed: Int): Int = {
+ val normalizedValue = if (collation.supportsSpaceTrimming) {
+ CollationFactory.applyTrimmingPolicy(value, dataType.collationId)
+ } else {
+ value
+ }
+ val key = collation.getCollator.getRawCollationKey(
+ normalizedValue.toValidString, hashScratch.rawCollationKey)
+ Murmur3_x86_32.hashUnsafeBytes(
+ key.bytes, Platform.BYTE_ARRAY_OFFSET, key.size, seed)
+ }
+
+ override protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ collation.equalsFunction.apply(
+ left.getUTF8String(leftOrdinal), right.getUTF8String(rightOrdinal))
+ }
+ }
+
+ private final class ArrayOperations(
+ elementType: DataType,
+ hashScratch: HashScratch) extends FieldOperations {
+ private val elementOperations = createFieldOperations(elementType, hashScratch)
+
+ override protected def hashNonNull(
+ input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ val array = input.getArray(ordinal)
Review Comment:
**Non-blocking:**
`getArray` allocates a wrapper on each call, so this creates one object per hash and two per equality probe. The same issue exists for `getStruct` in `StructOperations` (lines 267 and 275-277). Since these operations now run for every hash-aggregation key, please keep reusable hash/left/right views in each operations instance and repoint them before recursion.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847] Collation-aware hash aggregation [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on PR #58083:
URL: https://github.com/apache/spark/pull/58083#issuecomment-5386086534
LGTM, thank you @stefankandic
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #58083:
URL: https://github.com/apache/spark/pull/58083#discussion_r3882894255
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala:
##########
@@ -957,13 +959,21 @@ case class HashAggregateExec(
// per regular-map row (see below); emitting it in more than one runtime branch is unsafe
// because the projection's subexpression/writer state assigned in one branch would be read
// stale from another (e.g. the adaptive pass-through path would reuse the last probed key).
+ val (computeKeyHash, lookupBuffer) =
+ if (UnsafeRowUtils.isBinaryStable(groupingKeySchema)) {
+ val unsafeRowKeyHash = ctx.freshName("unsafeRowKeyHash")
+ (s"int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode();",
+ s"$hashMapTerm.getAggregationBufferFromUnsafeRow(" +
+ s"$unsafeRowKeys, $unsafeRowKeyHash)")
+ } else {
+ ("", s"$hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys)")
Review Comment:
**Non-blocking:** This adaptive partial-aggregation branch now uses the semantic-key lookup, but the changed collation tests never enable adaptive partial aggregation. Please add a focused `AdaptivePartialAggregationSuite` case with collation-equivalent spellings on opposite sides of the bypass cutoff, compare enabled output with the disabled reference, and assert `numBypassingRows > 0` so the branch cannot pass vacuously.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -557,31 +641,84 @@ public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location l
}
}
- /**
- * Handle returned by {@link BytesToBytesMap#lookup(Object, long, int)} function.
- */
+ private void safeLookupWithKeyOperations(
+ Object keyBase, long keyOffset, int keyLength, Location loc) {
+ assert(longArray != null);
+ assert(keyOperationsFactory != null);
+
+ final KeyOperations keyOperations = loc.getKeyOperations();
+ assert(keyOperations != null);
+ final int hash = keyOperations.hash(keyBase, keyOffset, keyLength);
+
+ numKeyLookups++;
+
+ int pos = hash & mask;
+ int step = 1;
+ while (true) {
+ numProbes++;
+ if (longArray.get(pos * 2) == 0) {
+ // This is a new key.
+ loc.with(pos, hash, false);
+ return;
+ } else {
+ long stored = longArray.get(pos * 2 + 1);
+ if ((int) (stored) == hash) {
+ // Full hash code matches. Let's compare the keys for equality.
+ loc.with(pos, hash, true);
+ if (loc.getKeyLength() == keyLength &&
+ ByteArrayMethods.arrayEquals(
+ keyBase,
+ keyOffset,
+ loc.getKeyBase(),
+ loc.getKeyOffset(),
+ keyLength)) {
+ return;
+ }
+ if (keyOperations.equals(
Review Comment:
**Non-blocking:** `KeyOperations` only requires equal keys to have equal hashes, so this branch must continue probing when the full hash matches but semantic equality is false. The new test covers only the `A`/`a` equivalence class. Please add two semantically unequal groups under a deterministic constant hash, retrieve equivalent spellings of both, verify their values and `numKeys == 2`; keeping this in `AbstractBytesToBytesMapSuite` will exercise both memory modes.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/UnsafeRowKeyOperations.scala:
##########
@@ -0,0 +1,383 @@
+/*
+ * 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.catalyst.util
+
+import com.ibm.icu.text.RawCollationKey
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{SpecializedGetters, UnsafeArrayData, UnsafeRow}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.unsafe.array.ByteArrayMethods
+import org.apache.spark.unsafe.hash.Murmur3_x86_32
+import org.apache.spark.unsafe.map.BytesToBytesMap
+import org.apache.spark.unsafe.types.UTF8String
+
+/** Hashes and compares unsafe grouping keys containing non-binary collated strings. */
+final class UnsafeRowKeyOperations(val schema: StructType)
+ extends BytesToBytesMap.KeyOperationsFactory {
+ import UnsafeRowKeyOperations._
+
+ require(
+ schema.forall(field => supportsDataType(field.dataType)),
+ s"Unsupported grouping key schema: $schema")
+
+ private val operations = new RowOperations(schema, new HashScratch)
+
+ def hash(row: UnsafeRow): Int = operations.hash(row, HASH_SEED)
+
+ def areEqual(left: UnsafeRow, right: UnsafeRow): Boolean = {
+ operations.areEqual(left, right)
+ }
+
+ override def create(): BytesToBytesMap.KeyOperations = new BytesToBytesMap.KeyOperations {
+ private val localOperations = new RowOperations(schema, new HashScratch)
+ private val left = new UnsafeRow(schema.length)
+ private val right = new UnsafeRow(schema.length)
+
+ override def hash(base: AnyRef, offset: Long, length: Int): Int = {
+ left.pointTo(base, offset, length)
+ localOperations.hash(left, HASH_SEED)
+ }
+
+ override def equals(
+ leftBase: AnyRef,
+ leftOffset: Long,
+ leftLength: Int,
+ rightBase: AnyRef,
+ rightOffset: Long,
+ rightLength: Int): Boolean = {
+ left.pointTo(leftBase, leftOffset, leftLength)
+ right.pointTo(rightBase, rightOffset, rightLength)
+ localOperations.areEqual(left, right)
+ }
+ }
+}
+
+object UnsafeRowKeyOperations {
+ private val HASH_SEED = 42
+
+ /** Returns whether this type can be hashed and compared as part of an unsafe row key. */
+ def supportsDataType(dataType: DataType): Boolean = {
+ if (UnsafeRowUtils.isBinaryStable(dataType)) {
+ true
+ } else {
+ dataType match {
+ case st: StringType => !st.supportsBinaryEquality
+ case ArrayType(elementType, _) => supportsDataType(elementType)
+ case StructType(fields) => fields.forall(field => supportsDataType(field.dataType))
+ case _ => false
+ }
+ }
+ }
+
+ private final class BinaryRegion {
+ var base: AnyRef = _
+ var offset: Long = 0L
+ var length: Int = 0
+ }
+
+ private final class HashScratch {
+ private var reusableRawCollationKey: RawCollationKey = _
+
+ def rawCollationKey: RawCollationKey = {
+ if (reusableRawCollationKey == null) {
+ reusableRawCollationKey = new RawCollationKey()
+ }
+ reusableRawCollationKey
+ }
+ }
+
+ private sealed trait FieldOperations {
+ final def hash(input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ if (input.isNullAt(ordinal)) seed else hashNonNull(input, ordinal, seed)
+ }
+
+ final def areEqual(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ val leftIsNull = left.isNullAt(leftOrdinal)
+ val rightIsNull = right.isNullAt(rightOrdinal)
+ if (leftIsNull || rightIsNull) {
+ leftIsNull == rightIsNull
+ } else {
+ areEqualNonNull(left, leftOrdinal, right, rightOrdinal)
+ }
+ }
+
+ protected def hashNonNull(input: SpecializedGetters, ordinal: Int, seed: Int): Int
+
+ protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean
+ }
+
+ private final class BinaryStableOperations(dataType: DataType) extends FieldOperations {
+ private val leftRegion = new BinaryRegion
+ private val rightRegion = new BinaryRegion
+
+ override protected def hashNonNull(
+ input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ setRegion(input, ordinal, leftRegion)
+ Murmur3_x86_32.hashUnsafeBytes(
+ leftRegion.base, leftRegion.offset, leftRegion.length, seed)
+ }
+
+ override protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ setRegion(left, leftOrdinal, leftRegion)
+ setRegion(right, rightOrdinal, rightRegion)
+ leftRegion.length == rightRegion.length && ByteArrayMethods.arrayEquals(
+ leftRegion.base,
+ leftRegion.offset,
+ rightRegion.base,
+ rightRegion.offset,
+ leftRegion.length)
+ }
+
+ private def setRegion(
+ input: SpecializedGetters, ordinal: Int, region: BinaryRegion): Unit = input match {
+ case row: UnsafeRow =>
+ region.base = row.getBaseObject
+ if (UnsafeRow.isFixedLength(dataType)) {
+ region.offset = row.getBaseOffset +
+ UnsafeRow.calculateBitSetWidthInBytes(row.numFields()) + ordinal * 8L
+ region.length = 8
+ } else {
+ setVariableLengthRegion(row.getBaseOffset, row.getLong(ordinal), region)
+ }
+
+ case other =>
+ throw SparkException.internalError(
+ s"Expected unsafe grouping-key storage, found ${other.getClass.getName}")
+ }
+ }
+
+ private final class CollatedStringOperations(
+ dataType: StringType,
+ hashScratch: HashScratch) extends FieldOperations {
+ private val collation = CollationFactory.fetchCollation(dataType.collationId)
+ private val useRawCollationKey = collation.provider == CollationFactory.PROVIDER_ICU
+
+ override protected def hashNonNull(
+ input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ val value = input.getUTF8String(ordinal)
+ if (useRawCollationKey) {
+ hashICUCollationKey(value, seed)
+ } else {
+ val key = collation.sortKeyFunction.apply(value)
+ Murmur3_x86_32.hashUnsafeBytes(key, Platform.BYTE_ARRAY_OFFSET, key.length, seed)
+ }
+ }
+
+ /**
+ * Hashes an ICU sort key from a reusable buffer after applying configured space trimming.
+ * `toValidString` applies Spark's replacement policy for malformed UTF-8 before calling ICU.
+ */
+ private def hashICUCollationKey(value: UTF8String, seed: Int): Int = {
+ val normalizedValue = if (collation.supportsSpaceTrimming) {
+ CollationFactory.applyTrimmingPolicy(value, dataType.collationId)
+ } else {
+ value
+ }
+ val key = collation.getCollator.getRawCollationKey(
+ normalizedValue.toValidString, hashScratch.rawCollationKey)
+ Murmur3_x86_32.hashUnsafeBytes(
+ key.bytes, Platform.BYTE_ARRAY_OFFSET, key.size, seed)
+ }
+
+ override protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ collation.equalsFunction.apply(
+ left.getUTF8String(leftOrdinal), right.getUTF8String(rightOrdinal))
+ }
+ }
+
+ private final class ArrayOperations(
+ elementType: DataType,
+ hashScratch: HashScratch) extends FieldOperations {
+ private val elementOperations = createFieldOperations(elementType, hashScratch)
+ private val hashArray = new UnsafeArrayData
+ private val leftArray = new UnsafeArrayData
+ private val rightArray = new UnsafeArrayData
+
+ override protected def hashNonNull(
+ input: SpecializedGetters, ordinal: Int, seed: Int): Int = {
+ pointToArray(input, ordinal, hashArray)
+ var result = seed
+ var index = 0
+ while (index < hashArray.numElements()) {
+ result = elementOperations.hash(hashArray, index, result)
+ index += 1
+ }
+ result
+ }
+
+ override protected def areEqualNonNull(
+ left: SpecializedGetters,
+ leftOrdinal: Int,
+ right: SpecializedGetters,
+ rightOrdinal: Int): Boolean = {
+ pointToArray(left, leftOrdinal, leftArray)
+ pointToArray(right, rightOrdinal, rightArray)
+ if (leftArray.numElements() != rightArray.numElements()) {
Review Comment:
**Test coverage:** All added array-equality cases use equal-length inputs, so this new cardinality guard is not exercised. Please add unequal-length and empty/non-empty nested collated arrays and assert `areEqual` is false in both operand orders under `CODEGEN_ONLY` and `NO_CODEGEN`.
##########
core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java:
##########
@@ -557,31 +641,84 @@ public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location l
}
}
- /**
- * Handle returned by {@link BytesToBytesMap#lookup(Object, long, int)} function.
- */
+ private void safeLookupWithKeyOperations(
+ Object keyBase, long keyOffset, int keyLength, Location loc) {
+ assert(longArray != null);
+ assert(keyOperationsFactory != null);
+
+ final KeyOperations keyOperations = loc.getKeyOperations();
+ assert(keyOperations != null);
+ final int hash = keyOperations.hash(keyBase, keyOffset, keyLength);
+
+ numKeyLookups++;
+
+ int pos = hash & mask;
+ int step = 1;
+ while (true) {
+ numProbes++;
+ if (longArray.get(pos * 2) == 0) {
+ // This is a new key.
+ loc.with(pos, hash, false);
+ return;
+ } else {
+ long stored = longArray.get(pos * 2 + 1);
+ if ((int) (stored) == hash) {
+ // Full hash code matches. Let's compare the keys for equality.
+ loc.with(pos, hash, true);
+ if (loc.getKeyLength() == keyLength &&
+ ByteArrayMethods.arrayEquals(
+ keyBase,
+ keyOffset,
+ loc.getKeyBase(),
+ loc.getKeyOffset(),
+ keyLength)) {
+ return;
+ }
+ if (keyOperations.equals(
+ keyBase,
+ keyOffset,
+ keyLength,
+ loc.getKeyBase(),
+ loc.getKeyOffset(),
+ loc.getKeyLength())) {
+ return;
+ }
+ }
+ }
+ pos = (pos + step) & mask;
+ step++;
+ }
+ }
+
+ /** Handle returned by this map's lookup methods. */
public final class Location {
/** An index into the hash map's Long array */
private int pos;
/** True if this location points to a position where a key is defined, false otherwise */
private boolean isDefined;
- /**
- * The hashcode of the most recent key passed to
- * {@link BytesToBytesMap#lookup(Object, long, int, int)}. Caching this hashcode here allows us
- * to avoid re-hashing the key when storing a value for that key.
- */
+ /** The hash code used by the most recent lookup. */
private int keyHashcode;
private Object baseObject; // the base object for key and value
private long keyOffset;
private int keyLength;
private long valueOffset;
private int valueLength;
+ @Nullable private KeyOperations keyOperations;
+
/**
* Memory page containing the record. Only set if created by {@link BytesToBytesMap#iterator()}.
*/
@Nullable private MemoryBlock memoryPage;
Review Comment:
**Non-blocking:** The `safeLookup` contract allows concurrent calls with distinct `Location`s, and each location can call the shared `KeyOperationsFactory.create()` here on first use. The factory contract requires independent results but does not say `create()` may be invoked concurrently. Please either document that requirement or serialize lazy creation; if serialization is chosen, a deliberately non-thread-safe factory would make the concurrency test cover 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58847][SQL] Collation-aware hash aggregation [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b closed pull request #58083: [SPARK-58847][SQL] Collation-aware hash aggregation
URL: https://github.com/apache/spark/pull/58083
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org