You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/03/19 20:16:37 UTC
(pinot) branch master updated: Add additional pruning step for Upsert Tables (#17898)
This is an automated email from the ASF dual-hosted git repository.
xiangfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new ad1cdb13af8 Add additional pruning step for Upsert Tables (#17898)
ad1cdb13af8 is described below
commit ad1cdb13af885c074458e76454b3d97c6055aa66
Author: Chaitanya Deepthi <45...@users.noreply.github.com>
AuthorDate: Thu Mar 19 13:16:30 2026 -0700
Add additional pruning step for Upsert Tables (#17898)
* Add additional pruning condition for Upsert Tables
* Fix the no broker query options
* Change the pruning to take in queryable docids as well
* Change the code to reuse from Utility class
* Checkstyle exceptions
---
.../core/query/pruner/SegmentPrunerService.java | 26 +++-
.../query/pruner/SegmentPrunerServiceTest.java | 145 +++++++++++++++++++++
2 files changed, 166 insertions(+), 5 deletions(-)
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/SegmentPrunerService.java b/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/SegmentPrunerService.java
index a426cb7abea..084a00f9351 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/SegmentPrunerService.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/SegmentPrunerService.java
@@ -25,11 +25,14 @@ import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.function.BiConsumer;
import javax.annotation.Nullable;
+import org.apache.pinot.common.utils.config.QueryOptionsUtils;
import org.apache.pinot.core.query.config.SegmentPrunerConfig;
import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.segment.local.upsert.UpsertUtils;
import org.apache.pinot.segment.spi.IndexSegment;
import org.apache.pinot.spi.trace.InvocationScope;
import org.apache.pinot.spi.trace.Tracing;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -105,7 +108,7 @@ public class SegmentPrunerService {
public List<IndexSegment> prune(List<IndexSegment> segments, QueryContext query, SegmentPrunerStatistics stats,
@Nullable ExecutorService executorService) {
try (InvocationScope scope = Tracing.getTracer().createScope(SegmentPrunerService.class)) {
- segments = removeEmptySegments(segments);
+ segments = removeEmptySegments(segments, query);
int invokedPrunersCount = 0;
for (SegmentPruner segmentPruner : _segmentPruners) {
if (segmentPruner.isApplicableTo(query)) {
@@ -134,19 +137,32 @@ public class SegmentPrunerService {
*
* @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.
* @return the new list with filtered elements. This is the list that have to be used.
*/
- private static List<IndexSegment> removeEmptySegments(List<IndexSegment> segments) {
+ private static List<IndexSegment> removeEmptySegments(List<IndexSegment> segments, @Nullable QueryContext query) {
int selected = 0;
for (IndexSegment segment : segments) {
- if (!isEmptySegment(segment)) {
+ if (!isEmptySegment(segment, query)) {
segments.set(selected++, segment);
}
}
return segments.subList(0, selected);
}
- private static boolean isEmptySegment(IndexSegment segment) {
- return segment.getSegmentMetadata().getTotalDocs() == 0;
+ private static boolean isEmptySegment(IndexSegment segment, @Nullable QueryContext query) {
+ if (segment.getSegmentMetadata().getTotalDocs() == 0) {
+ return true;
+ }
+ // For upsert tables, treat segments with 0 docs to query as empty only when the query does not skip upsert.
+ // When skipUpsert=true, the query returns all docs (including replaced), so the segment contributes rows.
+ // Use query options map directly: _skipUpsert is only set on the server; the broker has options in the map.
+ if (query != null && query.getQueryOptions() != null && QueryOptionsUtils.isSkipUpsert(query.getQueryOptions())) {
+ return false;
+ }
+ MutableRoaringBitmap queryableDocIds = UpsertUtils.getQueryableDocIdsSnapshotFromSegment(segment);
+ return queryableDocIds != null && queryableDocIds.isEmpty();
}
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/SegmentPrunerServiceTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/SegmentPrunerServiceTest.java
index c18c670c57c..93e382b4765 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/SegmentPrunerServiceTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/SegmentPrunerServiceTest.java
@@ -29,6 +29,7 @@ import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
import org.apache.pinot.segment.spi.IndexSegment;
import org.apache.pinot.segment.spi.SegmentMetadata;
+import org.apache.pinot.segment.spi.index.mutable.ThreadSafeMutableRoaringBitmap;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.utils.CommonConstants.Server;
import org.testng.Assert;
@@ -87,6 +88,135 @@ public class SegmentPrunerServiceTest {
Assert.assertEquals(stats.getInvalidSegments(), 0);
}
+ @Test
+ public void emptyValidPruned() {
+ 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");
+
+ List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, Collections.emptyList());
+ }
+
+ @Test
+ public void emptyValidRetainedWithSkipUpsert() {
+ 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(skipUpsert=true)");
+
+ List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, segments);
+ }
+
+ @Test
+ public void nonEmptyValidNotPrunedWhenQueryableNull() {
+ SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);
+ IndexSegment segment = mockUpsertIndexSegment(10, new ThreadSafeMutableRoaringBitmap(0), null);
+
+ List<IndexSegment> indexes = new ArrayList<>();
+ indexes.add(segment);
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("select col1 from t1");
+
+ List<IndexSegment> actual = service.prune(indexes, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, indexes);
+ }
+
+ @Test
+ public void nonUpsertSegmentWithNullDocIdBitmapsNotPruned() {
+ SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);
+ IndexSegment segment = mockUpsertIndexSegment(5, null, null);
+
+ List<IndexSegment> segments = new ArrayList<>();
+ segments.add(segment);
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("select col1 from t1");
+
+ List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, segments);
+ }
+
+ /**
+ * When queryable doc ids exist, a segment with no queryable rows is pruned even if valid doc ids still hold
+ * replaced rows (matches normal upsert query semantics).
+ */
+ @Test
+ public void emptyQueryablePruned() {
+ 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");
+
+ List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, Collections.emptyList());
+ }
+
+ @Test
+ public void emptyQueryableRetainedWithSkipUpsert() {
+ 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(skipUpsert=true)");
+
+ List<IndexSegment> actual = service.prune(segments, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, segments);
+ }
+
+ @Test
+ public void nonEmptyQueryableNotPruned() {
+ SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);
+ ThreadSafeMutableRoaringBitmap valid = new ThreadSafeMutableRoaringBitmap(0);
+ ThreadSafeMutableRoaringBitmap queryable = new ThreadSafeMutableRoaringBitmap(0);
+ IndexSegment segment = mockUpsertIndexSegment(10, valid, queryable);
+
+ List<IndexSegment> indexes = new ArrayList<>();
+ indexes.add(segment);
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("select col1 from t1");
+
+ List<IndexSegment> actual = service.prune(indexes, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, indexes);
+ }
+
+ /**
+ * Queryable bitmap takes precedence over valid when both are present.
+ */
+ @Test
+ public void nonEmptyQueryableOverridesEmptyValid() {
+ SegmentPrunerService service = new SegmentPrunerService(_emptyPrunerConf);
+ ThreadSafeMutableRoaringBitmap valid = new ThreadSafeMutableRoaringBitmap();
+ ThreadSafeMutableRoaringBitmap queryable = new ThreadSafeMutableRoaringBitmap(0);
+ IndexSegment segment = mockUpsertIndexSegment(10, valid, queryable);
+
+ List<IndexSegment> indexes = new ArrayList<>();
+ indexes.add(segment);
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("select col1 from t1");
+
+ List<IndexSegment> actual = service.prune(indexes, queryContext, new SegmentPrunerStatistics());
+
+ Assert.assertEquals(actual, indexes);
+ }
+
private IndexSegment mockIndexSegment(int totalDocs, String... columns) {
IndexSegment indexSegment = mock(IndexSegment.class);
when(indexSegment.getColumnNames()).thenReturn(new HashSet<>(Arrays.asList(columns)));
@@ -95,4 +225,19 @@ public class SegmentPrunerServiceTest {
when(indexSegment.getSegmentMetadata()).thenReturn(segmentMetadata);
return indexSegment;
}
+
+ /**
+ * Segment with upsert-style doc id bitmaps (valid and optional queryable).
+ */
+ private IndexSegment mockUpsertIndexSegment(int totalDocs,
+ ThreadSafeMutableRoaringBitmap validDocIds, ThreadSafeMutableRoaringBitmap queryableDocIds) {
+ IndexSegment indexSegment = mock(IndexSegment.class);
+ when(indexSegment.getColumnNames()).thenReturn(new HashSet<>(Arrays.asList("col1")));
+ SegmentMetadata segmentMetadata = mock(SegmentMetadata.class);
+ when(segmentMetadata.getTotalDocs()).thenReturn(totalDocs);
+ when(indexSegment.getSegmentMetadata()).thenReturn(segmentMetadata);
+ when(indexSegment.getValidDocIds()).thenReturn(validDocIds);
+ when(indexSegment.getQueryableDocIds()).thenReturn(queryableDocIds);
+ return indexSegment;
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org