[WIP] Optimize heap usage of dictionary build to support dict on high cardinality STRINGs (fast filtering)

pinot-perf/pom.xml

@@ -105,6 +105,11 @@
       <groupId>net.sf.jopt-simple</groupId>
       <artifactId>jopt-simple</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.openjdk.jol</groupId>
+      <artifactId>jol-core</artifactId>
+      <version>0.17</version>
+    </dependency>
   </dependencies>
   <build>
     <plugins>

pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkDictionaryHeapFootprint.java

@@ -0,0 +1,94 @@
+/**
+ * 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.pinot.perf;
+
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import java.util.Arrays;
+import java.util.Random;
+import org.openjdk.jol.info.GraphLayout;
+
+
+/**
+ * Measures, with JOL, the exact on-heap footprint of the two structures a STRING dictionary column holds during
+ * segment generation in map mode — the sorted {@code String[]} and the value->dictId {@code Object2IntOpenHashMap}
+ * (which together pin the distinct String objects) — across cardinality x average length. This is the heap that
+ * binary-search mode (3a + 3b) eliminates: in binary-search mode the map is never built and the array is released,
+ * so the on-heap footprint of these structures drops to ~0 (the dictionary lives off-heap in the mmap'd file).
+ *
+ * The printed "per-value" column validates the back-of-envelope model (~ 60 + L bytes/value with compressed oops).
+ *
+ * Run as a plain main (not JMH), e.g. with {@code -Xmx8g}:
+ *   java -cp pinot-perf/target/classes:... org.apache.pinot.perf.BenchmarkDictionaryHeapFootprint
+ */
+public final class BenchmarkDictionaryHeapFootprint {
+  private BenchmarkDictionaryHeapFootprint() {
+  }
+
+  private static final int[] CARDINALITIES = {100_000, 500_000, 1_000_000, 2_000_000, 3_000_000};
+  private static final int[] AVG_LENGTHS = {16, 32, 48, 100};
+
+  public static void main(String[] args) {
+    System.out.printf("%-10s %-8s %-18s %-14s %-16s%n", "card", "avgLen", "mapMode(MB)", "bytes/value",
+        "predicted(60+L)");
+    System.out.println("-----------------------------------------------------------------------");
+    for (int avgLength : AVG_LENGTHS) {
+      for (int cardinality : CARDINALITIES) {
+        String[] sortedValues = generateSortedStrings(cardinality, avgLength);
+        Object2IntOpenHashMap<String> valueToDictId = new Object2IntOpenHashMap<>(cardinality);
+        for (int i = 0; i < cardinality; i++) {
+          valueToDictId.put(sortedValues[i], i);
+        }
+
+        // Combined reachable graph of the sorted array + the map; shared String objects are counted once. This is the
+        // on-heap footprint map mode holds (and binary-search mode eliminates).
+        long mapModeBytes = GraphLayout.parseInstance(sortedValues, valueToDictId).totalSize();
+        double mapModeMb = mapModeBytes / (1024.0 * 1024.0);
+        double bytesPerValue = (double) mapModeBytes / cardinality;
+
+        System.out.printf("%-10d %-8d %-18.1f %-14.1f %-16d%n", cardinality, avgLength, mapModeMb, bytesPerValue,
+            60 + avgLength);
+
+        // Drop references so successive cells don't accumulate heap.
+        // CHECKSTYLE:OFF
+        sortedValues = null;
+        valueToDictId = null;
+        // CHECKSTYLE:ON
+        System.gc();
+      }
+    }
+    System.out.println();
+    System.out.println("binary-search mode (3a + 3b): on-heap footprint of these structures is ~0 "
+        + "(dictionary served from the off-heap mmap'd file).");
+  }
+
+  private static String[] generateSortedStrings(int cardinality, int avgLength) {
+    Random random = new Random(42);
+    String[] values = new String[cardinality];
+    for (int i = 0; i < cardinality; i++) {
+      StringBuilder sb = new StringBuilder();
+      sb.append(i).append('_');
+      while (sb.length() < avgLength) {
+        sb.append((char) ('a' + random.nextInt(26)));
+      }
+      values[i] = sb.toString();
+    }
+    Arrays.sort(values);
+    return values;
+  }
+}

pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkSegmentDictionaryCreation.java

@@ -0,0 +1,186 @@
+/**
+ * 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.pinot.perf;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Random;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+import org.openjdk.jmh.runner.options.TimeValue;
+
+
+/**
+ * Compares the two dictId-resolution strategies in {@link SegmentDictionaryCreator} for STRING columns:
+ *  - map mode: the on-heap value->dictId {@code Object2IntOpenHashMap} (default), O(1) lookup.
+ *  - binary-search mode: no map; dictIds resolved by binary-searching the off-heap mmap'd dictionary buffer (3a).
+ *
+ * {@code lookupMap} / {@code lookupBinarySearch} measure the per-row encode cost (the pass-2 hot path);
+ * {@code buildMap} / {@code buildBinarySearch} measure dictionary build cost (run with {@code -prof gc} to see the
+ * allocation a high-card column avoids by skipping the map). Swept over cardinality, average string length and
+ * fixed- vs var-length dictionary encoding.
+ *
+ * Run the heaviest cells with adequate heap, e.g. {@code -Xmx8g}.
+ */
+@State(Scope.Benchmark)
+@Fork(1)
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 5, time = 3)
+public class BenchmarkSegmentDictionaryCreation {
+  private static final FieldSpec STRING_FIELD = new DimensionFieldSpec("col", FieldSpec.DataType.STRING, true);
+
+  @Param({"100000", "1000000", "5000000"})
+  private int _cardinality;
+
+  @Param({"16", "64"})
+  private int _avgLength;
+
+  @Param({"true", "false"})
+  private boolean _useVarLength;
+
+  private File _baseDir;
+  private File _mapDir;
+  private File _bsDir;
+  private String[] _sortedValues;
+  private String[] _lookupKeys;
+  private SegmentDictionaryCreator _mapCreator;
+  private SegmentDictionaryCreator _bsCreator;
+
+  @Setup(Level.Trial)
+  public void setUp()
+      throws IOException {
+    _baseDir = new File(FileUtils.getTempDirectory(), "BenchmarkSegmentDictionaryCreation");
+    FileUtils.deleteQuietly(_baseDir);
+    _mapDir = new File(_baseDir, "map");
+    _bsDir = new File(_baseDir, "bs");
+    FileUtils.forceMkdir(_mapDir);
+    FileUtils.forceMkdir(_bsDir);
+
+    // Build a sorted set of distinct strings of roughly _avgLength bytes; prefix the index to guarantee uniqueness.
+    _sortedValues = new String[_cardinality];
+    Random random = new Random(42);
+    for (int i = 0; i < _cardinality; i++) {
+      StringBuilder sb = new StringBuilder();
+      sb.append(i).append('_');
+      while (sb.length() < _avgLength) {
+        sb.append((char) ('a' + random.nextInt(26)));
+      }
+      _sortedValues[i] = sb.toString();
+    }
+    Arrays.sort(_sortedValues);
+
+    // A bounded set of lookup keys in pseudo-random ("row") order.
+    int numKeys = Math.min(_cardinality, 200_000);
+    _lookupKeys = new String[numKeys];
+    Random keyRandom = new Random(7);
+    for (int i = 0; i < numKeys; i++) {
+      _lookupKeys[i] = _sortedValues[keyRandom.nextInt(_cardinality)];
+    }
+
+    _mapCreator = new SegmentDictionaryCreator(STRING_FIELD, _mapDir, _useVarLength);
+    _mapCreator.setOnHeapLookupMapMaxCardinality(Integer.MAX_VALUE);
+    _mapCreator.build(_sortedValues);
+
+    _bsCreator = new SegmentDictionaryCreator(STRING_FIELD, _bsDir, _useVarLength);
+    _bsCreator.setOnHeapLookupMapMaxCardinality(0);
+    _bsCreator.build(_sortedValues);
+  }
+
+  @TearDown(Level.Trial)
+  public void tearDown()
+      throws IOException {
+    _mapCreator.close();
+    _bsCreator.close();
+    FileUtils.deleteQuietly(_baseDir);
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.AverageTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public long lookupMap() {
+    long sum = 0;
+    for (String key : _lookupKeys) {
+      sum += _mapCreator.indexOfSV(key);
+    }
+    return sum;
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.AverageTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public long lookupBinarySearch() {
+    long sum = 0;
+    for (String key : _lookupKeys) {
+      sum += _bsCreator.indexOfSV(key);
+    }
+    return sum;
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.SingleShotTime)
+  @OutputTimeUnit(TimeUnit.MILLISECONDS)
+  public int buildMap()
+      throws IOException {
+    try (SegmentDictionaryCreator creator = new SegmentDictionaryCreator(STRING_FIELD, _mapDir, _useVarLength)) {
+      creator.setOnHeapLookupMapMaxCardinality(Integer.MAX_VALUE);
+      creator.build(_sortedValues);
+      return creator.indexOfSV(_sortedValues[0]);
+    }
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.SingleShotTime)
+  @OutputTimeUnit(TimeUnit.MILLISECONDS)
+  public int buildBinarySearch()
+      throws IOException {
+    try (SegmentDictionaryCreator creator = new SegmentDictionaryCreator(STRING_FIELD, _bsDir, _useVarLength)) {
+      creator.setOnHeapLookupMapMaxCardinality(0);
+      creator.build(_sortedValues);
+      return creator.indexOfSV(_sortedValues[0]);
+    }
+  }
+
+  public static void main(String[] args)
+      throws Exception {
+    Options opt = new OptionsBuilder().include(BenchmarkSegmentDictionaryCreation.class.getSimpleName())
+        .warmupTime(TimeValue.seconds(3)).warmupIterations(3).measurementTime(TimeValue.seconds(3))
+        .measurementIterations(5).forks(1).build();
+    new Runner(opt).run();
+  }
+}

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java

@@ -206,6 +206,7 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio
           .withImmutableToMutableIdMap(immutableToMutableIdMap)
           .withRealtimeConversion(segmentCreationSpec.isRealtimeConversion())
           .withConsumerDir(segmentCreationSpec.getConsumerDir())
+          .withDictionaryOnHeapLookupMapMaxBytes(_config.getDictionaryOnHeapLookupMapMaxBytes())
           .build();
       //@formatter:on
 
@@ -266,6 +267,19 @@ public void init(SegmentGeneratorConfig segmentCreationSpec, SegmentIndexCreatio
       _creatorsByColAndIndex.put(columnName, creatorsByIndex);
     }
 
+    // All dictionaries and value-array-consuming indexes (e.g. FST) are now built. For columns whose dictionary
+    // switched to binary-search dictId lookup, the on-heap value->dictId map was never built, so the sorted
+    // unique-values array is the only remaining reference pinning the distinct values. Release it now (before the
+    // row-indexing pass) so the GC can reclaim those values; min/max/cardinality are cached on the collector.
+    for (Map.Entry<String, SegmentDictionaryCreator> entry : _dictionaryCreatorMap.entrySet()) {
+      if (entry.getValue().isUsingBinarySearchLookup()) {
+        ColumnIndexCreationInfo columnIndexCreationInfo = _indexCreationInfoMap.get(entry.getKey());
+        if (columnIndexCreationInfo != null) {
+          columnIndexCreationInfo.getColumnStatistics().releaseSortedValues();
+        }
+      }
+    }
+
     // Although NullValueVector is implemented as an index, it needs to be treated in a different way than other indexes
     for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) {
       String columnName = fieldSpec.getName();