Count primary keys not replaced at removal, not from a stale bitmap clone

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<PrimaryKey> 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<PrimaryKey> 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<PrimaryKey> 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<PrimaryKey> primaryKeyIterator) {

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<PrimaryKey> primaryKeyIterator) {
+    removeSegmentAndGetNumKeysRemoved(segment, primaryKeyIterator, null);
+  }
+
+  protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, Iterator<PrimaryKey> primaryKeyIterator,
+      @Nullable List<PrimaryKey> 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<PrimaryKey> 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<PrimaryKey> 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();

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<PrimaryKey> 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,26 +311,33 @@ 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<PrimaryKey> 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();
     }
   }
 
   @Override
   protected void removeSegment(IndexSegment segment, Iterator<PrimaryKey> primaryKeyIterator) {
+    removeSegmentAndGetNumKeysRemoved(segment, primaryKeyIterator, null);
+  }
+
+  protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, Iterator<PrimaryKey> primaryKeyIterator,
+      @Nullable List<PrimaryKey> 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<PrimaryKey> 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<PrimaryKey> primaryK
                 RecordLocation.decrementSegmentCount(recordLocation.getDistinctSegmentCount()));
           });
     }
+    return numKeysRemoved.get();
   }
 
   @Override

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.<RecordInfo>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.<RecordInfo>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<RecordInfo> recordInfoIterator,
+        @Nullable IndexSegment oldSegment, @Nullable MutableRoaringBitmap validDocIdsForOldSegment) {
+      _duringSegmentReplacement.run();
+    }
+
+    @Override
+    protected int removeSegmentAndGetNumKeysRemoved(IndexSegment segment, MutableRoaringBitmap validDocIds,
+        @Nullable List<PrimaryKey> sampledKeysRemoved) {
+      _candidateValidDocIds = validDocIds.clone();
+      return _numKeysRemovedAtActionBoundary;
+    }
+
+    @Override
+    protected void updateInconsistentRowsMetric(String segmentName, int numKeysStillNotReplaced) {
+      _numInconsistentRows += numKeysStillNotReplaced;
+    }
+  }
 }

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<PrimaryKey> 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();
+  }
 }

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<PrimaryKey> 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 {