Skip to content

feat: reduce allocations in Spark Decimal access #9842

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
robert3005 merged 5 commits into from
Sep 24, 2026
Merged

feat: reduce allocations in Spark Decimal access #9842

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f3a1eb6
feat: reduce allocations in Spark Decimal access
xiaoh1024 Sep 11, 2026
3f9b5d3
Merge branch 'develop' into exp/decimal-accessor-pr
xiaoh1024 Sep 15, 2026
98a72d4
Merge branch 'develop' into exp/decimal-accessor-pr
xiaoh1024 Sep 21, 2026
f4185cd
Merge branch 'develop' into exp/decimal-accessor-pr
xiaoh1024 Sep 21, 2026
614ad61
test: fix Decimal accessor row round-trip checks
xiaoh1024 Sep 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 40 additions & 1 deletion java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexArrowColumnVector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit;
import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType;
import dev.vortex.spark.ArrowUtils;
import java.math.BigDecimal;
import java.nio.ByteOrder;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Decimal;
Expand Down Expand Up @@ -309,7 +311,12 @@ void initAccessor(ValueVector vector) {
} else if (vector instanceof Float8Vector float8Vector) {
accessor = new VortexArrowColumnVector.DoubleAccessor(float8Vector);
} else if (vector instanceof DecimalVector decimalVector) {
accessor = new VortexArrowColumnVector.DecimalAccessor(decimalVector);
int precision = decimalVector.getPrecision();
if (precision > 0 && precision <= Decimal.MAX_LONG_DIGITS()) {
accessor = new VortexArrowColumnVector.SmallDecimalAccessor(decimalVector, ByteOrder.nativeOrder());
} else {
accessor = new VortexArrowColumnVector.DecimalAccessor(decimalVector);
}
} else if (vector instanceof VarCharVector varCharVector) {
accessor = new VortexArrowColumnVector.StringAccessor(varCharVector);
} else if (vector instanceof LargeVarCharVector largeVarCharVector) {
Expand Down Expand Up @@ -556,6 +563,38 @@ final Decimal getDecimal(int rowId, int precision, int scale) {
}
}

static final class SmallDecimalAccessor extends VortexArrowColumnVector.ArrowVectorAccessor {

private final DecimalVector accessor;
private final int lowWordOffset;
private final int highWordOffset;

SmallDecimalAccessor(DecimalVector vector, ByteOrder byteOrder) {
super(vector);
this.accessor = vector;
// Arrow Java stores Decimal128 in native byte order, including the order of its two words.
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;
}
long offset = (long) rowId * DecimalVector.TYPE_WIDTH;
long unscaled = accessor.getDataBuffer().getLong(offset + lowWordOffset);
long high = accessor.getDataBuffer().getLong(offset + highWordOffset);
if (high == (unscaled >> 63)) {
// Decode with the source scale; Spark performs the requested rescaling and overflow checks.
// Keep the expanded Decimal representation for Spark's checked integer casts.
return Decimal.apply(BigDecimal.valueOf(unscaled, accessor.getScale()), precision, scale);

@robert3005 robert3005 Sep 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we use Decimal.apply(unscaled, p, s) ? What does BigDecimal.valueOf change here?

@xiaoh1024 xiaoh1024 Sep 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I initially tried Decimal.apply(unscaled, p, s), but it changes the internal representation and some downstream conversion behavior.

In Spark 3.5.9 and 4.1.2, for example, 127.999999999999999 throws an overflow exception on roundToByte() through the existing BigDecimal-based path, while the compact Decimal created by the long overload returns 127. The checked integer conversion test covers this difference.

Using BigDecimal.valueOf(unscaled, accessor.getScale()) preserves that behavior. It also reconstructs the value with the source scale, before Spark applies the requested precision and scale. Directly using the requested scale with the unscaled integer would change the value when the scales differ.

We still avoid the temporary byte array and BigInteger created by Arrow’s getObject(), while retaining the existing Spark Decimal behavior.

}
// Preserve all 128 bits when the stored integer does not fit in a long.
return Decimal.apply(accessor.getObject(rowId), precision, scale);
}
}

@Open
static class StringAccessor extends VortexArrowColumnVector.ArrowVectorAccessor {

Expand Down
315 changes: 315 additions & 0 deletions java/vortex-spark/src/test/java/dev/vortex/spark/read/VortexDecimalAccessorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,315 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

package dev.vortex.spark.read;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotSame;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;

import dev.vortex.relocated.org.apache.arrow.memory.RootAllocator;
import dev.vortex.relocated.org.apache.arrow.vector.DecimalVector;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.function.Supplier;
import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter;
import org.apache.spark.sql.internal.SQLConf;
import org.apache.spark.sql.types.Decimal;
import org.junit.jupiter.api.Test;

final class VortexDecimalAccessorTest {
@Test
void compactPrecisionsAndScalesMatchArrowValues() {
Random random = new Random(1024);
try (RootAllocator allocator = new RootAllocator()) {
long limit = 1;
for (int precision = 1; precision <= Decimal.MAX_LONG_DIGITS(); precision++) {
limit *= 10;
List<Long> values = new ArrayList<>(List.of(0L, 1L, -1L, limit - 1, 1 - limit));
for (long power = 10; power < limit; power *= 10) {
values.add(power - 1);
values.add(power);
values.add(power + 1);
values.add(1 - power);
values.add(-power);
values.add(-power - 1);
}
for (int i = 0; i < 32; i++) {
values.add(random.nextLong() % limit);
}
for (int scale = 0; scale <= precision; scale++) {
try (DecimalVector vector = new DecimalVector("decimal", allocator, precision, scale)) {
vector.allocateNew(values.size() + 1);
for (int row = 0; row < values.size(); row++) {
vector.set(row, BigDecimal.valueOf(values.get(row), scale));
}
vector.setNull(values.size());
vector.setValueCount(values.size() + 1);
VortexArrowColumnVector column = new VortexArrowColumnVector(vector);
for (int row = values.size() - 1; row >= 0; row--) {
BigDecimal expected = BigDecimal.valueOf(values.get(row), scale);
Decimal actual = column.getDecimal(row, precision, scale);
assertEquivalent(Decimal.apply(expected, precision, scale), actual, precision, scale);
}
assertNull(column.getDecimal(values.size(), precision, scale));
}
}
}
}
}

@Test
void smallDecimalReadsNativeWordsInEitherOrder() {
String[] inputs = {"0.00", "123.45", "-123.45", "9999999999999999.99", "-9999999999999999.99"};
try (RootAllocator allocator = new RootAllocator()) {
for (ByteOrder order : new ByteOrder[] {ByteOrder.LITTLE_ENDIAN, ByteOrder.BIG_ENDIAN}) {
try (DecimalVector vector = new DecimalVector("decimal", allocator, 18, 2)) {
vector.allocateNew(inputs.length + 1);
for (int row = 0; row < inputs.length; row++) {
BigDecimal value = new BigDecimal(inputs[row]);
vector.set(row, value);
long offset = (long) row * DecimalVector.TYPE_WIDTH;
long low = value.unscaledValue().longValueExact();
long high = value.signum() < 0 ? -1L : 0L;
// Model the two native-word positions; ArrowBuf supplies native-endian long reads.
vector.getDataBuffer().setLong(offset, order == ByteOrder.LITTLE_ENDIAN ? low : high);
vector.getDataBuffer()
.setLong(offset + Long.BYTES, order == ByteOrder.LITTLE_ENDIAN ? high : low);
}
vector.setNull(inputs.length);
vector.setValueCount(inputs.length + 1);
var accessor = new VortexArrowColumnVector.SmallDecimalAccessor(vector, order);
for (int row = 0; row < inputs.length; row++) {
BigDecimal expected = new BigDecimal(inputs[row]);
assertEquivalent(Decimal.apply(expected, 18, 2), accessor.getDecimal(row, 18, 2), 18, 2);
}
assertNull(accessor.getDecimal(inputs.length, 18, 2));
}
}
}
}

@Test
void requestedScaleAndPrecisionRetainRoundingAndOverflow() {
try (RootAllocator allocator = new RootAllocator();
DecimalVector vector = new DecimalVector("decimal", allocator, 5, 2)) {
vector.allocateNew(4);
vector.set(0, new BigDecimal("123.45"));
vector.set(1, new BigDecimal("-123.45"));
vector.set(2, new BigDecimal("999.95"));
vector.setNull(3);
vector.setValueCount(4);
VortexArrowColumnVector column = new VortexArrowColumnVector(vector);
assertEquals(new BigDecimal("123.5"), column.getDecimal(0, 5, 1).toJavaBigDecimal());
assertEquals(new BigDecimal("-123.5"), column.getDecimal(1, 5, 1).toJavaBigDecimal());
assertEquals(
new BigDecimal("123.450000000000000000"),
column.getDecimal(0, 38, 18).toJavaBigDecimal());
for (int row = 0; row < 3; row++) {
for (int precision : new int[] {2, 4, 5, 8, 19}) {
for (int scale : new int[] {0, 1, 2, 3}) {
assertMatchesReference(vector, column, row, precision, scale);
}
}
}
assertNull(column.getDecimal(3, 2, 0));
}
}

@Test
void wideVectorsRetainAllBitsEvenWhenRequestedPrecisionIsSmall() {
try (RootAllocator allocator = new RootAllocator()) {
for (int precision : new int[] {19, 20, 38}) {
BigInteger largest = BigInteger.TEN.pow(precision).subtract(BigInteger.ONE);
try (DecimalVector vector = new DecimalVector("decimal", allocator, precision, 2)) {
vector.allocateNew(4);
vector.set(0, new BigDecimal(largest, 2));
vector.set(1, new BigDecimal(largest.negate(), 2));
vector.set(2, new BigDecimal("1.23"));
vector.setNull(3);
vector.setValueCount(4);
VortexArrowColumnVector column = new VortexArrowColumnVector(vector);
for (int row = 0; row < 3; row++) {
assertMatchesReference(vector, column, row, precision, 2);
assertMatchesReference(vector, column, row, 18, 2);
}
assertNull(column.getDecimal(3, precision, 2));
}
}
}
}

@Test
void outOfPrecisionBufferValuesRetainReferenceResultsAndErrors() {
try (RootAllocator allocator = new RootAllocator();
DecimalVector vector = new DecimalVector("decimal", allocator, 18, 0)) {
// The byte setter permits values inconsistent with the field's precision.
BigInteger[] invalid = {
BigInteger.ONE.shiftLeft(64).add(BigInteger.valueOf(7)),
BigInteger.ONE.shiftLeft(64).negate().add(BigInteger.valueOf(7)),
BigInteger.TEN.pow(18),
BigInteger.TEN.pow(18).negate(),
BigInteger.valueOf(Long.MAX_VALUE),
BigInteger.valueOf(Long.MIN_VALUE),
BigInteger.valueOf(Long.MAX_VALUE).add(BigInteger.ONE),
BigInteger.valueOf(Long.MIN_VALUE).subtract(BigInteger.ONE)
};
vector.allocateNew(invalid.length);
for (int row = 0; row < invalid.length; row++) {
vector.setBigEndian(row, invalid[row].toByteArray());
}
vector.setValueCount(invalid.length);
VortexArrowColumnVector column = new VortexArrowColumnVector(vector);
for (int row = 0; row < invalid.length; row++) {
assertMatchesReference(vector, column, row, 18, 0);
assertMatchesReference(vector, column, row, 38, 0);
}
vector.setNull(0);
assertNull(column.getDecimal(0, 18, 0));
}
}

@Test
void slicedVectorsAndReturnedValuesRemainIndependent() {
try (RootAllocator allocator = new RootAllocator();
DecimalVector vector = new DecimalVector("decimal", allocator, 10, 2)) {
vector.allocateNew(4);
vector.set(0, new BigDecimal("999.99"));
vector.set(1, new BigDecimal("-12.34"));
vector.setNull(2);
vector.set(3, new BigDecimal("56.78"));
vector.setValueCount(4);
var transfer = vector.getTransferPair(allocator);
transfer.splitAndTransfer(1, 3);
try (DecimalVector slice = (DecimalVector) transfer.getTo()) {
VortexArrowColumnVector column = new VortexArrowColumnVector(slice);
Decimal first = column.getDecimal(0, 10, 2);
Decimal repeated = column.getDecimal(0, 10, 2);
Decimal last = column.getDecimal(2, 10, 2);
assertNotSame(first, repeated);
assertEquals(new BigDecimal("-12.34"), first.toJavaBigDecimal());
assertEquals(new BigDecimal("56.78"), last.toJavaBigDecimal());
first.set(0L);
assertEquals(new BigDecimal("-12.34"), repeated.toJavaBigDecimal());
assertEquals(new BigDecimal("56.78"), last.toJavaBigDecimal());
assertNull(column.getDecimal(1, 10, 2));
}
}
}

@Test
void negativeScaleRetainsReferenceConversionsAndValidation() {
SQLConf conf = SQLConf.get();
String key = "spark.sql.legacy.allowNegativeScaleOfDecimal";
String previous = conf.getConfString(key, "false");
boolean wasSet = conf.getAllConfs().contains(key);
conf.setConfString(key, "true");
try (RootAllocator allocator = new RootAllocator();
DecimalVector vector = new DecimalVector("decimal", allocator, 5, -2)) {
vector.allocateNew(1);
vector.set(0, new BigDecimal("1.23E+4"));
vector.setValueCount(1);
VortexArrowColumnVector column = new VortexArrowColumnVector(vector);
assertMatchesReference(vector, column, 0, 5, -2);
assertMatchesReference(vector, column, 0, 8, 0);
assertMatchesReference(vector, column, 0, 5, -3);
assertEquals(12300L, column.getDecimal(0, 5, -2).toLong());
conf.setConfString(key, "false");
assertMatchesReference(vector, column, 0, 5, -2);
} finally {
if (wasSet) {
conf.setConfString(key, previous);
} else {
conf.unsetConf(key);
}
}
}

@Test
void checkedIntegralConversionsRetainResultsAndOverflow() {
String[] values = {
"127.999999999999999",
"-128.999999999999999",
"32767.9999999999999",
"-32768.9999999999999",
"2147483647.99999999",
"-2147483648.99999999",
"123456789.999999999"
};
try (RootAllocator allocator = new RootAllocator()) {
for (String value : values) {
BigDecimal input = new BigDecimal(value);
int precision = input.precision();
int scale = input.scale();
try (DecimalVector vector = new DecimalVector("decimal", allocator, precision, scale)) {
vector.allocateNew(1);
vector.set(0, input);
vector.setValueCount(1);
VortexArrowColumnVector column = new VortexArrowColumnVector(vector);
Decimal expected = Decimal.apply(vector.getObject(0), precision, scale);
Decimal actual = column.getDecimal(0, precision, scale);
assertConversion(expected::roundToByte, actual::roundToByte);
assertConversion(expected::roundToShort, actual::roundToShort);
assertConversion(expected::roundToInt, actual::roundToInt);
assertConversion(expected::roundToLong, actual::roundToLong);
}
}
}
}

private static void assertConversion(Supplier<Object> expectedConversion, Supplier<Object> actualConversion) {
Object expected;
try {
expected = expectedConversion.get();
} catch (Exception expectedError) {
Exception actualError = assertThrows(expectedError.getClass(), actualConversion::get);
assertEquals(expectedError.getClass(), actualError.getClass());
assertEquals(expectedError.getMessage(), actualError.getMessage());
return;
}
assertEquals(expected, actualConversion.get());
}

private static void assertMatchesReference(
DecimalVector vector, VortexArrowColumnVector column, int row, int precision, int scale) {
Decimal expected;
try {
expected = Decimal.apply(vector.getObject(row), precision, scale);
} catch (Exception expectedError) {
Exception actualError =
assertThrows(expectedError.getClass(), () -> column.getDecimal(row, precision, scale));
assertEquals(expectedError.getClass(), actualError.getClass());
assertEquals(expectedError.getMessage(), actualError.getMessage());
return;
}
assertEquivalent(expected, column.getDecimal(row, precision, scale), precision, scale);
}

private static void assertEquivalent(Decimal expected, Decimal actual, int precision, int scale) {
assertEquals(expected.toJavaBigDecimal(), actual.toJavaBigDecimal());
assertEquals(expected.precision(), actual.precision());
assertEquals(expected.scale(), actual.scale());
assertEquals(expected.hashCode(), actual.hashCode());
assertEquals(0, expected.compare(actual));
assertEquals(expected.toLong(), actual.toLong());
assertEquals(expected, actual.clone());
Decimal expectedCopy = expected.clone();
Decimal actualCopy = actual.clone();
assertEquals(expectedCopy.changePrecision(precision, 0), actualCopy.changePrecision(precision, 0));
assertEquals(expectedCopy.toJavaBigDecimal(), actualCopy.toJavaBigDecimal());
UnsafeRowWriter expectedWriter = new UnsafeRowWriter(1);
Comment thread
xiaoh1024 marked this conversation as resolved.
UnsafeRowWriter actualWriter = new UnsafeRowWriter(1);
expectedWriter.resetRowWriter();
actualWriter.resetRowWriter();
expectedWriter.write(0, expected, precision, scale);
actualWriter.write(0, actual, precision, scale);
assertEquals(expected, expectedWriter.getRow().getDecimal(0, precision, scale));
assertEquals(actual, actualWriter.getRow().getDecimal(0, precision, scale));
assertEquals(expectedWriter.getRow(), actualWriter.getRow());
}
}
Loading