Add skipUpsertDelete query option to view valid docs instead of queryable docs (#19009)

Author: deepthi912

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -165,6 +165,10 @@ public static boolean isSkipUpsertView(Map<String, String> queryOptions) {
     return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SKIP_UPSERT_VIEW));
   }
 
+  public static boolean isSkipUpsertDelete(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SKIP_UPSERT_DELETE));
+  }
+
   public static boolean isTraceRuleProductions(Map<String, String> queryOptions) {
     return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.TRACE_RULE_PRODUCTIONS));
   }

pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java

@@ -97,20 +97,20 @@ public FilterPlanNode(SegmentContext segmentContext, QueryContext queryContext,
 
   @Override
   public BaseFilterOperator run() {
-    MutableRoaringBitmap queryableDocIdsSnapshot = _segmentContext.getQueryableDocIdsSnapshot();
+    MutableRoaringBitmap docIdsSnapshot = _segmentContext.getDocIdsSnapshot();
     int numDocs = _indexSegment.getSegmentMetadata().getTotalDocs();
 
     if (_filter != null) {
       BaseFilterOperator filterOperator = constructPhysicalOperator(_filter, numDocs);
-      if (queryableDocIdsSnapshot != null) {
-        BaseFilterOperator validDocFilter = new BitmapBasedFilterOperator(queryableDocIdsSnapshot, false, numDocs);
+      if (docIdsSnapshot != null) {
+        BaseFilterOperator validDocFilter = new BitmapBasedFilterOperator(docIdsSnapshot, false, numDocs);
         return FilterOperatorUtils.getAndFilterOperator(_queryContext, Arrays.asList(filterOperator, validDocFilter),
             numDocs);
       } else {
         return filterOperator;
       }
-    } else if (queryableDocIdsSnapshot != null) {
-      return new BitmapBasedFilterOperator(queryableDocIdsSnapshot, false, numDocs);
+    } else if (docIdsSnapshot != null) {
+      return new BitmapBasedFilterOperator(docIdsSnapshot, false, numDocs);
     } else {
       return new MatchAllFilterOperator(numDocs);
     }

pinot-core/src/main/java/org/apache/pinot/core/query/pruner/SegmentPrunerService.java

@@ -136,26 +136,31 @@ public List<IndexSegment> prune(List<IndexSegment> segments, QueryContext query,
    * @param segments the list of segments to be pruned. This is a destructive operation that may modify this list in an
    *                 undefined way. Therefore, this list should not be used after calling this method.
    * @param query    query context; when non-null and skipUpsert=true, segments with 0 queryable/valid docs are not
-   *                 treated as empty (they contribute replaced rows to the result). When queryable doc ids exist,
-   *                 emptiness is determined from them; otherwise {@link IndexSegment#getValidDocIds()} is used.
+   *                 treated as empty (they contribute replaced rows to the result). When skipUpsertDelete=true,
+   *                 emptiness is determined from valid docs (tombstones count as non-empty); otherwise from
+   *                 queryable docs.
    * @return the new list with filtered elements. This is the list that have to be used.
    */
   private static List<IndexSegment> removeEmptySegments(List<IndexSegment> segments, QueryContext query) {
     int selected = 0;
-    boolean skipUpsert = QueryOptionsUtils.isSkipUpsert(query.getQueryOptions());
+    Map<String, String> queryOptions = query.getQueryOptions();
+    boolean skipUpsert = QueryOptionsUtils.isSkipUpsert(queryOptions);
+    boolean skipUpsertDelete = QueryOptionsUtils.isSkipUpsertDelete(queryOptions);
     for (IndexSegment segment : segments) {
-      if (!isEmptySegment(segment, skipUpsert)) {
+      if (!isEmptySegment(segment, skipUpsert, skipUpsertDelete)) {
         segments.set(selected++, segment);
       }
     }
     return segments.subList(0, selected);
   }
 
-  private static boolean isEmptySegment(IndexSegment segment, boolean skipUpsert) {
+  private static boolean isEmptySegment(IndexSegment segment, boolean skipUpsert, boolean skipUpsertDelete) {
     if (segment.getSegmentMetadata().getTotalDocs() == 0) {
       return true;
     }
-    // Check if the segment has 0 queryable docIds while skipUpsert=false
-    return !skipUpsert && segment.hasNoQueryableDocs();
+    if (skipUpsert) {
+      return false;
+    }
+    return skipUpsertDelete ? segment.hasNoValidDocs() : segment.hasNoQueryableDocs();
   }
 }

pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeTest.java

@@ -94,7 +94,7 @@ public void testConsistentSnapshot()
     // Result should be invariant - always exactly 3 docs
     for (int i = 0; i < 10_000; i++) {
       SegmentContext segmentContext = new SegmentContext(segment);
-      segmentContext.setQueryableDocIdsSnapshot(UpsertUtils.getQueryableDocIdsSnapshotFromSegment(segment));
+      segmentContext.setDocIdsSnapshot(UpsertUtils.getQueryableDocIdsSnapshotFromSegment(segment));
       assertEquals(getNumberOfFilteredDocs(segmentContext, queryContext), 3);
     }
 

pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java

@@ -165,7 +165,7 @@ public void testPlanMaker(String query, Class<? extends Operator<?>> operatorCla
     assertTrue(operatorClass.isInstance(operator));
 
     SegmentContext segmentContext = new SegmentContext(_upsertIndexSegment);
-    segmentContext.setQueryableDocIdsSnapshot(UpsertUtils.getQueryableDocIdsSnapshotFromSegment(_upsertIndexSegment));
+    segmentContext.setDocIdsSnapshot(UpsertUtils.getQueryableDocIdsSnapshotFromSegment(_upsertIndexSegment));
     Operator<?> upsertOperator = PLAN_MAKER.makeSegmentPlanNode(segmentContext, queryContext).run();
     assertTrue(upsertOperatorClass.isInstance(upsertOperator));
   }

pinot-core/src/test/java/org/apache/pinot/core/query/pruner/SegmentPrunerServiceTest.java

@@ -186,6 +186,42 @@ public void emptyQueryableRetainedWithSkipUpsert() {
     Assert.assertEquals(actual, segments);
   }
 
+  @Test
+  public void emptyQueryableRetainedWithSkipUpsertDelete() {
+    SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);
+    ThreadSafeMutableRoaringBitmap valid = new ThreadSafeMutableRoaringBitmap(0);
+    ThreadSafeMutableRoaringBitmap queryable = new ThreadSafeMutableRoaringBitmap();
+    IndexSegment segment = mockUpsertIndexSegment(10, valid, queryable);
+
+    List<IndexSegment> segments = new ArrayList<>();
+    segments.add(segment);
+    QueryContext queryContext =
+        QueryContextConverterUtils.getQueryContext("select col1 from t1 option(skipUpsertDelete=true)");
+
+    List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+    Assert.assertEquals(actual, segments);
+  }
+
+  /**
+   * skipUpsertDelete checks valid-docs emptiness specifically, not skipUpsert's blanket bypass: a segment fully
+   * superseded elsewhere (0 valid docs) is still pruned.
+   */
+  @Test
+  public void emptyValidPrunedWithSkipUpsertDelete() {
+    SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);
+    IndexSegment segment = mockUpsertIndexSegment(10, new ThreadSafeMutableRoaringBitmap(), null);
+
+    List<IndexSegment> segments = new ArrayList<>();
+    segments.add(segment);
+    QueryContext queryContext =
+        QueryContextConverterUtils.getQueryContext("select col1 from t1 option(skipUpsertDelete=true)");
+
+    List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+    Assert.assertEquals(actual, List.of());
+  }
+
   @Test
   public void nonEmptyQueryableNotPruned() {
     SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java

@@ -40,6 +40,7 @@
 import org.apache.pinot.segment.local.segment.virtualcolumn.VirtualColumnContext;
 import org.apache.pinot.segment.local.startree.v2.store.StarTreeIndexContainer;
 import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager;
+import org.apache.pinot.segment.local.upsert.UpsertUtils;
 import org.apache.pinot.segment.local.upsert.UpsertViewManager;
 import org.apache.pinot.segment.spi.ColumnMetadata;
 import org.apache.pinot.segment.spi.FetchContext;
@@ -382,6 +383,11 @@ public boolean hasNoQueryableDocs() {
     return validDocIds != null && validDocIds.isEmpty();
   }
 
+  @Override
+  public boolean hasNoValidDocs() {
+    return UpsertUtils.hasNoValidDocs(_partitionUpsertMetadataManager, this);
+  }
+
   @Override
   public GenericRow getRecord(int docId, GenericRow reuse) {
     try (PinotSegmentRecordReader recordReader = new PinotSegmentRecordReader()) {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java

@@ -74,6 +74,7 @@
 import org.apache.pinot.segment.local.upsert.PartitionUpsertMetadataManager;
 import org.apache.pinot.segment.local.upsert.RecordInfo;
 import org.apache.pinot.segment.local.upsert.UpsertContext;
+import org.apache.pinot.segment.local.upsert.UpsertUtils;
 import org.apache.pinot.segment.local.upsert.UpsertViewManager;
 import org.apache.pinot.segment.local.utils.FixedIntArrayOffHeapIdMap;
 import org.apache.pinot.segment.local.utils.IdMap;
@@ -1242,6 +1243,11 @@ public boolean hasNoQueryableDocs() {
     return validDocIds != null && validDocIds.isEmpty();
   }
 
+  @Override
+  public boolean hasNoValidDocs() {
+    return UpsertUtils.hasNoValidDocs(_partitionUpsertMetadataManager, this);
+  }
+
   @Override
   public GenericRow getRecord(int docId, GenericRow reuse) {
     try (PinotSegmentRecordReader recordReader = new PinotSegmentRecordReader()) {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/ConcurrentMapTableUpsertMetadataManager.java

@@ -95,24 +95,28 @@ public void setSegmentContexts(List<SegmentContext> segmentContexts, Map<String,
     // Otherwise, get queryableDocIds bitmaps as kept by the segment objects directly as before.
     if (_context.getConsistencyMode() == UpsertConfig.ConsistencyMode.NONE || QueryOptionsUtils.isSkipUpsertView(
         queryOptions)) {
+      // No shared upsert-view lock exists in this branch, so a direct read is already safe here.
+      boolean skipUpsertDelete = QueryOptionsUtils.isSkipUpsertDelete(queryOptions);
       for (SegmentContext segmentContext : segmentContexts) {
         IndexSegment segment = segmentContext.getIndexSegment();
-        segmentContext.setQueryableDocIdsSnapshot(UpsertUtils.getQueryableDocIdsSnapshotFromSegment(segment));
+        segmentContext.setDocIdsSnapshot(skipUpsertDelete
+            ? UpsertUtils.getValidDocIdsSnapshotFromSegment(segment)
+            : UpsertUtils.getQueryableDocIdsSnapshotFromSegment(segment));
       }
       return;
     }
-    // All segments should have been tracked by partitionMetadataManagers to provide queries consistent upsert view.
+    // A consistency mode is active: UpsertViewManager knows the locking each mode requires for skipUpsertDelete too.
     _partitionMetadataManagerMap.forEach(
         (partitionID, upsertMetadataManager) -> upsertMetadataManager.getUpsertViewManager()
             .setSegmentContexts(segmentContexts, queryOptions));
     if (LOGGER.isDebugEnabled()) {
       for (SegmentContext segmentContext : segmentContexts) {
         IndexSegment segment = segmentContext.getIndexSegment();
-        if (segmentContext.getQueryableDocIdsSnapshot() == null) {
+        if (segmentContext.getDocIdsSnapshot() == null) {
           LOGGER.debug("No upsert view for segment: {}, type: {}, total: {}", segment.getSegmentName(),
               (segment instanceof ImmutableSegment ? "imm" : "mut"), segment.getSegmentMetadata().getTotalDocs());
         } else {
-          int cardCnt = segmentContext.getQueryableDocIdsSnapshot().getCardinality();
+          int cardCnt = segmentContext.getDocIdsSnapshot().getCardinality();
           LOGGER.debug("Got upsert view of segment: {}, type: {}, total: {}, valid: {}", segment.getSegmentName(),
               (segment instanceof ImmutableSegment ? "imm" : "mut"), segment.getSegmentMetadata().getTotalDocs(),
               cardCnt);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/UpsertUtils.java

@@ -57,6 +57,36 @@ public static MutableRoaringBitmap getQueryableDocIdsSnapshotFromSegment(IndexSe
         : (useEmptyForNull ? new MutableRoaringBitmap() : null);
   }
 
+  /// Unlike [#getQueryableDocIdsSnapshotFromSegment], never falls back to preferring queryable docs.
+  @Nullable
+  public static MutableRoaringBitmap getValidDocIdsSnapshotFromSegment(IndexSegment segment) {
+    return getValidDocIdsSnapshotFromSegment(segment, false);
+  }
+
+  /// Shared by {@code ImmutableSegmentImpl}/{@code MutableSegmentImpl}'s `hasNoValidDocs()`. Mirrors
+  /// `hasNoQueryableDocs()`'s consistency-mode-aware/live-fallback split, against the valid-docs cache instead.
+  public static boolean hasNoValidDocs(@Nullable PartitionUpsertMetadataManager partitionUpsertMetadataManager,
+      IndexSegment segment) {
+    if (partitionUpsertMetadataManager == null) {
+      return false;
+    }
+    UpsertViewManager viewManager = partitionUpsertMetadataManager.getUpsertViewManager();
+    if (viewManager != null) {
+      MutableRoaringBitmap validDocIdsSnapshot = viewManager.getValidDocIdsSnapshot(segment);
+      return validDocIdsSnapshot != null && validDocIdsSnapshot.isEmpty();
+    }
+    ThreadSafeMutableRoaringBitmap validDocIds = segment.getValidDocIds();
+    return validDocIds != null && validDocIds.isEmpty();
+  }
+
+  @Nullable
+  public static MutableRoaringBitmap getValidDocIdsSnapshotFromSegment(IndexSegment segment,
+      boolean useEmptyForNull) {
+    ThreadSafeMutableRoaringBitmap validDocIds = segment.getValidDocIds();
+    return validDocIds != null ? validDocIds.getMutableRoaringBitmap()
+        : (useEmptyForNull ? new MutableRoaringBitmap() : null);
+  }
+
   public static void doReplaceDocId(ThreadSafeMutableRoaringBitmap validDocIds,
       @Nullable ThreadSafeMutableRoaringBitmap queryableDocIds, int oldDocId, int newDocId, RecordInfo recordInfo) {
     validDocIds.replace(oldDocId, newDocId);