GH-1228: Arrow Format 1.5: Decimal32/Decimal64 support

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/ArrowToAvroUtils.java

@@ -50,6 +50,8 @@
 import org.apache.arrow.adapter.avro.producers.logical.AvroDateDayProducer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroDateMilliProducer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroDecimal256Producer;
+import org.apache.arrow.adapter.avro.producers.logical.AvroDecimal32Producer;
+import org.apache.arrow.adapter.avro.producers.logical.AvroDecimal64Producer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroDecimalProducer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroTimeMicroProducer;
 import org.apache.arrow.adapter.avro.producers.logical.AvroTimeMilliProducer;
@@ -70,6 +72,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.FieldVector;
 import org.apache.arrow.vector.FixedSizeBinaryVector;
@@ -603,6 +607,10 @@ private static BaseAvroProducer<?> createProducer(
 
         // Logical types
 
+      case DECIMAL32:
+        return new AvroDecimal32Producer((Decimal32Vector) vector);
+      case DECIMAL64:
+        return new AvroDecimal64Producer((Decimal64Vector) vector);
       case DECIMAL:
         return new AvroDecimalProducer((DecimalVector) vector);
       case DECIMAL256:

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimal32Producer.java

@@ -0,0 +1,46 @@
+/*
+ * 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.avro.producers.logical;
+
+import java.io.IOException;
+import java.math.BigDecimal;
+import org.apache.arrow.adapter.avro.producers.BaseAvroProducer;
+import org.apache.arrow.vector.Decimal32Vector;
+import org.apache.avro.io.Encoder;
+
+/**
+ * Producer that produces decimal values from a {@link Decimal32Vector}, writes data to an Avro
+ * encoder.
+ */
+public class AvroDecimal32Producer extends BaseAvroProducer<Decimal32Vector> {
+
+  // Logic is the same as for DecimalVector (128 bit)
+
+  byte[] encodedBytes = new byte[Decimal32Vector.TYPE_WIDTH];
+
+  /** Instantiate an AvroDecimal32Producer. */
+  public AvroDecimal32Producer(Decimal32Vector vector) {
+    super(vector);
+  }
+
+  @Override
+  public void produce(Encoder encoder) throws IOException {
+    BigDecimal value = vector.getObject(currentIndex++);
+    AvroDecimalProducer.encodeDecimal(value, encodedBytes);
+    encoder.writeFixed(encodedBytes);
+  }
+}

adapter/avro/src/main/java/org/apache/arrow/adapter/avro/producers/logical/AvroDecimal64Producer.java

@@ -0,0 +1,46 @@
+/*
+ * 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.avro.producers.logical;
+
+import java.io.IOException;
+import java.math.BigDecimal;
+import org.apache.arrow.adapter.avro.producers.BaseAvroProducer;
+import org.apache.arrow.vector.Decimal64Vector;
+import org.apache.avro.io.Encoder;
+
+/**
+ * Producer that produces decimal values from a {@link Decimal64Vector}, writes data to an Avro
+ * encoder.
+ */
+public class AvroDecimal64Producer extends BaseAvroProducer<Decimal64Vector> {
+
+  // Logic is the same as for DecimalVector (128 bit)
+
+  byte[] encodedBytes = new byte[Decimal64Vector.TYPE_WIDTH];
+
+  /** Instantiate an AvroDecimal64Producer. */
+  public AvroDecimal64Producer(Decimal64Vector vector) {
+    super(vector);
+  }
+
+  @Override
+  public void produce(Encoder encoder) throws IOException {
+    BigDecimal value = vector.getObject(currentIndex++);
+    AvroDecimalProducer.encodeDecimal(value, encodedBytes);
+    encoder.writeFixed(encodedBytes);
+  }
+}

adapter/avro/src/test/java/org/apache/arrow/adapter/avro/ArrowToAvroDataTest.java

@@ -41,6 +41,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.FieldVector;
 import org.apache.arrow.vector.FixedSizeBinaryVector;
@@ -1039,6 +1041,101 @@ record = datumReader.read(record, decoder);
     }
   }
 
+  @Test
+  public void testWriteNarrowDecimals() throws Exception {
+
+    // Field definitions
+    FieldType decimal32Field = new FieldType(false, new ArrowType.Decimal(9, 3, 32), null);
+    FieldType decimal64Field = new FieldType(false, new ArrowType.Decimal(18, 6, 64), null);
+    FieldType nullableDecimal32Field = new FieldType(true, new ArrowType.Decimal(9, 3, 32), null);
+    FieldType nullableDecimal64Field = new FieldType(true, new ArrowType.Decimal(18, 6, 64), null);
+
+    // Create empty vectors
+    BufferAllocator allocator = new RootAllocator();
+    Decimal32Vector decimal32Vector =
+        new Decimal32Vector(new Field("decimal32", decimal32Field, null), allocator);
+    Decimal64Vector decimal64Vector =
+        new Decimal64Vector(new Field("decimal64", decimal64Field, null), allocator);
+    Decimal32Vector nullableDecimal32Vector =
+        new Decimal32Vector(
+            new Field("nullableDecimal32", nullableDecimal32Field, null), allocator);
+    Decimal64Vector nullableDecimal64Vector =
+        new Decimal64Vector(
+            new Field("nullableDecimal64", nullableDecimal64Field, null), allocator);
+
+    // Set up VSR
+    List<FieldVector> vectors =
+        Arrays.asList(
+            decimal32Vector, decimal64Vector, nullableDecimal32Vector, nullableDecimal64Vector);
+    int rowCount = 3;
+
+    try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
+
+      root.setRowCount(rowCount);
+      root.allocate新建();
+
+      // Set test data
+      decimal32Vector.setSafe(0, new BigDecimal("123456.789"));
+      decimal32Vector.setSafe(1, new BigDecimal("-999999.999"));
+      decimal32Vector.setSafe(2, new BigDecimal("0.001"));
+
+      decimal64Vector.setSafe(0, new BigDecimal("123456789012.345678"));
+      decimal64Vector.setSafe(1, new BigDecimal("-999999999999.999999"));
+      decimal64Vector.setSafe(2, new BigDecimal("0.000001"));
+
+      nullableDecimal32Vector.setSafe(0, new BigDecimal("-1.500"));
+      nullableDecimal32Vector.setNull(1);
+      nullableDecimal32Vector.setSafe(2, new BigDecimal("42.000"));
+
+      nullableDecimal64Vector.setNull(0);
+      nullableDecimal64Vector.setSafe(1, new BigDecimal("-1.500000"));
+      nullableDecimal64Vector.setSafe(2, new BigDecimal("42.000000"));
+
+      File dataFile = new File(TMP, "testWriteNarrowDecimals.avro");
+
+      // Write an AVRO block using the producer classes
+      try (FileOutputStream fos = new FileOutputStream(dataFile)) {
+        BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null);
+        CompositeAvroProducer producer = ArrowToAvroUtils.createCompositeProducer(vectors);
+        for (int row = 0; row < rowCount; row++) {
+          producer.produce(encoder);
+        }
+        encoder.flush();
+      }
+
+      // Set up reading the AVRO block as a GenericRecord
+      Schema schema = ArrowToAvroUtils.createAvroSchema(root.getSchema().getFields());
+      GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>(schema);
+
+      try (InputStream inputStream = new FileInputStream(dataFile)) {
+
+        BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(inputStream, null);
+        GenericRecord record = null;
+
+        // Read and check values
+        for (int row = 0; row < rowCount; row++) {
+          record = datumReader.read(record, decoder);
+          assertEquals(decimal32Vector.getObject(row), decodeFixedDecimal(record, "decimal32"));
+          assertEquals(decimal64Vector.getObject(row), decodeFixedDecimal(record, "decimal64"));
+          if (nullableDecimal32Vector.isNull(row)) {
+            assertNull(record.get("nullableDecimal32"));
+          } else {
+            assertEquals(
+                nullableDecimal32Vector.getObject(row),
+                decodeFixedDecimal(record, "nullableDecimal32"));
+          }
+          if (nullableDecimal64Vector.isNull(row)) {
+            assertNull(record.get("nullableDecimal64"));
+          } else {
+            assertEquals(
+                nullableDecimal64Vector.getObject(row),
+                decodeFixedDecimal(record, "nullableDecimal64"));
+          }
+        }
+      }
+    }
+  }
+
   @Test
   public void testWriteNullableDecimals() throws Exception {
 

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);
         }

adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/ColumnBinderArrowTypeVisitor.java

@@ -25,6 +25,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.FieldVector;
 import org.apache.arrow.vector.FixedSizeBinaryVector;
@@ -211,7 +213,17 @@ public ColumnBinder visit(ArrowType.Bool type) {
 
   @Override
   public ColumnBinder visit(ArrowType.Decimal type) {
-    if (type.getBitWidth() == 128) {
+    if (type.getBitWidth() == 32) {
+      Decimal32Vector decimalVector = (Decimal32Vector) vector;
+      return jdbcType == null
+          ? new Decimal32Binder(decimalVector)
+          : new Decimal32Binder(decimalVector, jdbcType);
+    } else if (type.getBitWidth() == 64) {
+      Decimal64Vector decimalVector = (Decimal64Vector) vector;
+      return jdbcType == null
+          ? new Decimal64Binder(decimalVector)
+          : new Decimal64Binder(decimalVector, jdbcType);
+    } else if (type.getBitWidth() == 128) {
       DecimalVector decimalVector = (DecimalVector) vector;
       return jdbcType == null
           ? new Decimal128Binder(decimalVector)

adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/Decimal32Binder.java

@@ -0,0 +1,44 @@
+/*
+ * 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.binder;
+
+import java.math.BigDecimal;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.sql.Types;
+import org.apache.arrow.vector.Decimal32Vector;
+import org.apache.arrow.vector.util.DecimalUtility;
+
+/** A binder for 32-bit decimals. */
+public class Decimal32Binder extends BaseColumnBinder<Decimal32Vector> {
+  public Decimal32Binder(Decimal32Vector vector) {
+    this(vector, Types.DECIMAL);
+  }
+
+  public Decimal32Binder(Decimal32Vector vector, int jdbcType) {
+    super(vector, jdbcType);
+  }
+
+  @Override
+  public void bind(PreparedStatement statement, int parameterIndex, int rowIndex)
+      throws SQLException {
+    final BigDecimal value =
+        DecimalUtility.getBigDecimalFromArrowBuf(
+            vector.getDataBuffer(), rowIndex, vector.getScale(), Decimal32Vector.TYPE_WIDTH);
+    statement.setBigDecimal(parameterIndex, value);
+  }
+}

adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/binder/Decimal64Binder.java

@@ -0,0 +1,44 @@
+/*
+ * 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.binder;
+
+import java.math.BigDecimal;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.sql.Types;
+import org.apache.arrow.vector.Decimal64Vector;
+import org.apache.arrow.vector.util.DecimalUtility;
+
+/** A binder for 64-bit decimals. */
+public class Decimal64Binder extends BaseColumnBinder<Decimal64Vector> {
+  public Decimal64Binder(Decimal64Vector vector) {
+    this(vector, Types.DECIMAL);
+  }
+
+  public Decimal64Binder(Decimal64Vector vector, int jdbcType) {
+    super(vector, jdbcType);
+  }
+
+  @Override
+  public void bind(PreparedStatement statement, int parameterIndex, int rowIndex)
+      throws SQLException {
+    final BigDecimal value =
+        DecimalUtility.getBigDecimalFromArrowBuf(
+            vector.getDataBuffer(), rowIndex, vector.getScale(), Decimal64Vector.TYPE_WIDTH);
+    statement.setBigDecimal(parameterIndex, value);
+  }
+}