From 0710e002c403ab55079e721401c4ae2918c9d70a Mon Sep 17 00:00:00 2001 From: Peifeng Li Date: Fri, 11 Sep 2026 23:45:10 +0800 Subject: [PATCH] [SQL] Reduce allocations in ArrowColumnVector decimal reads Read native Decimal128 words for source precision 1-18 while preserving full-width fallback and Spark decimal conversion semantics. Add decimal regression coverage and a read benchmark. --- .../sql/vectorized/ArrowColumnVector.java | 41 +++- .../benchmark/ArrowDecimalReadBenchmark.scala | 94 ++++++++ .../vectorized/ArrowColumnVectorSuite.scala | 225 ++++++++++++++++++ 3 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowDecimalReadBenchmark.scala diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java index d6e2f3f2c12e6..e0cc93df5ea26 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java @@ -17,6 +17,9 @@ package org.apache.spark.sql.vectorized; +import java.math.BigDecimal; +import java.nio.ByteOrder; + import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.*; import org.apache.arrow.vector.complex.*; @@ -202,7 +205,12 @@ void initAccessor(ValueVector vector) { } else if (vector instanceof Float8Vector float8Vector) { accessor = new DoubleAccessor(float8Vector); } else if (vector instanceof DecimalVector decimalVector) { - accessor = new DecimalAccessor(decimalVector); + int precision = decimalVector.getPrecision(); + if (precision > 0 && precision <= Decimal.MAX_LONG_DIGITS()) { + accessor = new SmallDecimalAccessor(decimalVector, ByteOrder.nativeOrder()); + } else { + accessor = new DecimalAccessor(decimalVector); + } } else if (vector instanceof VarCharVector varCharVector) { accessor = new StringAccessor(varCharVector); } else if (vector instanceof LargeVarCharVector largeVarCharVector) { @@ -444,6 +452,37 @@ final double getDouble(int rowId) { } } + static class SmallDecimalAccessor extends ArrowVectorAccessor { + + private final DecimalVector accessor; + private final int lowWordOffset; + private final int highWordOffset; + + SmallDecimalAccessor(DecimalVector vector, ByteOrder byteOrder) { + super(vector); + this.accessor = vector; + this.lowWordOffset = byteOrder == ByteOrder.LITTLE_ENDIAN ? 0 : Long.BYTES; + this.highWordOffset = byteOrder == ByteOrder.LITTLE_ENDIAN ? Long.BYTES : 0; + } + + @Override + final Decimal getDecimal(int rowId, int precision, int scale) { + if (isNullAt(rowId)) return null; + // Arrow stores Decimal128 values in native byte order. + long offset = (long) rowId * DecimalVector.TYPE_WIDTH; + ArrowBuf data = accessor.getDataBuffer(); + long unscaled = data.getLong(offset + lowWordOffset); + long high = data.getLong(offset + highWordOffset); + // Preserve full-width values even when they exceed the declared source precision. + if (high == (unscaled >> 63)) { + // Keep BigDecimal-backed Decimal semantics, including checked integral conversions. + return Decimal.apply( + BigDecimal.valueOf(unscaled, accessor.getScale()), precision, scale); + } + return Decimal.apply(accessor.getObject(rowId), precision, scale); + } + } + static class DecimalAccessor extends ArrowVectorAccessor { private final DecimalVector accessor; diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowDecimalReadBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowDecimalReadBenchmark.scala new file mode 100644 index 0000000000000..f02cf8c496d62 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowDecimalReadBenchmark.scala @@ -0,0 +1,94 @@ +/* + * 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 java.math.{BigDecimal => JavaBigDecimal} + +import scala.util.Using + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.DecimalVector + +import org.apache.spark.benchmark.{Benchmark, BenchmarkBase} +import org.apache.spark.sql.types.Decimal +import org.apache.spark.sql.vectorized.ArrowColumnVector + +/** + * Measures Decimal128 reads through ArrowColumnVector, including scale conversion and wide values. + * To run this benchmark: + * {{{ + * build/sbt "sql/Test/runMain org.apache.spark.sql.execution.benchmark.ArrowDecimalReadBenchmark" + * }}} + * Set SPARK_GENERATE_BENCHMARK_FILES=1 to save results under benchmarks/. + */ +object ArrowDecimalReadBenchmark extends BenchmarkBase { + private val BatchSize = 4096 + private val BatchesPerIteration = 256 + private val results = new Array[Decimal](BatchSize) + + private def readDecimals(precision: Int, sourceScale: Int, targetScale: Int): Unit = { + Using.resource(new RootAllocator()) { allocator => + Using.resource(new DecimalVector("decimal", allocator, precision, sourceScale)) { vector => + vector.allocateNew(BatchSize) + val base = BigInt(10).pow(precision - 1) + for (row <- 0 until BatchSize) { + val unscaled = (base + row) * (if (row % 2 == 0) 1 else -1) + vector.set(row, new JavaBigDecimal(unscaled.bigInteger, sourceScale)) + } + vector.setValueCount(BatchSize) + val column = new ArrowColumnVector(vector) + val benchmark = new Benchmark( + s"decimal($precision,$sourceScale) to decimal($precision,$targetScale)", + BatchSize.toLong * BatchesPerIteration, output = output) + benchmark.addCase("Arrow getObject") { _ => + var batch = 0 + while (batch < BatchesPerIteration) { + var row = 0 + while (row < BatchSize) { + results(row) = if (vector.isNull(row)) null else { + Decimal(vector.getObject(row), precision, targetScale) + } + row += 1 + } + batch += 1 + } + } + benchmark.addCase("ArrowColumnVector") { _ => + var batch = 0 + while (batch < BatchesPerIteration) { + var row = 0 + while (row < BatchSize) { + results(row) = column.getDecimal(row, precision, targetScale) + row += 1 + } + batch += 1 + } + } + benchmark.run() + } + } + } + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + runBenchmark("Arrow decimal reads") { + readDecimals(18, 2, 2) + readDecimals(18, 2, 1) + readDecimals(38, 2, 2) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala index ea6bc4a78f02d..1e69555d0ede4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/vectorized/ArrowColumnVectorSuite.scala @@ -17,11 +17,20 @@ package org.apache.spark.sql.vectorized +import java.math.{BigDecimal => JavaBigDecimal} +import java.nio.ByteOrder + +import scala.util.{Failure, Random, Success, Try, Using} + +import org.apache.arrow.memory.RootAllocator import org.apache.arrow.vector._ import org.apache.arrow.vector.complex._ import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType} import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.expressions.UnsafeRow +import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.sql.util.ArrowUtils import org.apache.spark.unsafe.types.UTF8String @@ -224,6 +233,190 @@ class ArrowColumnVectorSuite extends SparkFunSuite { allocator.close() } + test("decimal precisions and scales") { + val random = new Random(1024) + for (precision <- 1 to Decimal.MAX_LONG_DIGITS) { + val limit = BigInt(10).pow(precision).toLong + val boundaries = (1 until precision).flatMap { digit => + val power = BigInt(10).pow(digit).toLong + Seq(power - 1, power, power + 1, 1 - power, -power, -power - 1) + } + val values = Seq(0L, 1L, -1L, limit - 1, 1 - limit) ++ boundaries ++ + Seq.fill(32)(random.nextLong() % limit) + for (scale <- 0 to precision) { + withDecimalVector(precision, scale) { vector => + vector.allocateNew(values.size + 1) + values.zipWithIndex.foreach { case (value, row) => + vector.set(row, JavaBigDecimal.valueOf(value, scale)) + } + vector.setNull(values.size) + vector.setValueCount(values.size + 1) + val column = new ArrowColumnVector(vector) + assert(column.dataType === DecimalType(precision, scale)) + assert(column.numNulls === 1) + for (row <- values.indices.reverse) { + checkDecimal(vector, column, row, precision, scale) + } + assert(column.getDecimal(values.size, precision, scale) === null) + } + } + } + } + + test("decimal native word order") { + val values = Seq(0L, 12345L, -12345L, 999999999999999999L, -999999999999999999L) + for (order <- Seq(ByteOrder.LITTLE_ENDIAN, ByteOrder.BIG_ENDIAN)) { + withDecimalVector(18, 2) { vector => + vector.allocateNew(values.size + 1) + values.zipWithIndex.foreach { case (low, row) => + vector.set(row, JavaBigDecimal.valueOf(low, 2)) + val offset = row.toLong * DecimalVector.TYPE_WIDTH + val high = low >> 63 + // Model native-word positions; ArrowBuf reads each long in native byte order. + vector.getDataBuffer.setLong(offset, if (order == ByteOrder.LITTLE_ENDIAN) low else high) + vector.getDataBuffer.setLong(offset + java.lang.Long.BYTES, + if (order == ByteOrder.LITTLE_ENDIAN) high else low) + } + vector.setNull(values.size) + vector.setValueCount(values.size + 1) + val accessor = new ArrowColumnVector.SmallDecimalAccessor(vector, order) + values.indices.foreach { row => + assert(accessor.getDecimal(row, 18, 2).toJavaBigDecimal === + JavaBigDecimal.valueOf(values(row), 2)) + } + assert(accessor.getDecimal(values.size, 18, 2) === null) + } + } + } + + test("decimal requested precision and scale") { + withDecimalVector(5, 2) { vector => + vector.allocateNew(4) + Seq("123.45", "-123.45", "999.95").zipWithIndex.foreach { case (value, row) => + vector.set(row, new JavaBigDecimal(value)) + } + vector.setNull(3) + vector.setValueCount(4) + val column = new ArrowColumnVector(vector) + assert(column.getDecimal(0, 5, 1).toJavaBigDecimal === new JavaBigDecimal("123.5")) + assert(column.getDecimal(1, 5, 1).toJavaBigDecimal === new JavaBigDecimal("-123.5")) + checkDecimal(vector, column, 0, 38, 18) + for (row <- 0 until 3; precision <- Seq(2, 4, 5, 8, 19); scale <- 0 to 3) { + checkDecimal(vector, column, row, precision, scale) + } + assert(column.getDecimal(3, 2, 0) === null) + } + } + + test("decimal wide source precision") { + for (precision <- Seq(19, 20, 38)) { + withDecimalVector(precision, 2) { vector => + val largest = BigInt(10).pow(precision) - 1 + vector.allocateNew(4) + vector.set(0, new JavaBigDecimal(largest.bigInteger, 2)) + vector.set(1, new JavaBigDecimal((-largest).bigInteger, 2)) + vector.set(2, new JavaBigDecimal("1.23")) + vector.setNull(3) + vector.setValueCount(4) + val column = new ArrowColumnVector(vector) + for (row <- 0 until 3; targetPrecision <- Seq(18, precision)) { + checkDecimal(vector, column, row, targetPrecision, 2) + } + assert(column.getDecimal(3, precision, 2) === null) + } + } + } + + test("decimal buffer values outside declared precision") { + withDecimalVector(18, 0) { vector => + val values = Seq( + (BigInt(1) << 64) + 7, -(BigInt(1) << 64) + 7, + BigInt(10).pow(18), -BigInt(10).pow(18), + BigInt(Long.MaxValue), BigInt(Long.MinValue), + BigInt(Long.MaxValue) + 1, BigInt(Long.MinValue) - 1) + vector.allocateNew(values.size) + values.zipWithIndex.foreach { case (value, row) => + // The byte setter permits values inconsistent with the declared precision. + vector.setBigEndian(row, value.toByteArray) + } + vector.setValueCount(values.size) + val column = new ArrowColumnVector(vector) + for (row <- values.indices; precision <- Seq(18, 38)) { + checkDecimal(vector, column, row, precision, 0) + } + vector.setNull(0) + assert(column.getDecimal(0, 18, 0) === null) + } + } + + test("decimal slices and returned values remain independent") { + withDecimalVector(10, 2) { vector => + vector.allocateNew(4) + vector.set(0, new JavaBigDecimal("999.99")) + vector.set(1, new JavaBigDecimal("-12.34")) + vector.setNull(2) + vector.set(3, new JavaBigDecimal("56.78")) + vector.setValueCount(4) + val transfer = vector.getTransferPair(vector.getAllocator) + transfer.splitAndTransfer(1, 3) + Using.resource(transfer.getTo.asInstanceOf[DecimalVector]) { slice => + val column = new ArrowColumnVector(slice) + val first = column.getDecimal(0, 10, 2) + val repeated = column.getDecimal(0, 10, 2) + val last = column.getDecimal(2, 10, 2) + assert(first ne repeated) + assert(first.toJavaBigDecimal === new JavaBigDecimal("-12.34")) + first.set(0L) + assert(repeated.toJavaBigDecimal === new JavaBigDecimal("-12.34")) + assert(last.toJavaBigDecimal === new JavaBigDecimal("56.78")) + assert(column.getDecimal(1, 10, 2) === null) + } + } + } + + test("decimal negative scale") { + val conf = new SQLConf + val key = SQLConf.LEGACY_ALLOW_NEGATIVE_SCALE_OF_DECIMAL_ENABLED.key + SQLConf.withExistingConf(conf) { + conf.setConfString(key, "true") + withDecimalVector(5, -2) { vector => + vector.allocateNew(1) + vector.set(0, new JavaBigDecimal("1.23E+4")) + vector.setValueCount(1) + val column = new ArrowColumnVector(vector) + Seq((5, -2), (8, 0), (5, -3)).foreach { case (precision, scale) => + checkDecimal(vector, column, 0, precision, scale) + } + assert(column.getDecimal(0, 5, -2).toLong === 12300L) + conf.setConfString(key, "false") + checkDecimal(vector, column, 0, 5, -2) + } + } + } + + test("decimal checked integral conversions") { + val values = Seq( + "127.999999999999999", "-128.999999999999999", + "32767.9999999999999", "-32768.9999999999999", + "2147483647.99999999", "-2147483648.99999999", "123456789.999999999") + values.foreach { value => + val input = new JavaBigDecimal(value) + val precision = input.precision() + val scale = input.scale() + withDecimalVector(precision, scale) { vector => + vector.allocateNew(1) + vector.set(0, input) + vector.setValueCount(1) + val expected = Decimal(vector.getObject(0), precision, scale) + val actual = new ArrowColumnVector(vector).getDecimal(0, precision, scale) + checkDecimalResult(expected.roundToByte(), actual.roundToByte()) + checkDecimalResult(expected.roundToShort(), actual.roundToShort()) + checkDecimalResult(expected.roundToInt(), actual.roundToInt()) + checkDecimalResult(expected.roundToLong(), actual.roundToLong()) + } + } + } + test("string") { val allocator = ArrowUtils.rootAllocator.newChildAllocator("string", 0, Long.MaxValue) val vector = ArrowUtils.toArrowField("string", StringType, nullable = true, null) @@ -684,4 +877,36 @@ class ArrowColumnVectorSuite extends SparkFunSuite { columnVector.close() allocator.close() } + + private def withDecimalVector(precision: Int, scale: Int)(f: DecimalVector => Unit): Unit = { + Using.resource(new RootAllocator()) { allocator => + Using.resource(new DecimalVector("decimal", allocator, precision, scale))(f) + } + } + + private def checkDecimal( + vector: DecimalVector, + column: ArrowColumnVector, + row: Int, + precision: Int, + scale: Int): Unit = { + def result(decimal: Decimal): (JavaBigDecimal, Int, Int, Long, UnsafeRow) = { + val writer = new UnsafeRowWriter(1) + writer.write(0, decimal, precision, scale) + (decimal.toJavaBigDecimal, decimal.precision, decimal.scale, decimal.toLong, writer.getRow) + } + checkDecimalResult( + result(Decimal(vector.getObject(row), precision, scale)), + result(column.getDecimal(row, precision, scale))) + } + + private def checkDecimalResult[T](expected: => T, actual: => T): Unit = { + Try(expected) match { + case Success(value) => assert(actual === value) + case Failure(error) => + val actualError = intercept[Exception](actual) + assert(actualError.getClass === error.getClass) + assert(actualError.getMessage === error.getMessage) + } + } }