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: [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]
