[SPARK-37037][SQL] Improve byte array sort by unify compareTo function of UTF8String and ByteArray

common/unsafe/src/main/java/org/apache/spark/unsafe/types/ByteArray.java

@@ -75,6 +75,42 @@ static long getPrefix(Object base, long offset, int numBytes) {
  return (IS_LITTLE_ENDIAN ? java.lang.Long.reverseBytes(p) : p) & ~mask;
  }
 
+ public static int compareBinary(byte[] leftBase, byte[] rightBase) {
+ return compareBinary(leftBase, Platform.BYTE_ARRAY_OFFSET, leftBase.length,
+ rightBase, Platform.BYTE_ARRAY_OFFSET, rightBase.length);
+ }
+
+ static int compareBinary(
+ Object leftBase,
+ long leftOffset,
+ int leftNumBytes,
+ Object rightBase,
+ long rightOffset,
+ int rightNumBytes) {
+ int len = Math.min(leftNumBytes, rightNumBytes);
+ int wordMax = (len / 8) * 8;
+ for (int i = 0; i < wordMax; i += 8) {
+ long left = Platform.getLong(leftBase, leftOffset + i);
+ long right = Platform.getLong(rightBase, rightOffset + i);
+ if (left != right) {
+ if (IS_LITTLE_ENDIAN) {
+ return Long.compareUnsigned(Long.reverseBytes(left), Long.reverseBytes(right));
+ } else {
+ return Long.compareUnsigned(left, right);
+ }
+ }
+ }
+ for (int i = wordMax; i < len; i++) {
+ // Both UTF-8 and byte array should be compared as unsigned int.
+ int res = (Platform.getByte(leftBase, leftOffset + i) & 0xFF) -
+ (Platform.getByte(rightBase, rightOffset + i) & 0xFF);
+ if (res != 0) {
+ return res;
+ }
+ }
+ return leftNumBytes - rightNumBytes;
+ }
+
  public static byte[] subStringSQL(byte[] bytes, int pos, int len) {
  // This pos calculation is according to UTF8String#subStringSQL
  if (pos > bytes.length) {

common/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java

@@ -1353,29 +1353,8 @@ public UTF8String copy() {
 
  @Override
  public int compareTo(@Nonnull final UTF8String other) {
- int len = Math.min(numBytes, other.numBytes);
- int wordMax = (len / 8) * 8;
- long roffset = other.offset;
- Object rbase = other.base;
- for (int i = 0; i < wordMax; i += 8) {
- long left = getLong(base, offset + i);
- long right = getLong(rbase, roffset + i);
- if (left != right) {
- if (IS_LITTLE_ENDIAN) {
- return Long.compareUnsigned(Long.reverseBytes(left), Long.reverseBytes(right));
- } else {
- return Long.compareUnsigned(left, right);
- }
- }
- }
- for (int i = wordMax; i < len; i++) {
- // In UTF-8, the byte should be unsigned, so we should compare them as unsigned int.
- int res = (getByte(i) & 0xFF) - (Platform.getByte(rbase, roffset + i) & 0xFF);
- if (res != 0) {
- return res;
- }
- }
- return numBytes - other.numBytes;
+ return ByteArray.compareBinary(
+ base, offset, numBytes, other.base, other.offset, other.numBytes);
  }
 
  public int compare(final UTF8String other) {

common/unsafe/src/test/java/org/apache/spark/unsafe/array/ByteArraySuite.java

@@ -48,4 +48,23 @@ public void testGetPrefix() {
  Assert.assertEquals(result, expected);
  }
  }
+
+ @Test
+ public void testCompareBinary() {
+ byte[] x1 = new byte[0];
+ byte[] y1 = new byte[]{(byte) 1, (byte) 2, (byte) 3};
+ assert(ByteArray.compareBinary(x1, y1) < 0);
+
+ byte[] x2 = new byte[]{(byte) 200, (byte) 100};
+ byte[] y2 = new byte[]{(byte) 100, (byte) 100};
+ assert(ByteArray.compareBinary(x2, y2) > 0);
+
+ byte[] x3 = new byte[]{(byte) 100, (byte) 200, (byte) 12};
+ byte[] y3 = new byte[]{(byte) 100, (byte) 200};
+ assert(ByteArray.compareBinary(x3, y3) > 0);
+
+ byte[] x4 = new byte[]{(byte) 100, (byte) 200};
+ byte[] y4 = new byte[]{(byte) 100, (byte) 200};
+ assert(ByteArray.compareBinary(x4, y4) == 0);
+ }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -650,7 +650,7 @@ class CodegenContext extends Logging {
  s"$clsName.compareFloats($c1, $c2)"
  // use c1 - c2 may overflow
  case dt: DataType if isPrimitiveType(dt) => s"($c1 > $c2 ? 1 : $c1 < $c2 ? -1 : 0)"
- case BinaryType => s"org.apache.spark.sql.catalyst.util.TypeUtils.compareBinary($c1, $c2)"
+ case BinaryType => s"org.apache.spark.unsafe.types.ByteArray.compareBinary($c1, $c2)"
  case NullType => "0"
  case array: ArrayType =>
  val elementType = array.elementType

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala

@@ -87,17 +87,6 @@ object TypeUtils {
  }
  }
 
- def compareBinary(x: Array[Byte], y: Array[Byte]): Int = {
- val limit = if (x.length <= y.length) x.length else y.length
- var i = 0
- while (i < limit) {
- val res = (x(i) & 0xff) - (y(i) & 0xff)
- if (res != 0) return res
- i += 1
- }
- x.length - y.length
- }
-
  /**
  * Returns true if the equals method of the elements of the data type is implemented properly.
  * This also means that they can be safely used in collections relying on the equals method,

sql/catalyst/src/main/scala/org/apache/spark/sql/types/BinaryType.scala

@@ -20,7 +20,7 @@ package org.apache.spark.sql.types
 import scala.reflect.runtime.universe.typeTag
 
 import org.apache.spark.annotation.Stable
-import org.apache.spark.sql.catalyst.util.TypeUtils
+import org.apache.spark.unsafe.types.ByteArray
 
 /**
  * The data type representing `Array[Byte]` values.
@@ -37,7 +37,7 @@ class BinaryType private() extends AtomicType {
  @transient private[sql] lazy val tag = typeTag[InternalType]
 
  private[sql] val ordering =
- (x: Array[Byte], y: Array[Byte]) => TypeUtils.compareBinary(x, y)
+ (x: Array[Byte], y: Array[Byte]) => ByteArray.compareBinary(x, y)
 
  /**
  * The default size of a value of the BinaryType is 100 bytes.

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/TypeUtilsSuite.scala

@@ -43,22 +43,4 @@ class TypeUtilsSuite extends SparkFunSuite {
  typeCheckPass(ArrayType(StringType, containsNull = true) ::
  ArrayType(StringType, containsNull = false) :: Nil)
  }
-
- test("compareBinary") {
- val x1 = Array[Byte]()
- val y1 = Array(1, 2, 3).map(_.toByte)
- assert(TypeUtils.compareBinary(x1, y1) < 0)
-
- val x2 = Array(200, 100).map(_.toByte)
- val y2 = Array(100, 100).map(_.toByte)
- assert(TypeUtils.compareBinary(x2, y2) > 0)
-
- val x3 = Array(100, 200, 12).map(_.toByte)
- val y3 = Array(100, 200).map(_.toByte)
- assert(TypeUtils.compareBinary(x3, y3) > 0)
-
- val x4 = Array(100, 200).map(_.toByte)
- val y4 = Array(100, 200).map(_.toByte)
- assert(TypeUtils.compareBinary(x4, y4) == 0)
- }
 }

sql/core/benchmarks/ByteArrayBenchmark-jdk11-results.txt

@@ -0,0 +1,16 @@
+================================================================================================
+byte array comparisons
+================================================================================================
+
+OpenJDK 64-Bit Server VM 11.0.13+8-LTS on Linux 5.8.0-1042-azure
+Intel(R) Xeon(R) Platinum 8171M CPU @ 2.60GHz
+Byte Array compareTo: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+2-7 byte 501 514 14 130.9 7.6 1.0X
+8-16 byte 976 993 10 67.1 14.9 0.5X
+16-32 byte 985 995 6 66.5 15.0 0.5X
+512-1024 byte 1260 1282 13 52.0 19.2 0.4X
+512 byte slow 3114 3193 46 21.0 47.5 0.2X
+2-7 byte 572 578 7 114.5 8.7 0.9X
+
+

sql/core/benchmarks/ByteArrayBenchmark-results.txt

@@ -0,0 +1,16 @@
+================================================================================================
+byte array comparisons
+================================================================================================
+
+OpenJDK 64-Bit Server VM 1.8.0_312-b07 on Linux 5.8.0-1042-azure
+Intel(R) Xeon(R) Platinum 8171M CPU @ 2.60GHz
+Byte Array compareTo: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+2-7 byte 407 418 9 161.1 6.2 1.0X
+8-16 byte 867 919 30 75.6 13.2 0.5X
+16-32 byte 882 916 23 74.3 13.5 0.5X
+512-1024 byte 1123 1167 31 58.4 17.1 0.4X
+512 byte slow 4054 4611 506 16.2 61.9 0.1X
+2-7 byte 430 450 16 152.4 6.6 0.9X
+
+

sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ByteArrayBenchmark.scala

@@ -0,0 +1,86 @@
+/*
+ * 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.execution.benchmark
+
+import scala.util.Random
+
+import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
+import org.apache.spark.unsafe.types.ByteArray
+
+/**
+ * Benchmark to measure performance for byte array comparisons.
+ * {{{
+ * To run this benchmark:
+ * 1. without sbt:
+ * bin/spark-submit --class <this class> --jars <spark core test jar> <sql core test jar>
+ * 2. build/sbt "sql/test:runMain <this class>"
+ * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/test:runMain <this class>"
+ * Results will be written to "benchmarks/<this class>-results.txt".
+ * }}}
+ */
+object ByteArrayBenchmark extends BenchmarkBase {
+
+ def byteArrayComparisons(iters: Long): Unit = {
+ val chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+ val random = new Random(0)
+ def randomBytes(min: Int, max: Int): Array[Byte] = {
+ val len = random.nextInt(max - min) + min
+ val bytes = new Array[Byte](len)
+ var i = 0
+ while (i < len) {
+ bytes(i) = chars.charAt(random.nextInt(chars.length())).toByte
+ i += 1
+ }
+ bytes
+ }
+
+ val count = 16 * 1000
+ val dataTiny = Seq.fill(count)(randomBytes(2, 7)).toArray
+ val dataSmall = Seq.fill(count)(randomBytes(8, 16)).toArray
+ val dataMedium = Seq.fill(count)(randomBytes(16, 32)).toArray
+ val dataLarge = Seq.fill(count)(randomBytes(512, 1024)).toArray
+ val dataLargeSlow = Seq.fill(count)(
+ Array.tabulate(512) {i => if (i < 511) 0.toByte else 1.toByte}).toArray
+
+ def compareBinary(data: Array[Array[Byte]]) = { _: Int =>
+ var sum = 0L
+ for (_ <- 0L until iters) {
+ var i = 0
+ while (i < count) {
+ sum += ByteArray.compareBinary(data(i), data((i + 1) % count))
+ i += 1
+ }
+ }
+ }
+
+ val benchmark = new Benchmark("Byte Array compareTo", count * iters, 25, output = output)
+ benchmark.addCase("2-7 byte")(compareBinary(dataTiny))
+ benchmark.addCase("8-16 byte")(compareBinary(dataSmall))
+ benchmark.addCase("16-32 byte")(compareBinary(dataMedium))
+ benchmark.addCase("512-1024 byte")(compareBinary(dataLarge))
+ benchmark.addCase("512 byte slow")(compareBinary(dataLargeSlow))
+ benchmark.addCase("2-7 byte")(compareBinary(dataTiny))
+ benchmark.run()
+ }
+
+ override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+ runBenchmark("byte array comparisons") {
+ byteArrayComparisons(1024 * 4)
+ }
+ }
+}