Kyligence#6 Configurable record counts for block size checks

parquet-arrow/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-avro/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-benchmarks/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-cascading/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-cascading3/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-cli/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-column/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

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

@@ -49,6 +49,8 @@ public class ParquetProperties {
   public static final int DEFAULT_MAXIMUM_RECORD_COUNT_FOR_CHECK = 10000;
   public static final int DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH = 64;
   public static final int DEFAULT_PAGE_ROW_COUNT_LIMIT = 20_000;
+  public static final boolean DEFAULT_PAGE_ROW_COUNT_CHECK_ENABLE = false;
+  public static final long DEFAULT_SINGLE_ROW_SIZE_LIMIT = 512 * 1024;
 
   public static final ValuesWriterFactory DEFAULT_VALUES_WRITER_FACTORY = new DefaultValuesWriterFactory();
 
@@ -87,10 +89,13 @@ public static WriterVersion fromString(String name) {
   private final ValuesWriterFactory valuesWriterFactory;
   private final int columnIndexTruncateLength;
   private final int pageRowCountLimit;
+  private final boolean pageSizeCheckEnable;
+  private final long singleRowSizeLimit;
 
   private ParquetProperties(WriterVersion writerVersion, int pageSize, int dictPageSize, boolean enableDict, int minRowCountForPageSizeCheck,
                             int maxRowCountForPageSizeCheck, boolean estimateNextSizeCheck, ByteBufferAllocator allocator,
-                            ValuesWriterFactory writerFactory, int columnIndexMinMaxTruncateLength, int pageRowCountLimit) {
+                            ValuesWriterFactory writerFactory, int columnIndexMinMaxTruncateLength, int pageRowCountLimit,
+                            boolean pageSizeCheckEnable, long singleRowSizeLimit) {
     this.pageSizeThreshold = pageSize;
     this.initialSlabSize = CapacityByteArrayOutputStream
       .initialSlabSizeHeuristic(MIN_SLAB_SIZE, pageSizeThreshold, 10);
@@ -105,6 +110,8 @@ private ParquetProperties(WriterVersion writerVersion, int pageSize, int dictPag
     this.valuesWriterFactory = writerFactory;
     this.columnIndexTruncateLength = columnIndexMinMaxTruncateLength;
     this.pageRowCountLimit = pageRowCountLimit;
+    this.pageSizeCheckEnable = pageSizeCheckEnable;
+    this.singleRowSizeLimit = singleRowSizeLimit;
   }
 
   public ValuesWriter newRepetitionLevelWriter(ColumnDescriptor path) {
@@ -145,6 +152,14 @@ public int getPageSizeThreshold() {
     return pageSizeThreshold;
   }
 
+  public static boolean isDefaultPageRowCountCheckEnable() {
+    return DEFAULT_PAGE_ROW_COUNT_CHECK_ENABLE;
+  }
+
+  public static long getDefaultSingleRowSizeLimit() {
+    return DEFAULT_SINGLE_ROW_SIZE_LIMIT;
+  }
+
   public int getInitialSlabSize() {
     return initialSlabSize;
   }
@@ -181,6 +196,10 @@ public int getMinRowCountForPageSizeCheck() {
     return minRowCountForPageSizeCheck;
   }
 
+  public boolean isPageSizeCheckEnable() {
+    return pageSizeCheckEnable;
+  }
+
   public int getMaxRowCountForPageSizeCheck() {
     return maxRowCountForPageSizeCheck;
   }
@@ -201,6 +220,10 @@ public int getPageRowCountLimit() {
     return pageRowCountLimit;
   }
 
+  public long getSingleRowSizeLimit() {
+    return singleRowSizeLimit;
+  }
+
   public static Builder builder() {
     return new Builder();
   }
@@ -221,6 +244,8 @@ public static class Builder {
     private ValuesWriterFactory valuesWriterFactory = DEFAULT_VALUES_WRITER_FACTORY;
     private int columnIndexTruncateLength = DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH;
     private int pageRowCountLimit = DEFAULT_PAGE_ROW_COUNT_LIMIT;
+    private boolean pageSizeCheckEnable = DEFAULT_PAGE_ROW_COUNT_CHECK_ENABLE;
+    private long singleRowSizeLimit = DEFAULT_SINGLE_ROW_SIZE_LIMIT;
 
     private Builder() {
     }
@@ -236,6 +261,7 @@ private Builder(ParquetProperties toCopy) {
       this.valuesWriterFactory = toCopy.valuesWriterFactory;
       this.allocator = toCopy.allocator;
       this.pageRowCountLimit = toCopy.pageRowCountLimit;
+      this.singleRowSizeLimit = toCopy.singleRowSizeLimit;
     }
 
     /**
@@ -330,11 +356,22 @@ public Builder withPageRowCountLimit(int rowCount) {
       return this;
     }
 
+    public Builder withPageRowCheckEnable(boolean enable) {
+      pageSizeCheckEnable = enable;
+      return this;
+    }
+
+    public Builder withSingleRowSizeLimit(long size) {
+      singleRowSizeLimit = size;
+      return this;
+    }
+
     public ParquetProperties build() {
       ParquetProperties properties =
         new ParquetProperties(writerVersion, pageSize, dictPageSize,
           enableDict, minRowCountForPageSizeCheck, maxRowCountForPageSizeCheck,
-          estimateNextSizeCheck, allocator, valuesWriterFactory, columnIndexTruncateLength, pageRowCountLimit);
+          estimateNextSizeCheck, allocator, valuesWriterFactory, columnIndexTruncateLength,
+          pageRowCountLimit, pageSizeCheckEnable, singleRowSizeLimit);
       // we pass a constructed but uninitialized factory to ParquetProperties above as currently
       // creation of ValuesWriters is invoked from within ParquetProperties. In the future
       // we'd like to decouple that and won't need to pass an object to properties and then pass the

parquet-common/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-encoding/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-format-structures/pom.xml

@@ -24,7 +24,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <artifactId>parquet-format-structures</artifactId>

parquet-generator/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-hadoop-bundle/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

parquet-hadoop/pom.xml

@@ -21,7 +21,7 @@
     <groupId>org.apache.parquet</groupId>
     <artifactId>parquet</artifactId>
     <relativePath>../pom.xml</relativePath>
-    <version>1.12.0-kylin-r2</version>
+    <version>1.12.0-kylin-r3</version>
   </parent>
 
   <modelVersion>4.0.0</modelVersion>

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

@@ -18,14 +18,6 @@
  */
 package org.apache.parquet.hadoop;
 
-import static java.lang.Math.max;
-import static java.lang.Math.min;
-import static org.apache.parquet.Preconditions.checkNotNull;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
 import org.apache.parquet.column.ColumnWriteStore;
 import org.apache.parquet.column.ParquetProperties;
 import org.apache.parquet.hadoop.CodecFactory.BytesCompressor;
@@ -39,6 +31,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import static java.lang.Math.max;
+import static java.lang.Math.min;
+import static org.apache.parquet.Preconditions.checkNotNull;
+
 class InternalParquetRecordWriter<T> {
   private static final Logger LOG = LoggerFactory.getLogger(InternalParquetRecordWriter.class);
 
@@ -66,23 +66,25 @@ class InternalParquetRecordWriter<T> {
   private ColumnChunkPageWriteStore pageStore;
   private RecordConsumer recordConsumer;
 
+  private boolean pageSizeCheckEnable;
+
   /**
    * @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
+   * @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,
-      BytesCompressor compressor,
-      boolean validating,
-      ParquetProperties props) {
+    ParquetFileWriter parquetFileWriter,
+    WriteSupport<T> writeSupport,
+    MessageType schema,
+    Map<String, String> extraMetaData,
+    long rowGroupSize,
+    BytesCompressor compressor,
+    boolean validating,
+    ParquetProperties props) {
     this.parquetFileWriter = parquetFileWriter;
     this.writeSupport = checkNotNull(writeSupport, "writeSupport");
     this.schema = schema;
@@ -92,6 +94,7 @@ public InternalParquetRecordWriter(
     this.nextRowGroupSize = rowGroupSizeThreshold;
     this.compressor = compressor;
     this.validating = validating;
+    this.pageSizeCheckEnable = props.isPageSizeCheckEnable();
     this.props = props;
     initStore();
   }
@@ -102,7 +105,7 @@ public ParquetMetadata getFooter() {
 
   private void initStore() {
     pageStore = new ColumnChunkPageWriteStore(compressor, schema, props.getAllocator(),
-        props.getColumnIndexTruncateLength());
+      props.getColumnIndexTruncateLength());
     columnStore = props.newColumnWriteStore(schema, pageStore);
     MessageColumnIO columnIO = new ColumnIOFactory(validating).getColumnIO(schema);
     this.recordConsumer = columnIO.getRecordWriter(columnStore);
@@ -125,9 +128,14 @@ public void close() throws IOException, InterruptedException {
   }
 
   public void write(T value) throws IOException, InterruptedException {
-    writeSupport.write(value);
-    ++ recordCount;
-    checkBlockSizeReached();
+    Long singleRowLimit = props.getSingleRowSizeLimit();
+    Long recordSize = writeSupport.writeAndFetchRecordSize(value);
+    ++recordCount;
+    if (singleRowLimit <= recordSize) {
+      checkBlockSizeReachedManual();
+    } else {
+      checkBlockSizeReached();
+    }
   }
 
   /**
@@ -144,25 +152,38 @@ private void checkBlockSizeReached() throws IOException {
       // flush the row group if it is within ~2 records of the limit
       // it is much better to be slightly under size than to be over at all
       if (memSize > (nextRowGroupSize - 2 * recordSize)) {
-        LOG.debug("mem size {} > {}: flushing {} records to disk.", memSize, nextRowGroupSize, recordCount);
+        LOG.info("mem size {} > {}: flushing {} records to disk.", memSize, nextRowGroupSize, recordCount);
         flushRowGroupToStore();
         initStore();
         recordCountForNextMemCheck = min(max(MINIMUM_RECORD_COUNT_FOR_CHECK, recordCount / 2), MAXIMUM_RECORD_COUNT_FOR_CHECK);
         this.lastRowGroupEndPos = parquetFileWriter.getPos();
       } else {
         recordCountForNextMemCheck = min(
-            max(MINIMUM_RECORD_COUNT_FOR_CHECK, (recordCount + (long)(nextRowGroupSize / ((float)recordSize))) / 2), // will check halfway
-            recordCount + MAXIMUM_RECORD_COUNT_FOR_CHECK // will not look more than max records ahead
-            );
-        LOG.debug("Checked mem at {} will check again at: {}", recordCount, recordCountForNextMemCheck);
+          max(MINIMUM_RECORD_COUNT_FOR_CHECK, (recordCount + (long)(nextRowGroupSize / ((float)recordSize))) / 2), // will check halfway
+          recordCount + MAXIMUM_RECORD_COUNT_FOR_CHECK // will not look more than max records ahead
+        );
+        LOG.info("Checked mem at {} will check again at: {}", recordCount, recordCountForNextMemCheck);
       }
     }
   }
 
+  private void checkBlockSizeReachedManual() throws IOException {
+    long memSize = columnStore.getBufferedSize();
+    long recordSize = memSize / recordCount;
+    if (memSize > (nextRowGroupSize - 2 * recordSize)) {
+      LOG.info("mem size {} > {}: flushing {} records to disk.", memSize, nextRowGroupSize, recordCount);
+      flushRowGroupToStore();
+      initStore();
+      this.lastRowGroupEndPos = parquetFileWriter.getPos();
+    } else {
+      LOG.info("Checked mem at {}", recordCount);
+    }
+  }
+
   private void flushRowGroupToStore()
-      throws IOException {
+    throws IOException {
     recordConsumer.flush();
-    LOG.debug("Flushing mem columnStore to file. allocated memory: {}", columnStore.getAllocatedSize());
+    LOG.info("Flushing mem columnStore to file. allocated memory: {}", columnStore.getAllocatedSize());
     if (columnStore.getAllocatedSize() > (3 * rowGroupSizeThreshold)) {
       LOG.warn("Too much memory used: {}", columnStore.memUsageString());
     }
@@ -174,8 +195,8 @@ private void flushRowGroupToStore()
       recordCount = 0;
       parquetFileWriter.endBlock();
       this.nextRowGroupSize = Math.min(
-          parquetFileWriter.getNextRowGroupSize(),
-          rowGroupSizeThreshold);
+        parquetFileWriter.getNextRowGroupSize(),
+        rowGroupSizeThreshold);
     }
 
     columnStore = null;