PARQUET-374: Add api to read dictionary from each column chunk for predicate pushdown

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

@@ -0,0 +1,34 @@
+/* 
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * 
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.parquet.column.page;
+
+import org.apache.parquet.column.ColumnDescriptor;
+
+/**
+ * contains all the dictionary readers for all the columns of the corresponding row group
+ */
+public interface DictionaryPageReadStore {
+
+  /**
+   *
+   * @param descriptor the descriptor of the column
+   * @return the dictionary page reader for that column
+   */
+  DictionaryPageReader getDictionaryPageReader(ColumnDescriptor descriptor);
+}

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

@@ -0,0 +1,35 @@
+/* 
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * 
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.parquet.column.page;
+
+/**
+ * Reader for a dictionary page from a given column chunk
+ */
+public interface DictionaryPageReader {
+
+ /**
+  * @return the dictionary page in that chunk or null if none
+  */
+  DictionaryPage readDictionaryPage();
+
+  /**
+   * @return the dictionary size
+   */
+  int getDictionarySize();
+}

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

@@ -0,0 +1,46 @@
+/* 
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * 
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.parquet.column.page;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.Dictionary;
+
+import java.io.IOException;
+
+/**
+ * Read a dictionary from dictionary page
+ */
+public class DictionaryReader {
+  private final ColumnDescriptor column;
+  private final DictionaryPageReader pageReader;
+
+  public DictionaryReader(ColumnDescriptor column, DictionaryPageReadStore dictionaryPageReadStore) {
+    this.column = column;
+    this.pageReader = dictionaryPageReadStore.getDictionaryPageReader(column);
+  }
+
+  public Dictionary readDictionary() throws IOException {
+    DictionaryPage dictionaryPage = pageReader.readDictionaryPage();
+    return dictionaryPage.getEncoding().initDictionary(column, dictionaryPage);
+  }
+
+  public int getDictionarySize() {
+    return pageReader.getDictionarySize();
+  }
+}

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkDictionaryPageReadStore.java

@@ -0,0 +1,73 @@
+/* 
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * 
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.parquet.hadoop;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.parquet.Ints;
+import org.apache.parquet.Log;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.DictionaryPageReadStore;
+import org.apache.parquet.column.page.DictionaryPageReader;
+import org.apache.parquet.hadoop.CodecFactory.BytesDecompressor;
+import org.apache.parquet.hadoop.ColumnChunkPageReadStore.ColumnChunkPageReader;
+
+class ColumnChunkDictionaryPageReadStore implements DictionaryPageReadStore {
+  /**
+   * DictionaryPageReader for a single column chunk. A column chunk contains
+   * several pages, first of which could be a dictionary page.
+   *
+   * This implementation is provided with a ColumnChunkPageReader delegate
+   */
+  static final class ColumnChunkDictionaryPageReader implements DictionaryPageReader {
+
+    private final ColumnChunkPageReader columnChunkPageReader;
+
+    ColumnChunkDictionaryPageReader(ColumnChunkPageReader columnChunkPageReader) {
+      this.columnChunkPageReader = columnChunkPageReader;
+    }
+
+    @Override
+    public int getDictionarySize() {
+      return columnChunkPageReader.getCompressedDictionaryPage().getDictionarySize();
+    }
+
+    @Override
+    public DictionaryPage readDictionaryPage() {
+      return columnChunkPageReader.readDictionaryPage();
+    }
+  }
+
+  private final Map<ColumnDescriptor, ColumnChunkDictionaryPageReader> readers;
+
+  ColumnChunkDictionaryPageReadStore(Map<ColumnDescriptor, ColumnChunkDictionaryPageReader> readers) {
+    this.readers = readers;
+  }
+
+  @Override
+  public DictionaryPageReader getDictionaryPageReader(ColumnDescriptor column) {
+    if (!readers.containsKey(column)) {
+      throw new IllegalArgumentException(column + " is not in the store: " + readers.keySet());
+    }
+    return readers.get(column);
+  }
+}

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageReadStore.java

@@ -139,6 +139,10 @@ public DictionaryPage readDictionaryPage() {
         throw new RuntimeException(e); // TODO: cleanup
       }
     }
+
+    public DictionaryPage getCompressedDictionaryPage() {
+      return compressedDictionaryPage;
+    }
   }
 
   private final Map<ColumnDescriptor, ColumnChunkPageReader> readers = new HashMap<ColumnDescriptor, ColumnChunkPageReader>();

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java

@@ -56,20 +56,26 @@
 import org.apache.parquet.Log;
 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.column.page.DataPage;
 import org.apache.parquet.column.page.DataPageV1;
 import org.apache.parquet.column.page.DataPageV2;
+import org.apache.parquet.column.page.DictionaryPageReadStore;
+import org.apache.parquet.column.page.DictionaryPageReader;
 import org.apache.parquet.column.page.DictionaryPage;
 import org.apache.parquet.column.page.PageReadStore;
 import org.apache.parquet.hadoop.metadata.ColumnPath;
 import org.apache.parquet.format.DataPageHeader;
 import org.apache.parquet.format.DataPageHeaderV2;
 import org.apache.parquet.format.DictionaryPageHeader;
 import org.apache.parquet.format.PageHeader;
+import org.apache.parquet.format.PageType;
 import org.apache.parquet.format.Util;
 import org.apache.parquet.format.converter.ParquetMetadataConverter;
 import org.apache.parquet.format.converter.ParquetMetadataConverter.MetadataFilter;
 import org.apache.parquet.hadoop.CodecFactory.BytesDecompressor;
+import org.apache.parquet.hadoop.ColumnChunkDictionaryPageReadStore.ColumnChunkDictionaryPageReader;
 import org.apache.parquet.hadoop.ColumnChunkPageReadStore.ColumnChunkPageReader;
 import org.apache.parquet.hadoop.metadata.BlockMetaData;
 import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
@@ -78,6 +84,7 @@
 import org.apache.parquet.hadoop.util.HiddenFileFilter;
 import org.apache.parquet.hadoop.util.counters.BenchmarkCounter;
 import org.apache.parquet.io.ParquetDecodingException;
+import org.apache.parquet.schema.MessageType;
 
 /**
  * Internal implementation of the Parquet file reader as a block container
@@ -517,7 +524,74 @@ public PageReadStore readNextRowGroup() throws IOException {
     return columnChunkPageReadStore;
   }
 
+  /**
+   * Reads the requested predicate columns dictionaries from the row group at the current file position.
+   * Not reading all the requested columns, only read predicate columns.
+   * eg. select columnA, columnB, columnC from table where columnC > 10
+   *     only reads columnC's dictionary
+   * @param requestedSchema the predicate columns to read
+   * @throws IOException if an error occurs while reading
+   * @return the DictionaryPageReadStore which can provide DictionaryPageReaders for each column.
+   */
+  public DictionaryPageReadStore getCurrentRowGroupDictionaries(MessageType requestedSchema) throws IOException {
+    if (currentBlock == blocks.size()) {
+      return null;
+    }
+    BlockMetaData block = blocks.get(currentBlock);
+    if (block.getRowCount() == 0) {
+      throw new RuntimeException("Illegal row group of 0 rows");
+    }
 
+    Map<ColumnDescriptor, ColumnChunkDictionaryPageReader> dictionaryReaders = new HashMap<ColumnDescriptor, ColumnChunkDictionaryPageReader>();
+    for (ColumnChunkMetaData columnChunkMetaData : block.getColumns()) {
+      for (ColumnDescriptor columnDescriptor : requestedSchema.getColumns()) {
+        if (columnChunkMetaData.getPath().equals(ColumnPath.get(columnDescriptor.getPath())) &&
+            isOnlyDictionaryEncodingPages(columnChunkMetaData.getEncodings())) {
+          long startingPosition = columnChunkMetaData.getStartingPos();
+          this.f.seek(startingPosition);
+
+          int totalSize = (int) columnChunkMetaData.getTotalSize();
+          byte[] buffer = new byte[totalSize];
+          // TODO: set read size to be (page header size + compressed dictionary page size)
+          this.f.readFully(buffer);
+
+          ChunkDescriptor chunkDescriptor = new ChunkDescriptor(columnDescriptor,
+                                                                columnChunkMetaData,
+                                                                startingPosition,
+                                                                totalSize);
+          Chunk chunk = new Chunk(chunkDescriptor, buffer, 0);
+          DictionaryPage compressedPage = chunk.readDictionaryPage();
+          if (compressedPage == null) {
+            dictionaryReaders.put(columnDescriptor, null);
+            break;
+          }
+
+          BytesDecompressor decompressor = codecFactory.getDecompressor(columnChunkMetaData.getCodec());
+          ColumnChunkPageReader columnChunkPageReader = new ColumnChunkPageReader(decompressor, new ArrayList<DataPage>(), compressedPage);
+          ColumnChunkDictionaryPageReader dictionaryPageReader = new ColumnChunkDictionaryPageReader(columnChunkPageReader);
+          dictionaryReaders.put(columnDescriptor, dictionaryPageReader);
+        }
+      }
+    }
+    return new ColumnChunkDictionaryPageReadStore(dictionaryReaders);
+  }
+
+  // are all the value pages dictionary encoded
+  private boolean isOnlyDictionaryEncodingPages(Set<Encoding> encodings) {
+    // if values have more than one encodings, definitely not dictionary encoding only
+    if (encodings.size() > 3) {
+      return false;
+    }
+    // heuristics: definition level, repetition level are not dictionary encoded
+    // if there is dictionary encoding, it must be values
+    // TODO: use PageEncodingStats to determine whether all value pages are dictionary encoded
+    for (Encoding encoding : encodings) {
+      if (encoding.usesDictionary()) {
+        return true;
+      }
+    }
+    return false;
+  }
 
   @Override
   public void close() throws IOException {
@@ -535,6 +609,8 @@ private class Chunk extends ByteArrayInputStream {
 
     private final ChunkDescriptor descriptor;
 
+    private PageHeader pageHeader;
+
     /**
      *
      * @param descriptor descriptor for the chunk
@@ -551,6 +627,26 @@ protected PageHeader readPageHeader() throws IOException {
       return Util.readPageHeader(this);
     }
 
+    /**
+     * Read only dictionary page in a given column chunk.
+     * readDictionaryPage should not be used interleaving with readAllPages
+     * either read only dictionary page, or read all pages
+     * @return dictionary page
+     */
+    public DictionaryPage readDictionaryPage() throws IOException {
+      this.pageHeader = readPageHeader();
+      int uncompressedPageSize = pageHeader.getUncompressed_page_size();
+      int compressedPageSize = pageHeader.getCompressed_page_size();
+      if (pageHeader.type == PageType.DICTIONARY_PAGE) {
+        DictionaryPageHeader dicHeader = pageHeader.getDictionary_page_header();
+        return new DictionaryPage(this.readAsBytesInput(compressedPageSize),
+                                uncompressedPageSize,
+                                dicHeader.getNum_values(),
+                                converter.getEncoding(dicHeader.getEncoding()));
+      }
+      return null;
+    }
+
     /**
      * Read all of the pages in a given column chunk.
      * @return the list of pages
@@ -560,24 +656,14 @@ public ColumnChunkPageReader readAllPages() throws IOException {
       DictionaryPage dictionaryPage = null;
       long valuesCountReadSoFar = 0;
       while (valuesCountReadSoFar < descriptor.metadata.getValueCount()) {
-        PageHeader pageHeader = readPageHeader();
-        int uncompressedPageSize = pageHeader.getUncompressed_page_size();
-        int compressedPageSize = pageHeader.getCompressed_page_size();
+        DictionaryPage nextDictionaryPage = readDictionaryPage();
+        if (dictionaryPage != null && nextDictionaryPage != null) {
+          throw new ParquetDecodingException("more than one dictionary page in column " + descriptor.col);
+        }
+        dictionaryPage = nextDictionaryPage;
+        int uncompressedPageSize = this.pageHeader.getUncompressed_page_size();
+        int compressedPageSize = this.pageHeader.getCompressed_page_size();
         switch (pageHeader.type) {
-          case DICTIONARY_PAGE:
-            // there is only one dictionary page per column chunk
-            if (dictionaryPage != null) {
-              throw new ParquetDecodingException("more than one dictionary page in column " + descriptor.col);
-            }
-          DictionaryPageHeader dicHeader = pageHeader.getDictionary_page_header();
-          dictionaryPage =
-                new DictionaryPage(
-                    this.readAsBytesInput(compressedPageSize),
-                    uncompressedPageSize,
-                    dicHeader.getNum_values(),
-                    converter.getEncoding(dicHeader.getEncoding())
-                    );
-            break;
           case DATA_PAGE:
             DataPageHeader dataHeaderV1 = pageHeader.getData_page_header();
             pagesInChunk.add(
@@ -616,6 +702,8 @@ public ColumnChunkPageReader readAllPages() throws IOException {
                     ));
             valuesCountReadSoFar += dataHeaderV2.getNum_values();
             break;
+          case DICTIONARY_PAGE:
+            throw new ParquetDecodingException("more than one dictionary page in column " + descriptor.col);
           default:
             if (DEBUG) LOG.debug("skipping page of type " + pageHeader.getType() + " of size " + compressedPageSize);
             this.skip(compressedPageSize);
@@ -801,5 +889,4 @@ public long endPos() {
     }
 
   }
-
 }