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 f0c8f850b89a..d78a23641864 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 @@ -494,17 +494,7 @@ public void addSegment(ImmutableSegmentImpl segment, @Nullable ThreadSafeMutable if (queryableDocIds == null && _deleteRecordColumn != null) { queryableDocIds = new ThreadSafeMutableRoaringBitmap(); } - addOrReplaceSegment(segment, validDocIds, queryableDocIds, recordInfoIterator, null, null); - } - - protected void addOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds, - @Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator recordInfoIterator, - @Nullable IndexSegment oldSegment, @Nullable MutableRoaringBitmap validDocIdsForOldSegment) { - if (_partialUpsertHandler != null) { - recordInfoIterator = resolveComparisonTies(recordInfoIterator, _hashFunction); - } - doAddOrReplaceSegment(segment, validDocIds, queryableDocIds, recordInfoIterator, oldSegment, - validDocIdsForOldSegment); + doAddOrReplaceSegment(segment, validDocIds, queryableDocIds, recordInfoIterator, null, null); } protected abstract void doAddOrReplaceSegment(ImmutableSegmentImpl segment, @@ -514,7 +504,7 @@ protected abstract void doAddOrReplaceSegment(ImmutableSegmentImpl segment, protected void addSegmentWithoutUpsert(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds, @Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator recordInfoIterator) { - addOrReplaceSegment(segment, validDocIds, queryableDocIds, recordInfoIterator, null, null); + doAddOrReplaceSegment(segment, validDocIds, queryableDocIds, recordInfoIterator, null, null); } /** @@ -668,8 +658,8 @@ public void replaceSegment(ImmutableSegment segment, @Nullable ThreadSafeMutable if (queryableDocIds == null && _deleteRecordColumn != null) { queryableDocIds = new ThreadSafeMutableRoaringBitmap(); } - addOrReplaceSegment((ImmutableSegmentImpl) segment, validDocIds, queryableDocIds, recordInfoIterator, oldSegment, - validDocIdsForOldSegment); + doAddOrReplaceSegment((ImmutableSegmentImpl) segment, validDocIds, queryableDocIds, recordInfoIterator, + oldSegment, validDocIdsForOldSegment); } if (_upsertViewManager != null) { // When using consistency mode, the old segment's bitmap is updated in place, so we get the validDocIds after 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 a06bd82395d9..f580712576f5 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 @@ -71,6 +71,9 @@ protected long getNumPrimaryKeys() { protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds, @Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator recordInfoIterator, @Nullable IndexSegment oldSegment, @Nullable MutableRoaringBitmap validDocIdsForOldSegment) { + if (_partialUpsertHandler != null) { + recordInfoIterator = resolveComparisonTies(recordInfoIterator, _hashFunction); + } String segmentName = segment.getSegmentName(); segment.enableUpsert(this, validDocIds, queryableDocIds); 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 9a1c4f133cb8..929df2ac3ac3 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 @@ -90,7 +90,9 @@ protected long getNumPrimaryKeys() { protected void doAddOrReplaceSegment(ImmutableSegmentImpl segment, ThreadSafeMutableRoaringBitmap validDocIds, @Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, Iterator recordInfoIterator, @Nullable IndexSegment oldSegment, @Nullable MutableRoaringBitmap validDocIdsForOldSegment) { - if (_partialUpsertHandler == null) { + if (_partialUpsertHandler != null) { + recordInfoIterator = resolveComparisonTies(recordInfoIterator, _hashFunction); + } else { // for full upsert, we are de-duping primary key once here to make sure that we are not adding // primary-key multiple times and subtracting just once in removeSegment. // for partial-upsert, we call this method in base class. @@ -298,7 +300,7 @@ public void replaceSegment(ImmutableSegment segment, @Nullable ThreadSafeMutable if (queryableDocIds == null && _deleteRecordColumn != null) { queryableDocIds = new ThreadSafeMutableRoaringBitmap(); } - addOrReplaceSegment((ImmutableSegmentImpl) segment, validDocIds, queryableDocIds, recordInfoIterator, + doAddOrReplaceSegment((ImmutableSegmentImpl) segment, validDocIds, queryableDocIds, recordInfoIterator, oldSegment, validDocIdsForOldSegment); } if (validDocIdsForOldSegment != null && !validDocIdsForOldSegment.isEmpty()) {