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

[KKcorps]

TL;DR

The "Found N primary keys not replaced" warning and the upsertInconsistentRows /
partialUpsertKeysNotReplaced metrics are computed from a validDocIds bitmap cloned before
segment replacement began, and never reconciled against the live state afterwards. Keys that
ingestion correctly moved onto a newer segment mid-replacement get reported as not replaced. This
counts at the point where the code already decides, per key, whether the key still belongs to the
segment being removed.

The problem

doAddOrReplaceSegment takes validDocIdsForOldSegment as a snapshot of the old segment's valid
docs. The replacement then runs, which takes time on a large segment. While it runs, ingestion keeps
going: a new record for one of those primary keys updates the live record-location map and points
the key at the newer consuming segment. That is correct upsert behaviour, and it is exactly what
should happen.

At the end, the code reports validDocIdsForOldSegment.getCardinality() — the size of the stale
clone. Every key that moved during the window is counted as "not replaced" even though it was
replaced properly.

sequenceDiagram
  participant I as Ingestion
  participant R as Replacement
  participant M as Record location map

  R->>M: clone validDocIds of old segment
  Note over R: clone says keys {A, B} still valid
  I->>M: new record for key A
  M-->>I: A now points at the consuming segment
  Note over R: replacement finishes
  R->>R: report clone cardinality = 2
  Note over R: ❌ A is counted, but A was replaced correctly

Loading

The consequence is not just a noisy log. UpsertInconsistentReplicas is built on these metrics and
is muted across a large number of production environments because of this false positive, so a real
divergence on those tables now goes unreported.

The approach

removeSegment already answers the question correctly. For each candidate key it does a
computeIfPresent and only acts when recordLocation.getSegment() == segment, so a key ingestion
moved away fails that check and is skipped. That per-key check is the authoritative answer to "was
this key actually still owned by the old segment".

1. Add BasePartitionUpsertMetadataManager.removeSegmentAndGetNumKeysRemoved(IndexSegment, MutableRoaringBitmap), returning how many keys passed that ownership check.
2. Default it to removeSegment(...) followed by the bitmap cardinality, so any metadata-manager
   implementation that does not override keeps its current behaviour exactly.
3. Override it in both ConcurrentMap managers, counting inside the computeIfPresent where the
   ownership check happens.
4. Report the warning and the metric from that count, and only when it is above zero.

Key components

Class / file	Change
BasePartitionUpsertMetadataManager	New removeSegmentAndGetNumKeysRemoved with a compatible default; reporting moved to use its return value
ConcurrentMapPartitionUpsertMetadataManager	Overrides it; removeSegment(segment, iterator) refactored to share one private body with an optional counter
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes	Same, plus doRemoveSegmentAndGetNumKeysRemoved, since that path walks every doc rather than the valid ones

Flow after the change

sequenceDiagram
  participant R as doAddOrReplaceSegment
  participant Rm as removeSegmentAndGetNumKeysRemoved
  participant M as Record location map

  R->>Rm: candidate keys from the (stale) clone
  loop per candidate key
    Rm->>M: computeIfPresent(key)
    alt still points at the old segment
      M-->>Rm: removed, count it
    else already moved by ingestion
      M-->>Rm: untouched, not counted
    end
  end
  Rm-->>R: numKeysStillNotReplaced
  alt count > 0
    R->>R: warn + updateInconsistentRowsMetric
  else count == 0
    R->>R: silent, the replacement was clean
  end

Loading

Why the reporting moved after removal in the consistent-deletes manager

ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletes.doRemoveSegment deliberately walks
every doc in the segment, not just the valid ones, so it can decrement distinctSegmentCount for
every key that was ever there. That means the count only becomes knowable once the removal has run,
so the warning and metric move below it. The shouldRevertMetadataOnInconsistency early return is
unchanged, and the removal itself was already unconditional on that path.

Behaviour changes

• The warning and the metric no longer fire when every candidate key had already been replaced.
  Previously they fired with an inflated count. A table with a clean replacement goes from a
  spurious warning to silence.
• When there is genuine inconsistency the count is now the real number of unreplaced keys, which
  will generally be smaller than what was reported before.
• No config, no metric names, no wire format changed. removeSegment keeps its signature and
  behaviour for every existing caller.

Naming the keys

The count tells you a table is inconsistent but not which rows, so the next step was always a hunt.
The warning now carries a bounded sample of the keys:

Found 2 primary keys not replaced for segment: X, first 2: [[200], [300]]

The sample is collected into a caller-supplied list rather than logged where it is gathered, and
that placement is the whole design. removeSegmentAndGetNumKeysRemoved is also on the plain
segment-removal path, where removing every key the segment owns is entirely normal. Logging keys
there would fire a warning on every retention deletion and name customer data for a healthy
operation. So the replacement path passes a list and reports it, while the plain removal path passes
null and neither samples nor logs.

The base class default cannot tell which of its candidates were actually removed, so it names no
keys, and logKeysNotReplaced falls back to the count-only message for any implementation that does
not override the hook.

Capped at 8 keys (NUM_SAMPLED_KEYS_NOT_REPLACED), because the count itself can be large and a
primary key is customer data. Getting the keys is free: they are already in hand at the increment,
and UpsertUtils.getPrimaryKeyIterator hands out a fresh PrimaryKey per doc
(PrimaryKeyReader.getPrimaryKey(int) allocates a new one), so a retained reference is safe and
needs no defensive copy.

Performance considerations

• No extra pass over anything. The counting is an increment inside the computeIfPresent the code
  already performs per key.
• The private removeSegment(segment, iterator, @Nullable int[] counter) body is shared by the
  counting and non-counting paths, so the non-counting path is unchanged apart from a null check per
  key.
• The removal was already happening on this path. Only the reporting moved.
• The key sampling adds one null check per removed key, and one bounded ArrayList per replacement.
  The plain removal path allocates nothing.

Testing

Test	Covers
BasePartitionUpsertMetadataManagerTest.testReplaceSegmentDoesNotReportKeyMovedBeforeRemoval	A key moved onto a newer segment during replacement is not reported
BasePartitionUpsertMetadataManagerTest.testReplaceSegmentReportsOnlyKeysStillOwnedAtRemoval	A genuinely unreplaced key is still reported, with the right count
Both removeSegmentAndGetNumKeysRemoved tests	The sample names only the key still owned by the old segment, not the one ingestion moved
ConcurrentMapPartitionUpsertMetadataManagerTest.testRemoveSegmentCountsOnlyKeysStillOwnedBySegment	The override counts only keys passing the ownership check
ConcurrentMapPartitionUpsertMetadataManagerForConsistentDeletesTest.testRemoveSegmentCountsOnlyKeysStillOwnedBySegment	Same for the consistent-deletes path

Every test in org.apache.pinot.segment.local.upsert passes: 104 tests, 0 failures. spotless,
checkstyle, license:format and license:check are clean on pinot-segment-local.

References

• Postmortem: ZD#7481, RCA-299
• StarTree's off-heap RocksDB metadata manager needs the same treatment on its async-removal path,
  where the count is only knowable during the background RocksDB scan. That is a separate change in
  the StarTree fork and does not affect this PR.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 71.42857% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.71%. Comparing base (6aaa040) to head (612883b).
⚠️ Report is 8 commits behind head on master.

Files with missing lines	Patch %	Lines
...tionUpsertMetadataManagerForConsistentDeletes.java	68.18%	4 Missing and 3 partials ⚠️
...t/ConcurrentMapPartitionUpsertMetadataManager.java	64.28%	4 Missing and 1 partial ⚠️
...cal/upsert/BasePartitionUpsertMetadataManager.java	84.61%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19504      +/-   ##
============================================
- Coverage     67.72%   67.71%   -0.01%     
  Complexity     1430     1430              
============================================
  Files          3489     3489              
  Lines        224634   224666      +32     
  Branches      35468    35473       +5     
============================================
+ Hits         152130   152136       +6     
- Misses        60485    60499      +14     
- Partials      12019    12031      +12     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	67.71% <71.42%> (-0.01%)	⬇️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	67.71% <71.42%> (-0.01%)	⬇️
unittests	67.71% <71.42%> (-0.01%)	⬇️
unittests1	57.76% <0.00%> (-0.03%)	⬇️
unittests2	39.45% <71.42%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.