From 8b144b6be0fa9359ca5643d72a787c2b2f095e0b Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Tue, 30 Jun 2026 00:52:35 -0700 Subject: [PATCH 1/5] Reduce dict-build heap: binary-search dictId lookup for high-card STRING/BYTES (3a) For high-cardinality STRING/BYTES columns, the on-heap value->dictId map (Object2IntOpenHashMap) is a dominant heap cost during offline segment generation and a frequent OOM source. Above a configurable cardinality threshold, skip building that map and resolve dictIds by binary-searching the just-written, mmap-backed dictionary buffer (via the same immutable dictionary reader used at load time) instead. Trades an O(1) on-heap lookup for an O(log N) off-heap read per value, removing the map from heap entirely. - Opt-in: threshold defaults to Integer.MAX_VALUE, so default behavior is unchanged. Decision is per-column, against that column's own cardinality. - Scoped to STRING/BYTES; numeric dictionaries keep their cheap primitive maps. - Lookup reader + mmap buffer are released in postIndexingCleanup()/close(). - 22 functional tests prove indexOfSV/indexOfMV return identical dictIds in map mode vs binary-search mode (STRING/BYTES x var/fixed x SV/MV), and the PinotBuffersAfterMethodCheckRule confirms no buffer leaks. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../impl/SegmentDictionaryCreator.java | 135 +++++++++- ...ctionaryCreatorBinarySearchLookupTest.java | 255 ++++++++++++++++++ 2 files changed, 377 insertions(+), 13 deletions(-) create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreatorBinarySearchLookupTest.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java index 2e4db2a6e6..84f2a2f69c 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreator.java @@ -35,6 +35,9 @@ import org.apache.pinot.segment.local.io.util.FixedByteValueReaderWriter; import org.apache.pinot.segment.local.io.util.VarLengthValueWriter; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; +import org.apache.pinot.segment.local.segment.index.readers.BaseImmutableDictionary; +import org.apache.pinot.segment.local.segment.index.readers.BytesDictionary; +import org.apache.pinot.segment.local.segment.index.readers.StringDictionary; import org.apache.pinot.segment.spi.index.IndexCreator; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.apache.pinot.spi.data.FieldSpec; @@ -71,6 +74,19 @@ public class SegmentDictionaryCreator implements IndexCreator { private Object2IntOpenHashMap _objectValueToIndexMap; private int _numBytesPerEntry = 0; + // For high-cardinality STRING/BYTES columns the on-heap value->dictId map (_objectValueToIndexMap) is the dominant + // heap cost during segment creation and a frequent OOM source. When the cardinality exceeds + // _onHeapLookupMapMaxCardinality, we skip building that map and instead resolve dictIds by binary-searching the + // just-written (off-heap, mmap-backed) dictionary buffer via an immutable dictionary reader. This trades an O(1) + // on-heap lookup for an O(log N) buffer read per value, removing the map from heap entirely. + // + // Defaults to Integer.MAX_VALUE (binary-search lookup OFF), so behavior is unchanged unless explicitly lowered. + // Must be set before build(). + private int _onHeapLookupMapMaxCardinality = Integer.MAX_VALUE; + private boolean _useBinarySearchLookup; + private PinotDataBuffer _lookupDictionaryBuffer; + private BaseImmutableDictionary _lookupDictionary; + public SegmentDictionaryCreator(String columnName, DataType storedType, File indexFile, boolean useVarLengthDictionary) { _columnName = columnName; @@ -102,6 +118,20 @@ public SegmentDictionaryCreator(FieldSpec fieldSpec, File indexDir) { this(fieldSpec, indexDir, false); } + /** + * Sets the cardinality above which STRING/BYTES dictionaries skip the on-heap value->dictId map and resolve + * dictIds by binary-searching the written dictionary buffer instead. Must be called before {@link #build(Object)}. + * The default ({@link Integer#MAX_VALUE}) keeps the on-heap map for all cardinalities (no behavior change). + */ + public void setOnHeapLookupMapMaxCardinality(int onHeapLookupMapMaxCardinality) { + _onHeapLookupMapMaxCardinality = onHeapLookupMapMaxCardinality; + } + + /** Whether {@link #build(Object)} skipped the on-heap value->dictId map in favor of binary-search lookups. */ + public boolean isUsingBinarySearchLookup() { + return _useBinarySearchLookup; + } + public void build(Object sortedValues) throws IOException { FileUtils.touch(_dictionaryFile); @@ -213,43 +243,61 @@ public void build(Object sortedValues) String[] sortedStrings = (String[]) sortedValues; numValues = sortedStrings.length; Preconditions.checkState(numValues > 0); - _objectValueToIndexMap = new Object2IntOpenHashMap<>(numValues); + _useBinarySearchLookup = numValues > _onHeapLookupMapMaxCardinality; + if (!_useBinarySearchLookup) { + _objectValueToIndexMap = new Object2IntOpenHashMap<>(numValues); + } // Get the maximum length of all entries byte[][] sortedStringBytes = new byte[numValues][]; for (int i = 0; i < numValues; i++) { String value = sortedStrings[i]; - _objectValueToIndexMap.put(value, i); + if (!_useBinarySearchLookup) { + _objectValueToIndexMap.put(value, i); + } byte[] valueBytes = value.getBytes(UTF_8); sortedStringBytes[i] = valueBytes; _numBytesPerEntry = Math.max(_numBytesPerEntry, valueBytes.length); } writeBytesValueDictionary(sortedStringBytes); + if (_useBinarySearchLookup) { + _lookupDictionary = openLookupDictionary(numValues); + } LOGGER.info( - "Created dictionary for STRING column: {} with cardinality: {}, max length in bytes: {}, range: {} to {}", - _columnName, numValues, _numBytesPerEntry, sortedStrings[0], sortedStrings[numValues - 1]); + "Created dictionary for STRING column: {} with cardinality: {}, max length in bytes: {}, range: {} to {}, " + + "binarySearchLookup: {}", _columnName, numValues, _numBytesPerEntry, sortedStrings[0], + sortedStrings[numValues - 1], _useBinarySearchLookup); return; case BYTES: ByteArray[] sortedBytes = (ByteArray[]) sortedValues; numValues = sortedBytes.length; Preconditions.checkState(numValues > 0); - _objectValueToIndexMap = new Object2IntOpenHashMap<>(numValues); + _useBinarySearchLookup = numValues > _onHeapLookupMapMaxCardinality; + if (!_useBinarySearchLookup) { + _objectValueToIndexMap = new Object2IntOpenHashMap<>(numValues); + } // Get the maximum length of all entries byte[][] sortedByteArrays = new byte[numValues][]; for (int i = 0; i < numValues; i++) { ByteArray value = sortedBytes[i]; sortedByteArrays[i] = value.getBytes(); - _objectValueToIndexMap.put(value, i); + if (!_useBinarySearchLookup) { + _objectValueToIndexMap.put(value, i); + } _numBytesPerEntry = Math.max(_numBytesPerEntry, value.getBytes().length); } writeBytesValueDictionary(sortedByteArrays); + if (_useBinarySearchLookup) { + _lookupDictionary = openLookupDictionary(numValues); + } LOGGER.info( - "Created dictionary for BYTES column: {} with cardinality: {}, max length in bytes: {}, range: {} to {}", - _columnName, numValues, _numBytesPerEntry, sortedBytes[0], sortedBytes[numValues - 1]); + "Created dictionary for BYTES column: {} with cardinality: {}, max length in bytes: {}, range: {} to {}, " + + "binarySearchLookup: {}", _columnName, numValues, _numBytesPerEntry, sortedBytes[0], + sortedBytes[numValues - 1], _useBinarySearchLookup); return; default: @@ -308,6 +356,29 @@ private void writeBytesValueDictionary(byte[][] bytesValues) } } + /** + * Opens an immutable dictionary reader over the dictionary buffer just written to {@link #_dictionaryFile}, used to + * resolve dictIds by binary search instead of the on-heap value->dictId map. The reader (and the mmap-backed + * buffer it owns) is released in {@link #closeLookupResources()}. Mirrors the load-time reader construction in + * {@code DictionaryIndexType.read}: dictionary buffers are big-endian, length is the cardinality, and + * numBytesPerValue is the max entry length (ignored for the var-length format, which is auto-detected from the + * buffer header). + */ + private BaseImmutableDictionary openLookupDictionary(int numValues) + throws IOException { + _lookupDictionaryBuffer = PinotDataBuffer.mapFile(_dictionaryFile, true, 0, _dictionaryFile.length(), + ByteOrder.BIG_ENDIAN, getClass().getSimpleName() + ".lookup"); + switch (_storedType) { + case STRING: + return new StringDictionary(_lookupDictionaryBuffer, numValues, _numBytesPerEntry); + case BYTES: + return new BytesDictionary(_lookupDictionaryBuffer, numValues, _numBytesPerEntry); + default: + throw new UnsupportedOperationException( + "Binary-search dictId lookup is not supported for data type: " + _storedType); + } + } + public int getNumBytesPerEntry() { return _numBytesPerEntry; } @@ -323,10 +394,13 @@ public int indexOfSV(Object value) { case DOUBLE: return _doubleValueToIndexMap.get((double) value); case STRING: + return _useBinarySearchLookup ? _lookupDictionary.indexOf((String) value) + : _objectValueToIndexMap.getInt(value); case BIG_DECIMAL: return _objectValueToIndexMap.getInt(value); case BYTES: - return _objectValueToIndexMap.getInt(new ByteArray((byte[]) value)); + return _useBinarySearchLookup ? ((BytesDictionary) _lookupDictionary).indexOf(new ByteArray((byte[]) value)) + : _objectValueToIndexMap.getInt(new ByteArray((byte[]) value)); default: throw new UnsupportedOperationException("Unsupported data type : " + _storedType); } @@ -358,13 +432,26 @@ public int[] indexOfMV(Object value) { } break; case STRING: - for (int i = 0; i < multiValues.length; i++) { - indexes[i] = _objectValueToIndexMap.getInt(multiValues[i]); + if (_useBinarySearchLookup) { + for (int i = 0; i < multiValues.length; i++) { + indexes[i] = _lookupDictionary.indexOf((String) multiValues[i]); + } + } else { + for (int i = 0; i < multiValues.length; i++) { + indexes[i] = _objectValueToIndexMap.getInt(multiValues[i]); + } } break; case BYTES: - for (int i = 0; i < multiValues.length; i++) { - indexes[i] = _objectValueToIndexMap.getInt(new ByteArray((byte[]) multiValues[i])); + if (_useBinarySearchLookup) { + BytesDictionary lookupDictionary = (BytesDictionary) _lookupDictionary; + for (int i = 0; i < multiValues.length; i++) { + indexes[i] = lookupDictionary.indexOf(new ByteArray((byte[]) multiValues[i])); + } + } else { + for (int i = 0; i < multiValues.length; i++) { + indexes[i] = _objectValueToIndexMap.getInt(new ByteArray((byte[]) multiValues[i])); + } } break; default: @@ -382,6 +469,27 @@ public void postIndexingCleanup() { _floatValueToIndexMap = null; _doubleValueToIndexMap = null; _objectValueToIndexMap = null; + closeLookupResources(); + } + + /** Releases the binary-search lookup dictionary reader and the mmap-backed buffer it owns. Idempotent. */ + private void closeLookupResources() { + if (_lookupDictionary != null) { + try { + _lookupDictionary.close(); + } catch (IOException e) { + LOGGER.warn("Caught exception while closing lookup dictionary for column: {}", _columnName, e); + } + _lookupDictionary = null; + } + if (_lookupDictionaryBuffer != null) { + try { + _lookupDictionaryBuffer.close(); + } catch (IOException e) { + LOGGER.warn("Caught exception while closing lookup dictionary buffer for column: {}", _columnName, e); + } + _lookupDictionaryBuffer = null; + } } @Override @@ -391,5 +499,6 @@ public void seal() { @Override public void close() { + closeLookupResources(); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreatorBinarySearchLookupTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreatorBinarySearchLookupTest.java new file mode 100644 index 0000000000..28cbfee189 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentDictionaryCreatorBinarySearchLookupTest.java @@ -0,0 +1,255 @@ +/** + * 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.segment.local.segment.creator.impl; + +import java.io.File; +import java.util.Arrays; +import java.util.Random; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.ByteArray; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/** + * Verifies that the binary-search dictId lookup path of {@link SegmentDictionaryCreator} (used for high-cardinality + * STRING/BYTES columns above the configured threshold) returns dictIds identical to the on-heap map path. The dictId of + * a value equals its position in the sorted dictionary, so both paths must agree with each other and with the index. + */ +public class SegmentDictionaryCreatorBinarySearchLookupTest implements PinotBuffersAfterMethodCheckRule { + private static final File TEMP_DIR = + new File(FileUtils.getTempDirectory(), "SegmentDictionaryCreatorBinarySearchLookupTest"); + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteQuietly(TEMP_DIR); + FileUtils.forceMkdir(TEMP_DIR); + } + + @AfterClass + public void tearDown() { + FileUtils.deleteQuietly(TEMP_DIR); + } + + @DataProvider(name = "stringConfigs") + public Object[][] stringConfigs() { + // {cardinality, avgLength, useVarLengthDictionary} + return new Object[][]{ + {1, 8, false}, {1, 8, true}, + {2, 8, false}, {2, 8, true}, + {16, 16, false}, {16, 16, true}, + {1000, 32, false}, {1000, 32, true}, + {5000, 64, false}, {5000, 64, true}, + {10000, 4, false}, {10000, 4, true} + }; + } + + @DataProvider(name = "bytesConfigs") + public Object[][] bytesConfigs() { + return new Object[][]{ + {1, 8, false}, {1, 8, true}, + {16, 16, false}, {16, 16, true}, + {1000, 32, false}, {1000, 32, true}, + {5000, 12, false}, {5000, 12, true} + }; + } + + /** + * For STRING: builds the dictionary twice from the same sorted values — once in on-heap map mode, once forced into + * binary-search mode — and asserts indexOfSV/indexOfMV agree with each other and with the sorted position for every + * value. + */ + @Test(dataProvider = "stringConfigs") + public void testStringMapVsBinarySearchEquivalence(int cardinality, int avgLength, boolean useVarLength) + throws Exception { + String[] sortedValues = generateSortedStrings(cardinality, avgLength); + FieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.STRING, true); + + // Map mode: threshold above cardinality keeps the on-heap map. + try (SegmentDictionaryCreator mapCreator = + new SegmentDictionaryCreator(fieldSpec, dir("str_map"), useVarLength)) { + mapCreator.setOnHeapLookupMapMaxCardinality(Integer.MAX_VALUE); + mapCreator.build(sortedValues); + assertFalse(mapCreator.isUsingBinarySearchLookup()); + + // Binary-search mode: threshold 0 forces binary search (cardinality > 0 always). + try (SegmentDictionaryCreator bsCreator = + new SegmentDictionaryCreator(fieldSpec, dir("str_bs"), useVarLength)) { + bsCreator.setOnHeapLookupMapMaxCardinality(0); + bsCreator.build(sortedValues); + assertTrue(bsCreator.isUsingBinarySearchLookup()); + + for (int i = 0; i < cardinality; i++) { + String value = sortedValues[i]; + assertEquals(mapCreator.indexOfSV(value), i, "map dictId mismatch at " + i); + assertEquals(bsCreator.indexOfSV(value), i, "binary-search dictId mismatch at " + i); + } + // Multi-value lookup over a shuffled subset. + Object[] mvValues = shuffledSubset(sortedValues); + int[] mapMv = mapCreator.indexOfMV(mvValues); + int[] bsMv = bsCreator.indexOfMV(mvValues); + assertEquals(bsMv, mapMv, "indexOfMV mismatch between map and binary-search modes"); + } + } + } + + /** + * For BYTES: same equivalence check as STRING. + */ + @Test(dataProvider = "bytesConfigs") + public void testBytesMapVsBinarySearchEquivalence(int cardinality, int avgLength, boolean useVarLength) + throws Exception { + ByteArray[] sortedValues = generateSortedByteArrays(cardinality, avgLength); + FieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.BYTES, true); + + try (SegmentDictionaryCreator mapCreator = + new SegmentDictionaryCreator(fieldSpec, dir("bytes_map"), useVarLength)) { + mapCreator.setOnHeapLookupMapMaxCardinality(Integer.MAX_VALUE); + mapCreator.build(sortedValues); + assertFalse(mapCreator.isUsingBinarySearchLookup()); + + try (SegmentDictionaryCreator bsCreator = + new SegmentDictionaryCreator(fieldSpec, dir("bytes_bs"), useVarLength)) { + bsCreator.setOnHeapLookupMapMaxCardinality(0); + bsCreator.build(sortedValues); + assertTrue(bsCreator.isUsingBinarySearchLookup()); + + for (int i = 0; i < cardinality; i++) { + byte[] value = sortedValues[i].getBytes(); + assertEquals(mapCreator.indexOfSV(value), i, "map dictId mismatch at " + i); + assertEquals(bsCreator.indexOfSV(value), i, "binary-search dictId mismatch at " + i); + } + Object[] mvValues = new Object[Math.min(cardinality, 32)]; + for (int i = 0; i < mvValues.length; i++) { + mvValues[i] = sortedValues[(i * 7) % cardinality].getBytes(); + } + int[] mapMv = mapCreator.indexOfMV(mvValues); + int[] bsMv = bsCreator.indexOfMV(mvValues); + assertEquals(bsMv, mapMv, "indexOfMV mismatch between map and binary-search modes"); + } + } + } + + /** + * The mode is decided strictly by {@code cardinality > threshold}, evaluated per column against that column's own + * cardinality. + */ + @Test + public void testThresholdBoundary() + throws Exception { + String[] sortedValues = generateSortedStrings(100, 16); + FieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.STRING, true); + + // threshold == cardinality -> NOT greater -> map mode. + try (SegmentDictionaryCreator creator = new SegmentDictionaryCreator(fieldSpec, dir("boundary_eq"), false)) { + creator.setOnHeapLookupMapMaxCardinality(100); + creator.build(sortedValues); + assertFalse(creator.isUsingBinarySearchLookup(), "cardinality == threshold should stay on map"); + } + + // threshold == cardinality - 1 -> greater -> binary search. + try (SegmentDictionaryCreator creator = new SegmentDictionaryCreator(fieldSpec, dir("boundary_lt"), false)) { + creator.setOnHeapLookupMapMaxCardinality(99); + creator.build(sortedValues); + assertTrue(creator.isUsingBinarySearchLookup(), "cardinality > threshold should use binary search"); + for (int i = 0; i < sortedValues.length; i++) { + assertEquals(creator.indexOfSV(sortedValues[i]), i); + } + } + } + + /** + * Default threshold keeps existing behavior (on-heap map) regardless of cardinality. + */ + @Test + public void testDefaultKeepsMapMode() + throws Exception { + String[] sortedValues = generateSortedStrings(1000, 16); + FieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.STRING, true); + try (SegmentDictionaryCreator creator = new SegmentDictionaryCreator(fieldSpec, dir("default_mode"), false)) { + creator.build(sortedValues); + assertFalse(creator.isUsingBinarySearchLookup(), "default must keep the on-heap map (no behavior change)"); + for (int i = 0; i < sortedValues.length; i++) { + assertEquals(creator.indexOfSV(sortedValues[i]), i); + } + } + } + + private File dir(String name) + throws Exception { + File dir = new File(TEMP_DIR, name); + FileUtils.deleteQuietly(dir); + FileUtils.forceMkdir(dir); + return dir; + } + + private static String[] generateSortedStrings(int cardinality, int avgLength) { + Random random = new Random(cardinality * 31L + avgLength); + String[] values = new String[cardinality]; + for (int i = 0; i < cardinality; i++) { + StringBuilder sb = new StringBuilder(); + // Prefix with the index to guarantee uniqueness, then pad with random chars to reach ~avgLength. + sb.append(i).append('_'); + while (sb.length() < avgLength) { + sb.append((char) ('a' + random.nextInt(26))); + } + values[i] = sb.toString(); + } + Arrays.sort(values); + return values; + } + + private static ByteArray[] generateSortedByteArrays(int cardinality, int avgLength) { + Random random = new Random(cardinality * 17L + avgLength); + ByteArray[] values = new ByteArray[cardinality]; + for (int i = 0; i < cardinality; i++) { + int len = Math.max(4, avgLength); + byte[] bytes = new byte[len]; + random.nextBytes(bytes); + // Encode the index into the leading bytes to guarantee uniqueness. + bytes[0] = (byte) (i >>> 24); + bytes[1] = (byte) (i >>> 16); + bytes[2] = (byte) (i >>> 8); + bytes[3] = (byte) i; + values[i] = new ByteArray(bytes); + } + Arrays.sort(values); + return values; + } + + private static Object[] shuffledSubset(String[] sortedValues) { + int size = Math.min(sortedValues.length, 32); + Object[] subset = new Object[size]; + for (int i = 0; i < size; i++) { + subset[i] = sortedValues[(i * 7) % sortedValues.length]; + } + return subset; + } +} From 84da1414739cbcabe17d40003a1f8ec2a8ffb8e5 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Tue, 30 Jun 2026 06:12:22 -0700 Subject: [PATCH 2/5] Release sorted unique-values array after dict build for binary-search columns (3b) 3a removes the on-heap value->dictId map for high-cardinality STRING/BYTES columns, but the distinct values stay pinned by the sorted unique-values array held in the stats collector through the whole creation. This frees that array right after init() (once the dictionary and any value-array-consuming indexes such as FST are built, before the row-indexing pass), so the GC can reclaim the distinct values. Together with 3a this drops a high-card column's dictionary heap to ~0. - ColumnStatistics gains a default no-op releaseSortedValues(); String/Bytes collectors override it to cache min/max/cardinality, then null _sortedValues. Getters return the cached scalars once released; numeric collectors keep the default no-op (3b is scoped to STRING/BYTES, the OOM-driving types). - ColumnIndexCreationInfo.getDistinctValueCount() falls back to the cached cardinality when the array is null (previously a dead null-branch). - SegmentColumnarIndexCreator.init() releases the array only for columns in binary-search mode (creator.isUsingBinarySearchLookup()). - Like 3a, dormant until a threshold config flips columns to binary search. - Tests: collector release keeps min/max/cardinality and ColumnIndexCreationInfo reporting; numeric release is a no-op; DictionariesTest regression green. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../impl/SegmentColumnarIndexCreator.java | 13 ++ .../BytesColumnPredIndexStatsCollector.java | 29 +++- .../StringColumnPreIndexStatsCollector.java | 29 +++- ...lumnStatisticsReleaseSortedValuesTest.java | 141 ++++++++++++++++++ .../spi/creator/ColumnIndexCreationInfo.java | 5 +- .../segment/spi/creator/ColumnStatistics.java | 12 +- 6 files changed, 221 insertions(+), 8 deletions(-) create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnStatisticsReleaseSortedValuesTest.java diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java index a55ef41dfd..5ab2eef9b9 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java @@ -266,6 +266,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 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(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPredIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPredIndexStatsCollector.java index 0fca8a3c77..4dcddbf057 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPredIndexStatsCollector.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/BytesColumnPredIndexStatsCollector.java @@ -36,6 +36,13 @@ public class BytesColumnPredIndexStatsCollector extends AbstractColumnStatistics private ByteArray[] _sortedValues; private boolean _sealed = false; + // Set when _sortedValues is released early (see releaseSortedValues()) to free heap. Once released, min/max/ + // cardinality are served from these cached scalars instead of the (now null) array. + private boolean _sortedValuesReleased = false; + private ByteArray _cachedMinValue; + private ByteArray _cachedMaxValue; + private int _cachedCardinality; + public BytesColumnPredIndexStatsCollector(String column, StatsCollectorConfig statsCollectorConfig) { super(column, statsCollectorConfig); } @@ -77,7 +84,7 @@ public void collect(Object entry) { @Override public ByteArray getMinValue() { if (_sealed) { - return _sortedValues[0]; + return _sortedValuesReleased ? _cachedMinValue : _sortedValues[0]; } throw new IllegalStateException("you must seal the collector first before asking for min value"); } @@ -85,7 +92,7 @@ public ByteArray getMinValue() { @Override public ByteArray getMaxValue() { if (_sealed) { - return _sortedValues[_sortedValues.length - 1]; + return _sortedValuesReleased ? _cachedMaxValue : _sortedValues[_sortedValues.length - 1]; } throw new IllegalStateException("you must seal the collector first before asking for max value"); } @@ -93,6 +100,7 @@ public ByteArray getMaxValue() { @Override public ByteArray[] getUniqueValuesSet() { if (_sealed) { + // null after releaseSortedValues(); callers needing only the count use getCardinality() instead. return _sortedValues; } throw new IllegalStateException("you must seal the collector first before asking for unique values set"); @@ -115,7 +123,10 @@ public int getMaxRowLengthInBytes() { @Override public int getCardinality() { - return _sealed ? _sortedValues.length : _values.size(); + if (!_sealed) { + return _values.size(); + } + return _sortedValuesReleased ? _cachedCardinality : _sortedValues.length; } @Override @@ -127,4 +138,16 @@ public void seal() { _sealed = true; } } + + @Override + public void releaseSortedValues() { + if (_sealed && !_sortedValuesReleased) { + // Cache the scalars still derived from the array, then drop the array so its ByteArray objects can be reclaimed. + _cachedCardinality = _sortedValues.length; + _cachedMinValue = _sortedValues[0]; + _cachedMaxValue = _sortedValues[_sortedValues.length - 1]; + _sortedValues = null; + _sortedValuesReleased = true; + } + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StringColumnPreIndexStatsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StringColumnPreIndexStatsCollector.java index d43adbe3c7..304dfab8c9 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StringColumnPreIndexStatsCollector.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/StringColumnPreIndexStatsCollector.java @@ -42,6 +42,13 @@ public class StringColumnPreIndexStatsCollector extends AbstractColumnStatistics private boolean _sealed = false; private CLPStatsCollector _clpStatsCollector; + // Set when _sortedValues is released early (see releaseSortedValues()) to free heap. Once released, min/max/ + // cardinality are served from these cached scalars instead of the (now null) array. + private boolean _sortedValuesReleased = false; + private String _cachedMinValue; + private String _cachedMaxValue; + private int _cachedCardinality; + public StringColumnPreIndexStatsCollector(String column, StatsCollectorConfig statsCollectorConfig) { super(column, statsCollectorConfig); if (_fieldConfig != null && _fieldConfig.getCompressionCodec() == FieldConfig.CompressionCodec.CLP) { @@ -102,7 +109,7 @@ public CLPStats getCLPStats() { @Override public String getMinValue() { if (_sealed) { - return _sortedValues[0]; + return _sortedValuesReleased ? _cachedMinValue : _sortedValues[0]; } throw new IllegalStateException("you must seal the collector first before asking for min value"); } @@ -110,7 +117,7 @@ public String getMinValue() { @Override public String getMaxValue() { if (_sealed) { - return _sortedValues[_sortedValues.length - 1]; + return _sortedValuesReleased ? _cachedMaxValue : _sortedValues[_sortedValues.length - 1]; } throw new IllegalStateException("you must seal the collector first before asking for max value"); } @@ -118,6 +125,7 @@ public String getMaxValue() { @Override public Object[] getUniqueValuesSet() { if (_sealed) { + // null after releaseSortedValues(); callers needing only the count use getCardinality() instead. return _sortedValues; } throw new IllegalStateException("you must seal the collector first before asking for unique values set"); @@ -135,7 +143,10 @@ public int getMaxRowLengthInBytes() { @Override public int getCardinality() { - return _sealed ? _sortedValues.length : _values.size(); + if (!_sealed) { + return _values.size(); + } + return _sortedValuesReleased ? _cachedCardinality : _sortedValues.length; } @Override @@ -151,6 +162,18 @@ public void seal() { } } + @Override + public void releaseSortedValues() { + if (_sealed && !_sortedValuesReleased) { + // Cache the scalars still derived from the array, then drop the array so its String objects can be reclaimed. + _cachedCardinality = _sortedValues.length; + _cachedMinValue = _sortedValues[0]; + _cachedMaxValue = _sortedValues[_sortedValues.length - 1]; + _sortedValues = null; + _sortedValuesReleased = true; + } + } + @VisibleForTesting public static class CLPStatsCollector { private final EncodedMessage _clpEncodedMessage; diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnStatisticsReleaseSortedValuesTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnStatisticsReleaseSortedValuesTest.java new file mode 100644 index 0000000000..2a34b7c9dd --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/stats/ColumnStatisticsReleaseSortedValuesTest.java @@ -0,0 +1,141 @@ +/** + * 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.segment.local.segment.creator.impl.stats; + +import org.apache.pinot.segment.spi.creator.ColumnIndexCreationInfo; +import org.apache.pinot.segment.spi.creator.StatsCollectorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; + + +/** + * Verifies {@link org.apache.pinot.segment.spi.creator.ColumnStatistics#releaseSortedValues()} (3b): after the sorted + * unique-values array is released to free heap, min/max/cardinality are still served from cached scalars, the array + * accessor returns null, and {@link ColumnIndexCreationInfo} keeps reporting the correct cardinality/min/max. + */ +public class ColumnStatisticsReleaseSortedValuesTest { + private static final TableConfig TABLE_CONFIG = + new TableConfigBuilder(TableType.OFFLINE).setTableName("test").build(); + + private StatsCollectorConfig statsConfig(String column, DataType dataType) { + Schema schema = new Schema(); + schema.addField(new DimensionFieldSpec(column, dataType, true)); + return new StatsCollectorConfig(TABLE_CONFIG, schema, null); + } + + @Test + public void testStringReleaseKeepsMinMaxCardinality() { + StringColumnPreIndexStatsCollector collector = + new StringColumnPreIndexStatsCollector("col", statsConfig("col", DataType.STRING)); + for (String v : new String[]{"banana", "apple", "cherry", "apple"}) { + collector.collect(v); + } + collector.seal(); + + // Before release: sorted array present, scalars derived from it. + assertNotNull(collector.getUniqueValuesSet()); + assertEquals(collector.getCardinality(), 3); + assertEquals(collector.getMinValue(), "apple"); + assertEquals(collector.getMaxValue(), "cherry"); + + collector.releaseSortedValues(); + + // After release: array gone, scalars served from cache. + assertNull(collector.getUniqueValuesSet(), "sorted array should be released"); + assertEquals(collector.getCardinality(), 3); + assertEquals(collector.getMinValue(), "apple"); + assertEquals(collector.getMaxValue(), "cherry"); + + // ColumnIndexCreationInfo must keep reporting the right values via the cached cardinality fallback. + ColumnIndexCreationInfo info = new ColumnIndexCreationInfo(collector, true, false, false, "null"); + assertEquals(info.getDistinctValueCount(), 3); + assertEquals(info.getMin(), "apple"); + assertEquals(info.getMax(), "cherry"); + assertNull(info.getSortedUniqueElementsArray()); + + // Idempotent. + collector.releaseSortedValues(); + assertEquals(collector.getCardinality(), 3); + } + + @Test + public void testBytesReleaseKeepsMinMaxCardinality() { + BytesColumnPredIndexStatsCollector collector = + new BytesColumnPredIndexStatsCollector("col", statsConfig("col", DataType.BYTES)); + byte[] a = {1, 2}; + byte[] b = {1, 3}; + byte[] c = {9, 0}; + for (byte[] v : new byte[][]{b, a, c, a}) { + collector.collect(v); + } + collector.seal(); + + assertNotNull(collector.getUniqueValuesSet()); + assertEquals(collector.getCardinality(), 3); + ByteArray min = collector.getMinValue(); + ByteArray max = collector.getMaxValue(); + assertEquals(min, new ByteArray(a)); + assertEquals(max, new ByteArray(c)); + + collector.releaseSortedValues(); + + assertNull(collector.getUniqueValuesSet(), "sorted array should be released"); + assertEquals(collector.getCardinality(), 3); + assertEquals(collector.getMinValue(), min); + assertEquals(collector.getMaxValue(), max); + + ColumnIndexCreationInfo info = new ColumnIndexCreationInfo(collector, true, true, false, new byte[0]); + assertEquals(info.getDistinctValueCount(), 3); + assertEquals(info.getMin(), min); + assertEquals(info.getMax(), max); + assertNull(info.getSortedUniqueElementsArray()); + } + + /** + * Numeric collectors inherit the default no-op {@code releaseSortedValues()} (3b is scoped to STRING/BYTES), so their + * sorted array must remain intact after the call. + */ + @Test + public void testNumericReleaseIsNoOp() { + IntColumnPreIndexStatsCollector collector = + new IntColumnPreIndexStatsCollector("col", statsConfig("col", DataType.INT)); + for (int v : new int[]{5, 1, 9, 1}) { + collector.collect(v); + } + collector.seal(); + + collector.releaseSortedValues(); + + // Default no-op: array still present and intact. + assertNotNull(collector.getUniqueValuesSet(), "numeric collector should not release its array"); + assertEquals(collector.getCardinality(), 3); + assertEquals(collector.getMinValue(), 1); + assertEquals(collector.getMaxValue(), 9); + } +} diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnIndexCreationInfo.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnIndexCreationInfo.java index 9e47502098..ff613fd4bb 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnIndexCreationInfo.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnIndexCreationInfo.java @@ -70,7 +70,10 @@ public Object getSortedUniqueElementsArray() { public int getDistinctValueCount() { Object uniqueValArray = _columnStatistics.getUniqueValuesSet(); if (uniqueValArray == null) { - return Constants.UNKNOWN_CARDINALITY; + // The sorted unique-values array may have been released after dictionary build to free heap + // (see ColumnStatistics#releaseSortedValues()); in that case the cardinality is cached on the collector. + int cardinality = _columnStatistics.getCardinality(); + return cardinality > 0 ? cardinality : Constants.UNKNOWN_CARDINALITY; } return ArrayUtils.getLength(uniqueValArray); } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnStatistics.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnStatistics.java index e904fd0a9a..b4159f8325 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnStatistics.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/ColumnStatistics.java @@ -40,10 +40,20 @@ public interface ColumnStatistics extends Serializable { /** * - * @return An array of elements that has the unique values for this column, sorted order. + * @return An array of elements that has the unique values for this column, sorted order. May return {@code null} + * after {@link #releaseSortedValues()} has been called to free heap during segment creation. */ Object getUniqueValuesSet(); + /** + * Releases the on-heap sorted unique-values array (the structure returned by {@link #getUniqueValuesSet()}) to free + * heap once it is no longer needed (i.e. after the dictionary and any value-array-consuming indexes are built). + * Implementations that support this must first cache anything still derived from the array — {@link #getMinValue()}, + * {@link #getMaxValue()}, {@link #getCardinality()} — so those keep working afterwards. Default is a no-op. + */ + default void releaseSortedValues() { + } + /** * * @return The number of unique values of this column. From 54a642592345ede08c32dd93eb3826d59850a32b Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Tue, 30 Jun 2026 07:38:04 -0700 Subject: [PATCH 3/5] Add dictionary heap/CPU benchmarks for binary-search dictId lookup (3a/3b) - BenchmarkDictionaryHeapFootprint: JOL GraphLayout retained size of the map-mode structures (sorted String[] + Object2IntOpenHashMap) across cardinality x avg length. Confirms the on-heap footprint binary-search mode eliminates, and validates the ~(60+L) bytes/value model. - BenchmarkSegmentDictionaryCreation: JMH lookupMap vs lookupBinarySearch per-row dictId cost, plus buildMap/buildBinarySearch, swept over cardinality x avg length x var/fixed encoding. - Adds jol-core (0.17) to pinot-perf. Measured (JOL): bytes/value tracks 60+L within a few percent (e.g. 1M x L=32 = 88 MB, 3M x L=100 = 455 MB), all eliminated in binary-search mode. Measured (JMH, indicative): binary search is a ~30-55x per-lookup penalty vs the hash map (constant factor, no cardinality crossover), i.e. ~2-4 s extra encode per fat column on a 5M-row segment -- the bounded CPU price for the heap saved. Co-Authored-By: Claude Opus 4.8 (1M context) --- pinot-perf/pom.xml | 5 + .../BenchmarkDictionaryHeapFootprint.java | 94 +++++++++ .../BenchmarkSegmentDictionaryCreation.java | 186 ++++++++++++++++++ 3 files changed, 285 insertions(+) create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkDictionaryHeapFootprint.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkSegmentDictionaryCreation.java diff --git a/pinot-perf/pom.xml b/pinot-perf/pom.xml index 61c6bec92e..0c2d8f287a 100644 --- a/pinot-perf/pom.xml +++ b/pinot-perf/pom.xml @@ -105,6 +105,11 @@ net.sf.jopt-simple jopt-simple + + org.openjdk.jol + jol-core + 0.17 + diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkDictionaryHeapFootprint.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkDictionaryHeapFootprint.java new file mode 100644 index 0000000000..4852b4a792 --- /dev/null +++ b/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 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; + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkSegmentDictionaryCreation.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkSegmentDictionaryCreation.java new file mode 100644 index 0000000000..bc1f2b9a98 --- /dev/null +++ b/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(); + } +} From 91331762a900d9562164cabb0ed445ae90ca233d Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Wed, 1 Jul 2026 02:15:09 -0700 Subject: [PATCH 4/5] Wire dict-build low-heap path to a per-table byte budget (Layer 1 config) 3a/3b added a per-column binary-search dictId lookup (skipping the on-heap value->dictId map) plus early release of the sorted unique-values array for high-cardinality STRING/BYTES columns, but the flip was driven only by SegmentDictionaryCreator._onHeapLookupMapMaxCardinality (default Integer.MAX_VALUE), which nothing in production sets -- so the feature was dormant. This wires it to a real, per-table config so it can fire end-to-end, gated on an estimated byte budget rather than a flat cardinality count. New knob IndexingConfig.dictionaryOnHeapLookupMapMaxBytes (a size in bytes; default Long.MAX_VALUE = effectively disabled, so behavior is unchanged unless set). A dict-enabled STRING/BYTES column flips to the low-heap path when its estimated on-heap footprint (60 + maxLen) * cardinality >= the budget, where cardinality = numValues and maxLen = the column's longest value in bytes (the L proxy from the (60 + L) per-value model measured in the design doc). A byte budget (vs a flat cardinality) correctly targets long-string columns, which a flat count under-targets. Plumbing: IndexingConfig -> SegmentGeneratorConfig -> IndexCreationContext -> DictionaryIndexType.createIndexCreator. The gate is applied by converting the byte budget to the equivalent per-column cardinality threshold (ceil(budget / (60 + maxLen)) - 1) and feeding the existing, already-tested setOnHeapLookupMapMaxCardinality(int) setter, so the tested lookup path in SegmentDictionaryCreator is untouched. No-op for numeric columns (cheap primitive maps, never the OOM driver) and when the budget is disabled. Compiles clean (pinot-spi, pinot-segment-spi, pinot-segment-local). Tests (gate-wiring unit test + end-to-end build/load test) to follow in a subsequent commit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../impl/SegmentColumnarIndexCreator.java | 1 + .../index/dictionary/DictionaryIndexType.java | 51 ++++++++++++++++++- .../spi/creator/IndexCreationContext.java | 26 +++++++++- .../spi/creator/SegmentGeneratorConfig.java | 10 ++++ .../spi/config/table/IndexingConfig.java | 30 +++++++++++ 5 files changed, 115 insertions(+), 3 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java index 5ab2eef9b9..356d87e157 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/SegmentColumnarIndexCreator.java +++ b/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 diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java index ed04808e14..02c65ee51a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java @@ -89,6 +89,13 @@ public class DictionaryIndexType private static final Logger LOGGER = LoggerFactory.getLogger(DictionaryIndexType.class); private static final List EXTENSIONS = Collections.singletonList(V1Constants.Dict.FILE_EXTENSION); + // Approximate on-heap overhead (bytes) per distinct value of the structures that the low-heap dictionary-creation + // path eliminates for high-cardinality STRING/BYTES columns: the value->dictId map entry, the sorted-array slot, and + // the String/ByteArray object headers -- everything except the value's own text bytes. Together with the column's + // max value length L this gives the measured per-value footprint (60 + L); see design doc "offheap-dictionary- + // creation" (sections 11 and 12). Used to convert the byte budget into a per-column cardinality threshold. + private static final long APPROX_ON_HEAP_LOOKUP_BYTES_PER_VALUE_OVERHEAD = 60L; + protected DictionaryIndexType() { super(StandardIndexes.DICTIONARY_ID); } @@ -157,7 +164,49 @@ protected ColumnConfigDeserializer createDeserializerForL @Override public SegmentDictionaryCreator createIndexCreator(IndexCreationContext context, DictionaryIndexConfig indexConfig) { boolean useVarLengthDictionary = shouldUseVarLengthDictionary(context, indexConfig); - return new SegmentDictionaryCreator(context.getFieldSpec(), context.getIndexDir(), useVarLengthDictionary); + SegmentDictionaryCreator creator = + new SegmentDictionaryCreator(context.getFieldSpec(), context.getIndexDir(), useVarLengthDictionary); + configureOnHeapLookupMap(creator, context); + return creator; + } + + /** + * For high-cardinality STRING/BYTES columns, configures the dictionary creator to skip the on-heap value->dictId + * map (and let the sorted unique-values array be released early) when the column's estimated on-heap footprint + * {@code (60 + maxLen) * cardinality} reaches the configured byte budget + * ({@link IndexCreationContext#getDictionaryOnHeapLookupMapMaxBytes()}); above the budget dictIds are resolved by + * binary-searching the mmap-backed dictionary buffer instead. No-op (default behavior) for numeric columns or when + * the budget is disabled ({@link Long#MAX_VALUE} or non-positive). + */ + private static void configureOnHeapLookupMap(SegmentDictionaryCreator creator, IndexCreationContext context) { + long maxBytesBudget = context.getDictionaryOnHeapLookupMapMaxBytes(); + if (maxBytesBudget <= 0 || maxBytesBudget == Long.MAX_VALUE) { + // Disabled: keep the on-heap map for all cardinalities (no behavior change). + return; + } + DataType storedType = context.getFieldSpec().getDataType().getStoredType(); + if (storedType != DataType.STRING && storedType != DataType.BYTES) { + // Numeric dictionaries use cheap primitive maps and are never the OOM driver; leave them untouched. + return; + } + // Per-value on-heap footprint of the eliminated structures: 60 bytes overhead + the value's max length in bytes. + long perValueBytes = + APPROX_ON_HEAP_LOOKUP_BYTES_PER_VALUE_OVERHEAD + Math.max(context.getLengthOfLongestEntry(), 0); + // We want binary-search lookup exactly when perValueBytes * cardinality >= budget, i.e. + // cardinality >= ceil(budget / perValueBytes). The creator flips on a strict cardinality > threshold, so + // threshold = ceil(budget / perValueBytes) - 1. Computed without overflow (budget < Long.MAX_VALUE here). + long ceilCardinality = maxBytesBudget / perValueBytes + (maxBytesBudget % perValueBytes == 0 ? 0 : 1); + long thresholdCardinality = ceilCardinality - 1; + int clampedThreshold; + if (thresholdCardinality >= Integer.MAX_VALUE) { + // No int cardinality can exceed this, so the map is always kept (effectively disabled for this column). + clampedThreshold = Integer.MAX_VALUE; + } else if (thresholdCardinality < 0) { + clampedThreshold = 0; + } else { + clampedThreshold = (int) thresholdCardinality; + } + creator.setOnHeapLookupMapMaxCardinality(clampedThreshold); } public boolean shouldUseVarLengthDictionary(IndexCreationContext context, DictionaryIndexConfig indexConfig) { diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/IndexCreationContext.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/IndexCreationContext.java index 5d3a1a78ba..932f756521 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/IndexCreationContext.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/IndexCreationContext.java @@ -55,6 +55,15 @@ public interface IndexCreationContext { int getTotalDocs(); + /** + * Byte budget gating the low-heap dictionary-creation path for high-cardinality STRING/BYTES columns (see + * {@link IndexingConfig#getDictionaryOnHeapLookupMapMaxBytes()}). When a column's estimated on-heap footprint + * {@code (60 + maxLen) * cardinality} is {@code >=} this value, the dictionary creator skips the on-heap + * value->dictId map and resolves dictIds by binary-searching the mmap-backed dictionary buffer instead. Defaults + * to {@link Long#MAX_VALUE} (optimization disabled). + */ + long getDictionaryOnHeapLookupMapMaxBytes(); + boolean hasDictionary(); Comparable getMinValue(); @@ -123,6 +132,7 @@ final class Builder { private int _cardinality; private int _totalNumberOfEntries; private int _totalDocs; + private long _dictionaryOnHeapLookupMapMaxBytes = IndexingConfig.DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES; private boolean _hasDictionary = true; private Comparable _minValue; private Comparable _maxValue; @@ -226,6 +236,11 @@ public Builder withTotalDocs(int totalDocs) { return this; } + public Builder withDictionaryOnHeapLookupMapMaxBytes(long dictionaryOnHeapLookupMapMaxBytes) { + _dictionaryOnHeapLookupMapMaxBytes = dictionaryOnHeapLookupMapMaxBytes; + return this; + } + public Builder withDictionary(boolean hasDictionary) { _hasDictionary = hasDictionary; return this; @@ -271,7 +286,7 @@ public Common build() { _maxRowLengthInBytes, _onHeap, Objects.requireNonNull(_fieldSpec), _sorted, _cardinality, _totalNumberOfEntries, _totalDocs, _hasDictionary, _minValue, _maxValue, _forwardIndexDisabled, _sortedUniqueElementsArray, _optimizedDictionary, _fixedLength, _textCommitOnClose, _columnStatistics, - _realtimeConversion, _consumerDir, _immutableToMutableIdMap); + _realtimeConversion, _consumerDir, _immutableToMutableIdMap, _dictionaryOnHeapLookupMapMaxBytes); } public Builder withSortedUniqueElementsArray(Object sortedUniqueElementsArray) { @@ -308,6 +323,7 @@ final class Common implements IndexCreationContext { private final boolean _realtimeConversion; private final File _consumerDir; private final int[] _immutableToMutableIdMap; + private final long _dictionaryOnHeapLookupMapMaxBytes; public Common(File indexDir, int lengthOfLongestEntry, int maxNumberOfMultiValueElements, int maxRowLengthInBytes, boolean onHeap, @@ -315,7 +331,7 @@ public Common(File indexDir, int lengthOfLongestEntry, int totalDocs, boolean hasDictionary, Comparable minValue, Comparable maxValue, boolean forwardIndexDisabled, Object sortedUniqueElementsArray, boolean optimizeDictionary, boolean fixedLength, boolean textCommitOnClose, ColumnStatistics columnStatistics, boolean realtimeConversion, File consumerDir, - int[] immutableToMutableIdMap) { + int[] immutableToMutableIdMap, long dictionaryOnHeapLookupMapMaxBytes) { _indexDir = indexDir; _lengthOfLongestEntry = lengthOfLongestEntry; _maxNumberOfMultiValueElements = maxNumberOfMultiValueElements; @@ -338,6 +354,7 @@ public Common(File indexDir, int lengthOfLongestEntry, _realtimeConversion = realtimeConversion; _consumerDir = consumerDir; _immutableToMutableIdMap = immutableToMutableIdMap; + _dictionaryOnHeapLookupMapMaxBytes = dictionaryOnHeapLookupMapMaxBytes; } public FieldSpec getFieldSpec() { @@ -380,6 +397,11 @@ public int getTotalDocs() { return _totalDocs; } + @Override + public long getDictionaryOnHeapLookupMapMaxBytes() { + return _dictionaryOnHeapLookupMapMaxBytes; + } + public boolean hasDictionary() { return _hasDictionary; } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java index 84614a52c4..1145fc87cd 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java @@ -111,6 +111,7 @@ public enum TimeColumnType { private boolean _optimizeDictionaryType = false; private double _noDictionarySizeRatioThreshold = IndexingConfig.DEFAULT_NO_DICTIONARY_SIZE_RATIO_THRESHOLD; private Double _noDictionaryCardinalityRatioThreshold; + private long _dictionaryOnHeapLookupMapMaxBytes = IndexingConfig.DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES; private boolean _realtimeConversion = false; // consumerDir contains data from the consuming segment, and is used during _realtimeConversion optimization private File _consumerDir; @@ -159,6 +160,7 @@ public SegmentGeneratorConfig(TableConfig tableConfig, Schema schema) { _optimizeDictionaryType = indexingConfig.isOptimizeDictionaryType(); _noDictionarySizeRatioThreshold = indexingConfig.getNoDictionarySizeRatioThreshold(); _noDictionaryCardinalityRatioThreshold = indexingConfig.getNoDictionaryCardinalityRatioThreshold(); + _dictionaryOnHeapLookupMapMaxBytes = indexingConfig.getDictionaryOnHeapLookupMapMaxBytes(); // Star-tree configs setStarTreeIndexConfigs(indexingConfig.getStarTreeIndexConfigs()); @@ -663,6 +665,14 @@ public void setNoDictionaryCardinalityRatioThreshold(@Nullable Double noDictiona _noDictionaryCardinalityRatioThreshold = noDictionaryCardinalityRatioThreshold; } + public long getDictionaryOnHeapLookupMapMaxBytes() { + return _dictionaryOnHeapLookupMapMaxBytes; + } + + public void setDictionaryOnHeapLookupMapMaxBytes(long dictionaryOnHeapLookupMapMaxBytes) { + _dictionaryOnHeapLookupMapMaxBytes = dictionaryOnHeapLookupMapMaxBytes; + } + public boolean isFailOnEmptySegment() { return _failOnEmptySegment; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/IndexingConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/IndexingConfig.java index 9c5a8c3454..51c55472f4 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/IndexingConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/IndexingConfig.java @@ -33,6 +33,12 @@ public class IndexingConfig extends BaseJsonConfig { // Default ratio for overriding dictionary for fixed width columns public static final double DEFAULT_NO_DICTIONARY_SIZE_RATIO_THRESHOLD = 0.85d; + // Default byte budget for the on-heap dictionary value->dictId lookup map during segment creation. Long.MAX_VALUE + // effectively disables the optimization: the estimated on-heap footprint (60 + maxLen) * cardinality can never reach + // it, so no column flips to binary-search lookup (default behavior unchanged). See + // _dictionaryOnHeapLookupMapMaxBytes. + public static final long DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES = Long.MAX_VALUE; + /** * This should be equal to the one specified in RangeIndexType. */ @@ -97,6 +103,22 @@ public class IndexingConfig extends BaseJsonConfig { // then create a dictionary for the column. A value around 0.1 (10%) is a reasonable starting point private Double _noDictionaryCardinalityRatioThreshold; + /** + * Byte budget that gates the low-heap dictionary-creation path for high-cardinality STRING/BYTES columns. During + * offline segment generation, a dictionary-encoded column normally keeps an on-heap value->dictId map plus the + * sorted unique-values array alive through the whole build; for high-cardinality columns this is a dominant heap + * cost and a frequent OOM source. When a column's estimated on-heap footprint {@code (60 + maxLen) * cardinality} + * (bytes; {@code maxLen} = the column's longest value in bytes, {@code cardinality} = number of distinct values) is + * {@code >=} this budget, the map is skipped and dictIds are resolved by binary-searching the (already off-heap, + * mmap-backed) dictionary buffer instead, and the sorted array is released early. Applies only to STRING/BYTES + * columns (numeric dictionaries use cheap primitive maps and are never the OOM driver). The value is a size in + * bytes; both sides of the comparison are byte sizes (e.g. cardinality 2,000,000 with maxLen 30 estimates + * {@code (60 + 30) * 2,000,000 ~= 172 MB}). Defaults to {@link Long#MAX_VALUE}, which effectively disables the + * optimization (no behavior change); set it to a smaller byte value to enable -- e.g. 134217728 for a 128 MB budget. + * Note a tiny value such as {@code 100} means 100 bytes and would flip essentially every column. + */ + private long _dictionaryOnHeapLookupMapMaxBytes = DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES; + // TODO: Add a new configuration related to the segment generation private boolean _autoGeneratedInvertedIndex; private String _segmentNameGeneratorType; @@ -413,6 +435,14 @@ public void setNoDictionaryCardinalityRatioThreshold(Double noDictionaryCardinal _noDictionaryCardinalityRatioThreshold = noDictionaryCardinalityRatioThreshold; } + public long getDictionaryOnHeapLookupMapMaxBytes() { + return _dictionaryOnHeapLookupMapMaxBytes; + } + + public void setDictionaryOnHeapLookupMapMaxBytes(long dictionaryOnHeapLookupMapMaxBytes) { + _dictionaryOnHeapLookupMapMaxBytes = dictionaryOnHeapLookupMapMaxBytes; + } + public String getSegmentNameGeneratorType() { return _segmentNameGeneratorType; } From c9c6e093e21e26996afd53c7e22aa7e390515414 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Sat, 4 Jul 2026 10:56:53 -0700 Subject: [PATCH 5/5] Add tests for the dict-build byte-budget gate (Layer 1) Covers the byte-budget wiring added in the previous commit, which flips a high-cardinality STRING/BYTES column to the low-heap binary-search dictId lookup (3a) + early sorted-array release (3b) when its estimated on-heap footprint (60 + maxLen) * cardinality >= the configured budget. DictionaryIndexTypeOnHeapLookupBudgetTest (gate-wiring unit test): asserts the per-column flip decision made by DictionaryIndexType.createIndexCreator matches the byte-budget formula exactly at the boundary (one below stays on the on-heap map, at/above flips), across several (budget, maxLen) points, for STRING and BYTES and for both single-value (SV) and multi-value (MV) columns -- the switch is scoped to the stored type, so it must flip identically for SV and MV. Also verifies that every non-STRING/BYTES stored type (INT, LONG, FLOAT, DOUBLE, BIG_DECIMAL), SV and MV, never flips even at a tiny budget, and that the default/disabled budget (Long.MAX_VALUE, <= 0) keeps map mode -- including that a fresh, unconfigured IndexingConfig defaults to disabled (no behavior change). This nails the ceil/off-by-one conversion from byte budget to the creator's strict cardinality > threshold. SegmentGenerationWithDictionaryOnHeapLookupBudgetTest (end-to-end): builds a real segment carrying all four in-scope column types -- SV STRING, MV STRING, SV BYTES, MV BYTES (plus a LONG metric) -- and, because the on-disk segment is identical whether or not a column flips, directly confirms which path each column took by reading SegmentDictionaryCreator.isUsingBinarySearchLookup() off the driver's columnar index creator. It asserts the switch fired for every STRING/BYTES column when enabled, fired for none under the default/disabled config, and that in both cases every row round-trips through the forward-index -> dictionary path (incl. MV element order and BYTES values) with correct min/max/cardinality metadata. A parity build (ON vs OFF) asserts the two paths produce identical row values and metadata, so a bug in the binary-search dictId encoding for any SV/MV STRING/BYTES column would surface. The direct fired-flag observation is what makes these correctness checks non-vacuous. All new tests pass (18 gate-wiring + 4 end-to-end); existing dictionary/context tests (DictionariesTest, SegmentDictionaryCreatorBinarySearchLookupTest, ColumnStatisticsReleaseSortedValuesTest, DictionaryIndexTypeTest, InvertedIndexVersionConfigTest) remain green. Spotless + checkstyle clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...nWithDictionaryOnHeapLookupBudgetTest.java | 359 ++++++++++++++++++ ...ionaryIndexTypeOnHeapLookupBudgetTest.java | 305 +++++++++++++++ 2 files changed, 664 insertions(+) create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithDictionaryOnHeapLookupBudgetTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexTypeOnHeapLookupBudgetTest.java diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithDictionaryOnHeapLookupBudgetTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithDictionaryOnHeapLookupBudgetTest.java new file mode 100644 index 0000000000..51633c7173 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/SegmentGenerationWithDictionaryOnHeapLookupBudgetTest.java @@ -0,0 +1,359 @@ +/** + * 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.segment.local.segment.index.creator; + +import java.io.File; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentColumnarIndexCreator; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader; +import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/** + * End-to-end test for the low-heap dictionary-creation path (3a binary-search dictId lookup + 3b early sorted-array + * release) driven by the {@code dictionaryOnHeapLookupMapMaxBytes} byte budget. + * + *

The switch is scoped to STRING/BYTES and must apply to both single-value (SV) and multi-value (MV) columns, so + * this segment carries all four: SV STRING, MV STRING, SV BYTES, MV BYTES (plus a LONG metric that must be untouched). + * + *

The optimization is behavior-preserving by design: whether or not a column flips to binary-search lookup, the + * on-disk segment is identical -- so a plain "build a correct segment" check cannot tell the flipped path even ran. + * This test therefore builds the same data twice -- once with the budget set low enough to flip every high-cardinality + * STRING/BYTES column, once with the feature disabled (default) -- loads both, and asserts (a) every row round-trips + * correctly through the forward-index -> dictionary path for all column types, (b) min/max/cardinality metadata are + * correct, and (c) the two builds are identical (parity). {@link org.apache.pinot.segment.local.segment.index + * .dictionary.DictionaryIndexTypeOnHeapLookupBudgetTest} separately proves the chosen budget/cardinality actually flips + * STRING/BYTES SV and MV columns (and never numeric), so this correctness check is not vacuous. + */ +public class SegmentGenerationWithDictionaryOnHeapLookupBudgetTest implements PinotBuffersAfterMethodCheckRule { + private static final File TEMP_DIR = + new File(FileUtils.getTempDirectory(), "SegmentGenerationWithDictionaryOnHeapLookupBudgetTest"); + + private static final String TABLE_NAME = "testTable"; + private static final String SV_STRING_COLUMN = "svString"; + private static final String MV_STRING_COLUMN = "mvString"; + private static final String SV_BYTES_COLUMN = "svBytes"; + private static final String MV_BYTES_COLUMN = "mvBytes"; + private static final String LONG_COLUMN = "longMetric"; + + // The in-scope (STRING/BYTES) dictionary-encoded columns that the switch may flip. + private static final String[] STRING_BYTES_COLUMNS = + {SV_STRING_COLUMN, MV_STRING_COLUMN, SV_BYTES_COLUMN, MV_BYTES_COLUMN}; + // All dictionary-encoded columns plus the LONG metric, for full row round-trip. + private static final String[] ALL_COLUMNS = + {SV_STRING_COLUMN, MV_STRING_COLUMN, SV_BYTES_COLUMN, MV_BYTES_COLUMN, LONG_COLUMN}; + + private static final int NUM_ROWS = 5000; + private static final int BYTES_LENGTH = 16; + // All values distinct: SV columns -> cardinality NUM_ROWS; MV columns (2 distinct values/row) -> cardinality 2*ROWS. + private static final int SV_CARDINALITY = NUM_ROWS; + private static final int MV_CARDINALITY = 2 * NUM_ROWS; + + // Must match DictionaryIndexType.APPROX_ON_HEAP_LOOKUP_BYTES_PER_VALUE_OVERHEAD. + private static final long PER_VALUE_OVERHEAD = 60L; + // Budget small enough that even the smallest cardinality column (60 * NUM_ROWS bytes) comfortably exceeds it, so + // every STRING/BYTES column -- SV and MV -- flips. + private static final long FLIP_BUDGET_BYTES = 100_000L; + + private Schema _schema; + private List _rows; + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteQuietly(TEMP_DIR); + FileUtils.forceMkdir(TEMP_DIR); + _schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME) + .addSingleValueDimension(SV_STRING_COLUMN, DataType.STRING) + .addMultiValueDimension(MV_STRING_COLUMN, DataType.STRING) + .addSingleValueDimension(SV_BYTES_COLUMN, DataType.BYTES) + .addMultiValueDimension(MV_BYTES_COLUMN, DataType.BYTES) + .addMetric(LONG_COLUMN, DataType.LONG) + .build(); + _rows = generateRows(); + } + + @AfterClass + public void tearDown() { + FileUtils.deleteQuietly(TEMP_DIR); + } + + /** + * Guards the premise of this test: with the chosen parameters every high-cardinality STRING/BYTES column's estimated + * on-heap footprint must exceed {@link #FLIP_BUDGET_BYTES}, so the low-heap path is actually exercised by the ON + * build. Uses the 60-byte per-value overhead as a lower bound (ignoring maxLen), so this holds regardless of the + * exact value lengths. + */ + @Test + public void testChosenParametersExceedGate() { + assertTrue(PER_VALUE_OVERHEAD * SV_CARDINALITY >= FLIP_BUDGET_BYTES, + "SV columns must flip: 60 * " + SV_CARDINALITY + " < budget " + FLIP_BUDGET_BYTES); + assertTrue(PER_VALUE_OVERHEAD * MV_CARDINALITY >= FLIP_BUDGET_BYTES, + "MV columns must flip: 60 * " + MV_CARDINALITY + " < budget " + FLIP_BUDGET_BYTES); + } + + /** + * Builds a real segment with the budget set low enough to flip every high-cardinality STRING/BYTES column (SV and MV) + * to binary-search lookup, loads it, and asserts every row's values (all column types) and the SV STRING column's + * min/max/cardinality metadata are correct, plus the cardinality of the other dictionary-encoded columns. + */ + @Test + public void testHighCardColumnsBinarySearchPathProducesCorrectSegment() + throws Exception { + SegmentIndexCreationDriverImpl driver = buildSegment("segment_on", FLIP_BUDGET_BYTES); + // Directly confirm the switch fired for every in-scope STRING/BYTES column in the real build. + assertSwitchFired(driver, true); + ImmutableSegment segment = ImmutableSegmentLoader.load(driver.getOutputDirectory(), ReadMode.mmap); + try { + assertRowsRoundTrip(segment); + assertMetadata(segment); + } finally { + segment.destroy(); + } + } + + /** + * The out-of-the-box path (knob left at its {@code Long.MAX_VALUE} default) is a complete no-op: the switch fires for + * no column, and the segment is still built correctly (row round-trip + metadata). Confirms disabling means "no + * behavior change", observed directly via the fired flag rather than merely inferred. + */ + @Test + public void testDefaultConfigIsNoOp() + throws Exception { + SegmentIndexCreationDriverImpl driver = buildSegmentWithDefaultConfig("segment_default"); + assertSwitchFired(driver, false); + ImmutableSegment segment = ImmutableSegmentLoader.load(driver.getOutputDirectory(), ReadMode.mmap); + try { + assertRowsRoundTrip(segment); + assertMetadata(segment); + } finally { + segment.destroy(); + } + } + + /** + * Parity: the segment built with the feature ON (STRING/BYTES columns flipped) must be identical -- row values and + * metadata -- to the segment built with the feature disabled (default budget). This is the core proof that the + * binary-search dictId encoding is correct for both SV and MV STRING/BYTES: any divergence in the forward index or + * dictionary would surface here. + */ + @Test + public void testBinarySearchPathMatchesOnHeapMapPath() + throws Exception { + SegmentIndexCreationDriverImpl onDriver = buildSegment("parity_on", FLIP_BUDGET_BYTES); + SegmentIndexCreationDriverImpl offDriver = buildSegment("parity_off", /*disabled*/ Long.MAX_VALUE); + // Confirm the two builds genuinely took different paths: ON flipped every in-scope column, OFF flipped none. + assertSwitchFired(onDriver, true); + assertSwitchFired(offDriver, false); + + ImmutableSegment onSegment = ImmutableSegmentLoader.load(onDriver.getOutputDirectory(), ReadMode.mmap); + ImmutableSegment offSegment = ImmutableSegmentLoader.load(offDriver.getOutputDirectory(), ReadMode.mmap); + try { + try (PinotSegmentRecordReader onReader = new PinotSegmentRecordReader(); + PinotSegmentRecordReader offReader = new PinotSegmentRecordReader()) { + onReader.init(onSegment); + offReader.init(offSegment); + for (int i = 0; i < NUM_ROWS; i++) { + GenericRow onRow = onReader.next(); + GenericRow offRow = offReader.next(); + for (String column : ALL_COLUMNS) { + assertValueEquals(onRow.getValue(column), offRow.getValue(column), + column + " mismatch between ON and OFF builds at row " + i); + } + } + } + + for (String column : STRING_BYTES_COLUMNS) { + ColumnMetadata onMeta = onSegment.getSegmentMetadata().getColumnMetadataFor(column); + ColumnMetadata offMeta = offSegment.getSegmentMetadata().getColumnMetadataFor(column); + assertEquals(onMeta.getCardinality(), offMeta.getCardinality(), column + " cardinality mismatch"); + assertEquals(onMeta.getMinValue(), offMeta.getMinValue(), column + " min mismatch"); + assertEquals(onMeta.getMaxValue(), offMeta.getMaxValue(), column + " max mismatch"); + } + } finally { + onSegment.destroy(); + offSegment.destroy(); + } + } + + private void assertRowsRoundTrip(ImmutableSegment segment) + throws Exception { + try (PinotSegmentRecordReader recordReader = new PinotSegmentRecordReader()) { + recordReader.init(segment); + for (int i = 0; i < NUM_ROWS; i++) { + GenericRow actual = recordReader.next(); + GenericRow expected = _rows.get(i); + for (String column : ALL_COLUMNS) { + assertValueEquals(actual.getValue(column), expected.getValue(column), + column + " value mismatch at row " + i); + } + } + } + } + + private void assertMetadata(ImmutableSegment segment) { + assertEquals(segment.getSegmentMetadata().getTotalDocs(), NUM_ROWS, "totalDocs mismatch"); + + ColumnMetadata svString = segment.getSegmentMetadata().getColumnMetadataFor(SV_STRING_COLUMN); + assertEquals(svString.getCardinality(), SV_CARDINALITY, "SV STRING cardinality mismatch"); + assertEquals(svString.getMinValue(), svString(0), "SV STRING min mismatch"); + assertEquals(svString.getMaxValue(), svString(NUM_ROWS - 1), "SV STRING max mismatch"); + + assertEquals(segment.getSegmentMetadata().getColumnMetadataFor(MV_STRING_COLUMN).getCardinality(), MV_CARDINALITY, + "MV STRING cardinality mismatch"); + assertEquals(segment.getSegmentMetadata().getColumnMetadataFor(SV_BYTES_COLUMN).getCardinality(), SV_CARDINALITY, + "SV BYTES cardinality mismatch"); + assertEquals(segment.getSegmentMetadata().getColumnMetadataFor(MV_BYTES_COLUMN).getCardinality(), MV_CARDINALITY, + "MV BYTES cardinality mismatch"); + } + + /** + * Value equality that handles multi-value columns (read back as {@code Object[]}) and BYTES columns (read back as + * {@code byte[]}), recursing into MV elements. + */ + private static void assertValueEquals(Object actual, Object expected, String message) { + if (expected instanceof Object[]) { + Object[] expectedArray = (Object[]) expected; + Object[] actualArray = (Object[]) actual; + assertEquals(actualArray.length, expectedArray.length, message + " (MV length)"); + for (int i = 0; i < expectedArray.length; i++) { + assertValueEquals(actualArray[i], expectedArray[i], message + " [" + i + "]"); + } + } else if (expected instanceof byte[]) { + assertTrue(Arrays.equals((byte[]) actual, (byte[]) expected), message + " (bytes)"); + } else { + assertEquals(actual, expected, message); + } + } + + private SegmentIndexCreationDriverImpl buildSegment(String segmentName, long dictionaryOnHeapLookupMapMaxBytes) + throws Exception { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build(); + tableConfig.getIndexingConfig().setDictionaryOnHeapLookupMapMaxBytes(dictionaryOnHeapLookupMapMaxBytes); + return buildSegment(segmentName, tableConfig); + } + + /** Builds with a completely unconfigured table (knob left at its default) -- the true out-of-the-box path. */ + private SegmentIndexCreationDriverImpl buildSegmentWithDefaultConfig(String segmentName) + throws Exception { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build(); + return buildSegment(segmentName, tableConfig); + } + + private SegmentIndexCreationDriverImpl buildSegment(String segmentName, TableConfig tableConfig) + throws Exception { + SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, _schema); + config.setOutDir(new File(TEMP_DIR, segmentName + "_out").getAbsolutePath()); + config.setSegmentName(segmentName); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config, new GenericRowRecordReader(_rows)); + driver.build(); + return driver; + } + + /** + * Directly observes, through the real build pipeline, whether each in-scope STRING/BYTES column's dictionary was + * built with the binary-search switch fired. The segment is byte-identical either way, so this reads the decision + * from the retained {@link SegmentDictionaryCreator#isUsingBinarySearchLookup()} on the driver's columnar index + * creator (reflection, since the driver doesn't expose it). This is what makes the correctness/parity checks + * non-vacuous: it proves the ON build actually flipped and the disabled build actually did not. + */ + private static void assertSwitchFired(SegmentIndexCreationDriverImpl driver, boolean expectedFired) + throws Exception { + Field indexCreatorField = SegmentIndexCreationDriverImpl.class.getDeclaredField("_indexCreator"); + indexCreatorField.setAccessible(true); + SegmentColumnarIndexCreator indexCreator = (SegmentColumnarIndexCreator) indexCreatorField.get(driver); + + Field mapField = SegmentColumnarIndexCreator.class.getDeclaredField("_dictionaryCreatorMap"); + mapField.setAccessible(true); + @SuppressWarnings("unchecked") + Map creators = + (Map) mapField.get(indexCreator); + + for (String column : STRING_BYTES_COLUMNS) { + SegmentDictionaryCreator creator = creators.get(column); + assertNotNull(creator, "no dictionary creator retained for column " + column); + assertEquals(creator.isUsingBinarySearchLookup(), expectedFired, + column + " isUsingBinarySearchLookup expected " + expectedFired); + } + } + + private List generateRows() { + List rows = new ArrayList<>(NUM_ROWS); + for (int i = 0; i < NUM_ROWS; i++) { + GenericRow row = new GenericRow(); + // All values distinct and, for the SV STRING column, fixed-width and lexicographically ordered by index, so + // cardinalities and SV STRING min/max are deterministic. Rows are emitted in index order (doc order for unsorted + // columns), so the record reader reads them back in the same order. + row.putValue(SV_STRING_COLUMN, svString(i)); + row.putValue(MV_STRING_COLUMN, new Object[]{mvString(i, 'a'), mvString(i, 'b')}); + row.putValue(SV_BYTES_COLUMN, bytesOf(i, (byte) 0)); + row.putValue(MV_BYTES_COLUMN, new Object[]{bytesOf(i, (byte) 1), bytesOf(i, (byte) 2)}); + row.putValue(LONG_COLUMN, (long) i); + rows.add(row); + } + return rows; + } + + private static String svString(int i) { + return String.format("val_%07d", i); + } + + private static String mvString(int i, char suffix) { + return String.format("mv%c_%07d", suffix, i); + } + + private static byte[] bytesOf(int i, byte marker) { + byte[] bytes = new byte[BYTES_LENGTH]; + // Leading marker distinguishes the SV column from the two MV elements; the encoded index makes every value unique. + bytes[0] = marker; + bytes[1] = (byte) (i >>> 24); + bytes[2] = (byte) (i >>> 16); + bytes[3] = (byte) (i >>> 8); + bytes[4] = (byte) i; + return bytes; + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexTypeOnHeapLookupBudgetTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexTypeOnHeapLookupBudgetTest.java new file mode 100644 index 0000000000..02c555c725 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexTypeOnHeapLookupBudgetTest.java @@ -0,0 +1,305 @@ +/** + * 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.segment.local.segment.index.dictionary; + +import java.io.File; +import java.math.BigDecimal; +import java.util.Arrays; +import java.util.Random; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentDictionaryCreator; +import org.apache.pinot.segment.spi.creator.IndexCreationContext; +import org.apache.pinot.segment.spi.index.DictionaryIndexConfig; +import org.apache.pinot.spi.config.table.IndexingConfig; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.ByteArray; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/** + * Verifies the byte-budget gate that {@link DictionaryIndexType#createIndexCreator(IndexCreationContext, + * DictionaryIndexConfig)} wires onto {@link SegmentDictionaryCreator} for high-cardinality STRING/BYTES columns. + * + *

The feature (3a binary-search dictId lookup + 3b early sorted-array release) fires when a column's estimated + * on-heap footprint {@code (60 + maxLen) * cardinality} reaches the configured budget + * ({@link IndexingConfig#getDictionaryOnHeapLookupMapMaxBytes()}). {@code createIndexCreator} converts that byte budget + * into a per-column cardinality threshold and feeds it to the creator; the creator then flips on + * {@code cardinality > threshold} during {@link SegmentDictionaryCreator#build(Object)}. These tests assert the flip + * decision (observed via {@link SegmentDictionaryCreator#isUsingBinarySearchLookup()}) matches the byte-budget formula + * exactly at the boundary, is a no-op for numeric columns, and stays off when the budget is disabled -- so + * {@link SegmentGenerationWithDictionaryOnHeapLookupBudgetTest}'s end-to-end correctness check is not vacuous. + */ +public class DictionaryIndexTypeOnHeapLookupBudgetTest implements PinotBuffersAfterMethodCheckRule { + private static final File TEMP_DIR = + new File(FileUtils.getTempDirectory(), "DictionaryIndexTypeOnHeapLookupBudgetTest"); + + // Must match DictionaryIndexType.APPROX_ON_HEAP_LOOKUP_BYTES_PER_VALUE_OVERHEAD. + private static final long PER_VALUE_OVERHEAD = 60L; + + private final DictionaryIndexType _indexType = new DictionaryIndexPlugin().getIndexType(); + + @BeforeClass + public void setUp() + throws Exception { + FileUtils.deleteQuietly(TEMP_DIR); + FileUtils.forceMkdir(TEMP_DIR); + } + + @AfterClass + public void tearDown() { + FileUtils.deleteQuietly(TEMP_DIR); + } + + /** + * At a fixed (budget, maxLen), the column flips exactly at the smallest cardinality whose estimated footprint + * {@code (60 + maxLen) * cardinality} reaches the budget: one below stays on the on-heap map, at/above flips to + * binary-search lookup. + */ + @DataProvider(name = "budgetConfigs") + public Object[][] budgetConfigs() { + // {budgetBytes, maxLen, singleValue}. boundaryCardinality is derived so both sides of the boundary are exercised. + // Each (budget, maxLen) is run for both single-value (SV) and multi-value (MV) columns: the gate keys off the + // stored type only, so STRING/BYTES must flip identically whether SV or MV. + long[][] budgets = { + {1000L, 40L}, // perValue 100 -> boundary cardinality 10 + {1000L, 15L}, // perValue 75 -> boundary cardinality ceil(1000/75) = 14 + {10_000L, 90L}, // perValue 150 -> boundary cardinality ceil(10000/150) = 67 + {512L, 4L} // perValue 64 -> boundary cardinality 8 + }; + Object[][] configs = new Object[budgets.length * 2][]; + int idx = 0; + for (long[] budget : budgets) { + configs[idx++] = new Object[]{budget[0], budget[1], true}; + configs[idx++] = new Object[]{budget[0], budget[1], false}; + } + return configs; + } + + @Test(dataProvider = "budgetConfigs") + public void testStringBoundary(long budgetBytes, long maxLen, boolean singleValue) + throws Exception { + long perValue = PER_VALUE_OVERHEAD + maxLen; + int boundaryCardinality = (int) ((budgetBytes + perValue - 1) / perValue); // ceil(budget / perValue) + + // One below the boundary: footprint < budget -> keep the on-heap map. + assertStringFlip((int) maxLen, budgetBytes, boundaryCardinality - 1, singleValue, false); + // At the boundary: footprint >= budget -> flip to binary-search lookup. + assertStringFlip((int) maxLen, budgetBytes, boundaryCardinality, singleValue, true); + // Comfortably above the boundary: still flipped. + assertStringFlip((int) maxLen, budgetBytes, boundaryCardinality + 5, singleValue, true); + } + + @Test(dataProvider = "budgetConfigs") + public void testBytesBoundary(long budgetBytes, long maxLen, boolean singleValue) + throws Exception { + long perValue = PER_VALUE_OVERHEAD + maxLen; + int boundaryCardinality = (int) ((budgetBytes + perValue - 1) / perValue); + + assertBytesFlip((int) maxLen, budgetBytes, boundaryCardinality - 1, singleValue, false); + assertBytesFlip((int) maxLen, budgetBytes, boundaryCardinality, singleValue, true); + } + + /** + * The switch is scoped to STRING/BYTES: every other stored type (numeric primitives and BIG_DECIMAL) uses cheap maps + * and must never flip, regardless of the budget or single-value/multi-value, even at a tiny budget that would flip an + * equivalent STRING/BYTES column. + */ + @Test + public void testNonStringBytesTypesNeverFlip() + throws Exception { + DataType[] types = {DataType.INT, DataType.LONG, DataType.FLOAT, DataType.DOUBLE, DataType.BIG_DECIMAL}; + for (DataType type : types) { + for (boolean singleValue : new boolean[]{true, false}) { + // BIG_DECIMAL is a single-value-only stored type; skip the MV case for it. + if (type == DataType.BIG_DECIMAL && !singleValue) { + continue; + } + FieldSpec fieldSpec = new DimensionFieldSpec("col", type, singleValue); + // Tiny budget (1 byte) that would flip an equivalent STRING/BYTES column of any cardinality. + IndexCreationContext context = + context(fieldSpec, dir("nonflip_" + type + "_" + (singleValue ? "sv" : "mv")), 8, /*budget*/ 1L, 1000); + try (SegmentDictionaryCreator creator = _indexType.createIndexCreator(context, DictionaryIndexConfig.DEFAULT)) { + creator.build(sortedValuesFor(type, 1000)); + assertFalse(creator.isUsingBinarySearchLookup(), + type + (singleValue ? " (SV)" : " (MV)") + " must never flip to binary-search lookup"); + } + } + } + } + + /** + * The default budget ({@link Long#MAX_VALUE}) disables the optimization: the estimated footprint can never reach it, + * so a STRING column stays on the on-heap map regardless of cardinality (no behavior change). + */ + @Test + public void testDisabledBudgetKeepsMapMode() + throws Exception { + // Out-of-the-box (unconfigured) tables must have the feature disabled -> no behavior change by default. + assertEquals(new IndexingConfig().getDictionaryOnHeapLookupMapMaxBytes(), + IndexingConfig.DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES); + assertEquals(IndexingConfig.DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES, Long.MAX_VALUE); + + long disabledBudget = IndexingConfig.DEFAULT_DICTIONARY_ON_HEAP_LOOKUP_MAP_MAX_BYTES; + assertStringFlip(32, disabledBudget, 100_000, true, false); + // A non-positive budget is likewise treated as disabled. + assertStringFlip(32, 0L, 100_000, true, false); + assertStringFlip(32, -1L, 100_000, true, false); + } + + private void assertStringFlip(int maxLen, long budgetBytes, int cardinality, boolean singleValue, + boolean expectedFlip) + throws Exception { + FieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.STRING, singleValue); + IndexCreationContext context = + context(fieldSpec, dir("str_" + budgetBytes + "_" + maxLen + "_" + cardinality + "_" + singleValue), maxLen, + budgetBytes, cardinality); + try (SegmentDictionaryCreator creator = _indexType.createIndexCreator(context, DictionaryIndexConfig.DEFAULT)) { + creator.build(generateSortedStrings(cardinality, maxLen)); + assertFlip(creator, expectedFlip, "STRING" + (singleValue ? " (SV)" : " (MV)"), budgetBytes, maxLen, cardinality); + } + } + + private void assertBytesFlip(int maxLen, long budgetBytes, int cardinality, boolean singleValue, boolean expectedFlip) + throws Exception { + FieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.BYTES, singleValue); + IndexCreationContext context = + context(fieldSpec, dir("bytes_" + budgetBytes + "_" + maxLen + "_" + cardinality + "_" + singleValue), maxLen, + budgetBytes, cardinality); + try (SegmentDictionaryCreator creator = _indexType.createIndexCreator(context, DictionaryIndexConfig.DEFAULT)) { + creator.build(generateSortedByteArrays(cardinality, maxLen)); + assertFlip(creator, expectedFlip, "BYTES" + (singleValue ? " (SV)" : " (MV)"), budgetBytes, maxLen, cardinality); + } + } + + private static void assertFlip(SegmentDictionaryCreator creator, boolean expectedFlip, String type, long budgetBytes, + int maxLen, int cardinality) { + String message = String.format( + "%s column with budget=%d, maxLen=%d, cardinality=%d: estimated footprint (60+%d)*%d=%d %s budget %d", type, + budgetBytes, maxLen, cardinality, maxLen, cardinality, (PER_VALUE_OVERHEAD + maxLen) * cardinality, + expectedFlip ? ">=" : "<", budgetBytes); + if (expectedFlip) { + assertTrue(creator.isUsingBinarySearchLookup(), "expected binary-search lookup: " + message); + } else { + assertFalse(creator.isUsingBinarySearchLookup(), "expected on-heap map: " + message); + } + } + + private IndexCreationContext context(FieldSpec fieldSpec, File indexDir, int maxLen, long budgetBytes, + int cardinality) { + return IndexCreationContext.builder() + .withIndexDir(indexDir) + .withFieldSpec(fieldSpec) + .withLengthOfLongestEntry(maxLen) + .withCardinality(cardinality) + .withDictionaryOnHeapLookupMapMaxBytes(budgetBytes) + .build(); + } + + private File dir(String name) + throws Exception { + File dir = new File(TEMP_DIR, name); + FileUtils.deleteQuietly(dir); + FileUtils.forceMkdir(dir); + return dir; + } + + private static Object sortedValuesFor(DataType type, int cardinality) { + switch (type) { + case INT: { + int[] values = new int[cardinality]; + for (int i = 0; i < cardinality; i++) { + values[i] = i; + } + return values; + } + case LONG: { + long[] values = new long[cardinality]; + for (int i = 0; i < cardinality; i++) { + values[i] = i; + } + return values; + } + case FLOAT: { + float[] values = new float[cardinality]; + for (int i = 0; i < cardinality; i++) { + values[i] = i; + } + return values; + } + case DOUBLE: { + double[] values = new double[cardinality]; + for (int i = 0; i < cardinality; i++) { + values[i] = i; + } + return values; + } + case BIG_DECIMAL: { + BigDecimal[] values = new BigDecimal[cardinality]; + for (int i = 0; i < cardinality; i++) { + values[i] = BigDecimal.valueOf(i); + } + return values; + } + default: + throw new IllegalArgumentException("Unsupported non-STRING/BYTES type: " + type); + } + } + + private static String[] generateSortedStrings(int cardinality, int maxLen) { + Random random = new Random(cardinality * 31L + maxLen); + String[] values = new String[cardinality]; + for (int i = 0; i < cardinality; i++) { + StringBuilder sb = new StringBuilder(); + sb.append(i).append('_'); + while (sb.length() < maxLen) { + sb.append((char) ('a' + random.nextInt(26))); + } + values[i] = sb.toString(); + } + Arrays.sort(values); + return values; + } + + private static ByteArray[] generateSortedByteArrays(int cardinality, int maxLen) { + Random random = new Random(cardinality * 17L + maxLen); + ByteArray[] values = new ByteArray[cardinality]; + for (int i = 0; i < cardinality; i++) { + int len = Math.max(4, maxLen); + byte[] bytes = new byte[len]; + random.nextBytes(bytes); + bytes[0] = (byte) (i >>> 24); + bytes[1] = (byte) (i >>> 16); + bytes[2] = (byte) (i >>> 8); + bytes[3] = (byte) i; + values[i] = new ByteArray(bytes); + } + Arrays.sort(values); + return values; + } +}