Skip to content

feat(upsert): add primary-key map size-in-bytes gauge #18980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tarun11Mavani wants to merge 2 commits into
base: master
Choose a base branch
from
Open

feat(upsert): add primary-key map size-in-bytes gauge #18980

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7c95ccd
feat(upsert): add primary-key map size-in-bytes gauge
tarun11Mavani Jul 13, 2026
b660f7c
Keep per-record primary-key gauge update count-only
tarun11Mavani Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public enum ServerGauge implements AbstractMetrics.Gauge {
PAUSELESS_CONSUMPTION_ENABLED("pauselessConsumptionEnabled", false),
// Upsert metrics
UPSERT_PRIMARY_KEYS_COUNT("upsertPrimaryKeysCount", false),
UPSERT_PRIMARY_KEY_MAP_SIZE_IN_BYTES("bytes", false),
// Dedup metrics
DEDUP_PRIMARY_KEYS_COUNT("dedupPrimaryKeysCount", false),
CONSUMPTION_QUOTA_UTILIZATION("ratio", false),
Expand Down
5 changes: 5 additions & 0 deletions .../test/java/org/apache/pinot/integration/tests/models/DummyTableUpsertMetadataManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,11 @@ protected long getNumPrimaryKeys() {
return 0;
}

@Override
protected long getPrimaryKeyMapSizeInBytes() {
return 0;
}

@Override
protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds,
@Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator<RecordInfo> recordInfoIterator,
Expand Down
23 changes: 17 additions & 6 deletions ...c/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -388,20 +388,31 @@ protected void doAddSegment(ImmutableSegmentImpl segment) {

// Update metrics
long numPrimaryKeys = getNumPrimaryKeys();
updatePrimaryKeyGauge(numPrimaryKeys);
updatePrimaryKeyGauge(numPrimaryKeys, getPrimaryKeyMapSizeInBytes());
_logger.info("Finished adding segment: {} in {}ms, current primary key count: {}", segmentName,
System.currentTimeMillis() - startTimeMs, numPrimaryKeys);
}

protected abstract long getNumPrimaryKeys();

/**
* Returns the estimated size in bytes of the primary-key-to-record-location map.
*/
protected abstract long getPrimaryKeyMapSizeInBytes();

protected void updatePrimaryKeyGauge(long numPrimaryKeys, long primaryKeySizeInBytes) {
updatePrimaryKeyGauge(numPrimaryKeys);
_serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId,
ServerGauge.UPSERT_PRIMARY_KEY_MAP_SIZE_IN_BYTES, primaryKeySizeInBytes);
}

protected void updatePrimaryKeyGauge(long numPrimaryKeys) {
_serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId, ServerGauge.UPSERT_PRIMARY_KEYS_COUNT,
numPrimaryKeys);
}

protected void updatePrimaryKeyGauge() {
updatePrimaryKeyGauge(getNumPrimaryKeys());
updatePrimaryKeyGauge(getNumPrimaryKeys(), getPrimaryKeyMapSizeInBytes());

@xiangfu0 xiangfu0 Jul 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updatePrimaryKeyGauge() is still called from the realtime doAddRecord() paths, so this now runs getPrimaryKeyMapSizeInBytes() for every ingested upsert record. That samples up to 1000 keys and calls PrimaryKey.asBytes() while Pinot is on the per-record ingestion path, adding bounded but repeated map iteration/allocation to a hot path. Please keep the per-record update count-only and refresh the byte-size gauge from segment lifecycle, a rate-limited path, or background sampling instead.

@tarun11Mavani tarun11Mavani Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks.
split the per-record path onto a count-only overload; segment-lifecycle events (add/preload/replace/remove/close) and the periodic TTL eviction path still refresh both gauges.

}

@Override
Expand Down Expand Up @@ -459,7 +470,7 @@ protected void doPreloadSegment(ImmutableSegmentImpl segment) {

// Update metrics
long numPrimaryKeys = getNumPrimaryKeys();
updatePrimaryKeyGauge(numPrimaryKeys);
updatePrimaryKeyGauge(numPrimaryKeys, getPrimaryKeyMapSizeInBytes());
_logger.info("Finished preloading segment: {} in {}ms, current primary key count: {}", segmentName,
System.currentTimeMillis() - startTimeMs, numPrimaryKeys);
}
Expand Down Expand Up @@ -645,7 +656,7 @@ protected void doReplaceSegment(ImmutableSegment segment, IndexSegment oldSegmen

// Update metrics
long numPrimaryKeys = getNumPrimaryKeys();
updatePrimaryKeyGauge(numPrimaryKeys);
updatePrimaryKeyGauge(numPrimaryKeys, getPrimaryKeyMapSizeInBytes());
_logger.info("Finished replacing segment: {} in {}ms, current primary key count: {}", segmentName,
System.currentTimeMillis() - startTimeMs, numPrimaryKeys);
}
Expand Down Expand Up @@ -812,7 +823,7 @@ protected void doRemoveSegment(IndexSegment segment) {

// Update metrics
long numPrimaryKeys = getNumPrimaryKeys();
updatePrimaryKeyGauge(numPrimaryKeys);
updatePrimaryKeyGauge(numPrimaryKeys, getPrimaryKeyMapSizeInBytes());
_logger.info("Finished removing segment: {} in {}ms, current primary key count: {}", segmentName,
System.currentTimeMillis() - startTimeMs, numPrimaryKeys);
}
Expand Down Expand Up @@ -1186,7 +1197,7 @@ public synchronized void close()
// We don't remove the segment from the metadata manager when
// it's closed. This was done to make table deletion faster. Since we don't remove the segment, we never decrease
// the primary key count. So, we set the primary key count to 0 here.
updatePrimaryKeyGauge(0);
updatePrimaryKeyGauge(0, 0);
_logger.info("Closed the metadata manager");
}

Expand Down
14 changes: 13 additions & 1 deletion ...va/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.util.concurrent.atomic.AtomicInteger;
import javax.annotation.Nullable;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.lucene.util.RamUsageEstimator;
import org.apache.pinot.common.metrics.ServerMeter;
import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl;
import org.apache.pinot.segment.local.segment.readers.LazyRow;
Expand Down Expand Up @@ -67,6 +68,15 @@ protected long getNumPrimaryKeys() {
return _primaryKeyToRecordLocationMap.size();
}

// Approximates ConcurrentHashMap.Node + RecordLocation object overhead per map entry.
private static final long PER_ENTRY_OVERHEAD_BYTES =
RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY + RamUsageEstimator.shallowSizeOfInstance(RecordLocation.class);

@Override
protected long getPrimaryKeyMapSizeInBytes() {
return UpsertUtils.estimatePrimaryKeyMapSizeInBytes(_primaryKeyToRecordLocationMap, PER_ENTRY_OVERHEAD_BYTES);
}

@Override
protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds,
@Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator<RecordInfo> recordInfoIterator,
Expand Down Expand Up @@ -412,7 +422,9 @@ protected boolean doAddRecord(MutableSegment segment, RecordInfo recordInfo) {
}
});

updatePrimaryKeyGauge();
// Per-record hot path: keep the update count-only. The byte-size gauge is refreshed at segment-lifecycle
// cadence instead (see BasePartitionUpsertMetadataManager#updatePrimaryKeyGauge()).
updatePrimaryKeyGauge(getNumPrimaryKeys());
return !isOutOfOrderRecord.get();
}

Expand Down
16 changes: 14 additions & 2 deletions ...segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import java.util.concurrent.locks.Lock;
import javax.annotation.Nullable;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.lucene.util.RamUsageEstimator;
import org.apache.pinot.common.metrics.ServerMeter;
import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl;
import org.apache.pinot.segment.local.segment.readers.LazyRow;
Expand Down Expand Up @@ -86,6 +87,15 @@ protected long getNumPrimaryKeys() {
return _primaryKeyToRecordLocationMap.size();
}

// Approximates ConcurrentHashMap.Node + RecordLocation object overhead per map entry.
private static final long PER_ENTRY_OVERHEAD_BYTES =
RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY + RamUsageEstimator.shallowSizeOfInstance(RecordLocation.class);

@Override
protected long getPrimaryKeyMapSizeInBytes() {
return UpsertUtils.estimatePrimaryKeyMapSizeInBytes(_primaryKeyToRecordLocationMap, PER_ENTRY_OVERHEAD_BYTES);
}

@Override
protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds,
@Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator<RecordInfo> recordInfoIterator,
Expand Down Expand Up @@ -261,7 +271,7 @@ protected void doRemoveSegment(IndexSegment segment) {
}
// Update metrics
long numPrimaryKeys = getNumPrimaryKeys();
updatePrimaryKeyGauge(numPrimaryKeys);
updatePrimaryKeyGauge(numPrimaryKeys, getPrimaryKeyMapSizeInBytes());
_logger.info("Finished removing segment: {} in {}ms, current primary key count: {}", segmentName,
System.currentTimeMillis() - startTimeMs, numPrimaryKeys);
}
Expand Down Expand Up @@ -544,7 +554,9 @@ protected boolean doAddRecord(MutableSegment segment, RecordInfo recordInfo) {
}
});

updatePrimaryKeyGauge();
// Per-record hot path: keep the update count-only. The byte-size gauge is refreshed at segment-lifecycle
// cadence instead (see BasePartitionUpsertMetadataManager#updatePrimaryKeyGauge()).
updatePrimaryKeyGauge(getNumPrimaryKeys());
return !isOutOfOrderRecord.get();
}

Expand Down
43 changes: 43 additions & 0 deletions pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@

@SuppressWarnings("rawtypes")
public class UpsertUtils {
// ponytail: bounded sample instead of a full map walk, so size estimation stays O(1) regardless of table size.
private static final int PRIMARY_KEY_SAMPLE_SIZE = 1000;

private UpsertUtils() {
}

Expand Down Expand Up @@ -86,6 +89,46 @@ public static void doRemoveDocId(IndexSegment segment, int docId) {
}
}

/**
* Estimates the total size in bytes of a primary-key-to-record-location map. Samples up to
* {@link #PRIMARY_KEY_SAMPLE_SIZE} keys to compute the average key content size, then extrapolates over all keys,
* so the cost stays bounded regardless of map size. Handles both unhashed ({@link PrimaryKey}) and hashed
* ({@link ByteArray}) keys.
*
* @param perEntryOverheadBytes estimated fixed per-entry overhead (e.g. map node + record location object)
*/
public static long estimatePrimaryKeyMapSizeInBytes(Map<Object, ?> primaryKeyToRecordLocationMap,
long perEntryOverheadBytes) {
int numKeys = primaryKeyToRecordLocationMap.size();
if (numKeys == 0) {
return 0;
}
long sampledKeyBytes = 0;
int numSampled = 0;
for (Object key : primaryKeyToRecordLocationMap.keySet()) {
sampledKeyBytes += getPrimaryKeyContentSizeInBytes(key);
if (++numSampled >= PRIMARY_KEY_SAMPLE_SIZE) {
break;
}
}
if (numSampled == 0) {
// ponytail: map drained concurrently between size() and keySet() iteration; nothing sampled.
return 0;
}
double avgKeyBytes = (double) sampledKeyBytes / numSampled;
return (long) (numKeys * (avgKeyBytes + perEntryOverheadBytes));
}

private static int getPrimaryKeyContentSizeInBytes(Object key) {
if (key instanceof PrimaryKey) {
return ((PrimaryKey) key).asBytes().length;
}
if (key instanceof ByteArray) {
return ((ByteArray) key).length();
}
return 0;
}

/**
* Returns an iterator of {@link RecordInfo} for all the documents from the segment.
*/
Expand Down
5 changes: 5 additions & 0 deletions ...st/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1035,6 +1035,11 @@ protected long getNumPrimaryKeys() {
return 0;
}

@Override
protected long getPrimaryKeyMapSizeInBytes() {
return 0;
}

@Override
protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds,
@Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator<RecordInfo> recordInfoIterator,
Expand Down
22 changes: 22 additions & 0 deletions ...ent/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,28 @@ public void testStartFinishOperation() {
TestUtils.waitForCondition(aVoid -> closed.get(), 10_000L, "Failed to close the metadata manager");
}

@Test
public void testGetPrimaryKeyMapSizeInBytes() {
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes upsertMetadataManager =
new ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes(REALTIME_TABLE_NAME, 0,
_contextBuilder.build());
assertEquals(upsertMetadataManager.getPrimaryKeyMapSizeInBytes(), 0);

Map<Object, ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.RecordLocation> recordLocationMap =
upsertMetadataManager._primaryKeyToRecordLocationMap;
IndexSegment mockSegment = mock(IndexSegment.class);
recordLocationMap.put(new PrimaryKey(new Object[]{"short"}),
new ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.RecordLocation(mockSegment, 0, 0, 1));
long shortKeyBytes = upsertMetadataManager.getPrimaryKeyMapSizeInBytes();
assertTrue(shortKeyBytes > 0);

recordLocationMap.clear();
recordLocationMap.put(new PrimaryKey(new Object[]{"a".repeat(1000)}),
new ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.RecordLocation(mockSegment, 0, 0, 1));
long longKeyBytes = upsertMetadataManager.getPrimaryKeyMapSizeInBytes();
assertTrue(longKeyBytes > shortKeyBytes);
}

@Test
public void testAddReplaceRemoveSegment()
throws IOException {
Expand Down
39 changes: 39 additions & 0 deletions ...rg/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
import org.apache.pinot.common.metrics.ServerGauge;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.common.utils.LLCSegmentName;
import org.apache.pinot.common.utils.UploadedRealtimeSegmentName;
Expand Down Expand Up @@ -80,11 +81,14 @@

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.*;
Expand Down Expand Up @@ -194,6 +198,41 @@ public void testStartFinishOperation() {
TestUtils.waitForCondition(aVoid -> closed.get(), 10_000L, "Failed to close the metadata manager");
}

@Test
public void testGetPrimaryKeyMapSizeInBytes() {
ConcurrentMapPartitionUpsertMetadataManager upsertMetadataManager =
new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME, 0, _contextBuilder.build());
assertEquals(upsertMetadataManager.getPrimaryKeyMapSizeInBytes(), 0);

Map<Object, RecordLocation> recordLocationMap = upsertMetadataManager._primaryKeyToRecordLocationMap;
IndexSegment mockSegment = mock(IndexSegment.class);
recordLocationMap.put(new PrimaryKey(new Object[]{"short"}), new RecordLocation(mockSegment, 0, 0));
long shortKeyBytes = upsertMetadataManager.getPrimaryKeyMapSizeInBytes();
assertTrue(shortKeyBytes > 0);

recordLocationMap.clear();
recordLocationMap.put(new PrimaryKey(new Object[]{"a".repeat(1000)}), new RecordLocation(mockSegment, 0, 0));
long longKeyBytes = upsertMetadataManager.getPrimaryKeyMapSizeInBytes();
assertTrue(longKeyBytes > shortKeyBytes);
}

@Test
public void testAddRecordUpdatesCountGaugeOnly() {
ServerMetrics serverMetrics = ServerMetrics.get();
reset(serverMetrics);
ConcurrentMapPartitionUpsertMetadataManager upsertMetadataManager =
new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME, 0, _contextBuilder.build());

ThreadSafeMutableRoaringBitmap validDocIds = new ThreadSafeMutableRoaringBitmap();
MutableSegment segment = mockMutableSegment(1, validDocIds, null);
upsertMetadataManager.addRecord(segment, new RecordInfo(makePrimaryKey(1), 0, 10, false));

// Per-record hot path only refreshes the count gauge, not the byte-size gauge
verify(serverMetrics).setValueOfPartitionGauge(REALTIME_TABLE_NAME, 0, ServerGauge.UPSERT_PRIMARY_KEYS_COUNT, 1);
verify(serverMetrics, never()).setValueOfPartitionGauge(eq(REALTIME_TABLE_NAME), eq(0),
eq(ServerGauge.UPSERT_PRIMARY_KEY_MAP_SIZE_IN_BYTES), anyLong());
}

@Test
public void testAddReplaceRemoveSegment()
throws IOException {
Expand Down
Loading
Loading