diff --git a/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java b/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java index f9d7e5cb9e..ddfc464152 100644 --- a/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java +++ b/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java @@ -23,6 +23,7 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.impl.Decimal256ReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.Decimal256Holder; @@ -205,8 +206,9 @@ public void set(int index, ArrowBuf buffer) { * @param value array of bytes containing decimal in big endian byte order. */ public void setBigEndian(int index, byte[] value) { - BitVectorHelper.setBit(validityBuffer, index); final int length = value.length; + checkValueLength(length); + BitVectorHelper.setBit(validityBuffer, index); // do the bound check. valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); @@ -222,29 +224,26 @@ public void setBigEndian(int index, byte[] value) { MemoryUtil.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]); } - if (length == TYPE_WIDTH) { - return; - } - if (length < TYPE_WIDTH) { // sign extend final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad); - return; } } else { - if (length <= TYPE_WIDTH) { - // copy data from value to outAddress - MemoryUtil.copyToMemory( - value, 0, outAddress + Decimal256Vector.TYPE_WIDTH - length, length); - // sign extend - final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); - MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad); - return; - } + // copy data from value to outAddress + MemoryUtil.copyToMemory(value, 0, outAddress + Decimal256Vector.TYPE_WIDTH - length, length); + // sign extend + final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad); } - throw new IllegalArgumentException( - "Invalid decimal value length. Valid length in [1 - 32], got " + length); + } + + private static void checkValueLength(int length) { + Preconditions.checkArgument( + length >= 0 && length <= TYPE_WIDTH, + "Invalid decimal value length. Valid length in [0 - %s], got %s", + TYPE_WIDTH, + length); } /** @@ -268,6 +267,7 @@ public void set(int index, long start, ArrowBuf buffer) { * @param length length of the value in the buffer */ public void setSafe(int index, long start, ArrowBuf buffer, int length) { + checkValueLength(length); handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); @@ -277,6 +277,10 @@ public void setSafe(int index, long start, ArrowBuf buffer, int length) { long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (length == 0) { + MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH, (byte) 0); + return; + } if (LITTLE_ENDIAN) { MemoryUtil.copyMemory(inAddress, outAddress, length); // sign extend @@ -305,6 +309,7 @@ public void setSafe(int index, long start, ArrowBuf buffer, int length) { * @param length length of the value in the buffer */ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) { + checkValueLength(length); handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); @@ -315,6 +320,10 @@ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) // not using buffer.getByte() to avoid boundary checks for every byte. long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (length == 0) { + MemoryUtil.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH, (byte) 0); + return; + } if (LITTLE_ENDIAN) { // swap bytes to convert BE to LE for (int byteIdx = 0; byteIdx < length; ++byteIdx) { @@ -322,7 +331,7 @@ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) MemoryUtil.putByte(outAddress + byteIdx, val); } // sign extend - if (length < 32) { + if (length < TYPE_WIDTH) { byte msb = MemoryUtil.getByte(inAddress); final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad); diff --git a/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java b/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java index 9bf1812cc6..927fe94fa1 100644 --- a/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java +++ b/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java @@ -23,6 +23,7 @@ import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.complex.impl.DecimalReaderImpl; import org.apache.arrow.vector.complex.reader.FieldReader; import org.apache.arrow.vector.holders.DecimalHolder; @@ -204,8 +205,9 @@ public void set(int index, ArrowBuf buffer) { * @param value array of bytes containing decimal in big endian byte order. */ public void setBigEndian(int index, byte[] value) { - BitVectorHelper.setBit(validityBuffer, index); final int length = value.length; + checkValueLength(length); + BitVectorHelper.setBit(validityBuffer, index); // do the bound check. valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); @@ -221,28 +223,26 @@ public void setBigEndian(int index, byte[] value) { MemoryUtil.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]); } - if (length == TYPE_WIDTH) { - return; - } - if (length < TYPE_WIDTH) { // sign extend final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); MemoryUtil.setMemory(outAddress + length, DecimalVector.TYPE_WIDTH - length, pad); - return; } } else { - if (length <= TYPE_WIDTH) { - // copy data from value to outAddress - MemoryUtil.copyToMemory(value, 0, outAddress + DecimalVector.TYPE_WIDTH - length, length); - // sign extend - final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); - MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH - length, pad); - return; - } + // copy data from value to outAddress + MemoryUtil.copyToMemory(value, 0, outAddress + DecimalVector.TYPE_WIDTH - length, length); + // sign extend + final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH - length, pad); } - throw new IllegalArgumentException( - "Invalid decimal value length. Valid length in [1 - 16], got " + length); + } + + private static void checkValueLength(int length) { + Preconditions.checkArgument( + length >= 0 && length <= TYPE_WIDTH, + "Invalid decimal value length. Valid length in [0 - %s], got %s", + TYPE_WIDTH, + length); } /** @@ -266,6 +266,7 @@ public void set(int index, long start, ArrowBuf buffer) { * @param length length of the value in the buffer */ public void setSafe(int index, long start, ArrowBuf buffer, int length) { + checkValueLength(length); handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); @@ -275,6 +276,10 @@ public void setSafe(int index, long start, ArrowBuf buffer, int length) { long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (length == 0) { + MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH, (byte) 0); + return; + } if (LITTLE_ENDIAN) { MemoryUtil.copyMemory(inAddress, outAddress, length); // sign extend @@ -303,6 +308,7 @@ public void setSafe(int index, long start, ArrowBuf buffer, int length) { * @param length length of the value in the buffer */ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) { + checkValueLength(length); handleSafe(index); BitVectorHelper.setBit(validityBuffer, index); @@ -313,6 +319,10 @@ public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) // not using buffer.getByte() to avoid boundary checks for every byte. long inAddress = buffer.memoryAddress() + start; long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (length == 0) { + MemoryUtil.setMemory(outAddress, DecimalVector.TYPE_WIDTH, (byte) 0); + return; + } if (LITTLE_ENDIAN) { // swap bytes to convert BE to LE for (int byteIdx = 0; byteIdx < length; ++byteIdx) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java b/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java index b995dc5d92..7de109f1e5 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDecimal256Vector.java @@ -17,12 +17,14 @@ package org.apache.arrow.vector; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.math.BigDecimal; import java.math.BigInteger; +import java.util.function.BiConsumer; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -268,7 +270,7 @@ public void decimalBE2NE() { IllegalArgumentException.class, () -> decimalVector.setBigEndian(insertionIdxCapture, new byte[33])); assertTrue( - ex.getMessage().equals("Invalid decimal value length. Valid length in [1 - 32], got 33")); + ex.getMessage().equals("Invalid decimal value length. Valid length in [0 - 32], got 33")); decimalVector.setValueCount(insertionIdx); // retrieve values and check if they are correct @@ -419,4 +421,77 @@ private void verifyWritingArrowBufWithBigEndianBytes( assertEquals(expectedValues[i], value); } } + + private static final int OVERSIZED_LENGTH = Decimal256Vector.TYPE_WIDTH + 1; + + private Decimal256Vector newVectorWithFirstEightValid() { + Decimal256Vector vector = new Decimal256Vector("decimal", allocator, 38, 0); + vector.allocateNew(64); + for (int i = 0; i < 8; i++) { + vector.setSafe(i, BigDecimal.valueOf(i)); + } + return vector; + } + + // Validity bits sit right after the value buffer, so an overrun of the last slot clobbers them. + private static void assertFirstEightValid(Decimal256Vector vector) { + for (int i = 0; i < 8; i++) { + assertFalse(vector.isNull(i), "validity bit " + i + " was overwritten by an oversized write"); + assertEquals(BigDecimal.valueOf(i), vector.getObject(i)); + } + } + + private void assertOversizedWriteRejected(BiConsumer write) { + try (Decimal256Vector vector = newVectorWithFirstEightValid(); + ArrowBuf src = allocator.buffer(OVERSIZED_LENGTH)) { + src.setZero(0, OVERSIZED_LENGTH); + assertThrows(IllegalArgumentException.class, () -> write.accept(vector, src)); + assertFirstEightValid(vector); + } + } + + @Test + public void setBigEndianSafeRejectsOversizedByteArray() { + assertOversizedWriteRejected( + (vector, src) -> + vector.setBigEndianSafe(vector.getValueCapacity() - 1, new byte[OVERSIZED_LENGTH])); + } + + @Test + public void setSafeFromArrowBufRejectsOversizedLength() { + assertOversizedWriteRejected( + (vector, src) -> vector.setSafe(vector.getValueCapacity() - 1, 0, src, OVERSIZED_LENGTH)); + } + + @Test + public void setBigEndianSafeFromArrowBufRejectsOversizedLength() { + assertOversizedWriteRejected( + (vector, src) -> + vector.setBigEndianSafe(vector.getValueCapacity() - 1, 0, src, OVERSIZED_LENGTH)); + } + + @Test + public void setSafeFromArrowBufWithZeroLengthDoesNotReadBeforeStart() { + try (Decimal256Vector vector = new Decimal256Vector("decimal", allocator, 38, 0); + ArrowBuf src = allocator.buffer(2)) { + vector.allocateNew(1); + // A negative byte immediately before `start` must not be used for sign extension. + src.setByte(0, (byte) 0xFF); + src.setByte(1, (byte) 0x00); + vector.setSafe(0, 1, src, 0); + assertEquals(BigDecimal.ZERO, vector.getObject(0)); + } + } + + @Test + public void setBigEndianSafeFromArrowBufWithZeroLengthDoesNotReadAtStart() { + try (Decimal256Vector vector = new Decimal256Vector("decimal", allocator, 38, 0); + ArrowBuf src = allocator.buffer(1)) { + vector.allocateNew(1); + // A negative byte at `start` must not be used for sign extension of an empty value. + src.setByte(0, (byte) 0xFF); + vector.setBigEndianSafe(0, 0, src, 0); + assertEquals(BigDecimal.ZERO, vector.getObject(0)); + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java b/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java index 85c11e8f3d..cca1bfe854 100644 --- a/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java +++ b/vector/src/test/java/org/apache/arrow/vector/TestDecimalVector.java @@ -17,12 +17,15 @@ package org.apache.arrow.vector; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import java.math.BigDecimal; import java.math.BigInteger; +import java.util.function.BiConsumer; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.types.pojo.ArrowType; @@ -263,7 +266,7 @@ public void decimalBE2NE() { } catch (IllegalArgumentException ex) { assertTrue( ex.getMessage() - .equals("Invalid decimal value length. Valid length in [1 - 16], got 17")); + .equals("Invalid decimal value length. Valid length in [0 - 16], got 17")); } decimalVector.setValueCount(insertionIdx); @@ -413,4 +416,77 @@ private void verifyWritingArrowBufWithBigEndianBytes( assertEquals(expectedValues[i], value); } } + + private static final int OVERSIZED_LENGTH = DecimalVector.TYPE_WIDTH + 1; + + private DecimalVector newVectorWithFirstEightValid() { + DecimalVector vector = new DecimalVector("decimal", allocator, 38, 0); + vector.allocateNew(64); + for (int i = 0; i < 8; i++) { + vector.setSafe(i, BigDecimal.valueOf(i)); + } + return vector; + } + + // Validity bits sit right after the value buffer, so an overrun of the last slot clobbers them. + private static void assertFirstEightValid(DecimalVector vector) { + for (int i = 0; i < 8; i++) { + assertFalse(vector.isNull(i), "validity bit " + i + " was overwritten by an oversized write"); + assertEquals(BigDecimal.valueOf(i), vector.getObject(i)); + } + } + + private void assertOversizedWriteRejected(BiConsumer write) { + try (DecimalVector vector = newVectorWithFirstEightValid(); + ArrowBuf src = allocator.buffer(OVERSIZED_LENGTH)) { + src.setZero(0, OVERSIZED_LENGTH); + assertThrows(IllegalArgumentException.class, () -> write.accept(vector, src)); + assertFirstEightValid(vector); + } + } + + @Test + public void setBigEndianSafeRejectsOversizedByteArray() { + assertOversizedWriteRejected( + (vector, src) -> + vector.setBigEndianSafe(vector.getValueCapacity() - 1, new byte[OVERSIZED_LENGTH])); + } + + @Test + public void setSafeFromArrowBufRejectsOversizedLength() { + assertOversizedWriteRejected( + (vector, src) -> vector.setSafe(vector.getValueCapacity() - 1, 0, src, OVERSIZED_LENGTH)); + } + + @Test + public void setBigEndianSafeFromArrowBufRejectsOversizedLength() { + assertOversizedWriteRejected( + (vector, src) -> + vector.setBigEndianSafe(vector.getValueCapacity() - 1, 0, src, OVERSIZED_LENGTH)); + } + + @Test + public void setSafeFromArrowBufWithZeroLengthDoesNotReadBeforeStart() { + try (DecimalVector vector = new DecimalVector("decimal", allocator, 38, 0); + ArrowBuf src = allocator.buffer(2)) { + vector.allocateNew(1); + // A negative byte immediately before `start` must not be used for sign extension. + src.setByte(0, (byte) 0xFF); + src.setByte(1, (byte) 0x00); + vector.setSafe(0, 1, src, 0); + assertEquals(BigDecimal.ZERO, vector.getObject(0)); + } + } + + @Test + public void setBigEndianSafeFromArrowBufWithZeroLengthDoesNotReadAtStart() { + try (DecimalVector vector = new DecimalVector("decimal", allocator, 38, 0); + ArrowBuf src = allocator.buffer(1)) { + vector.allocateNew(1); + // A negative byte at `start` must not be used for sign extension of an empty value. + src.setByte(0, (byte) 0xFF); + vector.setBigEndianSafe(0, 0, src, 0); + assertEquals(BigDecimal.ZERO, vector.getObject(0)); + } + } }