fix: ColumnReader.loadVector should initiate CometDictionary afte…

…r re-import arrays (#473)
apache · May 25, 2024 · 5bfe46d · 5bfe46d
1 parent 79431f8
commit 5bfe46d
Showing 1 changed file with 7 additions and 7 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
@@ -230,15 +230,15 @@ public CometDecodedVector loadVector() {
         // return plain vector.
         currentVector = cometVector;
         return currentVector;
-      } else if (dictionary == null) {
-        // There is dictionary from native side but the Java side dictionary hasn't been
-        // initialized yet.
-        Dictionary arrowDictionary = dictionaryProvider.lookup(dictionaryEncoding.getId());
-        CometPlainVector dictionaryVector =
-            new CometPlainVector(arrowDictionary.getVector(), useDecimal128, isUuid);
-        dictionary = new CometDictionary(dictionaryVector);
       }
 
+      // We should already re-initiate `CometDictionary` here because `Data.importVector` API will
+      // release the previous dictionary vector and create a new one.
+      Dictionary arrowDictionary = dictionaryProvider.lookup(dictionaryEncoding.getId());
+      CometPlainVector dictionaryVector =
+          new CometPlainVector(arrowDictionary.getVector(), useDecimal128, isUuid);
+      dictionary = new CometDictionary(dictionaryVector);
+
       currentVector =
           new CometDictionaryVector(
               cometVector, dictionary, dictionaryProvider, useDecimal128, false, isUuid);