support multiple byteWidth in writeLongToArrowBuf()

apache · Nov 2, 2020 · 89ee07a · 89ee07a
1 parent a8c17fd
commit 89ee07a
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 16 deletions.
diff --git a/java/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java b/java/vector/src/main/java/org/apache/arrow/vector/DecimalVector.java
@@ -364,7 +364,7 @@ public void set(int index, BigDecimal value) {
    */
   public void set(int index, long value) {
     BitVectorHelper.setBit(validityBuffer, index);
-    DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index);
+    DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index, TYPE_WIDTH);
   }
 
   /**

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java b/java/vector/src/main/java/org/apache/arrow/vector/util/DecimalUtility.java
@@ -132,17 +132,25 @@ public static void writeBigDecimalToArrowBuf(BigDecimal value, ArrowBuf bytebuf,
 
   /**
    * Write the given long to the ArrowBuf at the given value index.
-   * This routine extends the original sign bit to a new upper 64-bit in 128-bit.
+   * This routine extends the original sign bit to a new upper area in 128-bit or 256-bit.
    */
-  public static void writeLongToArrowBuf(long value, ArrowBuf bytebuf, int index) {
+  public static void writeLongToArrowBuf(long value, ArrowBuf bytebuf, int index, int byteWidth) {
+    if (byteWidth != 16 && byteWidth != 32) {
+      throw new UnsupportedOperationException("DeciimalUtility.writeLongToArrowBuf() currently supports " +
+          "128-bit or 256-bit width data");
+    }
     final long addressOfValue = bytebuf.memoryAddress() + (long) index * DECIMAL_BYTE_LENGTH;
     final long padValue = Long.signum(value) == -1 ? -1L : 0L;
     if (LITTLE_ENDIAN) {
       PlatformDependent.putLong(addressOfValue, value);
-      PlatformDependent.putLong(addressOfValue + Long.BYTES, padValue);
+      for (int i = 1; i <= (byteWidth - 8) / 8; i++) {
+        PlatformDependent.putLong(addressOfValue + Long.BYTES * i, padValue);
+      }
     } else {
-      PlatformDependent.putLong(addressOfValue, padValue);
-      PlatformDependent.putLong(addressOfValue + Long.BYTES, value);
+      for (int i = 0; i < (byteWidth - 8) / 8; i++) {
+        PlatformDependent.putLong(addressOfValue + Long.BYTES * i, padValue);
+      }
+      PlatformDependent.putLong(addressOfValue + Long.BYTES * (byteWidth - 8) / 8, value);
     }
   }
 

diff --git a/java/vector/src/test/java/org/apache/arrow/vector/util/DecimalUtilityTest.java b/java/vector/src/test/java/org/apache/arrow/vector/util/DecimalUtilityTest.java
@@ -34,16 +34,19 @@ public class DecimalUtilityTest {
 
   @Test
   public void testSetLongInDecimalArrowBuf() {
-    try (BufferAllocator allocator = new RootAllocator(128);
-         ArrowBuf buf = allocator.buffer(16);
-    ) {
-      int [] intValues = new int [] {Integer.MAX_VALUE, Integer.MIN_VALUE, 0};
-      for (int val : intValues) {
-        buf.clear();
-        DecimalUtility.writeLongToArrowBuf((long) val, buf, 0);
-        BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0);
-        BigDecimal expected = BigDecimal.valueOf(val);
-        Assert.assertEquals(expected, actual);
+    int[] byteLengths = new int[]{16, 32};
+    for (int x = 0; x < 2; x++) {
+      try (BufferAllocator allocator = new RootAllocator(128);
+           ArrowBuf buf = allocator.buffer(byteLengths[x]);
+      ) {
+        int [] intValues = new int [] {Integer.MAX_VALUE, Integer.MIN_VALUE, 0};
+        for (int val : intValues) {
+          buf.clear();
+          DecimalUtility.writeLongToArrowBuf((long) val, buf, 0, byteLengths[x]);
+          BigDecimal actual = DecimalUtility.getBigDecimalFromArrowBuf(buf, 0, 0, byteLengths[x]);
+          BigDecimal expected = BigDecimal.valueOf(val);
+          Assert.assertEquals(expected, actual);
+        }
       }
     }
   }