PARQUET-548: Add EncodingStats.

parquet-column/src/main/java/org/apache/parquet/column/EncodingStats.java

@@ -0,0 +1,162 @@
+/*
+ *  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;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Set;
+
+import static org.apache.parquet.column.Encoding.PLAIN_DICTIONARY;
+import static org.apache.parquet.column.Encoding.RLE_DICTIONARY;
+
+/**
+ * EncodingStats track dictionary and data page encodings for a single column within a row group.
+ * These are used when filtering row groups. For example, to filter a row group based on a column's
+ * dictionary, all of the data pages in that column must be dictionary-encoded. This class provides
+ * convenience methods for those checks, like {@link #hasNonDictionaryEncodedPages()}.
+ */
+public class EncodingStats {
+  final Map<Encoding, Integer> dictStats;
+  final Map<Encoding, Integer> dataStats;
+  private final boolean usesV2Pages;
+
+  private EncodingStats(Map<Encoding, Integer> dictStats,
+                        Map<Encoding, Integer> dataStats,
+                        boolean usesV2Pages) {
+    this.dictStats = dictStats;
+    this.dataStats = dataStats;
+    this.usesV2Pages = usesV2Pages;
+  }
+
+  public Set<Encoding> getDictionaryEncodings() {
+    return dictStats.keySet();
+  }
+
+  public Set<Encoding> getDataEncodings() {
+    return dataStats.keySet();
+  }
+
+  public int getNumDictionaryPagesEncodedAs(Encoding enc) {
+    if (dictStats.containsKey(enc)) {
+      return dictStats.get(enc);
+    } else {
+      return 0;
+    }
+  }
+
+  public int getNumDataPagesEncodedAs(Encoding enc) {
+    if (dataStats.containsKey(enc)) {
+      return dataStats.get(enc);
+    } else {
+      return 0;
+    }
+  }
+
+  public boolean hasDictionaryPages() {
+    return !dictStats.isEmpty();
+  }
+
+  public boolean hasDictionaryEncodedPages() {
+    Set<Encoding> encodings = dataStats.keySet();
+    return (encodings.contains(RLE_DICTIONARY) || encodings.contains(PLAIN_DICTIONARY));
+  }
+
+  public boolean hasNonDictionaryEncodedPages() {
+    if (dataStats.isEmpty()) {
+      return false; // no pages
+    }
+
+    // this modifies the set, so copy it
+    Set<Encoding> encodings = new HashSet<Encoding>(dataStats.keySet());
+    if (!encodings.remove(RLE_DICTIONARY) &&
+        !encodings.remove(PLAIN_DICTIONARY)) {
+      return true; // not dictionary encoded
+    }
+
+    if (encodings.isEmpty()) {
+      return false;
+    }
+
+    // at least one non-dictionary encoding is present
+    return true;
+  }
+
+  public boolean usesV2Pages() {
+    return usesV2Pages;
+  }
+
+  /**
+   * Used to build {@link EncodingStats} from metadata or to accumulate stats as pages are written.
+   */
+  public static class Builder {
+    private final Map<Encoding, Integer> dictStats = new LinkedHashMap<Encoding, Integer>();
+    private final Map<Encoding, Integer> dataStats = new LinkedHashMap<Encoding, Integer>();
+    private boolean usesV2Pages = false;
+
+    public Builder clear() {
+      this.usesV2Pages = false;
+      dictStats.clear();
+      dataStats.clear();
+      return this;
+    }
+
+    public Builder withV2Pages() {
+      this.usesV2Pages = true;
+      return this;
+    }
+
+    public Builder addDictEncoding(Encoding encoding) {
+      return addDictEncoding(encoding, 1);
+    }
+
+    public Builder addDictEncoding(Encoding encoding, int numPages) {
+      Integer pages = dictStats.get(encoding);
+      dictStats.put(encoding, numPages + (pages != null ? pages : 0));
+      return this;
+    }
+
+    public Builder addDataEncodings(Collection<Encoding> encodings) {
+      for (Encoding encoding : encodings) {
+        addDataEncoding(encoding);
+      }
+      return this;
+    }
+
+    public Builder addDataEncoding(Encoding encoding) {
+      return addDataEncoding(encoding, 1);
+    }
+
+    public Builder addDataEncoding(Encoding encoding, int numPages) {
+      Integer pages = dataStats.get(encoding);
+      dataStats.put(encoding, numPages + (pages != null ? pages : 0));
+      return this;
+    }
+
+    public EncodingStats build() {
+      return new EncodingStats(
+          Collections.unmodifiableMap(new LinkedHashMap<Encoding, Integer>(dictStats)),
+          Collections.unmodifiableMap(new LinkedHashMap<Encoding, Integer>(dataStats)),
+          usesV2Pages);
+    }
+  }
+}

parquet-column/src/test/java/org/apache/parquet/column/TestEncodingStats.java

@@ -0,0 +1,202 @@
+/*
+ *  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;
+
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class TestEncodingStats {
+  @Test
+  public void testReusedBuilder() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.withV2Pages();
+    builder.addDictEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.RLE_DICTIONARY, 3);
+    builder.addDataEncoding(Encoding.DELTA_BYTE_ARRAY);
+    builder.addDataEncoding(Encoding.DELTA_BYTE_ARRAY);
+    EncodingStats stats1 = builder.build();
+
+    Map<Encoding, Integer> expectedDictStats1 = new HashMap<Encoding, Integer>();
+    expectedDictStats1.put(Encoding.PLAIN, 1);
+    Map<Encoding, Integer> expectedDataStats1 = new HashMap<Encoding, Integer>();
+    expectedDataStats1.put(Encoding.RLE_DICTIONARY, 3);
+    expectedDataStats1.put(Encoding.DELTA_BYTE_ARRAY, 2);
+
+    builder.clear();
+    builder.addDataEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.PLAIN);
+    EncodingStats stats2 = builder.build();
+
+    Map<Encoding, Integer> expectedDictStats2 = new HashMap<Encoding, Integer>();
+    Map<Encoding, Integer> expectedDataStats2 = new HashMap<Encoding, Integer>();
+    expectedDataStats2.put(Encoding.PLAIN, 4);
+
+    assertEquals("Dictionary stats should be correct", expectedDictStats2, stats2.dictStats);
+    assertEquals("Data stats should be correct", expectedDataStats2, stats2.dataStats);
+
+    assertEquals("Dictionary stats should be correct after reuse", expectedDictStats1, stats1.dictStats);
+    assertEquals("Data stats should be correct after reuse", expectedDataStats1, stats1.dataStats);
+  }
+
+  @Test
+  public void testNoPages() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    EncodingStats stats = builder.build();
+
+    assertFalse(stats.usesV2Pages());
+    assertFalse("Should not have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertFalse("Should not have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertFalse("Should not have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testNoDataPages() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.addDictEncoding(Encoding.PLAIN_DICTIONARY);
+    EncodingStats stats = builder.build();
+
+    assertFalse(stats.usesV2Pages());
+    assertFalse("Should not have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertFalse("Should not have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertTrue("Should have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testV1AllDictionary() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.addDictEncoding(Encoding.PLAIN_DICTIONARY);
+    builder.addDataEncoding(Encoding.PLAIN_DICTIONARY);
+    builder.addDataEncoding(Encoding.PLAIN_DICTIONARY);
+    EncodingStats stats = builder.build();
+
+    assertFalse(stats.usesV2Pages());
+    assertTrue("Should have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertFalse("Should not have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertTrue("Should have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testV1NoDictionary() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.addDataEncoding(Encoding.PLAIN);
+    EncodingStats stats = builder.build();
+
+    assertFalse(stats.usesV2Pages());
+    assertFalse("Should not have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertTrue("Should have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertFalse("Should not have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testV1Fallback() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.addDictEncoding(Encoding.PLAIN_DICTIONARY);
+    builder.addDataEncoding(Encoding.PLAIN_DICTIONARY);
+    builder.addDataEncoding(Encoding.PLAIN_DICTIONARY);
+    builder.addDataEncoding(Encoding.PLAIN);
+    EncodingStats stats = builder.build();
+
+    assertFalse(stats.usesV2Pages());
+    assertTrue("Should have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertTrue("Should have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertTrue("Should have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testV2AllDictionary() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.withV2Pages();
+    builder.addDictEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.RLE_DICTIONARY);
+    EncodingStats stats = builder.build();
+
+    assertTrue(stats.usesV2Pages());
+    assertTrue("Should have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertFalse("Should not have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertTrue("Should have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testV2NoDictionary() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.withV2Pages();
+    builder.addDataEncoding(Encoding.DELTA_BINARY_PACKED);
+    builder.addDataEncoding(Encoding.DELTA_BINARY_PACKED);
+    EncodingStats stats = builder.build();
+
+    assertTrue(stats.usesV2Pages());
+    assertFalse("Should not have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertTrue("Should have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertFalse("Should not have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testV2Fallback() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.withV2Pages();
+    builder.addDictEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.RLE_DICTIONARY);
+    builder.addDataEncoding(Encoding.DELTA_BYTE_ARRAY);
+    builder.addDataEncoding(Encoding.DELTA_BYTE_ARRAY);
+    EncodingStats stats = builder.build();
+
+    assertTrue(stats.usesV2Pages());
+    assertTrue("Should have dictionary-encoded pages", stats.hasDictionaryEncodedPages());
+    assertTrue("Should have non-dictionary pages", stats.hasNonDictionaryEncodedPages());
+    assertTrue("Should have dictionary pages", stats.hasDictionaryPages());
+  }
+
+  @Test
+  public void testCounts() {
+    EncodingStats.Builder builder = new EncodingStats.Builder();
+    builder.withV2Pages();
+    builder.addDictEncoding(Encoding.PLAIN);
+    builder.addDataEncoding(Encoding.RLE_DICTIONARY, 4);
+    builder.addDataEncoding(Encoding.RLE_DICTIONARY);
+    builder.addDataEncoding(Encoding.DELTA_BYTE_ARRAY);
+    builder.addDataEncoding(Encoding.DELTA_BYTE_ARRAY);
+    EncodingStats stats = builder.build();
+
+    assertEquals("Count should match", 1, stats.getNumDictionaryPagesEncodedAs(Encoding.PLAIN));
+    assertEquals("Count should match", 0, stats.getNumDictionaryPagesEncodedAs(Encoding.PLAIN_DICTIONARY));
+    assertEquals("Count should match", 0, stats.getNumDictionaryPagesEncodedAs(Encoding.RLE));
+    assertEquals("Count should match", 0, stats.getNumDictionaryPagesEncodedAs(Encoding.BIT_PACKED));
+    assertEquals("Count should match", 0, stats.getNumDictionaryPagesEncodedAs(Encoding.DELTA_BYTE_ARRAY));
+    assertEquals("Count should match", 0, stats.getNumDictionaryPagesEncodedAs(Encoding.DELTA_BINARY_PACKED));
+    assertEquals("Count should match", 0, stats.getNumDictionaryPagesEncodedAs(Encoding.DELTA_LENGTH_BYTE_ARRAY));
+
+    assertEquals("Count should match", 5, stats.getNumDataPagesEncodedAs(Encoding.RLE_DICTIONARY));
+    assertEquals("Count should match", 2, stats.getNumDataPagesEncodedAs(Encoding.DELTA_BYTE_ARRAY));
+    assertEquals("Count should match", 0, stats.getNumDataPagesEncodedAs(Encoding.RLE));
+    assertEquals("Count should match", 0, stats.getNumDataPagesEncodedAs(Encoding.BIT_PACKED));
+    assertEquals("Count should match", 0, stats.getNumDataPagesEncodedAs(Encoding.PLAIN));
+    assertEquals("Count should match", 0, stats.getNumDataPagesEncodedAs(Encoding.PLAIN_DICTIONARY));
+    assertEquals("Count should match", 0, stats.getNumDataPagesEncodedAs(Encoding.DELTA_BINARY_PACKED));
+    assertEquals("Count should match", 0, stats.getNumDataPagesEncodedAs(Encoding.DELTA_LENGTH_BYTE_ARRAY));
+  }
+}

parquet-hadoop/src/main/java/org/apache/parquet/filter2/dictionarylevel/DictionaryFilter.java

@@ -22,6 +22,7 @@
 import org.apache.parquet.column.ColumnDescriptor;
 import org.apache.parquet.column.Dictionary;
 import org.apache.parquet.column.Encoding;
+import org.apache.parquet.column.EncodingStats;
 import org.apache.parquet.column.page.DictionaryPage;
 import org.apache.parquet.column.page.DictionaryPageReadStore;
 import org.apache.parquet.filter2.predicate.FilterPredicate;
@@ -329,6 +330,11 @@ public <T extends Comparable<T>, U extends UserDefinedPredicate<T>> Boolean visi
 
   @SuppressWarnings("deprecation")
   private static boolean hasNonDictionaryPages(ColumnChunkMetaData meta) {
+    EncodingStats stats = meta.getEncodingStats();
+    if (stats != null) {
+      return stats.hasNonDictionaryEncodedPages();
+    }
+
     // without EncodingStats, fall back to testing the encoding list
     Set<Encoding> encodings = new HashSet<Encoding>(meta.getEncodings());
     if (encodings.remove(Encoding.PLAIN_DICTIONARY)) {