GH-3263: Add DictionaryPage.decode to allow dictionary reuse in the ColumnReaderBase ctor (#3264)

Author: pyckle

parquet-column/src/main/java/org/apache/parquet/column/impl/ColumnReaderBase.java

@@ -453,17 +453,9 @@ void writeValue() {
     this.writerVersion = writerVersion;
     this.maxDefinitionLevel = path.getMaxDefinitionLevel();
     DictionaryPage dictionaryPage = pageReader.readDictionaryPage();
-    if (dictionaryPage != null) {
-      try {
-        this.dictionary = dictionaryPage.getEncoding().initDictionary(path, dictionaryPage);
-        if (converter.hasDictionarySupport()) {
-          converter.setDictionary(dictionary);
-        }
-      } catch (IOException e) {
-        throw new ParquetDecodingException("could not decode the dictionary for " + path, e);
-      }
-    } else {
-      this.dictionary = null;
+    this.dictionary = dictionaryPage == null ? null : dictionaryPage.decode(path);
+    if (dictionary != null && converter.hasDictionarySupport()) {
+      converter.setDictionary(dictionary);
     }
     this.totalValueCount = pageReader.getTotalValueCount();
     if (totalValueCount <= 0) {

parquet-column/src/main/java/org/apache/parquet/column/page/DictionaryPage.java

@@ -21,7 +21,10 @@
 import java.io.IOException;
 import java.util.Objects;
 import org.apache.parquet.bytes.BytesInput;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.Dictionary;
 import org.apache.parquet.column.Encoding;
+import org.apache.parquet.io.ParquetDecodingException;
 
 /**
  * Data for a dictionary page
@@ -74,6 +77,17 @@ public DictionaryPage copy() throws IOException {
     return new DictionaryPage(BytesInput.copy(bytes), getUncompressedSize(), dictionarySize, encoding);
   }
 
+  /**
+   * @return the decoded dictionary
+   */
+  public Dictionary decode(ColumnDescriptor path) {
+    try {
+      return getEncoding().initDictionary(path, this);
+    } catch (IOException e) {
+      throw new ParquetDecodingException("could not decode the dictionary for " + path, e);
+    }
+  }
+
   @Override
   public String toString() {
     return "Page [bytes.size=" + bytes.size() + ", entryCount=" + dictionarySize + ", uncompressedSize="

parquet-column/src/test/java/org/apache/parquet/column/impl/TestColumnReaderImpl.java

@@ -24,17 +24,21 @@
 import java.util.List;
 import org.apache.parquet.Version;
 import org.apache.parquet.VersionParser;
+import org.apache.parquet.bytes.BytesInput;
 import org.apache.parquet.column.ColumnDescriptor;
 import org.apache.parquet.column.ColumnReader;
+import org.apache.parquet.column.Dictionary;
 import org.apache.parquet.column.ParquetProperties;
 import org.apache.parquet.column.page.DataPage;
 import org.apache.parquet.column.page.DataPageV2;
+import org.apache.parquet.column.page.DictionaryPage;
 import org.apache.parquet.column.page.mem.MemPageReader;
 import org.apache.parquet.column.page.mem.MemPageWriter;
 import org.apache.parquet.io.api.Binary;
 import org.apache.parquet.io.api.PrimitiveConverter;
 import org.apache.parquet.schema.MessageType;
 import org.apache.parquet.schema.MessageTypeParser;
+import org.junit.Assert;
 import org.junit.Test;
 
 public class TestColumnReaderImpl {
@@ -53,8 +57,28 @@ public void addBinary(Binary value) {
 
   @Test
   public void test() throws Exception {
+    ColumnDescriptor col = requiredBinaryColumn();
+    MemPageWriter pageWriter = writeBinaryDictColumn(col);
+    List<DataPage> pages = pageWriter.getPages();
+    int valueCount = 0;
+    int rowCount = 0;
+    for (DataPage dataPage : pages) {
+      valueCount += dataPage.getValueCount();
+      rowCount += ((DataPageV2) dataPage).getRowCount();
+    }
+    assertEquals(rows, rowCount);
+    assertEquals(rows, valueCount);
+    MemPageReader pageReader = toReader(pageWriter);
+    validateExpectedValuesAndCount(col, pageReader);
+  }
+
+  private static ColumnDescriptor requiredBinaryColumn() {
     MessageType schema = MessageTypeParser.parseMessageType("message test { required binary foo; }");
     ColumnDescriptor col = schema.getColumns().get(0);
+    return col;
+  }
+
+  private MemPageWriter writeBinaryDictColumn(ColumnDescriptor col) {
     MemPageWriter pageWriter = new MemPageWriter();
     ColumnWriterV2 columnWriterV2 = new ColumnWriterV2(
         col,
@@ -72,16 +96,15 @@ public void test() throws Exception {
     }
     columnWriterV2.writePage();
     columnWriterV2.finalizeColumnChunk();
-    List<DataPage> pages = pageWriter.getPages();
-    int valueCount = 0;
-    int rowCount = 0;
-    for (DataPage dataPage : pages) {
-      valueCount += dataPage.getValueCount();
-      rowCount += ((DataPageV2) dataPage).getRowCount();
-    }
-    assertEquals(rows, rowCount);
-    assertEquals(rows, valueCount);
-    MemPageReader pageReader = new MemPageReader(rows, pages.iterator(), pageWriter.getDictionaryPage());
+    return pageWriter;
+  }
+
+  private MemPageReader toReader(MemPageWriter pageWriter) {
+    return new MemPageReader(rows, pageWriter.getPages().iterator(), pageWriter.getDictionaryPage());
+  }
+
+  private void validateExpectedValuesAndCount(ColumnDescriptor col, MemPageReader pageReader)
+      throws VersionParser.VersionParseException {
     ValidatingConverter converter = new ValidatingConverter();
     ColumnReader columnReader =
         new ColumnReaderImpl(col, pageReader, converter, VersionParser.parse(Version.FULL_VERSION));
@@ -124,7 +147,7 @@ public void testOptional() throws Exception {
     }
     assertEquals(rows, rowCount);
     assertEquals(rows, valueCount);
-    MemPageReader pageReader = new MemPageReader(rows, pages.iterator(), pageWriter.getDictionaryPage());
+    MemPageReader pageReader = toReader(pageWriter);
     ValidatingConverter converter = new ValidatingConverter();
     ColumnReader columnReader =
         new ColumnReaderImpl(col, pageReader, converter, VersionParser.parse(Version.FULL_VERSION));
@@ -135,4 +158,29 @@ public void testOptional() throws Exception {
     }
     assertEquals(0, converter.count);
   }
+
+  @Test
+  public void testDeduplicatedDecodedDictionary() throws Exception {
+    ColumnDescriptor col = requiredBinaryColumn();
+    MemPageWriter pageWriter = writeBinaryDictColumn(col);
+
+    DictionaryPage dictionaryPage = pageWriter.getDictionaryPage();
+    Assert.assertNotNull("Expected a dictionary", dictionaryPage);
+
+    Dictionary dict = dictionaryPage.decode(col);
+
+    // construct a page reader from a dictionary page that lacks bytes but stores the decoded data.
+    MemPageReader pageReader = new MemPageReader(
+        rows,
+        pageWriter.getPages().iterator(),
+        new DictionaryPage(
+            BytesInput.empty(), dictionaryPage.getDictionarySize(), dictionaryPage.getEncoding()) {
+          @Override
+          public Dictionary decode(ColumnDescriptor path) {
+            return dict;
+          }
+        });
+
+    validateExpectedValuesAndCount(col, pageReader);
+  }
 }