diff --git a/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java b/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java index e7f3e3935d2..013d60e1db7 100644 --- a/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java +++ b/java/vector/src/main/codegen/templates/UnionFixedSizeListWriter.java @@ -174,16 +174,25 @@ public void end() { @Override public void write(DecimalHolder holder) { + if (writer.idx() >= (idx() + 1) * listSize) { + throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); + } writer.write(holder); writer.setPosition(writer.idx() + 1); } public void writeDecimal(int start, ArrowBuf buffer) { + if (writer.idx() >= (idx() + 1) * listSize) { + throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); + } writer.writeDecimal(start, buffer); writer.setPosition(writer.idx() + 1); } public void writeDecimal(BigDecimal value) { + if (writer.idx() >= (idx() + 1) * listSize) { + throw new IllegalStateException(String.format("values at index %s is greater than listSize %s", idx(), listSize)); + } writer.writeDecimal(value); writer.setPosition(writer.idx() + 1); } diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java index 6843daf5cf0..df3ba85a100 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestFixedSizeListVector.java @@ -19,8 +19,11 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; +import java.math.BigDecimal; import java.util.Arrays; +import java.util.List; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.complex.FixedSizeListVector; @@ -265,6 +268,54 @@ public void testUnionFixedSizeListWriter() throws Exception { } } + @Test + public void testWriteDecimal() throws Exception { + try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /*listSize=*/3, allocator)) { + + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + final int valueCount = 100; + + for (int i = 0; i < valueCount; i++) { + writer.startList(); + writer.decimal().writeDecimal(new BigDecimal(i)); + writer.decimal().writeDecimal(new BigDecimal(i * 2)); + writer.decimal().writeDecimal(new BigDecimal(i * 3)); + writer.endList(); + } + vector.setValueCount(valueCount); + + for (int i = 0; i < valueCount; i++) { + List values = (List) vector.getObject(i); + assertEquals(3, values.size()); + assertEquals(new BigDecimal(i), values.get(0)); + assertEquals(new BigDecimal(i * 2), values.get(1)); + assertEquals(new BigDecimal(i * 3), values.get(2)); + } + } + } + + @Test + public void testDecimalIndexCheck() throws Exception { + try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /*listSize=*/3, allocator)) { + + UnionFixedSizeListWriter writer = vector.getWriter(); + writer.allocate(); + + IllegalStateException e = assertThrows(IllegalStateException.class, () -> { + writer.startList(); + writer.decimal().writeDecimal(new BigDecimal(1)); + writer.decimal().writeDecimal(new BigDecimal(2)); + writer.decimal().writeDecimal(new BigDecimal(3)); + writer.decimal().writeDecimal(new BigDecimal(4)); + writer.endList(); + }); + assertEquals("values at index 0 is greater than listSize 3", e.getMessage()); + } + } + + @Test(expected = IllegalStateException.class) public void testWriteIllegalData() throws Exception { try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 3, allocator)) {