From de55c0515f7e074e628a7792d8faf7311b2b19c2 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Wed, 29 May 2024 18:03:09 -0700
Subject: [PATCH] For review

---
 .../apache/comet/vector/CometDictionary.java  | 25 ++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/common/src/main/java/org/apache/comet/vector/CometDictionary.java b/common/src/main/java/org/apache/comet/vector/CometDictionary.java
index cf3beb7ee..24a6d6d8c 100644
--- a/common/src/main/java/org/apache/comet/vector/CometDictionary.java
+++ b/common/src/main/java/org/apache/comet/vector/CometDictionary.java
@@ -29,9 +29,13 @@ public class CometDictionary implements AutoCloseable {
   private CometPlainVector values;
   private final int numValues;
 
+  /** Decoded dictionary values. We only need to copy values for decimal type. */
+  private ByteArrayWrapper[] binaries;
+
   public CometDictionary(CometPlainVector values) {
     this.values = values;
     this.numValues = values.numValues();
+    initialize();
   }
 
   public void setDictionaryVector(CometPlainVector values) {
@@ -79,9 +83,7 @@ public byte[] decodeToBinary(int index) {
       case FIXEDSIZEBINARY:
         return values.getBinary(index);
       case DECIMAL:
-        byte[] bytes = new byte[DECIMAL_BYTE_WIDTH];
-        bytes = values.copyBinaryDecimal(index, bytes);
-        return bytes;
+        return binaries[index].bytes;
       default:
         throw new IllegalArgumentException(
             "Invalid Arrow minor type: " + values.getValueVector().getMinorType());
@@ -97,6 +99,23 @@ public void close() {
     values.close();
   }
 
+  private void initialize() {
+    switch (values.getValueVector().getMinorType()) {
+      case DECIMAL:
+        // We only need to copy values for decimal type as random access
+        // to the dictionary is not efficient for decimal (it needs to copy
+        // the value to a new byte array everytime).
+        binaries = new ByteArrayWrapper[numValues];
+        for (int i = 0; i < numValues; i++) {
+          // Need copying here since we re-use byte array for decimal
+          byte[] bytes = new byte[DECIMAL_BYTE_WIDTH];
+          bytes = values.copyBinaryDecimal(i, bytes);
+          binaries[i] = new ByteArrayWrapper(bytes);
+        }
+        break;
+    }
+  }
+
   private static class ByteArrayWrapper {
     private final byte[] bytes;