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 166ff0825006..9c641403ad67 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 @@ -685,10 +685,13 @@ public void replaceSegment(ImmutableSegment segment, @Nullable ThreadSafeMutable revertSegmentUpsertMetadata(oldSegment, segmentName, validDocIdsForOldSegment); return; } - _logger.warn("Found {} primary keys not replaced for segment: {}", - validDocIdsForOldSegment.getCardinality(), segmentName); - updateInconsistentRowsMetric(segmentName, validDocIdsForOldSegment.getCardinality()); - removeSegment(oldSegment, validDocIdsForOldSegment); + List sampledKeysNotReplaced = new ArrayList<>(NUM_SAMPLED_KEYS_NOT_REPLACED); + int numKeysStillNotReplaced = + removeSegmentAndGetNumKeysRemoved(oldSegment, validDocIdsForOldSegment, sampledKeysNotReplaced); + if (numKeysStillNotReplaced > 0) { + logKeysNotReplaced(segmentName, numKeysStillNotReplaced, sampledKeysNotReplaced); + updateInconsistentRowsMetric(segmentName, numKeysStillNotReplaced); + } } } @@ -740,6 +743,32 @@ private MutableRoaringBitmap getValidDocIdsForOldSegment(IndexSegment oldSegment return oldSegment.getValidDocIds() != null ? oldSegment.getValidDocIds().getMutableRoaringBitmap() : null; } + /// How many of the unreplaced primary keys to name in the log. Bounded because the count itself can be large, + /// and because a primary key is customer data. + protected static final int NUM_SAMPLED_KEYS_NOT_REPLACED = 8; + + /// Removes candidate keys and returns how many were still owned by the segment at removal time, appending up to + /// [#NUM_SAMPLED_KEYS_NOT_REPLACED] of them to `sampledKeysRemoved` when it is given. Implementations backed by + /// concurrent metadata should override this and count only removals that pass their authoritative ownership check. + /// The default preserves compatibility with existing metadata-manager implementations and names no keys, since it + /// cannot tell which of the candidates were actually removed. + protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, MutableRoaringBitmap validDocIds, + @Nullable List sampledKeysRemoved) { + removeSegment(segment, validDocIds); + return validDocIds.getCardinality(); + } + + /// Logs the unreplaced-key count, naming a few of the keys when the implementation could identify them. + protected void logKeysNotReplaced(String segmentName, int numKeysStillNotReplaced, + List sampledKeysNotReplaced) { + if (sampledKeysNotReplaced.isEmpty()) { + _logger.warn("Found {} primary keys not replaced for segment: {}", numKeysStillNotReplaced, segmentName); + } else { + _logger.warn("Found {} primary keys not replaced for segment: {}, first {}: {}", numKeysStillNotReplaced, + segmentName, sampledKeysNotReplaced.size(), sampledKeysNotReplaced); + } + } + protected abstract void removeSegment(IndexSegment segment, MutableRoaringBitmap validDocIds); protected void removeSegment(IndexSegment segment, Iterator primaryKeyIterator) { 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 57559b020694..3b992bc6db57 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 @@ -21,6 +21,7 @@ import com.google.common.annotations.VisibleForTesting; import java.util.HashMap; import java.util.Iterator; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; @@ -209,11 +210,21 @@ protected void addSegmentWithoutUpsert(ImmutableSegmentImpl segment, ThreadSafeM @Override protected void removeSegment(IndexSegment segment, Iterator primaryKeyIterator) { + removeSegmentAndGetNumKeysRemoved(segment, primaryKeyIterator, null); + } + + protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, Iterator primaryKeyIterator, + @Nullable List sampledKeysRemoved) { + AtomicInteger numKeysRemoved = new AtomicInteger(); while (primaryKeyIterator.hasNext()) { PrimaryKey primaryKey = primaryKeyIterator.next(); _primaryKeyToRecordLocationMap.computeIfPresent(HashUtils.hashPrimaryKey(primaryKey, _hashFunction), (pk, recordLocation) -> { if (recordLocation.getSegment() == segment) { + numKeysRemoved.getAndIncrement(); + if (sampledKeysRemoved != null && sampledKeysRemoved.size() < NUM_SAMPLED_KEYS_NOT_REPLACED) { + sampledKeysRemoved.add(primaryKey); + } if (_context.isTableTypeInconsistentDuringConsumption() && segment instanceof MutableSegment) { _previousKeyToRecordLocationMap.remove(pk); } @@ -222,6 +233,7 @@ protected void removeSegment(IndexSegment segment, Iterator primaryK return recordLocation; }); } + return numKeysRemoved.get(); } @Override @@ -291,6 +303,19 @@ protected void removeSegment(IndexSegment segment, MutableRoaringBitmap validDoc } } + @Override + protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, MutableRoaringBitmap validDocIds, + @Nullable List sampledKeysRemoved) { + try (PrimaryKeyReader primaryKeyReader = new PrimaryKeyReader(segment, _primaryKeyColumns)) { + return removeSegmentAndGetNumKeysRemoved(segment, + UpsertUtils.getPrimaryKeyIterator(primaryKeyReader, validDocIds), sampledKeysRemoved); + } catch (Exception e) { + throw new RuntimeException( + String.format("Caught exception while removing segment: %s, table: %s, message: %s", segment.getSegmentName(), + _tableNameWithType, e.getMessage()), e); + } + } + @Override public void doRemoveExpiredPrimaryKeys() { AtomicInteger numMetadataTTLKeysRemoved = new AtomicInteger(); 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 32352f7a3b0f..55c91271c9e3 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 @@ -20,9 +20,11 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; @@ -238,19 +240,26 @@ protected void addSegmentWithoutUpsert(ImmutableSegmentImpl segment, ThreadSafeM @Override protected void doRemoveSegment(IndexSegment segment) { + doRemoveSegmentAndGetNumKeysRemoved(segment, null); + } + + protected int doRemoveSegmentAndGetNumKeysRemoved(IndexSegment segment, + @Nullable List sampledKeysRemoved) { String segmentName = segment.getSegmentName(); _logger.info("Removing {} segment: {}, current primary key count: {}", segment instanceof ImmutableSegment ? "immutable" : "mutable", segmentName, getNumPrimaryKeys()); long startTimeMs = System.currentTimeMillis(); // For ConsistentDeletes, we need to iterate over ALL docs in the segment (not just valid ones) // to properly decrement distinctSegmentCount for every key that was ever in the segment + int numKeysRemoved = 0; try (PrimaryKeyReader primaryKeyReader = new PrimaryKeyReader(segment, _primaryKeyColumns)) { if (shouldRevertMetadataOnInconsistency(segment)) { revertAndRemoveSegment(segment, UpsertUtils.getRecordIterator(primaryKeyReader, segment.getSegmentMetadata().getTotalDocs())); } else { - removeSegment(segment, - UpsertUtils.getPrimaryKeyIterator(primaryKeyReader, segment.getSegmentMetadata().getTotalDocs())); + numKeysRemoved = removeSegmentAndGetNumKeysRemoved(segment, + UpsertUtils.getPrimaryKeyIterator(primaryKeyReader, segment.getSegmentMetadata().getTotalDocs()), + sampledKeysRemoved); } } catch (Exception e) { throw new RuntimeException( @@ -262,6 +271,7 @@ protected void doRemoveSegment(IndexSegment segment) { updatePrimaryKeyGauge(numPrimaryKeys); _logger.info("Finished removing segment: {} in {}ms, current primary key count: {}", segmentName, System.currentTimeMillis() - startTimeMs, numPrimaryKeys); + return numKeysRemoved; } protected void removeSegment(IndexSegment segment, MutableRoaringBitmap validDocIds) { @@ -301,19 +311,20 @@ public void replaceSegment(ImmutableSegment segment, @Nullable ThreadSafeMutable doAddOrReplaceSegment((ImmutableSegmentImpl) segment, validDocIds, queryableDocIds, recordInfoIterator, oldSegment, validDocIdsForOldSegment); } - if (validDocIdsForOldSegment != null && !validDocIdsForOldSegment.isEmpty()) { - if (shouldRevertMetadataOnInconsistency(oldSegment)) { - revertSegmentUpsertMetadata(oldSegment, segmentName, validDocIdsForOldSegment); - return; - } - _logger.warn("Found {} primary keys not replaced for segment: {}", - validDocIdsForOldSegment.getCardinality(), segmentName); - updateInconsistentRowsMetric(segmentName, validDocIdsForOldSegment.getCardinality()); + if (validDocIdsForOldSegment != null && !validDocIdsForOldSegment.isEmpty() + && shouldRevertMetadataOnInconsistency(oldSegment)) { + revertSegmentUpsertMetadata(oldSegment, segmentName, validDocIdsForOldSegment); + return; } // we want to always remove a segment in case of enableDeletedKeysCompactionConsistency = true // this is to account for the removal of primary-key in the to-be-removed segment and reduce // distinctSegmentCount by 1 - doRemoveSegment(oldSegment); + List sampledKeysNotReplaced = new ArrayList<>(NUM_SAMPLED_KEYS_NOT_REPLACED); + int numKeysStillNotReplaced = doRemoveSegmentAndGetNumKeysRemoved(oldSegment, sampledKeysNotReplaced); + if (numKeysStillNotReplaced > 0) { + logKeysNotReplaced(segmentName, numKeysStillNotReplaced, sampledKeysNotReplaced); + updateInconsistentRowsMetric(segmentName, numKeysStillNotReplaced); + } } finally { segmentLock.unlock(); } @@ -321,6 +332,12 @@ public void replaceSegment(ImmutableSegment segment, @Nullable ThreadSafeMutable @Override protected void removeSegment(IndexSegment segment, Iterator primaryKeyIterator) { + removeSegmentAndGetNumKeysRemoved(segment, primaryKeyIterator, null); + } + + protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, Iterator primaryKeyIterator, + @Nullable List sampledKeysRemoved) { + AtomicInteger numKeysRemoved = new AtomicInteger(); // We need to decrease the distinctSegmentCount for each unique primary key in this deleting segment by 1 // as the occurrence of the key in this segment is being removed. We are taking a set of unique primary keys // to avoid double counting the same key in the same segment. @@ -330,6 +347,10 @@ protected void removeSegment(IndexSegment segment, Iterator primaryK _primaryKeyToRecordLocationMap.computeIfPresent(HashUtils.hashPrimaryKey(primaryKey, _hashFunction), (pk, recordLocation) -> { if (recordLocation.getSegment() == segment) { + numKeysRemoved.getAndIncrement(); + if (sampledKeysRemoved != null && sampledKeysRemoved.size() < NUM_SAMPLED_KEYS_NOT_REPLACED) { + sampledKeysRemoved.add(primaryKey); + } if (_context.isTableTypeInconsistentDuringConsumption() && segment instanceof MutableSegment) { _previousKeyToRecordLocationMap.remove(pk); } @@ -343,6 +364,7 @@ protected void removeSegment(IndexSegment segment, Iterator primaryK RecordLocation.decrementSegmentCount(recordLocation.getDistinctSegmentCount())); }); } + return numKeysRemoved.get(); } @Override 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 ca3e2da05a55..f5b66f7b6440 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 @@ -998,6 +998,54 @@ public void testTrackNewlyAddedSegments() { "Failed to remove stale segments"); } + @Test + public void testReplaceSegmentDoesNotReportKeyMovedBeforeRemoval() + throws IOException { + UpsertContext upsertContext = mock(UpsertContext.class); + when(upsertContext.getConsistencyMode()).thenReturn(UpsertConfig.ConsistencyMode.NONE); + + ThreadSafeMutableRoaringBitmap liveValidDocIds = createDocIds(0); + IndexSegment oldSegment = mock(IndexSegment.class); + when(oldSegment.getSegmentName()).thenReturn("segment"); + when(oldSegment.getValidDocIds()).thenReturn(liveValidDocIds); + ImmutableSegmentImpl newSegment = mock(ImmutableSegmentImpl.class); + when(newSegment.getSegmentName()).thenReturn("segment"); + + try (ActionBoundaryMetadataManager upsertMetadataManager = + new ActionBoundaryMetadataManager("myTable", 0, upsertContext, () -> liveValidDocIds.remove(0), 0)) { + upsertMetadataManager.replaceSegment(newSegment, new ThreadSafeMutableRoaringBitmap(), null, + List.of().iterator(), oldSegment); + + assertEquals(upsertMetadataManager._numInconsistentRows, 0); + assertEquals(upsertMetadataManager._candidateValidDocIds.toArray(), new int[]{0}); + upsertMetadataManager.stop(); + } + } + + @Test + public void testReplaceSegmentReportsOnlyKeysStillOwnedAtRemoval() + throws IOException { + UpsertContext upsertContext = mock(UpsertContext.class); + when(upsertContext.getConsistencyMode()).thenReturn(UpsertConfig.ConsistencyMode.NONE); + + ThreadSafeMutableRoaringBitmap liveValidDocIds = createDocIds(0, 1); + IndexSegment oldSegment = mock(IndexSegment.class); + when(oldSegment.getSegmentName()).thenReturn("segment"); + when(oldSegment.getValidDocIds()).thenReturn(liveValidDocIds); + ImmutableSegmentImpl newSegment = mock(ImmutableSegmentImpl.class); + when(newSegment.getSegmentName()).thenReturn("segment"); + + try (ActionBoundaryMetadataManager upsertMetadataManager = + new ActionBoundaryMetadataManager("myTable", 0, upsertContext, () -> liveValidDocIds.remove(0), 1)) { + upsertMetadataManager.replaceSegment(newSegment, new ThreadSafeMutableRoaringBitmap(), null, + List.of().iterator(), oldSegment); + + assertEquals(upsertMetadataManager._numInconsistentRows, 1); + assertEquals(upsertMetadataManager._candidateValidDocIds.toArray(), new int[]{0, 1}); + upsertMetadataManager.stop(); + } + } + @Test public void testResolveComparisonTies() { // Build a record info list for testing @@ -1140,4 +1188,37 @@ protected int getPrevKeyToRecordLocationSize() { protected void clearPrevKeyToRecordLocation() { } } + + private static class ActionBoundaryMetadataManager extends DummyPartitionUpsertMetadataManager { + private final Runnable _duringSegmentReplacement; + private final int _numKeysRemovedAtActionBoundary; + private int _numInconsistentRows; + private MutableRoaringBitmap _candidateValidDocIds; + + private ActionBoundaryMetadataManager(String tableNameWithType, int partitionId, UpsertContext context, + Runnable duringSegmentReplacement, int numKeysRemovedAtActionBoundary) { + super(tableNameWithType, partitionId, context); + _duringSegmentReplacement = duringSegmentReplacement; + _numKeysRemovedAtActionBoundary = numKeysRemovedAtActionBoundary; + } + + @Override + protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds, + @Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator recordInfoIterator, + @Nullable IndexSegment oldSegment, @Nullable MutableRoaringBitmap validDocIdsForOldSegment) { + _duringSegmentReplacement.run(); + } + + @Override + protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, MutableRoaringBitmap validDocIds, + @Nullable List sampledKeysRemoved) { + _candidateValidDocIds = validDocIds.clone(); + return _numKeysRemovedAtActionBoundary; + } + + @Override + protected void updateInconsistentRowsMetric(String segmentName, int numKeysStillNotReplaced) { + _numInconsistentRows += numKeysStillNotReplaced; + } + } } 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 52a688907f89..292f40a70a22 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 @@ -1350,4 +1350,37 @@ public void testNoRevertForImmutableSegmentReplacement() upsertMetadataManager.stop(); upsertMetadataManager.close(); } + + @Test + public void testRemoveSegmentCountsOnlyKeysStillOwnedBySegment() + throws IOException { + ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes upsertMetadataManager = + new ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes(REALTIME_TABLE_NAME, 0, + _contextBuilder.build()); + IndexSegment oldSegment = mock(IndexSegment.class); + IndexSegment newerSegment = mock(IndexSegment.class); + PrimaryKey oldSegmentKey = makePrimaryKey(10); + PrimaryKey newerSegmentKey = makePrimaryKey(20); + Object oldSegmentMapKey = HashUtils.hashPrimaryKey(oldSegmentKey, HashFunction.NONE); + Object newerSegmentMapKey = HashUtils.hashPrimaryKey(newerSegmentKey, HashFunction.NONE); + upsertMetadataManager._primaryKeyToRecordLocationMap.put(oldSegmentMapKey, + new ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.RecordLocation(oldSegment, 0, 100, 1)); + upsertMetadataManager._primaryKeyToRecordLocationMap.put(newerSegmentMapKey, + new ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.RecordLocation(newerSegment, 1, 200, 2)); + + List sampledKeysRemoved = new ArrayList<>(); + int numKeysRemoved = upsertMetadataManager.removeSegmentAndGetNumKeysRemoved(oldSegment, + List.of(oldSegmentKey, newerSegmentKey).iterator(), sampledKeysRemoved); + + assertEquals(numKeysRemoved, 1); + // Only the key still owned by the old segment is named, not the one ingestion already moved. + assertEquals(sampledKeysRemoved, List.of(oldSegmentKey)); + assertFalse(upsertMetadataManager._primaryKeyToRecordLocationMap.containsKey(oldSegmentMapKey)); + ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.RecordLocation newerLocation = + upsertMetadataManager._primaryKeyToRecordLocationMap.get(newerSegmentMapKey); + assertSame(newerLocation.getSegment(), newerSegment); + assertEquals(newerLocation.getDistinctSegmentCount(), 1); + upsertMetadataManager.stop(); + upsertMetadataManager.close(); + } } 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 a3dcfb12821c..3a1125e1020e 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 @@ -206,6 +206,35 @@ public void testAddReplaceRemoveSegment() verifyAddReplaceRemoveSegment(HashFunction.MURMUR3, true); } + @Test + public void testRemoveSegmentCountsOnlyKeysStillOwnedBySegment() + throws IOException { + ConcurrentMapPartitionUpsertMetadataManager upsertMetadataManager = + new ConcurrentMapPartitionUpsertMetadataManager(REALTIME_TABLE_NAME, 0, _contextBuilder.build()); + IndexSegment oldSegment = mock(IndexSegment.class); + IndexSegment newerSegment = mock(IndexSegment.class); + PrimaryKey oldSegmentKey = makePrimaryKey(10); + PrimaryKey newerSegmentKey = makePrimaryKey(20); + Object oldSegmentMapKey = HashUtils.hashPrimaryKey(oldSegmentKey, HashFunction.NONE); + Object newerSegmentMapKey = HashUtils.hashPrimaryKey(newerSegmentKey, HashFunction.NONE); + upsertMetadataManager._primaryKeyToRecordLocationMap.put(oldSegmentMapKey, + new RecordLocation(oldSegment, 0, 100)); + upsertMetadataManager._primaryKeyToRecordLocationMap.put(newerSegmentMapKey, + new RecordLocation(newerSegment, 1, 200)); + + List sampledKeysRemoved = new ArrayList<>(); + int numKeysRemoved = upsertMetadataManager.removeSegmentAndGetNumKeysRemoved(oldSegment, + List.of(oldSegmentKey, newerSegmentKey).iterator(), sampledKeysRemoved); + + assertEquals(numKeysRemoved, 1); + // Only the key still owned by the old segment is named, not the one ingestion already moved. + assertEquals(sampledKeysRemoved, List.of(oldSegmentKey)); + assertFalse(upsertMetadataManager._primaryKeyToRecordLocationMap.containsKey(oldSegmentMapKey)); + assertSame(upsertMetadataManager._primaryKeyToRecordLocationMap.get(newerSegmentMapKey).getSegment(), newerSegment); + upsertMetadataManager.stop(); + upsertMetadataManager.close(); + } + @Test public void testRemoveExpiredPrimaryKeys() throws IOException {