PARQUET-1213: Column indexes: Limit index size (#480)

Author: gszadovszky

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/BinaryColumnIndexBuilder.java

@@ -102,4 +102,9 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
   int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
     return comparator.compare(maxValues.get(index1), maxValues.get(index2));
   }
+
+  @Override
+  int sizeOf(Object value) {
+    return ((Binary) value).length();
+  }
 }

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/BooleanColumnIndexBuilder.java

@@ -103,4 +103,9 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
   int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
     return comparator.compare(maxValues.get(index1), maxValues.get(index2));
   }
+
+  @Override
+  int sizeOf(Object value) {
+    return 1;
+  }
 }

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/ColumnIndexBuilder.java

@@ -195,13 +195,19 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
     int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
       return 0;
     }
+
+    @Override
+    int sizeOf(Object value) {
+      return 0;
+    }
   };
 
   private static final Map<PrimitiveTypeName, ColumnIndexBuilder> BUILDERS = new EnumMap<>(PrimitiveTypeName.class);
 
   private PrimitiveType type;
   private final BooleanList nullPages = new BooleanArrayList();
   private final LongList nullCounts = new LongArrayList();
+  private long minMaxSize;
 
   /**
    * @return a no-op builder that does not collect statistics objects and therefore returns {@code null} at
@@ -293,7 +299,11 @@ public static ColumnIndex build(
   public void add(Statistics<?> stats) {
     if (stats.hasNonNullValue()) {
       nullPages.add(false);
-      addMinMax(stats.genericGetMin(), stats.genericGetMax());
+      Object min = stats.genericGetMin();
+      Object max = stats.genericGetMax();
+      addMinMax(min, max);
+      minMaxSize += sizeOf(min);
+      minMaxSize += sizeOf(max);
     } else {
       nullPages.add(true);
       addMinMax(null, null);
@@ -316,7 +326,7 @@ private void fill(List<Boolean> nullPages, List<Long> nullCounts, List<ByteBuffe
               nullPages.size(), nullCounts == null ? "null" : nullCounts.size(), minValues.size(), maxValues.size()));
     }
     this.nullPages.addAll(nullPages);
-    // Null counts is optional in the format
+    // Nullcounts is optional in the format
     if (nullCounts != null) {
       this.nullCounts.addAll(nullCounts);
     }
@@ -325,7 +335,11 @@ private void fill(List<Boolean> nullPages, List<Long> nullCounts, List<ByteBuffe
       if (nullPages.get(i)) {
         addMinMaxFromBytes(null, null);
       } else {
-        addMinMaxFromBytes(minValues.get(i), maxValues.get(i));
+        ByteBuffer min = minValues.get(i);
+        ByteBuffer max = maxValues.get(i);
+        addMinMaxFromBytes(min, max);
+        minMaxSize += min.remaining();
+        minMaxSize += max.remaining();
       }
     }
   }
@@ -421,9 +435,26 @@ private void clear() {
     nullPages.clear();
     nullCounts.clear();
     clearMinMax();
+    minMaxSize = 0;
   }
 
   abstract void clearMinMax();
 
   abstract ColumnIndexBase createColumnIndex(PrimitiveType type);
+
+  abstract int sizeOf(Object value);
+
+  /**
+   * @return the number of pages added so far to this builder
+   */
+  public int getPageCount() {
+    return nullPages.size();
+  }
+
+  /**
+   * @return the sum of size in bytes of the min/max values added so far to this builder
+   */
+  public long getMinMaxSize() {
+    return minMaxSize;
+  }
 }

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/DoubleColumnIndexBuilder.java

@@ -67,7 +67,7 @@ private static double convert(ByteBuffer buffer) {
   }
 
   private static ByteBuffer convert(double value) {
-    return ByteBuffer.allocate(Double.SIZE / 8).order(LITTLE_ENDIAN).putDouble(0, value);
+    return ByteBuffer.allocate(Double.BYTES).order(LITTLE_ENDIAN).putDouble(0, value);
   }
 
   @Override
@@ -105,4 +105,9 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
   int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
     return comparator.compare(maxValues.get(index1), maxValues.get(index2));
   }
+
+  @Override
+  int sizeOf(Object value) {
+    return Double.BYTES;
+  }
 }

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/FloatColumnIndexBuilder.java

@@ -67,7 +67,7 @@ private static float convert(ByteBuffer buffer) {
   }
 
   private static ByteBuffer convert(float value) {
-    return ByteBuffer.allocate(Float.SIZE / 8).order(LITTLE_ENDIAN).putFloat(0, value);
+    return ByteBuffer.allocate(Float.BYTES).order(LITTLE_ENDIAN).putFloat(0, value);
   }
 
   @Override
@@ -105,4 +105,9 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
   int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
     return comparator.compare(maxValues.get(index1), maxValues.get(index2));
   }
+
+  @Override
+  int sizeOf(Object value) {
+    return Float.BYTES;
+  }
 }

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/IntColumnIndexBuilder.java

@@ -67,7 +67,7 @@ private static int convert(ByteBuffer buffer) {
   }
 
   private static ByteBuffer convert(int value) {
-    return ByteBuffer.allocate(Integer.SIZE / 8).order(LITTLE_ENDIAN).putInt(0, value);
+    return ByteBuffer.allocate(Integer.BYTES).order(LITTLE_ENDIAN).putInt(0, value);
   }
 
   @Override
@@ -105,4 +105,9 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
   int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
     return comparator.compare(maxValues.get(index1), maxValues.get(index2));
   }
+
+  @Override
+  int sizeOf(Object value) {
+    return Integer.BYTES;
+  }
 }

parquet-column/src/main/java/org/apache/parquet/internal/column/columnindex/LongColumnIndexBuilder.java

@@ -67,7 +67,7 @@ private static long convert(ByteBuffer buffer) {
   }
 
   private static ByteBuffer convert(long value) {
-    return ByteBuffer.allocate(Long.SIZE / 8).order(LITTLE_ENDIAN).putLong(0, value);
+    return ByteBuffer.allocate(Long.BYTES).order(LITTLE_ENDIAN).putLong(0, value);
   }
 
   @Override
@@ -105,4 +105,9 @@ int compareMinValues(PrimitiveComparator<Binary> comparator, int index1, int ind
   int compareMaxValues(PrimitiveComparator<Binary> comparator, int index1, int index2) {
     return comparator.compare(maxValues.get(index1), maxValues.get(index2));
   }
+
+  @Override
+  int sizeOf(Object value) {
+    return Long.BYTES;
+  }
 }