From 7c95ccdcacc016cd59583d91a9a76bf965a036e8 Mon Sep 17 00:00:00 2001 From: tarunm Date: Mon, 13 Jul 2026 10:03:41 +0000 Subject: [PATCH 1/2] feat(upsert): add primary-key map size-in-bytes gauge Add UPSERT_PRIMARY_KEY_MAP_SIZE_IN_BYTES ServerGauge to complement the existing key-count gauge. Tables with few large-string PKs and tables with many small UUID-like PKs can have very different memory footprints despite similar counts, so byte-size observability helps with capacity monitoring. UpsertUtils.estimatePrimaryKeyMapSizeInBytes samples up to 1000 keys to average PK content size (via PrimaryKey.asBytes() / ByteArray.length()), adds a per-entry overhead constant for the hashmap node + RecordLocation object, and extrapolates over the full key count, keeping the cost bounded regardless of map size. Computed at the same segment-lifecycle cadence as the existing count gauge (add/replace/remove/preload/close), not on the per-record ingest path. Wired into both ConcurrentMapPartitionUpsertMetadataManager and ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes via a new abstract getPrimaryKeyMapSizeInBytes() hook on BasePartitionUpsertMetadataManager. --- .../pinot/common/metrics/ServerGauge.java | 1 + .../DummyTableUpsertMetadataManager.java | 5 ++ .../BasePartitionUpsertMetadataManager.java | 21 ++++--- ...rentMapPartitionUpsertMetadataManager.java | 10 +++ ...rtMetadataManagerForConsistentDeletes.java | 12 +++- .../segment/local/upsert/UpsertUtils.java | 43 +++++++++++++ ...asePartitionUpsertMetadataManagerTest.java | 5 ++ ...tadataManagerForConsistentDeletesTest.java | 22 +++++++ ...MapPartitionUpsertMetadataManagerTest.java | 18 ++++++ .../segment/local/upsert/UpsertUtilsTest.java | 61 +++++++++++++++++++ 10 files changed, 190 insertions(+), 8 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java index 0a909ce43412..2334b0aa445d 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerGauge.java @@ -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), diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/models/DummyTableUpsertMetadataManager.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/models/DummyTableUpsertMetadataManager.java index 56ace9ff9c23..599eff0a0723 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/models/DummyTableUpsertMetadataManager.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/models/DummyTableUpsertMetadataManager.java @@ -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 recordInfoIterator, diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java index 733e4e24014e..88c828a82117 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java @@ -388,20 +388,27 @@ 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(); - protected void updatePrimaryKeyGauge(long numPrimaryKeys) { + /** + * 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) { _serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId, ServerGauge.UPSERT_PRIMARY_KEYS_COUNT, numPrimaryKeys); + _serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId, + ServerGauge.UPSERT_PRIMARY_KEY_MAP_SIZE_IN_BYTES, primaryKeySizeInBytes); } protected void updatePrimaryKeyGauge() { - updatePrimaryKeyGauge(getNumPrimaryKeys()); + updatePrimaryKeyGauge(getNumPrimaryKeys(), getPrimaryKeyMapSizeInBytes()); } @Override @@ -459,7 +466,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); } @@ -645,7 +652,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); } @@ -812,7 +819,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); } @@ -1186,7 +1193,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"); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java index f580712576f5..bb40d95af419 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java @@ -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; @@ -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 recordInfoIterator, diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java index 929df2ac3ac3..864d5bd7438a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java @@ -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; @@ -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 recordInfoIterator, @@ -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); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertUtils.java index 3978044b8785..21eb652b7ffb 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertUtils.java @@ -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() { } @@ -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 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. */ diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManagerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManagerTest.java index aefd1cfc8959..657b2ba946c5 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManagerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManagerTest.java @@ -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 recordInfoIterator, diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java index 7acfc0ace545..f4baf0fc731c 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.java @@ -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 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 { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java index dbba9b0e2eba..1757d9a0cc26 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java @@ -194,6 +194,24 @@ 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 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 testAddReplaceRemoveSegment() throws IOException { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertUtilsTest.java index de88a5b27d3b..bf5002b4cd42 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/UpsertUtilsTest.java @@ -19,16 +19,21 @@ package org.apache.pinot.segment.local.upsert; import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentImpl; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; import org.apache.pinot.segment.spi.index.mutable.ThreadSafeMutableRoaringBitmap; import org.apache.pinot.segment.spi.store.SegmentDirectory; +import org.apache.pinot.spi.data.readers.PrimaryKey; +import org.apache.pinot.spi.utils.ByteArray; import org.roaringbitmap.buffer.MutableRoaringBitmap; import org.testng.annotations.Test; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; @@ -135,6 +140,62 @@ public void testHasNoQueryableDocsConsistencyModeSnapshotAbsent() { assertFalse(segment.hasNoQueryableDocs()); } + // -------- estimatePrimaryKeyMapSizeInBytes -------- + + @Test + public void testEstimatePrimaryKeyMapSizeInBytesEmptyMap() { + assertEquals(UpsertUtils.estimatePrimaryKeyMapSizeInBytes(new ConcurrentHashMap<>(), 100), 0); + } + + @Test + public void testEstimatePrimaryKeyMapSizeInBytesScalesWithPrimaryKeyContentSize() { + Map shortKeyMap = new ConcurrentHashMap<>(); + Map longKeyMap = new ConcurrentHashMap<>(); + String longValuePrefix = "x".repeat(500); + for (int i = 0; i < 100; i++) { + shortKeyMap.put(new PrimaryKey(new Object[]{"k" + i}), new Object()); + longKeyMap.put(new PrimaryKey(new Object[]{longValuePrefix + i}), new Object()); + } + long shortKeyMapBytes = UpsertUtils.estimatePrimaryKeyMapSizeInBytes(shortKeyMap, 0); + long longKeyMapBytes = UpsertUtils.estimatePrimaryKeyMapSizeInBytes(longKeyMap, 0); + assertTrue(longKeyMapBytes > shortKeyMapBytes); + } + + @Test + public void testEstimatePrimaryKeyMapSizeInBytesWithHashedByteArrayKeys() { + // Hashed primary keys (e.g. MD5) are always the same fixed length regardless of PK content. + Map map = new ConcurrentHashMap<>(); + for (int i = 0; i < 50; i++) { + byte[] bytes = new byte[16]; + bytes[0] = (byte) i; + map.put(new ByteArray(bytes), new Object()); + } + assertEquals(UpsertUtils.estimatePrimaryKeyMapSizeInBytes(map, 0), 50L * 16); + } + + @Test + public void testEstimatePrimaryKeyMapSizeInBytesIncludesPerEntryOverhead() { + Map map = new ConcurrentHashMap<>(); + map.put(new ByteArray(new byte[16]), new Object()); + assertEquals(UpsertUtils.estimatePrimaryKeyMapSizeInBytes(map, 32), 16 + 32); + } + + @Test + public void testEstimatePrimaryKeyMapSizeInBytesExtrapolatesBeyondSampleSize() { + // With more keys than the sampling cap, only a bounded sample is scanned and the result is + // extrapolated over all keys. Uniform-length ByteArray keys make the result exact regardless + // of which keys happen to be sampled. + int numKeys = 2500; + Map map = new ConcurrentHashMap<>(); + for (int i = 0; i < numKeys; i++) { + byte[] bytes = new byte[16]; + bytes[0] = (byte) i; + bytes[1] = (byte) (i >> 8); + map.put(new ByteArray(bytes), new Object()); + } + assertEquals(UpsertUtils.estimatePrimaryKeyMapSizeInBytes(map, 8), (long) numKeys * (16 + 8)); + } + /// Returns a minimal {@link ImmutableSegmentImpl} usable for testing the upsert-aware methods. /// Mirrors the pattern used in {@code BasePartitionUpsertMetadataManagerTest#createImmutableSegment}. private static ImmutableSegmentImpl newSegment() { From b660f7cc5747a704f7858750c7843d3e9c1ddd2c Mon Sep 17 00:00:00 2001 From: tarunm Date: Tue, 14 Jul 2026 04:26:13 +0000 Subject: [PATCH 2/2] Keep per-record primary-key gauge update count-only updatePrimaryKeyGauge() was recomputing the byte-size gauge (bounded sample over the map, but still per-record overhead) on every ingested record via doAddRecord(), despite the intent of computing it only at segment-lifecycle cadence. Split it into a count-only overload used by the per-record hot path; segment lifecycle events (add/preload/replace /remove/close) and the periodic TTL eviction path continue to refresh both gauges. --- .../BasePartitionUpsertMetadataManager.java | 8 +++++-- ...rentMapPartitionUpsertMetadataManager.java | 4 +++- ...rtMetadataManagerForConsistentDeletes.java | 4 +++- ...MapPartitionUpsertMetadataManagerTest.java | 21 +++++++++++++++++++ 4 files changed, 33 insertions(+), 4 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java index 88c828a82117..dbca3ac8f520 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java @@ -401,12 +401,16 @@ protected void doAddSegment(ImmutableSegmentImpl segment) { protected abstract long getPrimaryKeyMapSizeInBytes(); protected void updatePrimaryKeyGauge(long numPrimaryKeys, long primaryKeySizeInBytes) { - _serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId, ServerGauge.UPSERT_PRIMARY_KEYS_COUNT, - numPrimaryKeys); + 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(), getPrimaryKeyMapSizeInBytes()); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java index bb40d95af419..49c96e620bff 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManager.java @@ -422,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(); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java index 864d5bd7438a..1675d89cd7b9 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.java @@ -554,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(); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java index 1757d9a0cc26..d87d9d044af0 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/upsert/ConcurrentMapPartitionUpsertMetadataManagerTest.java @@ -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; @@ -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.*; @@ -212,6 +216,23 @@ public void testGetPrimaryKeyMapSizeInBytes() { 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 {