fix: Reuse previous CometDictionary Java arrays (apache#489)

(cherry picked from commit ee15460)
himadripal · May 31, 2024 · ba16f40 · ba16f40
1 parent 3f20c24
commit ba16f40
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/common/src/main/java/org/apache/comet/parquet/ColumnReader.java b/common/src/main/java/org/apache/comet/parquet/ColumnReader.java
@@ -234,7 +234,11 @@ public CometDecodedVector loadVector() {
       Dictionary arrowDictionary = importer.getProvider().lookup(dictionaryEncoding.getId());
       CometPlainVector dictionaryVector =
           new CometPlainVector(arrowDictionary.getVector(), useDecimal128, isUuid);
-      dictionary = new CometDictionary(dictionaryVector);
+      if (dictionary != null) {
+        dictionary.setDictionaryVector(dictionaryVector);
+      } else {
+        dictionary = new CometDictionary(dictionaryVector);
+      }
 
       currentVector =
           new CometDictionaryVector(

diff --git a/common/src/main/java/org/apache/comet/vector/CometDictionary.java b/common/src/main/java/org/apache/comet/vector/CometDictionary.java
@@ -26,7 +26,7 @@
 public class CometDictionary implements AutoCloseable {
   private static final int DECIMAL_BYTE_WIDTH = 16;
 
-  private final CometPlainVector values;
+  private CometPlainVector values;
   private final int numValues;
 
   /** Decoded dictionary values. Only one of the following is set. */
@@ -47,6 +47,13 @@ public CometDictionary(CometPlainVector values) {
     initialize();
   }
 
+  public void setDictionaryVector(CometPlainVector values) {
+    this.values = values;
+    if (values.numValues() != numValues) {
+      throw new IllegalArgumentException("Mismatched dictionary size");
+    }
+  }
+
   public ValueVector getValueVector() {
     return values.getValueVector();
   }