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