From f86fe9c2df30e230f160a91ce99f924175e9c3a2 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 18 Jun 2026 00:39:02 +0100 Subject: [PATCH] Decimal32/Decimal64 support Signed-off-by: Robert Kruszewski --- .../arrow/adapter/jdbc/JdbcToArrowUtils.java | 17 +- .../jdbc/consumer/Decimal32Consumer.java | 114 ++++ .../jdbc/consumer/Decimal64Consumer.java | 114 ++++ .../consumer/NarrowDecimalConsumerTest.java | 188 ++++++ .../sort/DefaultVectorComparators.java | 77 +++ .../sort/TestDefaultVectorComparator.java | 62 ++ .../java/org/apache/arrow/c/FormatTest.java | 6 + .../org/apache/arrow/c/RoundtripTest.java | 18 + .../main/codegen/data/ValueVectorTypes.tdd | 32 + .../AbstractPromotableFieldWriter.java | 43 +- .../src/main/codegen/templates/ArrowType.java | 5 +- .../codegen/templates/HolderReaderImpl.java | 13 +- .../codegen/templates/PromotableWriter.java | 84 +-- .../codegen/templates/UnionMapWriter.java | 24 + .../main/codegen/templates/UnionReader.java | 4 +- .../apache/arrow/vector/Decimal32Vector.java | 608 ++++++++++++++++++ .../apache/arrow/vector/Decimal64Vector.java | 608 ++++++++++++++++++ .../arrow/vector/GenerateSampleData.java | 47 ++ .../arrow/vector/extension/OpaqueType.java | 8 +- .../arrow/vector/ipc/JsonFileReader.java | 57 +- .../arrow/vector/ipc/JsonFileWriter.java | 21 +- .../org/apache/arrow/vector/types/Types.java | 40 +- .../arrow/vector/util/DecimalUtility.java | 26 +- .../validate/ValidateVectorTypeVisitor.java | 40 +- .../arrow/vector/TestDecimal32Vector.java | 460 +++++++++++++ .../arrow/vector/TestDecimal64Vector.java | 463 +++++++++++++ .../arrow/vector/TestGenerateSampleData.java | 86 +++ .../apache/arrow/vector/ipc/TestJSONFile.java | 55 ++ .../testing/ValueVectorDataPopulator.java | 50 ++ .../validate/TestValidateVectorFull.java | 2 +- .../TestValidateVectorTypeVisitor.java | 18 + 31 files changed, 3234 insertions(+), 156 deletions(-) create mode 100644 adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal32Consumer.java create mode 100644 adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal64Consumer.java create mode 100644 adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/NarrowDecimalConsumerTest.java create mode 100644 vector/src/main/java/org/apache/arrow/vector/Decimal32Vector.java create mode 100644 vector/src/main/java/org/apache/arrow/vector/Decimal64Vector.java create mode 100644 vector/src/test/java/org/apache/arrow/vector/TestDecimal32Vector.java create mode 100644 vector/src/test/java/org/apache/arrow/vector/TestDecimal64Vector.java create mode 100644 vector/src/test/java/org/apache/arrow/vector/TestGenerateSampleData.java diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java index 1edb6261d6..da63e2e942 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java @@ -46,6 +46,8 @@ import org.apache.arrow.adapter.jdbc.consumer.CompositeJdbcConsumer; import org.apache.arrow.adapter.jdbc.consumer.DateConsumer; import org.apache.arrow.adapter.jdbc.consumer.Decimal256Consumer; +import org.apache.arrow.adapter.jdbc.consumer.Decimal32Consumer; +import org.apache.arrow.adapter.jdbc.consumer.Decimal64Consumer; import org.apache.arrow.adapter.jdbc.consumer.DecimalConsumer; import org.apache.arrow.adapter.jdbc.consumer.DoubleConsumer; import org.apache.arrow.adapter.jdbc.consumer.FloatConsumer; @@ -66,6 +68,8 @@ import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.Float4Vector; @@ -510,10 +514,21 @@ public static JdbcConsumer getConsumer( } case Decimal: final RoundingMode bigDecimalRoundingMode = config.getBigDecimalRoundingMode(); - if (((ArrowType.Decimal) arrowType).getBitWidth() == 256) { + final int decimalBitWidth = ((ArrowType.Decimal) arrowType).getBitWidth(); + if (decimalBitWidth == 256) { return Decimal256Consumer.createConsumer( (Decimal256Vector) vector, columnIndex, nullable, bigDecimalRoundingMode); + } else if (decimalBitWidth == 128) { + return DecimalConsumer.createConsumer( + (DecimalVector) vector, columnIndex, nullable, bigDecimalRoundingMode); + } else if (decimalBitWidth == 64) { + return Decimal64Consumer.createConsumer( + (Decimal64Vector) vector, columnIndex, nullable, bigDecimalRoundingMode); + } else if (decimalBitWidth == 32) { + return Decimal32Consumer.createConsumer( + (Decimal32Vector) vector, columnIndex, nullable, bigDecimalRoundingMode); } else { + // Any other bit width maps to MinorType.DECIMAL, so the root created a DecimalVector. return DecimalConsumer.createConsumer( (DecimalVector) vector, columnIndex, nullable, bigDecimalRoundingMode); } diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal32Consumer.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal32Consumer.java new file mode 100644 index 0000000000..6e89302931 --- /dev/null +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal32Consumer.java @@ -0,0 +1,114 @@ +/* + * 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.arrow.adapter.jdbc.consumer; + +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.sql.ResultSet; +import java.sql.SQLException; +import org.apache.arrow.vector.Decimal32Vector; + +/** + * Consumer which consume decimal type values from {@link ResultSet}. Write the data to {@link + * org.apache.arrow.vector.Decimal32Vector}. + */ +public abstract class Decimal32Consumer extends BaseConsumer { + private final RoundingMode bigDecimalRoundingMode; + private final int scale; + + /** + * Constructs a new consumer. + * + * @param vector the underlying vector for the consumer. + * @param index the column id for the consumer. + */ + public Decimal32Consumer(Decimal32Vector vector, int index) { + this(vector, index, null); + } + + /** + * Constructs a new consumer, with optional coercibility. + * + * @param vector the underlying vector for the consumer. + * @param index the column index for the consumer. + * @param bigDecimalRoundingMode java.math.RoundingMode to be applied if the BigDecimal scale does + * not match that of the target vector. Set to null to retain strict matching behavior (scale + * of source and target vector must match exactly). + */ + public Decimal32Consumer(Decimal32Vector vector, int index, RoundingMode bigDecimalRoundingMode) { + super(vector, index); + this.bigDecimalRoundingMode = bigDecimalRoundingMode; + this.scale = vector.getScale(); + } + + /** Creates a consumer for {@link Decimal32Vector}. */ + public static JdbcConsumer createConsumer( + Decimal32Vector vector, int index, boolean nullable, RoundingMode bigDecimalRoundingMode) { + if (nullable) { + return new NullableDecimal32Consumer(vector, index, bigDecimalRoundingMode); + } else { + return new NonNullableDecimal32Consumer(vector, index, bigDecimalRoundingMode); + } + } + + protected void set(BigDecimal value) { + if (bigDecimalRoundingMode != null && value.scale() != scale) { + value = value.setScale(scale, bigDecimalRoundingMode); + } + vector.set(currentIndex, value); + } + + /** Consumer for nullable decimal. */ + static class NullableDecimal32Consumer extends Decimal32Consumer { + + /** Instantiate a Decimal32Consumer. */ + public NullableDecimal32Consumer( + Decimal32Vector vector, int index, RoundingMode bigDecimalRoundingMode) { + super(vector, index, bigDecimalRoundingMode); + } + + @Override + public void consume(ResultSet resultSet) throws SQLException { + BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); + if (!resultSet.wasNull()) { + // for fixed width vectors, we have allocated enough memory proactively, + // so there is no need to call the setSafe method here. + set(value); + } + currentIndex++; + } + } + + /** Consumer for non-nullable decimal. */ + static class NonNullableDecimal32Consumer extends Decimal32Consumer { + + /** Instantiate a Decimal32Consumer. */ + public NonNullableDecimal32Consumer( + Decimal32Vector vector, int index, RoundingMode bigDecimalRoundingMode) { + super(vector, index, bigDecimalRoundingMode); + } + + @Override + public void consume(ResultSet resultSet) throws SQLException { + BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); + // for fixed width vectors, we have allocated enough memory proactively, + // so there is no need to call the setSafe method here. + set(value); + currentIndex++; + } + } +} diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal64Consumer.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal64Consumer.java new file mode 100644 index 0000000000..7db07ee2ac --- /dev/null +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal64Consumer.java @@ -0,0 +1,114 @@ +/* + * 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.arrow.adapter.jdbc.consumer; + +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.sql.ResultSet; +import java.sql.SQLException; +import org.apache.arrow.vector.Decimal64Vector; + +/** + * Consumer which consume decimal type values from {@link ResultSet}. Write the data to {@link + * org.apache.arrow.vector.Decimal64Vector}. + */ +public abstract class Decimal64Consumer extends BaseConsumer { + private final RoundingMode bigDecimalRoundingMode; + private final int scale; + + /** + * Constructs a new consumer. + * + * @param vector the underlying vector for the consumer. + * @param index the column id for the consumer. + */ + public Decimal64Consumer(Decimal64Vector vector, int index) { + this(vector, index, null); + } + + /** + * Constructs a new consumer, with optional coercibility. + * + * @param vector the underlying vector for the consumer. + * @param index the column index for the consumer. + * @param bigDecimalRoundingMode java.math.RoundingMode to be applied if the BigDecimal scale does + * not match that of the target vector. Set to null to retain strict matching behavior (scale + * of source and target vector must match exactly). + */ + public Decimal64Consumer(Decimal64Vector vector, int index, RoundingMode bigDecimalRoundingMode) { + super(vector, index); + this.bigDecimalRoundingMode = bigDecimalRoundingMode; + this.scale = vector.getScale(); + } + + /** Creates a consumer for {@link Decimal64Vector}. */ + public static JdbcConsumer createConsumer( + Decimal64Vector vector, int index, boolean nullable, RoundingMode bigDecimalRoundingMode) { + if (nullable) { + return new NullableDecimal64Consumer(vector, index, bigDecimalRoundingMode); + } else { + return new NonNullableDecimal64Consumer(vector, index, bigDecimalRoundingMode); + } + } + + protected void set(BigDecimal value) { + if (bigDecimalRoundingMode != null && value.scale() != scale) { + value = value.setScale(scale, bigDecimalRoundingMode); + } + vector.set(currentIndex, value); + } + + /** Consumer for nullable decimal. */ + static class NullableDecimal64Consumer extends Decimal64Consumer { + + /** Instantiate a Decimal64Consumer. */ + public NullableDecimal64Consumer( + Decimal64Vector vector, int index, RoundingMode bigDecimalRoundingMode) { + super(vector, index, bigDecimalRoundingMode); + } + + @Override + public void consume(ResultSet resultSet) throws SQLException { + BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); + if (!resultSet.wasNull()) { + // for fixed width vectors, we have allocated enough memory proactively, + // so there is no need to call the setSafe method here. + set(value); + } + currentIndex++; + } + } + + /** Consumer for non-nullable decimal. */ + static class NonNullableDecimal64Consumer extends Decimal64Consumer { + + /** Instantiate a Decimal64Consumer. */ + public NonNullableDecimal64Consumer( + Decimal64Vector vector, int index, RoundingMode bigDecimalRoundingMode) { + super(vector, index, bigDecimalRoundingMode); + } + + @Override + public void consume(ResultSet resultSet) throws SQLException { + BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet); + // for fixed width vectors, we have allocated enough memory proactively, + // so there is no need to call the setSafe method here. + set(value); + currentIndex++; + } + } +} diff --git a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/NarrowDecimalConsumerTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/NarrowDecimalConsumerTest.java new file mode 100644 index 0000000000..c4ca5715ea --- /dev/null +++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/NarrowDecimalConsumerTest.java @@ -0,0 +1,188 @@ +/* + * 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.arrow.adapter.jdbc.consumer; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Types; +import java.util.Arrays; +import java.util.List; +import org.apache.arrow.adapter.jdbc.JdbcToArrowConfig; +import org.apache.arrow.adapter.jdbc.JdbcToArrowConfigBuilder; +import org.apache.arrow.adapter.jdbc.JdbcToArrowUtils; +import org.apache.arrow.adapter.jdbc.ResultSetUtility; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; +import org.junit.jupiter.api.Test; + +/** + * Tests for {@link Decimal32Consumer} and {@link Decimal64Consumer}, the consumers for the narrow + * (32-bit and 64-bit) Arrow decimal vectors, including dispatch through {@link + * JdbcToArrowUtils#getConsumer}. + */ +public class NarrowDecimalConsumerTest extends AbstractConsumerTest { + + /** Builds a single-column {@link ResultSet} of decimals; the first value must be non-null. */ + private ResultSet decimalResultSet(List values) throws SQLException { + ResultSetUtility.MockResultSet.Builder builder = ResultSetUtility.MockResultSet.builder(); + for (BigDecimal value : values) { + builder.addDataElement(value, Types.DECIMAL).finishRow(); + } + return builder.build(); + } + + @Test + void decimal32Nullable() throws SQLException, IOException { + List values = + Arrays.asList( + new BigDecimal("123.45"), null, new BigDecimal("-67.89"), new BigDecimal("0.00")); + try (Decimal32Vector vector = new Decimal32Vector("decimal32", allocator, 9, 2)) { + vector.allocateNew(values.size()); + JdbcConsumer consumer = + Decimal32Consumer.createConsumer(vector, 1, /* nullable= */ true, null); + ResultSet rs = decimalResultSet(values); + for (int i = 0; i < values.size(); i++) { + assertTrue(rs.next()); + consumer.consume(rs); + } + vector.setValueCount(values.size()); + + for (int i = 0; i < values.size(); i++) { + if (values.get(i) == null) { + assertTrue(vector.isNull(i), "expected null at row " + i); + } else { + assertEquals(0, values.get(i).compareTo(vector.getObject(i)), "mismatch at row " + i); + } + } + } + } + + @Test + void decimal32NonNullable() throws SQLException, IOException { + List values = + Arrays.asList( + new BigDecimal("1.23"), new BigDecimal("-9999999.99"), new BigDecimal("0.00")); + try (Decimal32Vector vector = new Decimal32Vector("decimal32", allocator, 9, 2)) { + vector.allocateNew(values.size()); + JdbcConsumer consumer = + Decimal32Consumer.createConsumer(vector, 1, /* nullable= */ false, null); + ResultSet rs = decimalResultSet(values); + for (int i = 0; i < values.size(); i++) { + assertTrue(rs.next()); + consumer.consume(rs); + } + vector.setValueCount(values.size()); + + for (int i = 0; i < values.size(); i++) { + assertEquals(0, values.get(i).compareTo(vector.getObject(i)), "mismatch at row " + i); + } + } + } + + @Test + void decimal32RoundingMode() throws SQLException, IOException { + // ResultSet supplies scale-4 values; the vector is scale-2, so the consumer must coerce. + List input = Arrays.asList(new BigDecimal("1.2345"), new BigDecimal("6.7891")); + try (Decimal32Vector vector = new Decimal32Vector("decimal32", allocator, 9, 2)) { + vector.allocateNew(input.size()); + JdbcConsumer consumer = + Decimal32Consumer.createConsumer(vector, 1, /* nullable= */ false, RoundingMode.HALF_UP); + ResultSet rs = decimalResultSet(input); + for (int i = 0; i < input.size(); i++) { + assertTrue(rs.next()); + consumer.consume(rs); + } + vector.setValueCount(input.size()); + + assertEquals(0, new BigDecimal("1.23").compareTo(vector.getObject(0))); + assertEquals(0, new BigDecimal("6.79").compareTo(vector.getObject(1))); + } + } + + @Test + void decimal64Nullable() throws SQLException, IOException { + List values = + Arrays.asList( + new BigDecimal("123456789.0123"), + null, + new BigDecimal("-98765.4321"), + new BigDecimal("0.0000")); + try (Decimal64Vector vector = new Decimal64Vector("decimal64", allocator, 18, 4)) { + vector.allocateNew(values.size()); + JdbcConsumer consumer = + Decimal64Consumer.createConsumer(vector, 1, /* nullable= */ true, null); + ResultSet rs = decimalResultSet(values); + for (int i = 0; i < values.size(); i++) { + assertTrue(rs.next()); + consumer.consume(rs); + } + vector.setValueCount(values.size()); + + for (int i = 0; i < values.size(); i++) { + if (values.get(i) == null) { + assertTrue(vector.isNull(i), "expected null at row " + i); + } else { + assertEquals(0, values.get(i).compareTo(vector.getObject(i)), "mismatch at row " + i); + } + } + } + } + + @Test + void decimal64RoundingMode() throws SQLException, IOException { + // ResultSet supplies scale-6 values; the vector is scale-4, so the consumer must coerce. + List input = + Arrays.asList(new BigDecimal("12.345678"), new BigDecimal("-0.987654")); + try (Decimal64Vector vector = new Decimal64Vector("decimal64", allocator, 18, 4)) { + vector.allocateNew(input.size()); + JdbcConsumer consumer = + Decimal64Consumer.createConsumer(vector, 1, /* nullable= */ false, RoundingMode.HALF_UP); + ResultSet rs = decimalResultSet(input); + for (int i = 0; i < input.size(); i++) { + assertTrue(rs.next()); + consumer.consume(rs); + } + vector.setValueCount(input.size()); + + assertEquals(0, new BigDecimal("12.3457").compareTo(vector.getObject(0))); + assertEquals(0, new BigDecimal("-0.9877").compareTo(vector.getObject(1))); + } + } + + @Test + void getConsumerDispatchesNarrowDecimals() { + JdbcToArrowConfig config = + new JdbcToArrowConfigBuilder(allocator, JdbcToArrowUtils.getUtcCalendar()).build(); + try (Decimal32Vector v32 = new Decimal32Vector("d32", allocator, 9, 2); + Decimal64Vector v64 = new Decimal64Vector("d64", allocator, 18, 4)) { + JdbcConsumer c32 = + JdbcToArrowUtils.getConsumer( + v32.getField().getType(), 1, /* nullable= */ true, v32, config); + JdbcConsumer c64 = + JdbcToArrowUtils.getConsumer( + v64.getField().getType(), 1, /* nullable= */ true, v64, config); + assertTrue(c32 instanceof Decimal32Consumer, "expected a Decimal32Consumer"); + assertTrue(c64 instanceof Decimal64Consumer, "expected a Decimal64Consumer"); + } + } +} diff --git a/algorithm/src/main/java/org/apache/arrow/algorithm/sort/DefaultVectorComparators.java b/algorithm/src/main/java/org/apache/arrow/algorithm/sort/DefaultVectorComparators.java index ec650cd9dc..d3fcce3731 100644 --- a/algorithm/src/main/java/org/apache/arrow/algorithm/sort/DefaultVectorComparators.java +++ b/algorithm/src/main/java/org/apache/arrow/algorithm/sort/DefaultVectorComparators.java @@ -22,12 +22,15 @@ import java.time.Duration; import org.apache.arrow.memory.util.ArrowBufPointer; import org.apache.arrow.memory.util.ByteFunctionHelpers; +import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FixedSizeBinaryVector; @@ -92,6 +95,10 @@ public static VectorValueComparator createDefaultComp return (VectorValueComparator) new DateDayComparator(); } else if (vector instanceof DateMilliVector) { return (VectorValueComparator) new DateMilliComparator(); + } else if (vector instanceof Decimal32Vector) { + return (VectorValueComparator) new Decimal32Comparator(); + } else if (vector instanceof Decimal64Vector) { + return (VectorValueComparator) new Decimal64Comparator(); } else if (vector instanceof Decimal256Vector) { return (VectorValueComparator) new Decimal256Comparator(); } else if (vector instanceof DecimalVector) { @@ -459,6 +466,76 @@ public VectorValueComparator createNew() { } } + /** + * Default comparator for Decimal32 type. The comparison is based on values, with null comes + * first. Values are compared by their unscaled representation, so both vectors must have the same + * scale. + */ + public static class Decimal32Comparator extends VectorValueComparator { + + public Decimal32Comparator() { + super(Decimal32Vector.TYPE_WIDTH); + } + + @Override + public void attachVectors(Decimal32Vector vector1, Decimal32Vector vector2) { + Preconditions.checkArgument( + vector1.getScale() == vector2.getScale(), + "Cannot compare Decimal32 vectors with different scales: %s and %s", + vector1.getScale(), + vector2.getScale()); + super.attachVectors(vector1, vector2); + } + + @Override + public int compareNotNull(int index1, int index2) { + int value1 = vector1.getDataBuffer().getInt((long) index1 * Decimal32Vector.TYPE_WIDTH); + int value2 = vector2.getDataBuffer().getInt((long) index2 * Decimal32Vector.TYPE_WIDTH); + + return Integer.compare(value1, value2); + } + + @Override + public VectorValueComparator createNew() { + return new Decimal32Comparator(); + } + } + + /** + * Default comparator for Decimal64 type. The comparison is based on values, with null comes + * first. Values are compared by their unscaled representation, so both vectors must have the same + * scale. + */ + public static class Decimal64Comparator extends VectorValueComparator { + + public Decimal64Comparator() { + super(Decimal64Vector.TYPE_WIDTH); + } + + @Override + public void attachVectors(Decimal64Vector vector1, Decimal64Vector vector2) { + Preconditions.checkArgument( + vector1.getScale() == vector2.getScale(), + "Cannot compare Decimal64 vectors with different scales: %s and %s", + vector1.getScale(), + vector2.getScale()); + super.attachVectors(vector1, vector2); + } + + @Override + public int compareNotNull(int index1, int index2) { + long value1 = vector1.getDataBuffer().getLong((long) index1 * Decimal64Vector.TYPE_WIDTH); + long value2 = vector2.getDataBuffer().getLong((long) index2 * Decimal64Vector.TYPE_WIDTH); + + return Long.compare(value1, value2); + } + + @Override + public VectorValueComparator createNew() { + return new Decimal64Comparator(); + } + } + /** * Default comparator for Decimal256 type. The comparison is based on values, with null comes * first. diff --git a/algorithm/src/test/java/org/apache/arrow/algorithm/sort/TestDefaultVectorComparator.java b/algorithm/src/test/java/org/apache/arrow/algorithm/sort/TestDefaultVectorComparator.java index 2a046533e8..508125b70c 100644 --- a/algorithm/src/test/java/org/apache/arrow/algorithm/sort/TestDefaultVectorComparator.java +++ b/algorithm/src/test/java/org/apache/arrow/algorithm/sort/TestDefaultVectorComparator.java @@ -28,6 +28,8 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FixedSizeBinaryVector; @@ -722,6 +724,66 @@ public void testCompareDecimal256() { } } + @Test + public void testCompareDecimal32() { + try (Decimal32Vector vec = new Decimal32Vector("", allocator, 9, 1)) { + vec.allocateNew(8); + ValueVectorDataPopulator.setVector( + vec, -1L, 0L, 1L, null, 1L, 5L, (long) Integer.MIN_VALUE + 1L, (long) Integer.MAX_VALUE); + + VectorValueComparator comparator = + DefaultVectorComparators.createDefaultComparator(vec); + comparator.attachVector(vec); + + assertTrue(comparator.compare(0, 1) < 0); + assertTrue(comparator.compare(0, 2) < 0); + assertTrue(comparator.compare(2, 1) > 0); + + // test equality + assertTrue(comparator.compare(5, 5) == 0); + assertTrue(comparator.compare(2, 4) == 0); + + // null first + assertTrue(comparator.compare(3, 4) < 0); + assertTrue(comparator.compare(5, 3) > 0); + + // potential overflow + assertTrue(comparator.compare(6, 7) < 0); + assertTrue(comparator.compare(7, 6) > 0); + assertTrue(comparator.compare(7, 7) == 0); + } + } + + @Test + public void testCompareDecimal64() { + try (Decimal64Vector vec = new Decimal64Vector("", allocator, 18, 1)) { + vec.allocateNew(8); + ValueVectorDataPopulator.setVector( + vec, -1L, 0L, 1L, null, 1L, 5L, Long.MIN_VALUE + 1L, Long.MAX_VALUE); + + VectorValueComparator comparator = + DefaultVectorComparators.createDefaultComparator(vec); + comparator.attachVector(vec); + + assertTrue(comparator.compare(0, 1) < 0); + assertTrue(comparator.compare(0, 2) < 0); + assertTrue(comparator.compare(2, 1) > 0); + + // test equality + assertTrue(comparator.compare(5, 5) == 0); + assertTrue(comparator.compare(2, 4) == 0); + + // null first + assertTrue(comparator.compare(3, 4) < 0); + assertTrue(comparator.compare(5, 3) > 0); + + // potential overflow + assertTrue(comparator.compare(6, 7) < 0); + assertTrue(comparator.compare(7, 6) > 0); + assertTrue(comparator.compare(7, 7) == 0); + } + } + @Test public void testCompareDuration() { try (DurationVector vec = diff --git a/c/src/test/java/org/apache/arrow/c/FormatTest.java b/c/src/test/java/org/apache/arrow/c/FormatTest.java index c773324330..6f78943d84 100644 --- a/c/src/test/java/org/apache/arrow/c/FormatTest.java +++ b/c/src/test/java/org/apache/arrow/c/FormatTest.java @@ -38,6 +38,9 @@ public void testAsString() { assertEquals("d:1,1", Format.asString(new ArrowType.Decimal(1, 1, 128))); assertEquals("d:1,1,1", Format.asString(new ArrowType.Decimal(1, 1, 1))); assertEquals("d:9,1,1", Format.asString(new ArrowType.Decimal(9, 1, 1))); + assertEquals("d:5,2,32", Format.asString(new ArrowType.Decimal(5, 2, 32))); + assertEquals("d:5,2,64", Format.asString(new ArrowType.Decimal(5, 2, 64))); + assertEquals("d:5,2,256", Format.asString(new ArrowType.Decimal(5, 2, 256))); assertEquals("tDs", Format.asString(new ArrowType.Duration(TimeUnit.SECOND))); assertEquals("tDm", Format.asString(new ArrowType.Duration(TimeUnit.MILLISECOND))); assertEquals("tDu", Format.asString(new ArrowType.Duration(TimeUnit.MICROSECOND))); @@ -134,6 +137,9 @@ public void testAsType() assertEquals(new ArrowType.Decimal(1, 1, 128), Format.asType("d:1,1", 0L)); assertEquals(new ArrowType.Decimal(1, 1, 1), Format.asType("d:1,1,1", 0L)); assertEquals(new ArrowType.Decimal(9, 1, 1), Format.asType("d:9,1,1", 0L)); + assertEquals(new ArrowType.Decimal(5, 2, 32), Format.asType("d:5,2,32", 0L)); + assertEquals(new ArrowType.Decimal(5, 2, 64), Format.asType("d:5,2,64", 0L)); + assertEquals(new ArrowType.Decimal(5, 2, 256), Format.asType("d:5,2,256", 0L)); assertEquals(new ArrowType.FixedSizeBinary(1), Format.asType("w:1", 0L)); assertEquals(new ArrowType.FixedSizeList(3), Format.asType("+w:3", 0L)); assertEquals( diff --git a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java index f6ff88571e..8787cb460e 100644 --- a/c/src/test/java/org/apache/arrow/c/RoundtripTest.java +++ b/c/src/test/java/org/apache/arrow/c/RoundtripTest.java @@ -41,6 +41,8 @@ import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FieldVector; @@ -294,6 +296,22 @@ public void testDateMilliVector() { imported.close(); } + @Test + public void testDecimal32Vector() { + try (final Decimal32Vector vector = new Decimal32Vector("v", allocator, 1, 1)) { + setVector(vector, 1L, 2L, 3L, null); + assertTrue(roundtrip(vector, Decimal32Vector.class)); + } + } + + @Test + public void testDecimal64Vector() { + try (final Decimal64Vector vector = new Decimal64Vector("v", allocator, 1, 1)) { + setVector(vector, 1L, 2L, 3L, null); + assertTrue(roundtrip(vector, Decimal64Vector.class)); + } + } + @Test public void testDecimalVector() { try (final DecimalVector vector = new DecimalVector("v", allocator, 1, 1)) { diff --git a/vector/src/main/codegen/data/ValueVectorTypes.tdd b/vector/src/main/codegen/data/ValueVectorTypes.tdd index ad1f1b93bb..9c6c1a5402 100644 --- a/vector/src/main/codegen/data/ValueVectorTypes.tdd +++ b/vector/src/main/codegen/data/ValueVectorTypes.tdd @@ -165,6 +165,38 @@ } ] }, + { + major: "Fixed", + width: 8, + javaType: "ArrowBuf", + boxedType: "ArrowBuf", + + minor: [ + { + class: "Decimal64", + maxPrecisionDigits: 18, nDecimalDigits: 2, friendlyType: "BigDecimal", + typeParams: [ {name: "scale", type: "int"}, { name: "precision", type: "int"}], + arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Decimal", + fields: [{name: "start", type: "long"}, {name: "buffer", type: "ArrowBuf"}] + } + ] + }, + { + major: "Fixed", + width: 4, + javaType: "ArrowBuf", + boxedType: "ArrowBuf", + + minor: [ + { + class: "Decimal32", + maxPrecisionDigits: 9, nDecimalDigits: 1, friendlyType: "BigDecimal", + typeParams: [ {name: "scale", type: "int"}, { name: "precision", type: "int"}], + arrowType: "org.apache.arrow.vector.types.pojo.ArrowType.Decimal", + fields: [{name: "start", type: "long"}, {name: "buffer", type: "ArrowBuf"}] + } + ] + }, { major: "Fixed", diff --git a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java index 2e7792fcfe..fc33f00ae9 100644 --- a/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java +++ b/vector/src/main/codegen/templates/AbstractPromotableFieldWriter.java @@ -120,46 +120,27 @@ public void endEntry() { <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> <#assign fields = minor.fields!type.fields /> - <#if minor.class == "Decimal"> + <#if minor.class?starts_with("Decimal")> + <#if minor.class == "Decimal"><#assign startType = "int" /><#else><#assign startType = "long" /> @Override - public void write(DecimalHolder holder) { - getWriter(MinorType.DECIMAL).write(holder); - } - - public void writeDecimal(int start, ArrowBuf buffer, ArrowType arrowType) { - getWriter(MinorType.DECIMAL).writeDecimal(start, buffer, arrowType); - } - - public void writeDecimal(int start, ArrowBuf buffer) { - getWriter(MinorType.DECIMAL).writeDecimal(start, buffer); - } - - public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) { - getWriter(MinorType.DECIMAL).writeBigEndianBytesToDecimal(value, arrowType); + public void write(${name}Holder holder) { + getWriter(MinorType.${name?upper_case}).write(holder); } - public void writeBigEndianBytesToDecimal(byte[] value) { - getWriter(MinorType.DECIMAL).writeBigEndianBytesToDecimal(value); - } - <#elseif minor.class == "Decimal256"> - @Override - public void write(Decimal256Holder holder) { - getWriter(MinorType.DECIMAL256).write(holder); + public void write${name}(${startType} start, ArrowBuf buffer, ArrowType arrowType) { + getWriter(MinorType.${name?upper_case}).write${name}(start, buffer, arrowType); } - public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) { - getWriter(MinorType.DECIMAL256).writeDecimal256(start, buffer, arrowType); + public void write${name}(${startType} start, ArrowBuf buffer) { + getWriter(MinorType.${name?upper_case}).write${name}(start, buffer); } - public void writeDecimal256(long start, ArrowBuf buffer) { - getWriter(MinorType.DECIMAL256).writeDecimal256(start, buffer); - } - public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { - getWriter(MinorType.DECIMAL256).writeBigEndianBytesToDecimal256(value, arrowType); + public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType) { + getWriter(MinorType.${name?upper_case}).writeBigEndianBytesTo${name}(value, arrowType); } - public void writeBigEndianBytesToDecimal256(byte[] value) { - getWriter(MinorType.DECIMAL256).writeBigEndianBytesToDecimal256(value); + public void writeBigEndianBytesTo${name}(byte[] value) { + getWriter(MinorType.${name?upper_case}).writeBigEndianBytesTo${name}(value); } <#elseif is_timestamp_tz(minor.class)> @Override diff --git a/vector/src/main/codegen/templates/ArrowType.java b/vector/src/main/codegen/templates/ArrowType.java index b428f09155..6567757e98 100644 --- a/vector/src/main/codegen/templates/ArrowType.java +++ b/vector/src/main/codegen/templates/ArrowType.java @@ -360,8 +360,9 @@ public static org.apache.arrow.vector.types.pojo.ArrowType getTypeForField(org.a <#if type.name == "Decimal"> - if (bitWidth != defaultDecimalBitWidth && bitWidth != 256) { - throw new IllegalArgumentException("Library only supports 128-bit and 256-bit decimal values"); + if (bitWidth != 32 && bitWidth != 64 && bitWidth != defaultDecimalBitWidth && bitWidth != 256) { + throw new IllegalArgumentException( + "Library only supports 32-bit, 64-bit, 128-bit and 256-bit decimal values"); } return new ${name}(<#list type.fields as field><#if field.valueType??>${field.valueType}.fromFlatbufID(${field.name})<#else>${field.name}<#if field_has_next>, ); diff --git a/vector/src/main/codegen/templates/HolderReaderImpl.java b/vector/src/main/codegen/templates/HolderReaderImpl.java index cdbb65c4f6..b67c1b26be 100644 --- a/vector/src/main/codegen/templates/HolderReaderImpl.java +++ b/vector/src/main/codegen/templates/HolderReaderImpl.java @@ -127,16 +127,9 @@ public void read(Nullable${name}Holder h) { return DurationVector.toDuration(holder.value, holder.unit); <#elseif minor.class == "Bit" > return Boolean.valueOf(holder.value != 0); - <#elseif minor.class == "Decimal"> - byte[] bytes = new byte[${type.width}]; - holder.buffer.getBytes(holder.start, bytes, 0, ${type.width}); - ${friendlyType} value = new BigDecimal(new BigInteger(bytes), holder.scale); - return value; - <#elseif minor.class == "Decimal256"> - byte[] bytes = new byte[${type.width}]; - holder.buffer.getBytes(holder.start, bytes, 0, ${type.width}); - ${friendlyType} value = new BigDecimal(new BigInteger(bytes), holder.scale); - return value; + <#elseif minor.class?starts_with("Decimal")> + return DecimalUtility.getBigDecimalFromArrowBufAtOffset( + holder.buffer, holder.start, holder.scale, ${type.width}); <#elseif minor.class == "FixedSizeBinary"> byte[] value = new byte [holder.byteWidth]; holder.buffer.getBytes(0, value, 0, holder.byteWidth); diff --git a/vector/src/main/codegen/templates/PromotableWriter.java b/vector/src/main/codegen/templates/PromotableWriter.java index 11d34f72c9..20d1d47572 100644 --- a/vector/src/main/codegen/templates/PromotableWriter.java +++ b/vector/src/main/codegen/templates/PromotableWriter.java @@ -44,8 +44,10 @@ public class PromotableWriter extends AbstractPromotableFieldWriter { protected final LargeListViewVector largeListViewVector; protected final NullableStructWriterFactory nullableStructWriterFactory; protected int position; - protected static final int MAX_DECIMAL_PRECISION = 38; - protected static final int MAX_DECIMAL256_PRECISION = 76; + protected static final int MAX_DECIMAL32_PRECISION = Decimal32Vector.MAX_PRECISION; + protected static final int MAX_DECIMAL64_PRECISION = Decimal64Vector.MAX_PRECISION; + protected static final int MAX_DECIMAL_PRECISION = DecimalVector.MAX_PRECISION; + protected static final int MAX_DECIMAL256_PRECISION = Decimal256Vector.MAX_PRECISION; protected enum State { UNTYPED, @@ -316,6 +318,8 @@ public void setPosition(int index) { protected boolean requiresArrowType(MinorType type) { return type == MinorType.DECIMAL + || type == MinorType.DECIMAL32 + || type == MinorType.DECIMAL64 || type == MinorType.MAP || type == MinorType.DURATION || type == MinorType.FIXEDSIZEBINARY @@ -404,81 +408,47 @@ protected FieldWriter promoteToUnion() { return writer; } + <#list vv.types as type><#list type.minor as minor><#assign name = minor.class?cap_first /> + <#if minor.class?starts_with("Decimal")> @Override - public void write(DecimalHolder holder) { + public void write(${name}Holder holder) { getWriter( - MinorType.DECIMAL, - new ArrowType.Decimal(MAX_DECIMAL_PRECISION, holder.scale, /*bitWidth=*/ 128)) + MinorType.${name?upper_case}, + new ArrowType.Decimal(MAX_${name?upper_case}_PRECISION, holder.scale, /*bitWidth=*/ ${type.width * 8})) .write(holder); } @Override - public void writeDecimal(long start, ArrowBuf buffer, ArrowType arrowType) { + public void write${name}(long start, ArrowBuf buffer, ArrowType arrowType) { getWriter( - MinorType.DECIMAL, + MinorType.${name?upper_case}, new ArrowType.Decimal( - MAX_DECIMAL_PRECISION, + MAX_${name?upper_case}_PRECISION, ((ArrowType.Decimal) arrowType).getScale(), - /*bitWidth=*/ 128)) - .writeDecimal(start, buffer, arrowType); + /*bitWidth=*/ ${type.width * 8})) + .write${name}(start, buffer, arrowType); } @Override - public void writeDecimal(BigDecimal value) { + public void write${name}(BigDecimal value) { getWriter( - MinorType.DECIMAL, - new ArrowType.Decimal(MAX_DECIMAL_PRECISION, value.scale(), /*bitWidth=*/ 128)) - .writeDecimal(value); + MinorType.${name?upper_case}, + new ArrowType.Decimal(MAX_${name?upper_case}_PRECISION, value.scale(), /*bitWidth=*/ ${type.width * 8})) + .write${name}(value); } @Override - public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) { + public void writeBigEndianBytesTo${name}(byte[] value, ArrowType arrowType) { getWriter( - MinorType.DECIMAL, + MinorType.${name?upper_case}, new ArrowType.Decimal( - MAX_DECIMAL_PRECISION, + MAX_${name?upper_case}_PRECISION, ((ArrowType.Decimal) arrowType).getScale(), - /*bitWidth=*/ 128)) - .writeBigEndianBytesToDecimal(value, arrowType); - } - - @Override - public void write(Decimal256Holder holder) { - getWriter( - MinorType.DECIMAL256, - new ArrowType.Decimal(MAX_DECIMAL256_PRECISION, holder.scale, /*bitWidth=*/ 256)) - .write(holder); - } - - @Override - public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) { - getWriter( - MinorType.DECIMAL256, - new ArrowType.Decimal( - MAX_DECIMAL256_PRECISION, - ((ArrowType.Decimal) arrowType).getScale(), - /*bitWidth=*/ 256)) - .writeDecimal256(start, buffer, arrowType); - } - - @Override - public void writeDecimal256(BigDecimal value) { - getWriter( - MinorType.DECIMAL256, - new ArrowType.Decimal(MAX_DECIMAL256_PRECISION, value.scale(), /*bitWidth=*/ 256)) - .writeDecimal256(value); - } - - @Override - public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) { - getWriter( - MinorType.DECIMAL256, - new ArrowType.Decimal( - MAX_DECIMAL256_PRECISION, - ((ArrowType.Decimal) arrowType).getScale(), - /*bitWidth=*/ 256)) - .writeBigEndianBytesToDecimal256(value, arrowType); + /*bitWidth=*/ ${type.width * 8})) + .writeBigEndianBytesTo${name}(value, arrowType); } + + @Override public void writeVarBinary(byte[] value) { diff --git a/vector/src/main/codegen/templates/UnionMapWriter.java b/vector/src/main/codegen/templates/UnionMapWriter.java index 8bbf6ae0a4..e8446752ea 100644 --- a/vector/src/main/codegen/templates/UnionMapWriter.java +++ b/vector/src/main/codegen/templates/UnionMapWriter.java @@ -183,6 +183,30 @@ public Decimal256Writer decimal256() { } } + @Override + public Decimal32Writer decimal32() { + switch (mode) { + case KEY: + return entryWriter.decimal32(MapVector.KEY_NAME); + case VALUE: + return entryWriter.decimal32(MapVector.VALUE_NAME); + default: + return this; + } + } + + @Override + public Decimal64Writer decimal64() { + switch (mode) { + case KEY: + return entryWriter.decimal64(MapVector.KEY_NAME); + case VALUE: + return entryWriter.decimal64(MapVector.VALUE_NAME); + default: + return this; + } + } + @Override public StructWriter struct() { diff --git a/vector/src/main/codegen/templates/UnionReader.java b/vector/src/main/codegen/templates/UnionReader.java index 0edae7ade0..7c4a201072 100644 --- a/vector/src/main/codegen/templates/UnionReader.java +++ b/vector/src/main/codegen/templates/UnionReader.java @@ -39,7 +39,9 @@ @SuppressWarnings("unused") public class UnionReader extends AbstractFieldReader { - private static final int NUM_SUPPORTED_TYPES = 51; + // The reader/type arrays are indexed by MinorType.ordinal(), so they must be sized to the total + // number of minor types (not a hardcoded count, which would break whenever a type is added). + private static final int NUM_SUPPORTED_TYPES = MinorType.values().length; private BaseReader[] readers = new BaseReader[NUM_SUPPORTED_TYPES]; public UnionVector data; diff --git a/vector/src/main/java/org/apache/arrow/vector/Decimal32Vector.java b/vector/src/main/java/org/apache/arrow/vector/Decimal32Vector.java new file mode 100644 index 0000000000..a511f8926d --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/Decimal32Vector.java @@ -0,0 +1,608 @@ +/* + * 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.arrow.vector; + +import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED; + +import java.math.BigDecimal; +import java.nio.ByteOrder; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.vector.complex.impl.Decimal32ReaderImpl; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.holders.Decimal32Holder; +import org.apache.arrow.vector.holders.NullableDecimal32Holder; +import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.DecimalUtility; +import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.validate.ValidateUtil; + +/** + * Decimal32Vector implements a fixed width vector (4 bytes) of decimal values which could be null. + * A validity buffer (bit vector) is maintained to track which elements in the vector are null. + */ +public final class Decimal32Vector extends BaseFixedWidthVector + implements ValueIterableVector { + public static final int MAX_PRECISION = 9; + public static final byte TYPE_WIDTH = 4; + private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; + + private final int precision; + private final int scale; + + private static void checkByteLength(int length) { + if (length < 1 || length > TYPE_WIDTH) { + throw new IllegalArgumentException( + "Invalid decimal value length. Valid length in [1 - " + TYPE_WIDTH + "], got " + length); + } + } + + /** + * Instantiate a Decimal32Vector. This doesn't allocate any memory for the data in vector. + * + * @param name name of the vector + * @param allocator allocator for memory management. + */ + public Decimal32Vector(String name, BufferAllocator allocator, int precision, int scale) { + this( + name, + FieldType.nullable(new ArrowType.Decimal(precision, scale, /* bitWidth= */ TYPE_WIDTH * 8)), + allocator); + } + + /** + * Instantiate a Decimal32Vector. This doesn't allocate any memory for the data in vector. + * + * @param name name of the vector + * @param fieldType type of Field materialized by this vector + * @param allocator allocator for memory management. + */ + public Decimal32Vector(String name, FieldType fieldType, BufferAllocator allocator) { + this(new Field(name, fieldType, null), allocator); + } + + /** + * Instantiate a Decimal32Vector. This doesn't allocate any memory for the data in vector. + * + * @param field field materialized by this vector + * @param allocator allocator for memory management. + */ + public Decimal32Vector(Field field, BufferAllocator allocator) { + super(field, allocator, TYPE_WIDTH); + ArrowType.Decimal arrowType = (ArrowType.Decimal) field.getFieldType().getType(); + this.precision = arrowType.getPrecision(); + this.scale = arrowType.getScale(); + } + + @Override + protected FieldReader getReaderImpl() { + return new Decimal32ReaderImpl(Decimal32Vector.this); + } + + /** + * Get minor type for this vector. The vector holds values belonging to a particular type. + * + * @return {@link org.apache.arrow.vector.types.Types.MinorType} + */ + @Override + public MinorType getMinorType() { + return MinorType.DECIMAL32; + } + + /*----------------------------------------------------------------* + | | + | vector value retrieval methods | + | | + *----------------------------------------------------------------*/ + + /** + * Get the element at the given index from the vector. + * + * @param index position of element + * @return element at given index + */ + public ArrowBuf get(int index) throws IllegalStateException { + if (NULL_CHECKING_ENABLED && isSet(index) == 0) { + throw new IllegalStateException("Value at index is null"); + } + return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); + } + + /** + * Get the element at the given index from the vector and sets the state in holder. If element at + * given index is null, holder.isSet will be zero. + * + * @param index position of element + */ + public void get(int index, NullableDecimal32Holder holder) { + if (isSet(index) == 0) { + holder.isSet = 0; + return; + } + holder.isSet = 1; + holder.buffer = valueBuffer; + holder.precision = precision; + holder.scale = scale; + holder.start = (long) index * TYPE_WIDTH; + } + + /** + * Same as {@link #get(int)}. + * + * @param index position of element + * @return element at given index + */ + @Override + public BigDecimal getObject(int index) { + if (isSet(index) == 0) { + return null; + } else { + return getObjectNotNull(index); + } + } + + /** + * Same as {@link #getObject(int)} but does not check for null. + * + * @param index position of element + * @return element at given index + */ + public BigDecimal getObjectNotNull(int index) { + return DecimalUtility.getBigDecimalFromArrowBuf(valueBuffer, index, scale, TYPE_WIDTH); + } + + /** Return precision for the decimal value. */ + public int getPrecision() { + return precision; + } + + /** Return scale for the decimal value. */ + public int getScale() { + return scale; + } + + /*----------------------------------------------------------------* + | | + | vector value setter methods | + | | + *----------------------------------------------------------------*/ + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param buffer ArrowBuf containing decimal value. + */ + public void set(int index, ArrowBuf buffer) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, 0, TYPE_WIDTH); + } + + /** + * Set the decimal element at given index to the provided array of bytes. Decimal is now + * implemented as Native Endian. This API allows the user to pass a decimal value in the form of + * byte array in BE byte order. + * + *

Consumers of Arrow code can use this API instead of first swapping the source bytes (doing a + * write and read) and then finally writing to ArrowBuf of decimal vector. + * + *

This method takes care of adding the necessary padding if the length of byte array is less + * than 4 (length of decimal type). + * + * @param index position of element + * @param value array of bytes containing decimal in big endian byte order. + */ + public void setBigEndian(int index, byte[] value) { + final int length = value.length; + if (length > TYPE_WIDTH) { + throw new IllegalArgumentException( + "Invalid decimal value length. Valid length in [0 - " + TYPE_WIDTH + "], got " + length); + } + BitVectorHelper.setBit(validityBuffer, index); + + // do the bound check. + valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); + + long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (length == 0) { + MemoryUtil.setMemory(outAddress, Decimal32Vector.TYPE_WIDTH, (byte) 0); + return; + } + if (LITTLE_ENDIAN) { + // swap bytes to convert BE to LE + for (int byteIdx = 0; byteIdx < length; ++byteIdx) { + MemoryUtil.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]); + } + if (length < TYPE_WIDTH) { + // sign extend + final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress + length, Decimal32Vector.TYPE_WIDTH - length, pad); + } + } else { + // copy data from value to outAddress + MemoryUtil.copyToMemory(value, 0, outAddress + Decimal32Vector.TYPE_WIDTH - length, length); + if (length < TYPE_WIDTH) { + // sign extend + final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal32Vector.TYPE_WIDTH - length, pad); + } + } + } + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param start start index of data in the buffer + * @param buffer ArrowBuf containing decimal value. + */ + public void set(int index, long start, ArrowBuf buffer) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, start, TYPE_WIDTH); + } + + /** + * Sets the element at given index using the buffer whose size maybe <= 4 bytes. + * + * @param index index to write the decimal to + * @param start start of value in the buffer + * @param buffer contains the decimal in native endian bytes + * @param length length of the value in the buffer + */ + public void setSafe(int index, long start, ArrowBuf buffer, int length) { + checkByteLength(length); + handleSafe(index); + BitVectorHelper.setBit(validityBuffer, index); + + // do the bound checks. + buffer.checkBytes(start, start + length); + valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); + + long inAddress = buffer.memoryAddress() + start; + long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (LITTLE_ENDIAN) { + MemoryUtil.copyMemory(inAddress, outAddress, length); + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress + length - 1); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress + length, Decimal32Vector.TYPE_WIDTH - length, pad); + } + } else { + MemoryUtil.copyMemory(inAddress, outAddress + Decimal32Vector.TYPE_WIDTH - length, length); + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal32Vector.TYPE_WIDTH - length, pad); + } + } + } + + /** + * Sets the element at given index using the buffer whose size maybe <= 4 bytes. + * + * @param index index to write the decimal to + * @param start start of value in the buffer + * @param buffer contains the decimal in big endian bytes + * @param length length of the value in the buffer + */ + public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) { + checkByteLength(length); + handleSafe(index); + BitVectorHelper.setBit(validityBuffer, index); + + // do the bound checks. + buffer.checkBytes(start, start + length); + valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); + + // 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 (LITTLE_ENDIAN) { + // swap bytes to convert BE to LE + for (int byteIdx = 0; byteIdx < length; ++byteIdx) { + byte val = MemoryUtil.getByte((inAddress + length - 1) - byteIdx); + MemoryUtil.putByte(outAddress + byteIdx, val); + } + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress + length, Decimal32Vector.TYPE_WIDTH - length, pad); + } + } else { + MemoryUtil.copyMemory(inAddress, outAddress + Decimal32Vector.TYPE_WIDTH - length, length); + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal32Vector.TYPE_WIDTH - length, pad); + } + } + } + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param value BigDecimal containing decimal value. + */ + public void set(int index, BigDecimal value) { + DecimalUtility.checkPrecisionAndScale(value, precision, scale); + BitVectorHelper.setBit(validityBuffer, index); + DecimalUtility.writeBigDecimalToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); + } + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param value long value. + */ + public void set(int index, long value) { + // write the value first: it range-checks and may throw, and the slot must + // not be marked valid if the write does not happen. + DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); + BitVectorHelper.setBit(validityBuffer, index); + } + + /** + * Set the element at the given index to the value set in data holder. If the value in holder is + * not indicated as set, element in the at the given index will be null. + * + * @param index position of element + * @param holder nullable data holder for value of element + */ + public void set(int index, NullableDecimal32Holder holder) throws IllegalArgumentException { + if (holder.isSet < 0) { + throw new IllegalArgumentException(); + } else if (holder.isSet > 0) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); + } else { + BitVectorHelper.unsetBit(validityBuffer, index); + } + } + + /** + * Set the element at the given index to the value set in data holder. + * + * @param index position of element + * @param holder data holder for value of element + */ + public void set(int index, Decimal32Holder holder) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); + } + + /** + * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than + * or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param buffer ArrowBuf containing decimal value. + */ + public void setSafe(int index, ArrowBuf buffer) { + handleSafe(index); + set(index, buffer); + } + + /** + * Same as {@link #setBigEndian(int, byte[])} except that it handles the case when index is + * greater than or equal to existing value capacity {@link #getValueCapacity()}. + */ + public void setBigEndianSafe(int index, byte[] value) { + handleSafe(index); + setBigEndian(index, value); + } + + /** + * Same as {@link #set(int, long, ArrowBuf)} except that it handles the case when index is greater + * than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param start start index of data in the buffer + * @param buffer ArrowBuf containing decimal value. + */ + public void setSafe(int index, long start, ArrowBuf buffer) { + handleSafe(index); + set(index, start, buffer); + } + + /** + * Same as {@link #set(int, BigDecimal)} except that it handles the case when index is greater + * than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param value BigDecimal containing decimal value. + */ + public void setSafe(int index, BigDecimal value) { + handleSafe(index); + set(index, value); + } + + /** + * Same as {@link #set(int, long)} except that it handles the case when index is greater than or + * equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param value long value. + */ + public void setSafe(int index, long value) { + handleSafe(index); + set(index, value); + } + + /** + * Same as {@link #set(int, NullableDecimal32Holder)} except that it handles the case when index + * is greater than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param holder nullable data holder for value of element + */ + public void setSafe(int index, NullableDecimal32Holder holder) throws IllegalArgumentException { + handleSafe(index); + set(index, holder); + } + + /** + * Same as {@link #set(int, Decimal32Holder)} except that it handles the case when index is + * greater than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param holder data holder for value of element + */ + public void setSafe(int index, Decimal32Holder holder) { + handleSafe(index); + set(index, holder); + } + + /** + * Store the given value at a particular position in the vector. isSet indicates whether the value + * is NULL or not. + * + * @param index position of the new value + * @param isSet 0 for NULL value, 1 otherwise + * @param start start position of the value in the buffer + * @param buffer buffer containing the value to be stored in the vector + */ + public void set(int index, int isSet, long start, ArrowBuf buffer) { + if (isSet > 0) { + set(index, start, buffer); + } else { + BitVectorHelper.unsetBit(validityBuffer, index); + } + } + + /** + * Same as {@link #set(int, int, long, ArrowBuf)} except that it handles the case when the + * position of new value is beyond the current value capacity of the vector. + * + * @param index position of the new value + * @param isSet 0 for NULL value, 1 otherwise + * @param start start position of the value in the buffer + * @param buffer buffer containing the value to be stored in the vector + */ + public void setSafe(int index, int isSet, long start, ArrowBuf buffer) { + handleSafe(index); + set(index, isSet, start, buffer); + } + + @Override + public void validateScalars() { + for (int i = 0; i < getValueCount(); ++i) { + BigDecimal value = getObject(i); + if (value != null) { + ValidateUtil.validateOrThrow( + DecimalUtility.checkPrecisionAndScaleNoThrow(value, getPrecision(), getScale()), + "Invalid value for Decimal32Vector at position " + + i + + ". Value does not fit in precision " + + getPrecision() + + " and scale " + + getScale() + + "."); + } + } + } + + /*----------------------------------------------------------------* + | | + | vector transfer | + | | + *----------------------------------------------------------------*/ + + /** + * Construct a TransferPair comprising this and a target vector of the same type. + * + * @param ref name of the target vector + * @param allocator allocator for the target vector + * @return {@link TransferPair} + */ + @Override + public TransferPair getTransferPair(String ref, BufferAllocator allocator) { + return new TransferImpl(ref, allocator); + } + + /** + * Construct a TransferPair comprising this and a target vector of the same type. + * + * @param field Field object used by the target vector + * @param allocator allocator for the target vector + * @return {@link TransferPair} + */ + @Override + public TransferPair getTransferPair(Field field, BufferAllocator allocator) { + return new TransferImpl(field, allocator); + } + + /** + * Construct a TransferPair with a desired target vector of the same type. + * + * @param to target vector + * @return {@link TransferPair} + */ + @Override + public TransferPair makeTransferPair(ValueVector to) { + return new TransferImpl((Decimal32Vector) to); + } + + private class TransferImpl implements TransferPair { + Decimal32Vector to; + + public TransferImpl(String ref, BufferAllocator allocator) { + to = + (Decimal32Vector.this.field != null && Decimal32Vector.this.field.getFieldType() != null) + ? new Decimal32Vector(ref, Decimal32Vector.this.field.getFieldType(), allocator) + : new Decimal32Vector( + ref, allocator, Decimal32Vector.this.precision, Decimal32Vector.this.scale); + } + + public TransferImpl(Field field, BufferAllocator allocator) { + to = new Decimal32Vector(field, allocator); + } + + public TransferImpl(Decimal32Vector to) { + this.to = to; + } + + @Override + public Decimal32Vector getTo() { + return to; + } + + @Override + public void transfer() { + transferTo(to); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + splitAndTransferTo(startIndex, length, to); + } + + @Override + public void copyValueSafe(int fromIndex, int toIndex) { + to.copyFromSafe(fromIndex, toIndex, Decimal32Vector.this); + } + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/Decimal64Vector.java b/vector/src/main/java/org/apache/arrow/vector/Decimal64Vector.java new file mode 100644 index 0000000000..bb1f223cb9 --- /dev/null +++ b/vector/src/main/java/org/apache/arrow/vector/Decimal64Vector.java @@ -0,0 +1,608 @@ +/* + * 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.arrow.vector; + +import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED; + +import java.math.BigDecimal; +import java.nio.ByteOrder; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.vector.complex.impl.Decimal64ReaderImpl; +import org.apache.arrow.vector.complex.reader.FieldReader; +import org.apache.arrow.vector.holders.Decimal64Holder; +import org.apache.arrow.vector.holders.NullableDecimal64Holder; +import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.util.DecimalUtility; +import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.validate.ValidateUtil; + +/** + * Decimal64Vector implements a fixed width vector (8 bytes) of decimal values which could be null. + * A validity buffer (bit vector) is maintained to track which elements in the vector are null. + */ +public final class Decimal64Vector extends BaseFixedWidthVector + implements ValueIterableVector { + public static final int MAX_PRECISION = 18; + public static final byte TYPE_WIDTH = 8; + private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; + + private final int precision; + private final int scale; + + private static void checkByteLength(int length) { + if (length < 1 || length > TYPE_WIDTH) { + throw new IllegalArgumentException( + "Invalid decimal value length. Valid length in [1 - " + TYPE_WIDTH + "], got " + length); + } + } + + /** + * Instantiate a Decimal64Vector. This doesn't allocate any memory for the data in vector. + * + * @param name name of the vector + * @param allocator allocator for memory management. + */ + public Decimal64Vector(String name, BufferAllocator allocator, int precision, int scale) { + this( + name, + FieldType.nullable(new ArrowType.Decimal(precision, scale, /* bitWidth= */ TYPE_WIDTH * 8)), + allocator); + } + + /** + * Instantiate a Decimal64Vector. This doesn't allocate any memory for the data in vector. + * + * @param name name of the vector + * @param fieldType type of Field materialized by this vector + * @param allocator allocator for memory management. + */ + public Decimal64Vector(String name, FieldType fieldType, BufferAllocator allocator) { + this(new Field(name, fieldType, null), allocator); + } + + /** + * Instantiate a Decimal64Vector. This doesn't allocate any memory for the data in vector. + * + * @param field field materialized by this vector + * @param allocator allocator for memory management. + */ + public Decimal64Vector(Field field, BufferAllocator allocator) { + super(field, allocator, TYPE_WIDTH); + ArrowType.Decimal arrowType = (ArrowType.Decimal) field.getFieldType().getType(); + this.precision = arrowType.getPrecision(); + this.scale = arrowType.getScale(); + } + + @Override + protected FieldReader getReaderImpl() { + return new Decimal64ReaderImpl(Decimal64Vector.this); + } + + /** + * Get minor type for this vector. The vector holds values belonging to a particular type. + * + * @return {@link org.apache.arrow.vector.types.Types.MinorType} + */ + @Override + public MinorType getMinorType() { + return MinorType.DECIMAL64; + } + + /*----------------------------------------------------------------* + | | + | vector value retrieval methods | + | | + *----------------------------------------------------------------*/ + + /** + * Get the element at the given index from the vector. + * + * @param index position of element + * @return element at given index + */ + public ArrowBuf get(int index) throws IllegalStateException { + if (NULL_CHECKING_ENABLED && isSet(index) == 0) { + throw new IllegalStateException("Value at index is null"); + } + return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH); + } + + /** + * Get the element at the given index from the vector and sets the state in holder. If element at + * given index is null, holder.isSet will be zero. + * + * @param index position of element + */ + public void get(int index, NullableDecimal64Holder holder) { + if (isSet(index) == 0) { + holder.isSet = 0; + return; + } + holder.isSet = 1; + holder.buffer = valueBuffer; + holder.precision = precision; + holder.scale = scale; + holder.start = (long) index * TYPE_WIDTH; + } + + /** + * Same as {@link #get(int)}. + * + * @param index position of element + * @return element at given index + */ + @Override + public BigDecimal getObject(int index) { + if (isSet(index) == 0) { + return null; + } else { + return getObjectNotNull(index); + } + } + + /** + * Same as {@link #getObject(int)} but does not check for null. + * + * @param index position of element + * @return element at given index + */ + public BigDecimal getObjectNotNull(int index) { + return DecimalUtility.getBigDecimalFromArrowBuf(valueBuffer, index, scale, TYPE_WIDTH); + } + + /** Return precision for the decimal value. */ + public int getPrecision() { + return precision; + } + + /** Return scale for the decimal value. */ + public int getScale() { + return scale; + } + + /*----------------------------------------------------------------* + | | + | vector value setter methods | + | | + *----------------------------------------------------------------*/ + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param buffer ArrowBuf containing decimal value. + */ + public void set(int index, ArrowBuf buffer) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, 0, TYPE_WIDTH); + } + + /** + * Set the decimal element at given index to the provided array of bytes. Decimal is now + * implemented as Native Endian. This API allows the user to pass a decimal value in the form of + * byte array in BE byte order. + * + *

Consumers of Arrow code can use this API instead of first swapping the source bytes (doing a + * write and read) and then finally writing to ArrowBuf of decimal vector. + * + *

This method takes care of adding the necessary padding if the length of byte array is less + * than 8 (length of decimal type). + * + * @param index position of element + * @param value array of bytes containing decimal in big endian byte order. + */ + public void setBigEndian(int index, byte[] value) { + final int length = value.length; + if (length > TYPE_WIDTH) { + throw new IllegalArgumentException( + "Invalid decimal value length. Valid length in [0 - " + TYPE_WIDTH + "], got " + length); + } + BitVectorHelper.setBit(validityBuffer, index); + + // do the bound check. + valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); + + long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (length == 0) { + MemoryUtil.setMemory(outAddress, Decimal64Vector.TYPE_WIDTH, (byte) 0); + return; + } + if (LITTLE_ENDIAN) { + // swap bytes to convert BE to LE + for (int byteIdx = 0; byteIdx < length; ++byteIdx) { + MemoryUtil.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]); + } + if (length < TYPE_WIDTH) { + // sign extend + final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress + length, Decimal64Vector.TYPE_WIDTH - length, pad); + } + } else { + // copy data from value to outAddress + MemoryUtil.copyToMemory(value, 0, outAddress + Decimal64Vector.TYPE_WIDTH - length, length); + if (length < TYPE_WIDTH) { + // sign extend + final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal64Vector.TYPE_WIDTH - length, pad); + } + } + } + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param start start index of data in the buffer + * @param buffer ArrowBuf containing decimal value. + */ + public void set(int index, long start, ArrowBuf buffer) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, start, TYPE_WIDTH); + } + + /** + * Sets the element at given index using the buffer whose size maybe <= 8 bytes. + * + * @param index index to write the decimal to + * @param start start of value in the buffer + * @param buffer contains the decimal in native endian bytes + * @param length length of the value in the buffer + */ + public void setSafe(int index, long start, ArrowBuf buffer, int length) { + checkByteLength(length); + handleSafe(index); + BitVectorHelper.setBit(validityBuffer, index); + + // do the bound checks. + buffer.checkBytes(start, start + length); + valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); + + long inAddress = buffer.memoryAddress() + start; + long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH; + if (LITTLE_ENDIAN) { + MemoryUtil.copyMemory(inAddress, outAddress, length); + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress + length - 1); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress + length, Decimal64Vector.TYPE_WIDTH - length, pad); + } + } else { + MemoryUtil.copyMemory(inAddress, outAddress + Decimal64Vector.TYPE_WIDTH - length, length); + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal64Vector.TYPE_WIDTH - length, pad); + } + } + } + + /** + * Sets the element at given index using the buffer whose size maybe <= 8 bytes. + * + * @param index index to write the decimal to + * @param start start of value in the buffer + * @param buffer contains the decimal in big endian bytes + * @param length length of the value in the buffer + */ + public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) { + checkByteLength(length); + handleSafe(index); + BitVectorHelper.setBit(validityBuffer, index); + + // do the bound checks. + buffer.checkBytes(start, start + length); + valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH); + + // 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 (LITTLE_ENDIAN) { + // swap bytes to convert BE to LE + for (int byteIdx = 0; byteIdx < length; ++byteIdx) { + byte val = MemoryUtil.getByte((inAddress + length - 1) - byteIdx); + MemoryUtil.putByte(outAddress + byteIdx, val); + } + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress + length, Decimal64Vector.TYPE_WIDTH - length, pad); + } + } else { + MemoryUtil.copyMemory(inAddress, outAddress + Decimal64Vector.TYPE_WIDTH - length, length); + // sign extend + if (length < TYPE_WIDTH) { + byte msb = MemoryUtil.getByte(inAddress); + final byte pad = (byte) (msb < 0 ? 0xFF : 0x00); + MemoryUtil.setMemory(outAddress, Decimal64Vector.TYPE_WIDTH - length, pad); + } + } + } + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param value BigDecimal containing decimal value. + */ + public void set(int index, BigDecimal value) { + DecimalUtility.checkPrecisionAndScale(value, precision, scale); + BitVectorHelper.setBit(validityBuffer, index); + DecimalUtility.writeBigDecimalToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); + } + + /** + * Set the element at the given index to the given value. + * + * @param index position of element + * @param value long value. + */ + public void set(int index, long value) { + // write the value first: it range-checks and may throw, and the slot must + // not be marked valid if the write does not happen. + DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index, TYPE_WIDTH); + BitVectorHelper.setBit(validityBuffer, index); + } + + /** + * Set the element at the given index to the value set in data holder. If the value in holder is + * not indicated as set, element in the at the given index will be null. + * + * @param index position of element + * @param holder nullable data holder for value of element + */ + public void set(int index, NullableDecimal64Holder holder) throws IllegalArgumentException { + if (holder.isSet < 0) { + throw new IllegalArgumentException(); + } else if (holder.isSet > 0) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); + } else { + BitVectorHelper.unsetBit(validityBuffer, index); + } + } + + /** + * Set the element at the given index to the value set in data holder. + * + * @param index position of element + * @param holder data holder for value of element + */ + public void set(int index, Decimal64Holder holder) { + BitVectorHelper.setBit(validityBuffer, index); + valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH); + } + + /** + * Same as {@link #set(int, ArrowBuf)} except that it handles the case when index is greater than + * or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param buffer ArrowBuf containing decimal value. + */ + public void setSafe(int index, ArrowBuf buffer) { + handleSafe(index); + set(index, buffer); + } + + /** + * Same as {@link #setBigEndian(int, byte[])} except that it handles the case when index is + * greater than or equal to existing value capacity {@link #getValueCapacity()}. + */ + public void setBigEndianSafe(int index, byte[] value) { + handleSafe(index); + setBigEndian(index, value); + } + + /** + * Same as {@link #set(int, long, ArrowBuf)} except that it handles the case when index is greater + * than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param start start index of data in the buffer + * @param buffer ArrowBuf containing decimal value. + */ + public void setSafe(int index, long start, ArrowBuf buffer) { + handleSafe(index); + set(index, start, buffer); + } + + /** + * Same as {@link #set(int, BigDecimal)} except that it handles the case when index is greater + * than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param value BigDecimal containing decimal value. + */ + public void setSafe(int index, BigDecimal value) { + handleSafe(index); + set(index, value); + } + + /** + * Same as {@link #set(int, long)} except that it handles the case when index is greater than or + * equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param value long value. + */ + public void setSafe(int index, long value) { + handleSafe(index); + set(index, value); + } + + /** + * Same as {@link #set(int, NullableDecimal64Holder)} except that it handles the case when index + * is greater than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param holder nullable data holder for value of element + */ + public void setSafe(int index, NullableDecimal64Holder holder) throws IllegalArgumentException { + handleSafe(index); + set(index, holder); + } + + /** + * Same as {@link #set(int, Decimal64Holder)} except that it handles the case when index is + * greater than or equal to existing value capacity {@link #getValueCapacity()}. + * + * @param index position of element + * @param holder data holder for value of element + */ + public void setSafe(int index, Decimal64Holder holder) { + handleSafe(index); + set(index, holder); + } + + /** + * Store the given value at a particular position in the vector. isSet indicates whether the value + * is NULL or not. + * + * @param index position of the new value + * @param isSet 0 for NULL value, 1 otherwise + * @param start start position of the value in the buffer + * @param buffer buffer containing the value to be stored in the vector + */ + public void set(int index, int isSet, long start, ArrowBuf buffer) { + if (isSet > 0) { + set(index, start, buffer); + } else { + BitVectorHelper.unsetBit(validityBuffer, index); + } + } + + /** + * Same as {@link #set(int, int, long, ArrowBuf)} except that it handles the case when the + * position of new value is beyond the current value capacity of the vector. + * + * @param index position of the new value + * @param isSet 0 for NULL value, 1 otherwise + * @param start start position of the value in the buffer + * @param buffer buffer containing the value to be stored in the vector + */ + public void setSafe(int index, int isSet, long start, ArrowBuf buffer) { + handleSafe(index); + set(index, isSet, start, buffer); + } + + @Override + public void validateScalars() { + for (int i = 0; i < getValueCount(); ++i) { + BigDecimal value = getObject(i); + if (value != null) { + ValidateUtil.validateOrThrow( + DecimalUtility.checkPrecisionAndScaleNoThrow(value, getPrecision(), getScale()), + "Invalid value for Decimal64Vector at position " + + i + + ". Value does not fit in precision " + + getPrecision() + + " and scale " + + getScale() + + "."); + } + } + } + + /*----------------------------------------------------------------* + | | + | vector transfer | + | | + *----------------------------------------------------------------*/ + + /** + * Construct a TransferPair comprising this and a target vector of the same type. + * + * @param ref name of the target vector + * @param allocator allocator for the target vector + * @return {@link TransferPair} + */ + @Override + public TransferPair getTransferPair(String ref, BufferAllocator allocator) { + return new TransferImpl(ref, allocator); + } + + /** + * Construct a TransferPair comprising this and a target vector of the same type. + * + * @param field Field object used by the target vector + * @param allocator allocator for the target vector + * @return {@link TransferPair} + */ + @Override + public TransferPair getTransferPair(Field field, BufferAllocator allocator) { + return new TransferImpl(field, allocator); + } + + /** + * Construct a TransferPair with a desired target vector of the same type. + * + * @param to target vector + * @return {@link TransferPair} + */ + @Override + public TransferPair makeTransferPair(ValueVector to) { + return new TransferImpl((Decimal64Vector) to); + } + + private class TransferImpl implements TransferPair { + Decimal64Vector to; + + public TransferImpl(String ref, BufferAllocator allocator) { + to = + (Decimal64Vector.this.field != null && Decimal64Vector.this.field.getFieldType() != null) + ? new Decimal64Vector(ref, Decimal64Vector.this.field.getFieldType(), allocator) + : new Decimal64Vector( + ref, allocator, Decimal64Vector.this.precision, Decimal64Vector.this.scale); + } + + public TransferImpl(Field field, BufferAllocator allocator) { + to = new Decimal64Vector(field, allocator); + } + + public TransferImpl(Decimal64Vector to) { + this.to = to; + } + + @Override + public Decimal64Vector getTo() { + return to; + } + + @Override + public void transfer() { + transferTo(to); + } + + @Override + public void splitAndTransfer(int startIndex, int length) { + splitAndTransferTo(startIndex, length, to); + } + + @Override + public void copyValueSafe(int fromIndex, int toIndex) { + to.copyFromSafe(fromIndex, toIndex, Decimal64Vector.this); + } + } +} diff --git a/vector/src/main/java/org/apache/arrow/vector/GenerateSampleData.java b/vector/src/main/java/org/apache/arrow/vector/GenerateSampleData.java index e8250f9f57..3099e104a3 100644 --- a/vector/src/main/java/org/apache/arrow/vector/GenerateSampleData.java +++ b/vector/src/main/java/org/apache/arrow/vector/GenerateSampleData.java @@ -33,6 +33,12 @@ public static void generateTestData(final ValueVector vector, final int valueCou writeIntData((IntVector) vector, valueCount); } else if (vector instanceof DecimalVector) { writeDecimalData((DecimalVector) vector, valueCount); + } else if (vector instanceof Decimal32Vector) { + writeDecimal32Data((Decimal32Vector) vector, valueCount); + } else if (vector instanceof Decimal64Vector) { + writeDecimal64Data((Decimal64Vector) vector, valueCount); + } else if (vector instanceof Decimal256Vector) { + writeDecimal256Data((Decimal256Vector) vector, valueCount); } else if (vector instanceof BitVector) { writeBooleanData((BitVector) vector, valueCount); } else if (vector instanceof VarCharVector) { @@ -118,6 +124,47 @@ private static void writeDecimalData(DecimalVector vector, int valueCount) { vector.setValueCount(valueCount); } + // The narrow and wide decimal sample values are derived from the vector's own scale (rather than + // hardcoded) so they fit any precision/scale, including the precision-9 Decimal32 limit. + private static void writeDecimal32Data(Decimal32Vector vector, int valueCount) { + final BigDecimal even = BigDecimal.valueOf(1, vector.getScale()); + final BigDecimal odd = BigDecimal.valueOf(2, vector.getScale()); + for (int i = 0; i < valueCount; i++) { + if (i % 2 == 0) { + vector.setSafe(i, even); + } else { + vector.setSafe(i, odd); + } + } + vector.setValueCount(valueCount); + } + + private static void writeDecimal64Data(Decimal64Vector vector, int valueCount) { + final BigDecimal even = BigDecimal.valueOf(1, vector.getScale()); + final BigDecimal odd = BigDecimal.valueOf(2, vector.getScale()); + for (int i = 0; i < valueCount; i++) { + if (i % 2 == 0) { + vector.setSafe(i, even); + } else { + vector.setSafe(i, odd); + } + } + vector.setValueCount(valueCount); + } + + private static void writeDecimal256Data(Decimal256Vector vector, int valueCount) { + final BigDecimal even = BigDecimal.valueOf(1, vector.getScale()); + final BigDecimal odd = BigDecimal.valueOf(2, vector.getScale()); + for (int i = 0; i < valueCount; i++) { + if (i % 2 == 0) { + vector.setSafe(i, even); + } else { + vector.setSafe(i, odd); + } + } + vector.setValueCount(valueCount); + } + private static void writeIntData(IntVector vector, int valueCount) { final int even = 1000; final int odd = 2000; diff --git a/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java b/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java index 780a4ee659..2cc537964e 100644 --- a/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java +++ b/vector/src/main/java/org/apache/arrow/vector/extension/OpaqueType.java @@ -28,6 +28,8 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FieldVector; @@ -307,7 +309,11 @@ public FieldVector visit(Bool type) { @Override public FieldVector visit(Decimal type) { - if (type.getBitWidth() == 128) { + if (type.getBitWidth() == 32) { + return new Decimal32Vector(Field.nullable(name, type), allocator); + } else if (type.getBitWidth() == 64) { + return new Decimal64Vector(Field.nullable(name, type), allocator); + } else if (type.getBitWidth() == 128) { return new DecimalVector(Field.nullable(name, type), allocator); } else if (type.getBitWidth() == 256) { return new Decimal256Vector(Field.nullable(name, type), allocator); diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java index e4bab7eb80..cc82d9be2d 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java @@ -58,6 +58,8 @@ import org.apache.arrow.vector.BitVectorHelper; import org.apache.arrow.vector.BufferLayout.BufferType; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.Float4Vector; @@ -597,43 +599,36 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException } }; - BufferReader DECIMAL = - new BufferReader() { - @Override - protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { - final long size = (long) count * DecimalVector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + BufferReader DECIMAL32 = decimalReader(Decimal32Vector.TYPE_WIDTH); - for (int i = 0; i < count; i++) { - parser.nextToken(); - BigDecimal decimalValue = new BigDecimal(parser.readValueAs(String.class)); - DecimalUtility.writeBigDecimalToArrowBuf( - decimalValue, buf, i, DecimalVector.TYPE_WIDTH); - } + BufferReader DECIMAL64 = decimalReader(Decimal64Vector.TYPE_WIDTH); - buf.writerIndex(size); - return buf; - } - }; + BufferReader DECIMAL = decimalReader(DecimalVector.TYPE_WIDTH); - BufferReader DECIMAL256 = - new BufferReader() { - @Override - protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { - final long size = (long) count * Decimal256Vector.TYPE_WIDTH; - ArrowBuf buf = allocator.buffer(size); + BufferReader DECIMAL256 = decimalReader(Decimal256Vector.TYPE_WIDTH); + private BufferReader decimalReader(int typeWidth) { + return new BufferReader() { + @Override + protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException { + final long size = (long) count * typeWidth; + ArrowBuf buf = allocator.buffer(size); + try { for (int i = 0; i < count; i++) { parser.nextToken(); BigDecimal decimalValue = new BigDecimal(parser.readValueAs(String.class)); - DecimalUtility.writeBigDecimalToArrowBuf( - decimalValue, buf, i, Decimal256Vector.TYPE_WIDTH); + DecimalUtility.writeBigDecimalToArrowBuf(decimalValue, buf, i, typeWidth); } - - buf.writerIndex(size); - return buf; + } catch (Exception e) { + buf.close(); + throw e; } - }; + + buf.writerIndex(size); + return buf; + } + }; + } ArrowBuf readBinaryValues(BufferAllocator allocator, int count) throws IOException { ArrayList values = new ArrayList<>(count); @@ -774,6 +769,12 @@ private List readIntoBuffer( case FLOAT8: reader = helper.FLOAT8; break; + case DECIMAL32: + reader = helper.DECIMAL32; + break; + case DECIMAL64: + reader = helper.DECIMAL64; + break; case DECIMAL: reader = helper.DECIMAL; break; diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java index 68700fe6af..bbc44fdd9b 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileWriter.java @@ -34,6 +34,7 @@ import java.util.Set; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.BaseFixedWidthVector; import org.apache.arrow.vector.BaseLargeVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BaseVariableWidthViewVector; @@ -42,8 +43,6 @@ import org.apache.arrow.vector.BufferLayout.BufferType; import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; -import org.apache.arrow.vector.Decimal256Vector; -import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.FixedSizeBinaryVector; @@ -78,6 +77,7 @@ import org.apache.arrow.vector.dictionary.Dictionary; import org.apache.arrow.vector.dictionary.DictionaryProvider; import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.DecimalUtility; @@ -564,22 +564,15 @@ private void writeValueToGenerator( generator.writeString(new String(b, "UTF-8")); break; } + case DECIMAL32: + case DECIMAL64: case DECIMAL: - { - int scale = ((DecimalVector) vector).getScale(); - BigDecimal decimalValue = - DecimalUtility.getBigDecimalFromArrowBuf( - buffer, index, scale, DecimalVector.TYPE_WIDTH); - // We write the unscaled value, because the scale is stored in the type metadata. - generator.writeString(decimalValue.unscaledValue().toString()); - break; - } case DECIMAL256: { - int scale = ((Decimal256Vector) vector).getScale(); + int scale = ((ArrowType.Decimal) vector.getField().getType()).getScale(); + int typeWidth = ((BaseFixedWidthVector) vector).getTypeWidth(); BigDecimal decimalValue = - DecimalUtility.getBigDecimalFromArrowBuf( - buffer, index, scale, Decimal256Vector.TYPE_WIDTH); + DecimalUtility.getBigDecimalFromArrowBuf(buffer, index, scale, typeWidth); // We write the unscaled value, because the scale is stored in the type metadata. generator.writeString(decimalValue.unscaledValue().toString()); break; diff --git a/vector/src/main/java/org/apache/arrow/vector/types/Types.java b/vector/src/main/java/org/apache/arrow/vector/types/Types.java index 17503f98c8..ca74a0835c 100644 --- a/vector/src/main/java/org/apache/arrow/vector/types/Types.java +++ b/vector/src/main/java/org/apache/arrow/vector/types/Types.java @@ -28,6 +28,8 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.ExtensionTypeVector; @@ -81,6 +83,8 @@ import org.apache.arrow.vector.complex.impl.DateDayWriterImpl; import org.apache.arrow.vector.complex.impl.DateMilliWriterImpl; import org.apache.arrow.vector.complex.impl.Decimal256WriterImpl; +import org.apache.arrow.vector.complex.impl.Decimal32WriterImpl; +import org.apache.arrow.vector.complex.impl.Decimal64WriterImpl; import org.apache.arrow.vector.complex.impl.DecimalWriterImpl; import org.apache.arrow.vector.complex.impl.DenseUnionWriter; import org.apache.arrow.vector.complex.impl.DurationWriterImpl; @@ -802,6 +806,30 @@ public FieldWriter getNewFieldWriter(ValueVector vector) { "FieldWriter for run-end encoded vector is not implemented yet."); } }, + DECIMAL32(null) { + @Override + public FieldVector getNewVector( + Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { + return new Decimal32Vector(field, allocator); + } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + return new Decimal32WriterImpl((Decimal32Vector) vector); + } + }, + DECIMAL64(null) { + @Override + public FieldVector getNewVector( + Field field, BufferAllocator allocator, CallBack schemaChangeCallback) { + return new Decimal64Vector(field, allocator); + } + + @Override + public FieldWriter getNewFieldWriter(ValueVector vector) { + return new Decimal64WriterImpl((Decimal64Vector) vector); + } + }, ; private final ArrowType type; @@ -948,10 +976,16 @@ public MinorType visit(Bool type) { @Override public MinorType visit(Decimal type) { - if (type.getBitWidth() == 256) { - return MinorType.DECIMAL256; + switch (type.getBitWidth()) { + case 32: + return MinorType.DECIMAL32; + case 64: + return MinorType.DECIMAL64; + case 256: + return MinorType.DECIMAL256; + default: + return MinorType.DECIMAL; } - return MinorType.DECIMAL; } @Override diff --git a/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java b/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java index 31b79fe53a..6cc16219b0 100644 --- a/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java +++ b/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java @@ -45,9 +45,17 @@ private DecimalUtility() {} */ public static BigDecimal getBigDecimalFromArrowBuf( ArrowBuf bytebuf, int index, int scale, int byteWidth) { + return getBigDecimalFromArrowBufAtOffset(bytebuf, (long) index * byteWidth, scale, byteWidth); + } + + /** + * Read an ArrowType.Decimal at the given byte offset in the ArrowBuf and convert to a BigDecimal + * with the given scale. + */ + public static BigDecimal getBigDecimalFromArrowBufAtOffset( + ArrowBuf bytebuf, long startIndex, int scale, int byteWidth) { byte[] value = new byte[byteWidth]; byte temp; - final long startIndex = (long) index * byteWidth; bytebuf.getBytes(startIndex, value, 0, byteWidth); if (LITTLE_ENDIAN) { @@ -118,7 +126,7 @@ public static boolean checkPrecisionAndScale( */ public static boolean checkPrecisionAndScaleNoThrow( BigDecimal value, int vectorPrecision, int vectorScale) { - return value.scale() == vectorScale && value.precision() < vectorPrecision; + return value.scale() == vectorScale && value.precision() <= vectorPrecision; } /** @@ -159,15 +167,23 @@ public static void writeBigDecimalToArrowBuf( /** * Write the given long to the ArrowBuf at the given value index. This routine extends the - * original sign bit to a new upper area in 128-bit or 256-bit. + * original sign bit to a new upper area in 64-bit, 128-bit or 256-bit. */ public static void writeLongToArrowBuf(long value, ArrowBuf bytebuf, int index, int byteWidth) { - if (byteWidth != 16 && byteWidth != 32) { + if (byteWidth != 4 && byteWidth != 8 && byteWidth != 16 && byteWidth != 32) { throw new UnsupportedOperationException( "DecimalUtility.writeLongToArrowBuf() currently supports " - + "128-bit or 256-bit width data"); + + "32-bit, 64-bit, 128-bit or 256-bit width data"); } final long addressOfValue = bytebuf.memoryAddress() + (long) index * byteWidth; + if (byteWidth == 4) { + if (value < Integer.MIN_VALUE || value > Integer.MAX_VALUE) { + throw new UnsupportedOperationException( + "Decimal value does not fit in 32-bit width: " + value); + } + MemoryUtil.putInt(addressOfValue, (int) value); + return; + } final long padValue = Long.signum(value) == -1 ? -1L : 0L; if (LITTLE_ENDIAN) { MemoryUtil.putLong(addressOfValue, value); diff --git a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java index 395852ef79..df16dd176a 100644 --- a/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java +++ b/vector/src/main/java/org/apache/arrow/vector/validate/ValidateVectorTypeVisitor.java @@ -27,6 +27,8 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.ExtensionTypeVector; @@ -155,7 +157,33 @@ private void validateDecimalVector(ValueVector vector) { decimalType.getScale(), decimalType.getPrecision()); switch (decimalType.getBitWidth()) { + case Decimal32Vector.TYPE_WIDTH * 8: + validateOrThrow( + vector instanceof Decimal32Vector, + "Expected Decimal32Vector for decimal 32, actual %s.", + vector.getClass()); + validateOrThrow( + decimalType.getPrecision() >= 1 + && decimalType.getPrecision() <= Decimal32Vector.MAX_PRECISION, + "Invalid precision %s for decimal 32.", + decimalType.getPrecision()); + break; + case Decimal64Vector.TYPE_WIDTH * 8: + validateOrThrow( + vector instanceof Decimal64Vector, + "Expected Decimal64Vector for decimal 64, actual %s.", + vector.getClass()); + validateOrThrow( + decimalType.getPrecision() >= 1 + && decimalType.getPrecision() <= Decimal64Vector.MAX_PRECISION, + "Invalid precision %s for decimal 64.", + decimalType.getPrecision()); + break; case DecimalVector.TYPE_WIDTH * 8: + validateOrThrow( + vector instanceof DecimalVector, + "Expected DecimalVector for decimal 128, actual %s.", + vector.getClass()); validateOrThrow( decimalType.getPrecision() >= 1 && decimalType.getPrecision() <= DecimalVector.MAX_PRECISION, @@ -163,6 +191,10 @@ private void validateDecimalVector(ValueVector vector) { decimalType.getPrecision()); break; case Decimal256Vector.TYPE_WIDTH * 8: + validateOrThrow( + vector instanceof Decimal256Vector, + "Expected Decimal256Vector for decimal 256, actual %s.", + vector.getClass()); validateOrThrow( decimalType.getPrecision() >= 1 && decimalType.getPrecision() <= Decimal256Vector.MAX_PRECISION, @@ -171,7 +203,8 @@ private void validateDecimalVector(ValueVector vector) { break; default: throw new ValidateUtil.ValidateException( - "Only decimal 128 or decimal 256 are supported for decimal types"); + "Only decimal 32, decimal 64, decimal 128 or decimal 256 are supported for decimal" + + " types"); } } @@ -273,7 +306,10 @@ public Void visit(BaseFixedWidthVector vector, Void value) { validateIntVector(vector, 64, false); } else if (vector instanceof BitVector) { validateVectorCommon(vector, ArrowType.Bool.class); - } else if (vector instanceof DecimalVector || vector instanceof Decimal256Vector) { + } else if (vector instanceof Decimal32Vector + || vector instanceof Decimal64Vector + || vector instanceof DecimalVector + || vector instanceof Decimal256Vector) { validateVectorCommon(vector, ArrowType.Decimal.class); validateDecimalVector(vector); } else if (vector instanceof DateDayVector) { diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDecimal32Vector.java b/vector/src/test/java/org/apache/arrow/vector/TestDecimal32Vector.java new file mode 100644 index 0000000000..0a83c9a646 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestDecimal32Vector.java @@ -0,0 +1,460 @@ +/* + * 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.arrow.vector; + +import static org.junit.jupiter.api.Assertions.assertEquals; +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.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.nio.channels.Channels; +import java.util.Collections; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.complex.impl.Decimal32HolderReaderImpl; +import org.apache.arrow.vector.complex.impl.NullableDecimal32HolderReaderImpl; +import org.apache.arrow.vector.holders.Decimal32Holder; +import org.apache.arrow.vector.holders.NullableDecimal32Holder; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.TransferPair; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestDecimal32Vector { + + private static long[] intValues; + + static { + intValues = new long[40]; + for (int i = 0; i < intValues.length / 2; i++) { + intValues[i] = 1L << (i + 1); + intValues[2 * i] = -1L * (1 << (i + 1)); + } + } + + private int scale = 3; + + private BufferAllocator allocator; + + @BeforeEach + public void init() { + allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); + } + + @AfterEach + public void terminate() throws Exception { + allocator.close(); + } + + @Test + public void testValuesWriteRead() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, scale, 32), allocator)) { + + try (Decimal32Vector oldConstructor = new Decimal32Vector("decimal", allocator, 9, scale)) { + assertEquals(decimalVector.getField().getType(), oldConstructor.getField().getType()); + } + + decimalVector.allocateNew(); + BigDecimal[] values = new BigDecimal[intValues.length]; + for (int i = 0; i < intValues.length; i++) { + BigDecimal decimal = new BigDecimal(BigInteger.valueOf(intValues[i]), scale); + values[i] = decimal; + decimalVector.setSafe(i, decimal); + } + + decimalVector.setValueCount(intValues.length); + + for (int i = 0; i < intValues.length; i++) { + BigDecimal value = decimalVector.getObject(i); + assertEquals(values[i], value, "unexpected data at index: " + i); + } + } + } + + @Test + public void testDecimal32DifferentScaleAndPrecision() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(4, 2, 32), allocator)) { + decimalVector.allocateNew(); + + // test Decimal32 with different scale + { + BigDecimal decimal = new BigDecimal(BigInteger.valueOf(0), 3); + UnsupportedOperationException ue = + assertThrows( + UnsupportedOperationException.class, () -> decimalVector.setSafe(0, decimal)); + assertEquals( + "BigDecimal scale must equal that in the Arrow vector: 3 != 2", ue.getMessage()); + } + + // test BigDecimal with larger precision than initialized + { + BigDecimal decimal = new BigDecimal(BigInteger.valueOf(12345), 2); + UnsupportedOperationException ue = + assertThrows( + UnsupportedOperationException.class, () -> decimalVector.setSafe(0, decimal)); + assertEquals( + "BigDecimal precision cannot be greater than that in the Arrow vector: 5 > 4", + ue.getMessage()); + } + decimalVector.setValueCount(1); + assertTrue(decimalVector.isNull(0)); + } + } + + @Test + public void testWriteBigEndian() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, 4, 32), allocator)) { + decimalVector.allocateNew(); + BigDecimal decimal1 = new BigDecimal("12345.6789"); + BigDecimal decimal2 = new BigDecimal("1.1234"); + BigDecimal decimal3 = new BigDecimal("1.0000"); + BigDecimal decimal4 = new BigDecimal("0.1111"); + BigDecimal decimal5 = new BigDecimal("98765.4321"); + BigDecimal decimal6 = new BigDecimal("-12345.6789"); + BigDecimal decimal7 = new BigDecimal("-1.0001"); + BigDecimal decimal8 = new BigDecimal("5.3434"); + + decimalVector.setBigEndian(0, decimal1.unscaledValue().toByteArray()); + decimalVector.setBigEndian(1, decimal2.unscaledValue().toByteArray()); + decimalVector.setBigEndian(2, decimal3.unscaledValue().toByteArray()); + decimalVector.setBigEndian(3, decimal4.unscaledValue().toByteArray()); + decimalVector.setBigEndian(4, decimal5.unscaledValue().toByteArray()); + decimalVector.setBigEndian(5, decimal6.unscaledValue().toByteArray()); + decimalVector.setBigEndian(6, decimal7.unscaledValue().toByteArray()); + decimalVector.setBigEndian(7, decimal8.unscaledValue().toByteArray()); + + decimalVector.setValueCount(8); + assertEquals(8, decimalVector.getValueCount()); + assertEquals(decimal1, decimalVector.getObject(0)); + assertEquals(decimal2, decimalVector.getObject(1)); + assertEquals(decimal3, decimalVector.getObject(2)); + assertEquals(decimal4, decimalVector.getObject(3)); + assertEquals(decimal5, decimalVector.getObject(4)); + assertEquals(decimal6, decimalVector.getObject(5)); + assertEquals(decimal7, decimalVector.getObject(6)); + assertEquals(decimal8, decimalVector.getObject(7)); + } + } + + @Test + public void testLongReadWrite() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, 0, 32), allocator)) { + decimalVector.allocateNew(); + + // The full precision-9 range fits in a Decimal32's 32-bit word. + long[] longValues = {0L, -2L, 999999999L, -999999999L, 187L}; + + for (int i = 0; i < longValues.length; ++i) { + decimalVector.set(i, longValues[i]); + } + + decimalVector.setValueCount(longValues.length); + + for (int i = 0; i < longValues.length; ++i) { + assertEquals(new BigDecimal(longValues[i]), decimalVector.getObject(i)); + } + + assertThrows(UnsupportedOperationException.class, () -> decimalVector.set(0, Long.MAX_VALUE)); + + // a failed overflowing set must not mark the slot valid + final int overflowIdx = longValues.length; + assertThrows( + UnsupportedOperationException.class, + () -> decimalVector.set(overflowIdx, Long.MAX_VALUE)); + decimalVector.setValueCount(overflowIdx + 1); + assertTrue(decimalVector.isNull(overflowIdx)); + } + } + + @Test + public void testHolderReaderUsesDecimalByteOrder() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, 2, 32), allocator)) { + decimalVector.allocateNew(); + decimalVector.set(0, new BigDecimal("1.23")); + BigDecimal expected = new BigDecimal("1234567.89"); + decimalVector.set(1, expected); + decimalVector.setValueCount(2); + + NullableDecimal32Holder nullableHolder = new NullableDecimal32Holder(); + decimalVector.get(1, nullableHolder); + assertEquals( + expected, new NullableDecimal32HolderReaderImpl(nullableHolder).readBigDecimal()); + + Decimal32Holder holder = new Decimal32Holder(); + holder.buffer = nullableHolder.buffer; + holder.start = nullableHolder.start; + holder.precision = nullableHolder.precision; + holder.scale = nullableHolder.scale; + assertEquals(expected, new Decimal32HolderReaderImpl(holder).readBigDecimal()); + } + } + + @Test + public void testBigDecimalReadWrite() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, 4, 32), allocator)) { + decimalVector.allocateNew(); + BigDecimal decimal1 = new BigDecimal("12345.6789"); + BigDecimal decimal2 = new BigDecimal("1.1234"); + BigDecimal decimal3 = new BigDecimal("1.0000"); + BigDecimal decimal4 = new BigDecimal("-0.1111"); + BigDecimal decimal5 = new BigDecimal("-98765.4321"); + BigDecimal decimal6 = new BigDecimal("-2.2222"); + BigDecimal decimal7 = new BigDecimal("7.6666"); + BigDecimal decimal8 = new BigDecimal("12121.3434"); + + decimalVector.set(0, decimal1); + decimalVector.set(1, decimal2); + decimalVector.set(2, decimal3); + decimalVector.set(3, decimal4); + decimalVector.set(4, decimal5); + decimalVector.set(5, decimal6); + decimalVector.set(6, decimal7); + decimalVector.set(7, decimal8); + + decimalVector.setValueCount(8); + assertEquals(8, decimalVector.getValueCount()); + assertEquals(decimal1, decimalVector.getObject(0)); + assertEquals(decimal2, decimalVector.getObject(1)); + assertEquals(decimal3, decimalVector.getObject(2)); + assertEquals(decimal4, decimalVector.getObject(3)); + assertEquals(decimal5, decimalVector.getObject(4)); + assertEquals(decimal6, decimalVector.getObject(5)); + assertEquals(decimal7, decimalVector.getObject(6)); + assertEquals(decimal8, decimalVector.getObject(7)); + } + } + + /** + * Test {@link Decimal32Vector#setBigEndian(int, byte[])} which takes BE layout input and stores + * in native-endian (NE) layout. Cases to cover: input byte array in different lengths in range + * [1-4] and negative values. + */ + @Test + public void decimalBE2NE() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, 2, 32), allocator)) { + decimalVector.allocateNew(); + + BigInteger[] testBigInts = + new BigInteger[] { + new BigInteger("0"), + new BigInteger("-1"), + new BigInteger("23"), + new BigInteger("234234"), + new BigInteger("-234234"), + new BigInteger("12345"), + new BigInteger("-345345"), + new BigInteger("754533"), + new BigInteger("999999999"), // 9 nines, fits in 4 bytes + new BigInteger("-999999999") + }; + + int insertionIdx = 0; + insertionIdx++; // insert a null + for (BigInteger val : testBigInts) { + decimalVector.setBigEndian(insertionIdx++, val.toByteArray()); + } + insertionIdx++; // insert a null + // insert a zero length buffer + decimalVector.setBigEndian(insertionIdx++, new byte[0]); + + // Try inserting a buffer larger than 4 bytes and expect a failure + final int insertionIdxCapture = insertionIdx; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, + () -> decimalVector.setBigEndian(insertionIdxCapture, new byte[5])); + assertTrue( + ex.getMessage().equals("Invalid decimal value length. Valid length in [0 - 4], got 5")); + decimalVector.setValueCount(insertionIdx); + + // retrieve values and check if they are correct + int outputIdx = 0; + assertTrue(decimalVector.isNull(outputIdx++)); + for (BigInteger expected : testBigInts) { + final BigDecimal actual = decimalVector.getObject(outputIdx++); + assertEquals(expected, actual.unscaledValue()); + } + assertTrue(decimalVector.isNull(outputIdx++)); + assertEquals(BigInteger.valueOf(0), decimalVector.getObject(outputIdx).unscaledValue()); + } + } + + @Test + public void setUsingArrowBufOfLEInts() { + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(5, 2, 32), allocator); + ArrowBuf buf = allocator.buffer(8)) { + decimalVector.allocateNew(); + + // add a positive value equivalent to 705.32 + int val = 70532; + buf.setInt(0, val); + decimalVector.setSafe(0, 0, buf, 4); + + // add a -ve value equivalent to -705.32 + val = -70532; + buf.setInt(4, val); + decimalVector.setSafe(1, 4, buf, 4); + + decimalVector.setValueCount(2); + + BigDecimal[] expectedValues = + new BigDecimal[] {BigDecimal.valueOf(705.32), BigDecimal.valueOf(-705.32)}; + for (int i = 0; i < 2; i++) { + BigDecimal value = decimalVector.getObject(i); + assertEquals(expectedValues[i], value); + } + } + } + + /** Round-trip a Decimal32 vector through the Arrow IPC stream format. */ + @Test + public void testIpcRoundtrip() throws Exception { + Field field = new Field("decimal", FieldType.nullable(new ArrowType.Decimal(9, 4, 32)), null); + Schema schema = new Schema(Collections.singletonList(field)); + byte[] serialized; + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + Decimal32Vector vector = (Decimal32Vector) root.getVector("decimal"); + vector.allocateNew(); + vector.set(0, new BigDecimal("12345.6789")); + vector.setNull(1); + vector.set(2, new BigDecimal("-9876.5432")); + root.setRowCount(3); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(out))) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + serialized = out.toByteArray(); + } + + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(serialized), allocator)) { + VectorSchemaRoot readRoot = reader.getVectorSchemaRoot(); + ArrowType.Decimal readType = + (ArrowType.Decimal) readRoot.getSchema().getFields().get(0).getType(); + assertEquals(32, readType.getBitWidth()); + assertEquals(9, readType.getPrecision()); + assertEquals(4, readType.getScale()); + + assertTrue(reader.loadNextBatch()); + Decimal32Vector vector = (Decimal32Vector) readRoot.getVector("decimal"); + assertEquals(3, vector.getValueCount()); + assertEquals(new BigDecimal("12345.6789"), vector.getObject(0)); + assertTrue(vector.isNull(1)); + assertEquals(new BigDecimal("-9876.5432"), vector.getObject(2)); + } + } + + @Test + public void testGetTransferPairWithField() { + final Decimal32Vector fromVector = new Decimal32Vector("decimal", allocator, 9, scale); + final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); + final Decimal32Vector toVector = (Decimal32Vector) transferPair.getTo(); + // Field inside a new vector created by reusing a field should be the same in memory as the + // original field. + assertSame(fromVector.getField(), toVector.getField()); + fromVector.close(); + toVector.close(); + } + + @Test + public void testGetTransferPairWithoutField() { + final Decimal32Vector fromVector = new Decimal32Vector("decimal", allocator, 9, scale); + final TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField().getName(), allocator); + final Decimal32Vector toVector = (Decimal32Vector) transferPair.getTo(); + // A new Field created inside a new vector should reuse the field type (should be the same in + // memory as the original Field's field type). + assertSame(fromVector.getField().getFieldType(), toVector.getField().getFieldType()); + fromVector.close(); + toVector.close(); + } + + @Test + public void testValidateScalarsAtMaxPrecision() { + // A value whose precision exactly equals the vector's precision must validate cleanly. + // Regression: DecimalUtility.checkPrecisionAndScaleNoThrow used a strict '<' comparison, so a + // value that set() accepts was wrongly rejected by validateScalars() at the precision boundary. + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(5, 2, 32), allocator)) { + decimalVector.allocateNew(); + decimalVector.set(0, new BigDecimal("999.99")); // precision 5 == vector precision + decimalVector.setValueCount(1); + decimalVector.validateScalars(); // must not throw + } + } + + @Test + public void testSetSafeRejectsTooLongLength() { + // Regression: setSafe/setBigEndianSafe with length > TYPE_WIDTH used to write past the slot + // before any validation, silently corrupting neighbouring values. They must now reject the + // length up front without mutating the vector. + try (Decimal32Vector decimalVector = + TestUtils.newVector( + Decimal32Vector.class, "decimal", new ArrowType.Decimal(9, 0, 32), allocator); + ArrowBuf buf = allocator.buffer(16)) { + decimalVector.allocateNew(2); + decimalVector.set(1, 123L); // neighbour that must remain intact + buf.setLong(0, 0x0102030405060708L); + + final int tooLong = Decimal32Vector.TYPE_WIDTH + 1; + assertThrows(IllegalArgumentException.class, () -> decimalVector.setSafe(0, 0, buf, tooLong)); + assertThrows( + IllegalArgumentException.class, () -> decimalVector.setBigEndianSafe(0, 0, buf, tooLong)); + assertThrows(IllegalArgumentException.class, () -> decimalVector.setSafe(0, 0, buf, 0)); + assertThrows( + IllegalArgumentException.class, () -> decimalVector.setBigEndianSafe(0, 0, buf, 0)); + assertThrows(IllegalArgumentException.class, () -> decimalVector.setSafe(0, 0, buf, -1)); + assertThrows( + IllegalArgumentException.class, () -> decimalVector.setBigEndianSafe(0, 0, buf, -1)); + + decimalVector.setValueCount(2); + assertTrue(decimalVector.isNull(0)); // rejected write left index 0 untouched + assertEquals(new BigDecimal(123), decimalVector.getObject(1)); + } + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/TestDecimal64Vector.java b/vector/src/test/java/org/apache/arrow/vector/TestDecimal64Vector.java new file mode 100644 index 0000000000..a7358bc950 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestDecimal64Vector.java @@ -0,0 +1,463 @@ +/* + * 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.arrow.vector; + +import static org.junit.jupiter.api.Assertions.assertEquals; +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.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.nio.channels.Channels; +import java.util.Collections; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.complex.impl.Decimal64HolderReaderImpl; +import org.apache.arrow.vector.complex.impl.NullableDecimal64HolderReaderImpl; +import org.apache.arrow.vector.holders.Decimal64Holder; +import org.apache.arrow.vector.holders.NullableDecimal64Holder; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.TransferPair; +import org.apache.arrow.vector.validate.ValidateUtil; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestDecimal64Vector { + + private static long[] intValues; + + static { + intValues = new long[60]; + for (int i = 0; i < intValues.length / 2; i++) { + intValues[i] = 1L << (i + 1); + intValues[2 * i] = -1L * (1 << (i + 1)); + } + } + + private int scale = 3; + + private BufferAllocator allocator; + + @BeforeEach + public void init() { + allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100); + } + + @AfterEach + public void terminate() throws Exception { + allocator.close(); + } + + @Test + public void testValuesWriteRead() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, scale, 64), allocator)) { + + try (Decimal64Vector oldConstructor = new Decimal64Vector("decimal", allocator, 18, scale)) { + assertEquals(decimalVector.getField().getType(), oldConstructor.getField().getType()); + } + + decimalVector.allocateNew(); + BigDecimal[] values = new BigDecimal[intValues.length]; + for (int i = 0; i < intValues.length; i++) { + BigDecimal decimal = new BigDecimal(BigInteger.valueOf(intValues[i]), scale); + values[i] = decimal; + decimalVector.setSafe(i, decimal); + } + + decimalVector.setValueCount(intValues.length); + + for (int i = 0; i < intValues.length; i++) { + BigDecimal value = decimalVector.getObject(i); + assertEquals(values[i], value, "unexpected data at index: " + i); + } + } + } + + @Test + public void testDecimal64DifferentScaleAndPrecision() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(4, 2, 64), allocator)) { + decimalVector.allocateNew(); + + // test Decimal64 with different scale + { + BigDecimal decimal = new BigDecimal(BigInteger.valueOf(0), 3); + UnsupportedOperationException ue = + assertThrows( + UnsupportedOperationException.class, () -> decimalVector.setSafe(0, decimal)); + assertEquals( + "BigDecimal scale must equal that in the Arrow vector: 3 != 2", ue.getMessage()); + } + + // test BigDecimal with larger precision than initialized + { + BigDecimal decimal = new BigDecimal(BigInteger.valueOf(12345), 2); + UnsupportedOperationException ue = + assertThrows( + UnsupportedOperationException.class, () -> decimalVector.setSafe(0, decimal)); + assertEquals( + "BigDecimal precision cannot be greater than that in the Arrow vector: 5 > 4", + ue.getMessage()); + } + decimalVector.setValueCount(1); + assertTrue(decimalVector.isNull(0)); + } + } + + @Test + public void testWriteBigEndian() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 9, 64), allocator)) { + decimalVector.allocateNew(); + BigDecimal decimal1 = new BigDecimal("123456789.000000000"); + BigDecimal decimal2 = new BigDecimal("11.123456789"); + BigDecimal decimal3 = new BigDecimal("1.000000000"); + BigDecimal decimal4 = new BigDecimal("0.111111111"); + BigDecimal decimal5 = new BigDecimal("987654321.123456789"); + BigDecimal decimal6 = new BigDecimal("-123456789.123456789"); + BigDecimal decimal7 = new BigDecimal("-1.000000001"); + BigDecimal decimal8 = new BigDecimal("55.343434343"); + + byte[] decimalValue1 = decimal1.unscaledValue().toByteArray(); + byte[] decimalValue2 = decimal2.unscaledValue().toByteArray(); + byte[] decimalValue3 = decimal3.unscaledValue().toByteArray(); + byte[] decimalValue4 = decimal4.unscaledValue().toByteArray(); + byte[] decimalValue5 = decimal5.unscaledValue().toByteArray(); + byte[] decimalValue6 = decimal6.unscaledValue().toByteArray(); + byte[] decimalValue7 = decimal7.unscaledValue().toByteArray(); + byte[] decimalValue8 = decimal8.unscaledValue().toByteArray(); + + decimalVector.setBigEndian(0, decimalValue1); + decimalVector.setBigEndian(1, decimalValue2); + decimalVector.setBigEndian(2, decimalValue3); + decimalVector.setBigEndian(3, decimalValue4); + decimalVector.setBigEndian(4, decimalValue5); + decimalVector.setBigEndian(5, decimalValue6); + decimalVector.setBigEndian(6, decimalValue7); + decimalVector.setBigEndian(7, decimalValue8); + + decimalVector.setValueCount(8); + assertEquals(8, decimalVector.getValueCount()); + assertEquals(decimal1, decimalVector.getObject(0)); + assertEquals(decimal2, decimalVector.getObject(1)); + assertEquals(decimal3, decimalVector.getObject(2)); + assertEquals(decimal4, decimalVector.getObject(3)); + assertEquals(decimal5, decimalVector.getObject(4)); + assertEquals(decimal6, decimalVector.getObject(5)); + assertEquals(decimal7, decimalVector.getObject(6)); + assertEquals(decimal8, decimalVector.getObject(7)); + } + } + + @Test + public void testLongReadWrite() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 0, 64), allocator)) { + decimalVector.allocateNew(); + + long[] longValues = {0L, -2L, 999999999999999999L, -999999999999999999L, 187L}; + + for (int i = 0; i < longValues.length; ++i) { + decimalVector.set(i, longValues[i]); + } + + decimalVector.setValueCount(longValues.length); + + for (int i = 0; i < longValues.length; ++i) { + assertEquals(new BigDecimal(longValues[i]), decimalVector.getObject(i)); + } + + decimalVector.set(0, Long.MAX_VALUE); + decimalVector.set(1, Long.MIN_VALUE); + decimalVector.setValueCount(2); + assertThrows(ValidateUtil.ValidateException.class, decimalVector::validateScalars); + } + } + + @Test + public void testHolderReaderUsesDecimalByteOrder() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 2, 64), allocator)) { + decimalVector.allocateNew(); + decimalVector.set(0, new BigDecimal("1.23")); + BigDecimal expected = new BigDecimal("1234567890123456.78"); + decimalVector.set(1, expected); + decimalVector.setValueCount(2); + + NullableDecimal64Holder nullableHolder = new NullableDecimal64Holder(); + decimalVector.get(1, nullableHolder); + assertEquals( + expected, new NullableDecimal64HolderReaderImpl(nullableHolder).readBigDecimal()); + + Decimal64Holder holder = new Decimal64Holder(); + holder.buffer = nullableHolder.buffer; + holder.start = nullableHolder.start; + holder.precision = nullableHolder.precision; + holder.scale = nullableHolder.scale; + assertEquals(expected, new Decimal64HolderReaderImpl(holder).readBigDecimal()); + } + } + + @Test + public void testBigDecimalReadWrite() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 9, 64), allocator)) { + decimalVector.allocateNew(); + BigDecimal decimal1 = new BigDecimal("123456789.000000000"); + BigDecimal decimal2 = new BigDecimal("11.123456789"); + BigDecimal decimal3 = new BigDecimal("1.000000000"); + BigDecimal decimal4 = new BigDecimal("-0.111111111"); + BigDecimal decimal5 = new BigDecimal("-987654321.123456789"); + BigDecimal decimal6 = new BigDecimal("-2.222222222"); + BigDecimal decimal7 = new BigDecimal("7.666666667"); + BigDecimal decimal8 = new BigDecimal("121212121.343434343"); + + decimalVector.set(0, decimal1); + decimalVector.set(1, decimal2); + decimalVector.set(2, decimal3); + decimalVector.set(3, decimal4); + decimalVector.set(4, decimal5); + decimalVector.set(5, decimal6); + decimalVector.set(6, decimal7); + decimalVector.set(7, decimal8); + + decimalVector.setValueCount(8); + assertEquals(8, decimalVector.getValueCount()); + assertEquals(decimal1, decimalVector.getObject(0)); + assertEquals(decimal2, decimalVector.getObject(1)); + assertEquals(decimal3, decimalVector.getObject(2)); + assertEquals(decimal4, decimalVector.getObject(3)); + assertEquals(decimal5, decimalVector.getObject(4)); + assertEquals(decimal6, decimalVector.getObject(5)); + assertEquals(decimal7, decimalVector.getObject(6)); + assertEquals(decimal8, decimalVector.getObject(7)); + } + } + + /** + * Test {@link Decimal64Vector#setBigEndian(int, byte[])} which takes BE layout input and stores + * in native-endian (NE) layout. Cases to cover: input byte array in different lengths in range + * [1-8] and negative values. + */ + @Test + public void decimalBE2NE() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 2, 64), allocator)) { + decimalVector.allocateNew(); + + BigInteger[] testBigInts = + new BigInteger[] { + new BigInteger("0"), + new BigInteger("-1"), + new BigInteger("23"), + new BigInteger("234234"), + new BigInteger("-234234234"), + new BigInteger("234234234234"), + new BigInteger("-56345345345345"), + new BigInteger("999999999999999999"), // 18 nines, fits in 8 bytes + new BigInteger("-999999999999999999"), + new BigInteger("-345345"), + new BigInteger("754533") + }; + + int insertionIdx = 0; + insertionIdx++; // insert a null + for (BigInteger val : testBigInts) { + decimalVector.setBigEndian(insertionIdx++, val.toByteArray()); + } + insertionIdx++; // insert a null + // insert a zero length buffer + decimalVector.setBigEndian(insertionIdx++, new byte[0]); + + // Try inserting a buffer larger than 8 bytes and expect a failure + final int insertionIdxCapture = insertionIdx; + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, + () -> decimalVector.setBigEndian(insertionIdxCapture, new byte[9])); + assertTrue( + ex.getMessage().equals("Invalid decimal value length. Valid length in [0 - 8], got 9")); + decimalVector.setValueCount(insertionIdx); + + // retrieve values and check if they are correct + int outputIdx = 0; + assertTrue(decimalVector.isNull(outputIdx++)); + for (BigInteger expected : testBigInts) { + final BigDecimal actual = decimalVector.getObject(outputIdx++); + assertEquals(expected, actual.unscaledValue()); + } + assertTrue(decimalVector.isNull(outputIdx++)); + assertEquals(BigInteger.valueOf(0), decimalVector.getObject(outputIdx).unscaledValue()); + } + } + + @Test + public void setUsingArrowLongLEBytes() { + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 0, 64), allocator); + ArrowBuf buf = allocator.buffer(16)) { + decimalVector.allocateNew(); + + long val = Long.MAX_VALUE; + buf.setLong(0, val); + decimalVector.setSafe(0, 0, buf, 8); + + val = Long.MIN_VALUE; + buf.setLong(8, val); + decimalVector.setSafe(1, 8, buf, 8); + + decimalVector.setValueCount(2); + + BigDecimal[] expectedValues = + new BigDecimal[] {BigDecimal.valueOf(Long.MAX_VALUE), BigDecimal.valueOf(Long.MIN_VALUE)}; + for (int i = 0; i < 2; i++) { + BigDecimal value = decimalVector.getObject(i); + assertEquals(expectedValues[i], value); + } + } + } + + /** Round-trip a Decimal64 vector through the Arrow IPC stream format. */ + @Test + public void testIpcRoundtrip() throws Exception { + Field field = new Field("decimal", FieldType.nullable(new ArrowType.Decimal(18, 4, 64)), null); + Schema schema = new Schema(Collections.singletonList(field)); + byte[] serialized; + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + Decimal64Vector vector = (Decimal64Vector) root.getVector("decimal"); + vector.allocateNew(); + vector.set(0, new BigDecimal("12345.6789")); + vector.setNull(1); + vector.set(2, new BigDecimal("-98765.4321")); + root.setRowCount(3); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(out))) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + serialized = out.toByteArray(); + } + + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(serialized), allocator)) { + VectorSchemaRoot readRoot = reader.getVectorSchemaRoot(); + ArrowType.Decimal readType = + (ArrowType.Decimal) readRoot.getSchema().getFields().get(0).getType(); + assertEquals(64, readType.getBitWidth()); + assertEquals(18, readType.getPrecision()); + assertEquals(4, readType.getScale()); + + assertTrue(reader.loadNextBatch()); + Decimal64Vector vector = (Decimal64Vector) readRoot.getVector("decimal"); + assertEquals(3, vector.getValueCount()); + assertEquals(new BigDecimal("12345.6789"), vector.getObject(0)); + assertTrue(vector.isNull(1)); + assertEquals(new BigDecimal("-98765.4321"), vector.getObject(2)); + } + } + + @Test + public void testGetTransferPairWithField() { + final Decimal64Vector fromVector = new Decimal64Vector("decimal", allocator, 10, scale); + final TransferPair transferPair = fromVector.getTransferPair(fromVector.getField(), allocator); + final Decimal64Vector toVector = (Decimal64Vector) transferPair.getTo(); + // Field inside a new vector created by reusing a field should be the same in memory as the + // original field. + assertSame(fromVector.getField(), toVector.getField()); + fromVector.close(); + toVector.close(); + } + + @Test + public void testGetTransferPairWithoutField() { + final Decimal64Vector fromVector = new Decimal64Vector("decimal", allocator, 10, scale); + final TransferPair transferPair = + fromVector.getTransferPair(fromVector.getField().getName(), allocator); + final Decimal64Vector toVector = (Decimal64Vector) transferPair.getTo(); + // A new Field created inside a new vector should reuse the field type (should be the same in + // memory as the original Field's field type). + assertSame(fromVector.getField().getFieldType(), toVector.getField().getFieldType()); + fromVector.close(); + toVector.close(); + } + + @Test + public void testValidateScalarsAtMaxPrecision() { + // A value whose precision exactly equals the vector's precision must validate cleanly. + // Regression: DecimalUtility.checkPrecisionAndScaleNoThrow used a strict '<' comparison, so a + // value that set() accepts was wrongly rejected by validateScalars() at the precision boundary. + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 2, 64), allocator)) { + decimalVector.allocateNew(); + decimalVector.set(0, new BigDecimal("9999999999999999.99")); // precision 18 == vector prec + decimalVector.setValueCount(1); + decimalVector.validateScalars(); // must not throw + } + } + + @Test + public void testSetSafeRejectsTooLongLength() { + // Regression: setSafe/setBigEndianSafe with length > TYPE_WIDTH used to write past the slot + // before any validation, silently corrupting neighbouring values. They must now reject the + // length up front without mutating the vector. + try (Decimal64Vector decimalVector = + TestUtils.newVector( + Decimal64Vector.class, "decimal", new ArrowType.Decimal(18, 0, 64), allocator); + ArrowBuf buf = allocator.buffer(16)) { + decimalVector.allocateNew(2); + decimalVector.set(1, 123L); // neighbour that must remain intact + buf.setLong(0, 0x0102030405060708L); + + final int tooLong = Decimal64Vector.TYPE_WIDTH + 1; + assertThrows(IllegalArgumentException.class, () -> decimalVector.setSafe(0, 0, buf, tooLong)); + assertThrows( + IllegalArgumentException.class, () -> decimalVector.setBigEndianSafe(0, 0, buf, tooLong)); + assertThrows(IllegalArgumentException.class, () -> decimalVector.setSafe(0, 0, buf, 0)); + assertThrows( + IllegalArgumentException.class, () -> decimalVector.setBigEndianSafe(0, 0, buf, 0)); + assertThrows(IllegalArgumentException.class, () -> decimalVector.setSafe(0, 0, buf, -1)); + assertThrows( + IllegalArgumentException.class, () -> decimalVector.setBigEndianSafe(0, 0, buf, -1)); + + decimalVector.setValueCount(2); + assertTrue(decimalVector.isNull(0)); // rejected write left index 0 untouched + assertEquals(new BigDecimal(123), decimalVector.getObject(1)); + } + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/TestGenerateSampleData.java b/vector/src/test/java/org/apache/arrow/vector/TestGenerateSampleData.java new file mode 100644 index 0000000000..2bfa8333c0 --- /dev/null +++ b/vector/src/test/java/org/apache/arrow/vector/TestGenerateSampleData.java @@ -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.arrow.vector; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; + +import java.math.BigDecimal; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class TestGenerateSampleData { + + private BufferAllocator allocator; + + @BeforeEach + public void init() { + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void terminate() { + allocator.close(); + } + + // Sample values derive from the vector's scale (not hardcoded), so they fit any precision/scale, + // including Decimal32's precision-9 limit that the scale-10 DecimalVector values would exceed. + + @Test + public void testDecimal32() { + try (Decimal32Vector vector = + new Decimal32Vector( + "decimal32", FieldType.nullable(new ArrowType.Decimal(9, 2, 32)), allocator)) { + GenerateSampleData.generateTestData(vector, 10); + assertEquals(10, vector.getValueCount()); + assertEquals(new BigDecimal("0.01"), vector.getObject(0)); + assertEquals(new BigDecimal("0.02"), vector.getObject(1)); + assertFalse(vector.isNull(9)); + } + } + + @Test + public void testDecimal64() { + try (Decimal64Vector vector = + new Decimal64Vector( + "decimal64", FieldType.nullable(new ArrowType.Decimal(18, 4, 64)), allocator)) { + GenerateSampleData.generateTestData(vector, 10); + assertEquals(10, vector.getValueCount()); + assertEquals(new BigDecimal("0.0001"), vector.getObject(0)); + assertEquals(new BigDecimal("0.0002"), vector.getObject(1)); + assertFalse(vector.isNull(9)); + } + } + + @Test + public void testDecimal256() { + try (Decimal256Vector vector = + new Decimal256Vector( + "decimal256", FieldType.nullable(new ArrowType.Decimal(40, 6, 256)), allocator)) { + GenerateSampleData.generateTestData(vector, 10); + assertEquals(10, vector.getValueCount()); + assertEquals(new BigDecimal("0.000001"), vector.getObject(0)); + assertEquals(new BigDecimal("0.000002"), vector.getObject(1)); + assertFalse(vector.isNull(9)); + } + } +} diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java index b394a667d1..07501e9fcb 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java @@ -19,9 +19,11 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.File; import java.io.IOException; +import java.math.BigDecimal; import java.net.URL; import java.nio.file.Files; import java.nio.file.Path; @@ -29,6 +31,8 @@ import java.util.Collections; import java.util.List; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.complex.StructVector; @@ -431,6 +435,8 @@ public void testRoundtripEmptyVector() throws Exception { Field.nullable("binary", ArrowType.Binary.INSTANCE), Field.nullable("largebinary", ArrowType.LargeBinary.INSTANCE), Field.nullable("fixedsizebinary", new ArrowType.FixedSizeBinary(2)), + Field.nullable("decimal32", new ArrowType.Decimal(3, 2, 32)), + Field.nullable("decimal64", new ArrowType.Decimal(3, 2, 64)), Field.nullable("decimal128", new ArrowType.Decimal(3, 2, 128)), Field.nullable("decimal128", new ArrowType.Decimal(3, 2, 256)), new Field( @@ -521,4 +527,53 @@ public void testRoundtripEmptyVector() throws Exception { } } } + + @Test + public void testRoundtripDecimal32And64() throws Exception { + final Schema schema = + new Schema( + Arrays.asList( + Field.nullable("decimal32", new ArrowType.Decimal(9, 4, 32)), + Field.nullable("decimal64", new ArrowType.Decimal(18, 4, 64)))); + try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + Decimal32Vector v32 = (Decimal32Vector) root.getVector("decimal32"); + Decimal64Vector v64 = (Decimal64Vector) root.getVector("decimal64"); + root.allocateNew(); + v32.set(0, new BigDecimal("12345.6789")); + v32.setNull(1); + v32.set(2, new BigDecimal("-9876.5432")); + v64.set(0, new BigDecimal("12345678901234.5678")); + v64.setNull(1); + v64.set(2, new BigDecimal("-98765432109876.5432")); + root.setRowCount(3); + + Path outputPath = Files.createTempFile("arrow-decimal-", ".json"); + File outputFile = outputPath.toFile(); + outputFile.deleteOnExit(); + + try (final JsonFileWriter jsonWriter = + new JsonFileWriter(outputFile, JsonFileWriter.config().pretty(true))) { + jsonWriter.start(schema, null); + jsonWriter.write(root); + } + + try (JsonFileReader reader = new JsonFileReader(outputFile, allocator)) { + Schema readSchema = reader.start(); + assertEquals(schema, readSchema); + try (final VectorSchemaRoot data = reader.read()) { + assertNotNull(data); + assertEquals(3, data.getRowCount()); + Decimal32Vector r32 = (Decimal32Vector) data.getVector("decimal32"); + Decimal64Vector r64 = (Decimal64Vector) data.getVector("decimal64"); + assertEquals(new BigDecimal("12345.6789"), r32.getObject(0)); + assertTrue(r32.isNull(1)); + assertEquals(new BigDecimal("-9876.5432"), r32.getObject(2)); + assertEquals(new BigDecimal("12345678901234.5678"), r64.getObject(0)); + assertTrue(r64.isNull(1)); + assertEquals(new BigDecimal("-98765432109876.5432"), r64.getObject(2)); + } + assertNull(reader.read()); + } + } + } } diff --git a/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java b/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java index 849fe6d667..f52f84ca33 100644 --- a/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java +++ b/vector/src/test/java/org/apache/arrow/vector/testing/ValueVectorDataPopulator.java @@ -29,6 +29,8 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FixedSizeBinaryVector; @@ -174,6 +176,54 @@ public static void setVector(Decimal256Vector vector, BigDecimal... values) { vector.setValueCount(length); } + /** Populate values for Decimal32Vector. */ + public static void setVector(Decimal32Vector vector, Long... values) { + final int length = values.length; + vector.allocateNew(length); + for (int i = 0; i < length; i++) { + if (values[i] != null) { + vector.set(i, values[i]); + } + } + vector.setValueCount(length); + } + + /** Populate values for Decimal32Vector. */ + public static void setVector(Decimal32Vector vector, BigDecimal... values) { + final int length = values.length; + vector.allocateNew(length); + for (int i = 0; i < length; i++) { + if (values[i] != null) { + vector.set(i, values[i]); + } + } + vector.setValueCount(length); + } + + /** Populate values for Decimal64Vector. */ + public static void setVector(Decimal64Vector vector, Long... values) { + final int length = values.length; + vector.allocateNew(length); + for (int i = 0; i < length; i++) { + if (values[i] != null) { + vector.set(i, values[i]); + } + } + vector.setValueCount(length); + } + + /** Populate values for Decimal64Vector. */ + public static void setVector(Decimal64Vector vector, BigDecimal... values) { + final int length = values.length; + vector.allocateNew(length); + for (int i = 0; i < length; i++) { + if (values[i] != null) { + vector.set(i, values[i]); + } + } + vector.setValueCount(length); + } + /** * Populate values for DurationVector. * diff --git a/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorFull.java b/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorFull.java index 6993fde8fa..5f852389aa 100644 --- a/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorFull.java +++ b/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorFull.java @@ -318,7 +318,7 @@ public void testValidateDecimal() { public void testValidateDecimal256() { try (final Decimal256Vector vector = new Decimal256Vector( - Field.nullable("v", new ArrowType.Decimal(2, 0, DecimalVector.TYPE_WIDTH * 8)), + Field.nullable("v", new ArrowType.Decimal(2, 0, Decimal256Vector.TYPE_WIDTH * 8)), allocator)) { vector.validateFull(); setVector(vector, 1L); diff --git a/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorTypeVisitor.java b/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorTypeVisitor.java index 5454008364..01a11ae428 100644 --- a/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorTypeVisitor.java +++ b/vector/src/test/java/org/apache/arrow/vector/validate/TestValidateVectorTypeVisitor.java @@ -26,6 +26,8 @@ import org.apache.arrow.vector.DateDayVector; import org.apache.arrow.vector.DateMilliVector; import org.apache.arrow.vector.Decimal256Vector; +import org.apache.arrow.vector.Decimal32Vector; +import org.apache.arrow.vector.Decimal64Vector; import org.apache.arrow.vector.DecimalVector; import org.apache.arrow.vector.DurationVector; import org.apache.arrow.vector.FixedSizeBinaryVector; @@ -283,6 +285,14 @@ public void testDecimalVector() { "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(38, 10, 128)), allocator)); + testPositiveCase( + () -> + new Decimal32Vector( + "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(9, 2, 32)), allocator)); + testPositiveCase( + () -> + new Decimal64Vector( + "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(18, 2, 64)), allocator)); testPositiveCase( () -> new Decimal256Vector( @@ -321,6 +331,14 @@ public void testDecimalVector() { () -> new Decimal256Vector( "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(30, 10, 64)), allocator)); + testNegativeCase( + () -> + new DecimalVector( + "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(9, 2, 32)), allocator)); + testNegativeCase( + () -> + new Decimal64Vector( + "dec", FieldType.nullable(ArrowType.Decimal.createDecimal(18, 2, 128)), allocator)); testNegativeCase( () -> new Decimal256Vector(