[PARQUET-99] Page size check properties

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

@@ -55,6 +55,10 @@
  */
 public class ParquetProperties {
 
+  public static final boolean DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK = true;
+  public static final int DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK = 100;
+  public static final int DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK = 10000;
+
   public enum WriterVersion {
     PARQUET_1_0 ("v1"),
     PARQUET_2_0 ("v2");
@@ -78,16 +82,28 @@ public static WriterVersion fromString(String name) {
   private final int dictionaryPageSizeThreshold;
   private final WriterVersion writerVersion;
   private final boolean enableDictionary;
+  private final int minRowCountForPageSizeCheck;
+  private final int maxRowCountForPageSizeCheck;
+  private final boolean estimateNextSizeCheck;
   private final ByteBufferAllocator allocator;
 
   public ParquetProperties(int dictPageSize, WriterVersion writerVersion, boolean enableDict) {
     this(dictPageSize, writerVersion, enableDict, new HeapByteBufferAllocator());
   }
 
   public ParquetProperties(int dictPageSize, WriterVersion writerVersion, boolean enableDict, ByteBufferAllocator allocator) {
+    this(dictPageSize, writerVersion, enableDict, DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK, DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK,
+        DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK, allocator);
+  }
+
+  public ParquetProperties(int dictPageSize, WriterVersion writerVersion, boolean enableDict, int minRowCountForPageSizeCheck,
+                           int maxRowCountForPageSizeCheck, boolean estimateNextSizeCheck, ByteBufferAllocator allocator) {
     this.dictionaryPageSizeThreshold = dictPageSize;
     this.writerVersion = writerVersion;
     this.enableDictionary = enableDict;
+    this.minRowCountForPageSizeCheck = minRowCountForPageSizeCheck;
+    this.maxRowCountForPageSizeCheck = maxRowCountForPageSizeCheck;
+    this.estimateNextSizeCheck = estimateNextSizeCheck;
     Preconditions.checkNotNull(allocator, "ByteBufferAllocator");
     this.allocator = allocator;
   }
@@ -97,8 +113,7 @@ public ValuesWriter getColumnDescriptorValuesWriter(int maxLevel, int initialSiz
       return new DevNullValuesWriter();
     } else {
       return new RunLengthBitPackingHybridValuesWriter(
-          getWidthFromMaxInt(maxLevel), initialSizePerCol, pageSize, this.allocator
-      );
+          getWidthFromMaxInt(maxLevel), initialSizePerCol, pageSize, this.allocator);
     }
   }
 
@@ -245,15 +260,38 @@ public ColumnWriteStore newColumnWriteStore(
           pageStore,
           pageSize,
           dictionaryPageSizeThreshold,
-          enableDictionary, writerVersion, allocator);
+          enableDictionary,
+          minRowCountForPageSizeCheck,
+          estimateNextSizeCheck,
+          writerVersion, 
+          allocator);
     case PARQUET_2_0:
       return new ColumnWriteStoreV2(
           schema,
           pageStore,
           pageSize,
-          new ParquetProperties(dictionaryPageSizeThreshold, writerVersion, enableDictionary, allocator));
+          new ParquetProperties(
+              dictionaryPageSizeThreshold,
+              writerVersion,
+              enableDictionary,
+              minRowCountForPageSizeCheck,
+              maxRowCountForPageSizeCheck,
+              estimateNextSizeCheck,
+              allocator));
     default:
       throw new IllegalArgumentException("unknown version " + writerVersion);
     }
   }
+
+  public int getMinRowCountForPageSizeCheck() {
+    return minRowCountForPageSizeCheck;
+  }
+
+  public int getMaxRowCountForPageSizeCheck() {
+    return maxRowCountForPageSizeCheck;
+  }
+
+  public boolean isEstimateNextSizeCheck() {
+    return estimateNextSizeCheck;
+  }
 }

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

@@ -33,22 +33,33 @@
 import org.apache.parquet.column.page.PageWriteStore;
 import org.apache.parquet.column.page.PageWriter;
 
+import static org.apache.parquet.column.ParquetProperties.DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK;
+import static org.apache.parquet.column.ParquetProperties.DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK;
+
 public class ColumnWriteStoreV1 implements ColumnWriteStore {
 
   private final Map<ColumnDescriptor, ColumnWriterV1> columns = new TreeMap<ColumnDescriptor, ColumnWriterV1>();
   private final PageWriteStore pageWriteStore;
   private final int pageSizeThreshold;
   private final int dictionaryPageSizeThreshold;
   private final boolean enableDictionary;
+  private final int initialRowCountForPageSizeCheck;
+  private final boolean estimateNextSizeCheck;
   private final WriterVersion writerVersion;
   private final ByteBufferAllocator allocator;
 
   public ColumnWriteStoreV1(PageWriteStore pageWriteStore, int pageSizeThreshold, int dictionaryPageSizeThreshold, boolean enableDictionary, WriterVersion writerVersion, ByteBufferAllocator allocator) {
+    this(pageWriteStore, pageSizeThreshold, dictionaryPageSizeThreshold, enableDictionary, DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK, DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK, writerVersion, allocator);
+  }
+
+  public ColumnWriteStoreV1(PageWriteStore pageWriteStore, int pageSizeThreshold, int dictionaryPageSizeThreshold, boolean enableDictionary, int initialRowCountForPageSizeCheck, boolean estimateNextSizeCheck, WriterVersion writerVersion, ByteBufferAllocator allocator) {
     super();
     this.pageWriteStore = pageWriteStore;
     this.pageSizeThreshold = pageSizeThreshold;
     this.dictionaryPageSizeThreshold = dictionaryPageSizeThreshold;
     this.enableDictionary = enableDictionary;
+    this.initialRowCountForPageSizeCheck = initialRowCountForPageSizeCheck;
+    this.estimateNextSizeCheck = estimateNextSizeCheck;
     this.writerVersion = writerVersion;
     this.allocator = allocator;
   }
@@ -68,7 +79,7 @@ public Set<ColumnDescriptor> getColumnDescriptors() {
 
   private ColumnWriterV1 newMemColumn(ColumnDescriptor path) {
     PageWriter pageWriter = pageWriteStore.getPageWriter(path);
-    return new ColumnWriterV1(path, pageWriter, pageSizeThreshold, dictionaryPageSizeThreshold, enableDictionary, writerVersion, allocator);
+    return new ColumnWriterV1(path, pageWriter, pageSizeThreshold, dictionaryPageSizeThreshold, enableDictionary, initialRowCountForPageSizeCheck, estimateNextSizeCheck, writerVersion, allocator);
   }
 
   @Override

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

@@ -40,16 +40,14 @@
 
 public class ColumnWriteStoreV2 implements ColumnWriteStore {
 
-  // will wait for at least that many records before checking again
-  private static final int MINIMUM_RECORD_COUNT_FOR_CHECK = 100;
-  private static final int MAXIMUM_RECORD_COUNT_FOR_CHECK = 10000;
   // will flush even if size bellow the threshold by this much to facilitate page alignment
   private static final float THRESHOLD_TOLERANCE_RATIO = 0.1f; // 10 %
 
   private final Map<ColumnDescriptor, ColumnWriterV2> columns;
   private final Collection<ColumnWriterV2> writers;
+  private final ParquetProperties parquetProperties;
   private long rowCount;
-  private long rowCountForNextSizeCheck = MINIMUM_RECORD_COUNT_FOR_CHECK;
+  private long rowCountForNextSizeCheck;
   private final long thresholdTolerance;
   private final ByteBufferAllocator allocator;
 
@@ -61,6 +59,7 @@ public ColumnWriteStoreV2(
       int pageSizeThreshold,
       ParquetProperties parquetProps) {
     super();
+    this.parquetProperties = parquetProps;
     this.pageSizeThreshold = pageSizeThreshold;
     this.thresholdTolerance = (long)(pageSizeThreshold * THRESHOLD_TOLERANCE_RATIO);
     this.allocator = parquetProps.getAllocator();
@@ -71,6 +70,8 @@ public ColumnWriteStoreV2(
     }
     this.columns = unmodifiableMap(mcolumns);
     this.writers = this.columns.values();
+
+    this.rowCountForNextSizeCheck = parquetProperties.getMinRowCountForPageSizeCheck();
   }
 
   public ColumnWriter getColumnWriter(ColumnDescriptor path) {
@@ -158,20 +159,25 @@ private void sizeCheck() {
       }
       long rowsToFillPage =
           usedMem == 0 ?
-              MAXIMUM_RECORD_COUNT_FOR_CHECK
+              parquetProperties.getMaxRowCountForPageSizeCheck()
               : (long)((float)rows) / usedMem * remainingMem;
       if (rowsToFillPage < minRecordToWait) {
         minRecordToWait = rowsToFillPage;
       }
     }
     if (minRecordToWait == Long.MAX_VALUE) {
-      minRecordToWait = MINIMUM_RECORD_COUNT_FOR_CHECK;
+      minRecordToWait = parquetProperties.getMinRowCountForPageSizeCheck();
+    }
+
+    if(parquetProperties.isEstimateNextSizeCheck()) {
+      // will check again halfway
+      rowCountForNextSizeCheck = rowCount +
+          min(
+              max(minRecordToWait / 2, parquetProperties.getMinRowCountForPageSizeCheck()), // no less than MINIMUM_RECORD_COUNT_FOR_CHECK
+              parquetProperties.getMaxRowCountForPageSizeCheck()); // no more than MAXIMUM_RECORD_COUNT_FOR_CHECK
+    } else {
+      rowCountForNextSizeCheck = rowCount + parquetProperties.getMinRowCountForPageSizeCheck();
     }
-    // will check again halfway
-    rowCountForNextSizeCheck = rowCount +
-        min(
-            max(minRecordToWait / 2, MINIMUM_RECORD_COUNT_FOR_CHECK), // no less than MINIMUM_RECORD_COUNT_FOR_CHECK
-            MAXIMUM_RECORD_COUNT_FOR_CHECK); // no more than MAXIMUM_RECORD_COUNT_FOR_CHECK
   }
 
 }

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

@@ -47,7 +47,6 @@
 final class ColumnWriterV1 implements ColumnWriter {
   private static final Log LOG = Log.getLog(ColumnWriterV1.class);
   private static final boolean DEBUG = Log.DEBUG;
-  private static final int INITIAL_COUNT_FOR_SIZE_CHECK = 100;
   private static final int MIN_SLAB_SIZE = 64;
 
   private final ColumnDescriptor path;
@@ -57,7 +56,9 @@ final class ColumnWriterV1 implements ColumnWriter {
   private ValuesWriter definitionLevelColumn;
   private ValuesWriter dataColumn;
   private int valueCount;
+  private int initialRowCountForPageSizeCheck;
   private int valueCountForNextSizeCheck;
+  private boolean estimateNextSizeCheck;
 
   private Statistics statistics;
 
@@ -67,13 +68,19 @@ public ColumnWriterV1(
       int pageSizeThreshold,
       int dictionaryPageSizeThreshold,
       boolean enableDictionary,
+      int initialRowCountForPageSizeCheck,
+      boolean estimateNextSizeCheck,
       WriterVersion writerVersion,
       ByteBufferAllocator allocator) {
     this.path = path;
     this.pageWriter = pageWriter;
     this.pageSizeThreshold = pageSizeThreshold;
     // initial check of memory usage. So that we have enough data to make an initial prediction
-    this.valueCountForNextSizeCheck = INITIAL_COUNT_FOR_SIZE_CHECK;
+    this.initialRowCountForPageSizeCheck = initialRowCountForPageSizeCheck;
+    this.valueCountForNextSizeCheck = initialRowCountForPageSizeCheck;
+    // Do not attempt to predict next size check.  Prevents issues with rows that vary significantly in size.
+    this.estimateNextSizeCheck = estimateNextSizeCheck;
+
     resetStatistics();
 
     ParquetProperties parquetProps = new ParquetProperties(dictionaryPageSizeThreshold, writerVersion, enableDictionary, allocator);
@@ -111,11 +118,17 @@ private void accountForValueWritten() {
           + dataColumn.getBufferedSize();
       if (memSize > pageSizeThreshold) {
         // we will write the current page and check again the size at the predicted middle of next page
-        valueCountForNextSizeCheck = valueCount / 2;
+        if(estimateNextSizeCheck) {
+          valueCountForNextSizeCheck = valueCount / 2;
+        } else {
+          valueCountForNextSizeCheck = initialRowCountForPageSizeCheck;
+        }
         writePage();
-      } else {
+      } else if (estimateNextSizeCheck) {
         // not reached the threshold, will check again midway
         valueCountForNextSizeCheck = (int)(valueCount + ((float)valueCount * pageSizeThreshold / memSize)) / 2 + 1;
+      } else {
+        valueCountForNextSizeCheck += initialRowCountForPageSizeCheck;
       }
     }
   }

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

@@ -19,7 +19,10 @@
 package org.apache.parquet.column.impl;
 
 import static junit.framework.Assert.assertEquals;
+import static org.apache.parquet.column.ParquetProperties.DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK;
 import static org.apache.parquet.column.ParquetProperties.WriterVersion.PARQUET_2_0;
+import static org.apache.parquet.column.ParquetProperties.DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK;
+import static org.apache.parquet.column.ParquetProperties.DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK;
 
 import java.util.List;
 
@@ -59,7 +62,9 @@ public void test() throws Exception {
     MessageType schema = MessageTypeParser.parseMessageType("message test { required binary foo; }");
     ColumnDescriptor col = schema.getColumns().get(0);
     MemPageWriter pageWriter = new MemPageWriter();
-    ColumnWriterV2 columnWriterV2 = new ColumnWriterV2(col, pageWriter, new ParquetProperties(1024, PARQUET_2_0, true), 2048);
+    ColumnWriterV2 columnWriterV2 = new ColumnWriterV2(col, pageWriter, new ParquetProperties(1024, PARQUET_2_0, true,
+        DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK, DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK,
+        DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK, new HeapByteBufferAllocator()), 2048);
     for (int i = 0; i < rows; i++) {
       columnWriterV2.write(Binary.fromString("bar" + i % 10), 0, 0);
       if ((i + 1) % 1000 == 0) {
@@ -94,7 +99,9 @@ public void testOptional() throws Exception {
     MessageType schema = MessageTypeParser.parseMessageType("message test { optional binary foo; }");
     ColumnDescriptor col = schema.getColumns().get(0);
     MemPageWriter pageWriter = new MemPageWriter();
-    ColumnWriterV2 columnWriterV2 = new ColumnWriterV2(col, pageWriter, new ParquetProperties(1024, PARQUET_2_0, true), 2048);
+    ColumnWriterV2 columnWriterV2 = new ColumnWriterV2(col, pageWriter, new ParquetProperties(1024, PARQUET_2_0, true,
+        DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK, DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK,
+        DEFAULT_ESTIMATE_ROW_COUNT_FOR_PAGE_SIZE_CHECK, new HeapByteBufferAllocator()), 2048);
     for (int i = 0; i < rows; i++) {
       columnWriterV2.writeNull(0, 0);
       if ((i + 1) % 1000 == 0) {

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

@@ -68,6 +68,44 @@ class InternalParquetRecordWriter<T> {
   private RecordConsumer recordConsumer;
 
 
+  /**
+   * @param parquetFileWriter the file to write to
+   * @param writeSupport the class to convert incoming records
+   * @param schema the schema of the records
+   * @param extraMetaData extra meta data to write in the footer of the file
+   * @param rowGroupSize the size of a block in the file (this will be approximate)
+   * @param compressor the codec used to compress
+   */
+  public InternalParquetRecordWriter(
+          ParquetFileWriter parquetFileWriter,
+          WriteSupport<T> writeSupport,
+          MessageType schema,
+          Map<String, String> extraMetaData,
+          long rowGroupSize,
+          int pageSize,
+          BytesCompressor compressor,
+          int dictionaryPageSize,
+          boolean enableDictionary,
+          boolean validating,
+          WriterVersion writerVersion,
+          ByteBufferAllocator allocator) {
+    this(parquetFileWriter,
+         writeSupport,
+         schema,
+         extraMetaData,
+         rowGroupSize,
+         pageSize,
+         compressor,
+         dictionaryPageSize,
+         enableDictionary,
+         ParquetProperties.DEFAULT_MINIMUM_RECORD_COUNT_FOR_CHECK,
+         ParquetProperties.DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK,
+         false,
+         validating,
+         writerVersion,
+         allocator);
+  }
+
   /**
    * @param parquetFileWriter the file to write to
    * @param writeSupport the class to convert incoming records
@@ -86,6 +124,9 @@ public InternalParquetRecordWriter(
       BytesCompressor compressor,
       int dictionaryPageSize,
       boolean enableDictionary,
+      int minRowCountForSizeCheck,
+      int maxRowCountForSizeCheck,
+      boolean constantNextSizeCheck,
       boolean validating,
       WriterVersion writerVersion,
       ByteBufferAllocator allocator) {
@@ -99,7 +140,8 @@ public InternalParquetRecordWriter(
     this.pageSize = pageSize;
     this.compressor = compressor;
     this.validating = validating;
-    this.parquetProperties = new ParquetProperties(dictionaryPageSize, writerVersion, enableDictionary, allocator);
+    this.parquetProperties = new ParquetProperties(dictionaryPageSize, writerVersion, enableDictionary,
+        minRowCountForSizeCheck, maxRowCountForSizeCheck, constantNextSizeCheck, allocator);
     initStore();
   }