This is an automated email from the ASF dual-hosted git repository.
raghavyadav01 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 66eb74d69d1 Add requiresFetch flag to a segment to for prefetch
planning (#20068)
66eb74d69d1 is described below
commit 66eb74d69d185ae99772b5a41cb6acef2aa583b4
Author: Jhow <[email protected]>
AuthorDate: Sat Oct 31 09:22:41 2026 +0800
Add requiresFetch flag to a segment to for prefetch planning (#20068)
* add requiresFetch flag to a segment to skip prefetch planning
* do run planner for each segment, let the planner decide to skip
* rework
* nit
* verify fetch lifecycle in test, add test for streaming
* Address review nits in InstancePlanMakerPrefetchTest
---
.../core/plan/maker/InstancePlanMakerImplV2.java | 74 +++---
.../pinot/core/query/prefetch/FetchPlanner.java | 10 +-
.../query/pruner/BloomFilterSegmentPruner.java | 5 +-
.../plan/maker/InstancePlanMakerPrefetchTest.java | 283 +++++++++++++++++++++
.../query/pruner/BloomFilterSegmentPrunerTest.java | 50 ++++
.../immutable/ImmutableSegmentImpl.java | 5 +
.../org/apache/pinot/segment/spi/IndexSegment.java | 8 +
.../pinot/segment/spi/store/SegmentDirectory.java | 8 +
8 files changed, 402 insertions(+), 41 deletions(-)
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
index 464bbd77897..5f5911e51c4 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
@@ -30,6 +30,7 @@ import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
+import java.util.function.Function;
import javax.annotation.Nullable;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
@@ -188,6 +189,29 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
_groupByTrimThreshold = groupByTrimThreshold;
}
+ /// Adds one plan node per segment to `planNodes`. With prefetch enabled, a
segment for which the [FetchPlanner]
+ /// plans a [FetchContext] gets a plan node that acquires and releases its
buffers around execution, and is added to
+ /// `fetchSegmentContexts` with its context added to `fetchContexts` at the
same position. A segment for which the
+ /// planner returns null (it needs no fetch) gets the plain plan node, as
without prefetch. The fetch lists are only
+ /// written with prefetch enabled, so they can be immutable empty lists
otherwise.
+ private void addSegmentPlanNodes(List<SegmentContext> segmentContexts,
QueryContext queryContext,
+ Function<SegmentContext, PlanNode> segmentPlanNodeMaker, List<PlanNode>
planNodes,
+ List<SegmentContext> fetchSegmentContexts, List<FetchContext>
fetchContexts) {
+ boolean enablePrefetch = queryContext.isEnablePrefetch();
+ for (SegmentContext segmentContext : segmentContexts) {
+ FetchContext fetchContext =
+ enablePrefetch ?
_fetchPlanner.planFetchForProcessing(segmentContext.getIndexSegment(),
queryContext) : null;
+ if (fetchContext != null) {
+ fetchSegmentContexts.add(segmentContext);
+ fetchContexts.add(fetchContext);
+ planNodes.add(new
AcquireReleaseColumnsSegmentPlanNode(segmentPlanNodeMaker.apply(segmentContext),
+ segmentContext, fetchContext));
+ } else {
+ planNodes.add(segmentPlanNodeMaker.apply(segmentContext));
+ }
+ }
+ }
+
@Override
public Plan makeInstancePlan(List<SegmentContext> segmentContexts,
QueryContext queryContext,
ExecutorService executorService) {
@@ -197,27 +221,15 @@ public class InstancePlanMakerImplV2 implements PlanMaker
{
int numSegments = segmentContexts.size();
List<PlanNode> planNodes = new ArrayList<>(numSegments);
- List<FetchContext> fetchContexts;
- if (queryContext.isEnablePrefetch()) {
- fetchContexts = new ArrayList<>(numSegments);
- for (SegmentContext segmentContext : segmentContexts) {
- FetchContext fetchContext =
-
_fetchPlanner.planFetchForProcessing(segmentContext.getIndexSegment(),
queryContext);
- fetchContexts.add(fetchContext);
- planNodes.add(
- new
AcquireReleaseColumnsSegmentPlanNode(makeSegmentPlanNode(segmentContext,
queryContext), segmentContext,
- fetchContext));
- }
- } else {
- fetchContexts = List.of();
- for (SegmentContext segmentContext : segmentContexts) {
- planNodes.add(makeSegmentPlanNode(segmentContext, queryContext));
- }
- }
+ boolean enablePrefetch = queryContext.isEnablePrefetch();
+ List<SegmentContext> fetchSegmentContexts = enablePrefetch ? new
ArrayList<>(numSegments) : List.of();
+ List<FetchContext> fetchContexts = enablePrefetch ? new
ArrayList<>(numSegments) : List.of();
+ addSegmentPlanNodes(segmentContexts, queryContext, segmentContext ->
makeSegmentPlanNode(segmentContext,
+ queryContext), planNodes, fetchSegmentContexts, fetchContexts);
CombinePlanNode combinePlanNode = createCombinePlanNode(planNodes,
queryContext, executorService, null);
return new GlobalPlanImplV0(
- new InstanceResponsePlanNode(combinePlanNode, segmentContexts,
fetchContexts, queryContext));
+ new InstanceResponsePlanNode(combinePlanNode, fetchSegmentContexts,
fetchContexts, queryContext));
}
/// Resolves [AndRestrictionPushdownMode] for a query.
@@ -403,27 +415,15 @@ public class InstancePlanMakerImplV2 implements PlanMaker
{
int numSegments = segmentContexts.size();
List<PlanNode> planNodes = new ArrayList<>(numSegments);
- List<FetchContext> fetchContexts;
- if (queryContext.isEnablePrefetch()) {
- fetchContexts = new ArrayList<>(numSegments);
- for (SegmentContext segmentContext : segmentContexts) {
- FetchContext fetchContext =
-
_fetchPlanner.planFetchForProcessing(segmentContext.getIndexSegment(),
queryContext);
- fetchContexts.add(fetchContext);
- planNodes.add(
- new
AcquireReleaseColumnsSegmentPlanNode(makeStreamingSegmentPlanNode(segmentContext,
queryContext),
- segmentContext, fetchContext));
- }
- } else {
- fetchContexts = List.of();
- for (SegmentContext segmentContext : segmentContexts) {
- planNodes.add(makeStreamingSegmentPlanNode(segmentContext,
queryContext));
- }
- }
+ boolean enablePrefetch = queryContext.isEnablePrefetch();
+ List<SegmentContext> fetchSegmentContexts = enablePrefetch ? new
ArrayList<>(numSegments) : List.of();
+ List<FetchContext> fetchContexts = enablePrefetch ? new
ArrayList<>(numSegments) : List.of();
+ addSegmentPlanNodes(segmentContexts, queryContext, segmentContext ->
makeStreamingSegmentPlanNode(segmentContext,
+ queryContext), planNodes, fetchSegmentContexts, fetchContexts);
CombinePlanNode combinePlanNode = createCombinePlanNode(planNodes,
queryContext, executorService, streamer);
- return new GlobalPlanImplV0(
- new StreamingInstanceResponsePlanNode(combinePlanNode,
segmentContexts, fetchContexts, queryContext, streamer));
+ return new GlobalPlanImplV0(new
StreamingInstanceResponsePlanNode(combinePlanNode, fetchSegmentContexts,
+ fetchContexts, queryContext, streamer));
}
@Override
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/prefetch/FetchPlanner.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/prefetch/FetchPlanner.java
index e320f01812b..cbeb4f409da 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/prefetch/FetchPlanner.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/prefetch/FetchPlanner.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.core.query.prefetch;
+import javax.annotation.Nullable;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.segment.spi.FetchContext;
import org.apache.pinot.segment.spi.IndexSegment;
@@ -27,6 +28,9 @@ import org.apache.pinot.segment.spi.IndexSegment;
/// read from some slow storage backend, to reduce the query latency. Although
the plan methods take in an indexSegment
/// object, it should mainly access the segment metadata like what types of
index it has, as the index data may not be
/// available, but leaving this to the implementations to check and handle per
their own need.
+///
+/// A planner can return null for a segment that needs no fetch, e.g. one
whose data is all local, so its fetch calls
+/// would be no-ops. The caller then skips the fetch lifecycle (prefetch,
acquire, release) for that segment entirely.
/// TODO: this interface and the dependent classes like QueryContext should be
moved to query-spi pkg, yet to be added.
public interface FetchPlanner {
/// Plan what index data to prefetch to help prune the segment before
processing it. For example, one can fetch bloom
@@ -34,7 +38,8 @@ public interface FetchPlanner {
///
/// @param indexSegment segment to be pruned.
/// @param queryContext context extracted from the query.
- /// @return context to guide data prefetching.
+ /// @return context to guide data prefetching, or null if the segment needs
no fetch.
+ @Nullable
FetchContext planFetchForPruning(IndexSegment indexSegment, QueryContext
queryContext);
/// Plan what index data to prefetch to process it. For example, one can
fetch all types of index for columns tracked
@@ -42,6 +47,7 @@ public interface FetchPlanner {
///
/// @param indexSegment segment to be processed.
/// @param queryContext context extracted from the query.
- /// @return context to guide data prefetching.
+ /// @return context to guide data prefetching, or null if the segment needs
no fetch.
+ @Nullable
FetchContext planFetchForProcessing(IndexSegment indexSegment, QueryContext
queryContext);
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPruner.java
b/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPruner.java
index d3b6d00b363..388f7e9ec2d 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPruner.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPruner.java
@@ -89,11 +89,12 @@ public class BloomFilterSegmentPruner extends
ValueBasedSegmentPruner {
int numSegments = segments.size();
FetchContext[] fetchContexts = new FetchContext[numSegments];
try {
- // Prefetch bloom filter for columns within the EQ/IN predicate if exists
+ // Prefetch bloom filter for columns within the EQ/IN predicate if
exists. The planner returns null for a segment
+ // that needs no fetch, which is then pruned without the fetch lifecycle.
for (int i = 0; i < numSegments; i++) {
IndexSegment segment = segments.get(i);
FetchContext fetchContext = _fetchPlanner.planFetchForPruning(segment,
query);
- if (!fetchContext.isEmpty()) {
+ if (fetchContext != null && !fetchContext.isEmpty()) {
segment.prefetch(fetchContext);
fetchContexts[i] = fetchContext;
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerPrefetchTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerPrefetchTest.java
new file mode 100644
index 00000000000..e72da4411b9
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerPrefetchTest.java
@@ -0,0 +1,283 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.plan.maker;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
+import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
+import org.apache.pinot.core.query.prefetch.DefaultFetchPlanner;
+import org.apache.pinot.core.query.prefetch.FetchPlanner;
+import org.apache.pinot.core.query.prefetch.FetchPlannerRegistry;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.FetchContext;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.SegmentContext;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.query.QueryThreadContext;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.mockito.ArgumentCaptor;
+import org.mockito.MockedStatic;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// With prefetch enabled, [InstancePlanMakerImplV2] runs the fetch lifecycle
only for segments for which the
+/// [FetchPlanner] plans a fetch context. A segment for which the planner
returns null (it needs no fetch) is planned
+/// and executed as without prefetch.
+public class InstancePlanMakerPrefetchTest {
+ private static final String TABLE_NAME = "testTable";
+ private static final String SEGMENT_NAME = "testSegment";
+ private static final int NUM_DOCS = 100;
+ // Every tenth document matches
+ private static final String QUERY = "SELECT COUNT(*), SUM(m) FROM testTable
WHERE d = 'v3'";
+ // A selection-only query is streamed; the limit keeps every matching row of
both segments
+ private static final String STREAMING_QUERY = "SELECT m FROM testTable WHERE
d = 'v3' LIMIT 1000";
+ private static final long MATCHES_PER_SEGMENT = NUM_DOCS / 10;
+ private static final double SUM_PER_SEGMENT = 3 + 13 + 23 + 33 + 43 + 53 +
63 + 73 + 83 + 93;
+
+ private File _tempDir;
+ private ImmutableSegment _segment;
+ private ExecutorService _executorService;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ _tempDir =
Files.createTempDirectory("instance-plan-maker-prefetch-test").toFile();
+ Schema schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME)
+ .addSingleValueDimension("d", DataType.STRING)
+ .addMetric("m", DataType.LONG)
+ .build();
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build();
+ List<GenericRow> rows = new ArrayList<>(NUM_DOCS);
+ for (int i = 0; i < NUM_DOCS; i++) {
+ GenericRow row = new GenericRow();
+ row.putValue("d", "v" + (i % 10));
+ row.putValue("m", (long) i);
+ rows.add(row);
+ }
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig,
schema);
+ config.setTableName(TABLE_NAME);
+ config.setSegmentName(SEGMENT_NAME);
+ config.setOutDir(_tempDir.getAbsolutePath());
+ SegmentIndexCreationDriverImpl driver = new
SegmentIndexCreationDriverImpl();
+ driver.init(config, new GenericRowRecordReader(rows));
+ driver.build();
+ _segment =
+ ImmutableSegmentLoader.load(new File(_tempDir, SEGMENT_NAME), new
IndexLoadingConfig(tableConfig, schema));
+ _executorService =
QueryThreadContext.contextAwareExecutorService(Executors.newFixedThreadPool(2));
+ }
+
+ @AfterClass
+ public void tearDown() {
+ _executorService.shutdownNow();
+ _segment.destroy();
+ FileUtils.deleteQuietly(_tempDir);
+ }
+
+ @Test
+ public void testPrefetchRunsFetchLifecycleOnlyForPlannedSegments()
+ throws Exception {
+ IndexSegment local = spy(_segment);
+ IndexSegment remote = spy(_segment);
+
+ assertResult(execute(plannerSkipping(Set.of(local)), List.of(local,
remote), true), 2);
+
+ verifyFetchLifecycle(remote);
+ verifyNoFetch(local);
+ }
+
+ /// Streaming selection has its own start/stop and per-segment fetch
lifecycle; only the remote segment gets it.
+ @Test
+ public void testStreamingSelectionRunsFetchLifecycleOnlyForPlannedSegments()
+ throws Exception {
+ IndexSegment local = spy(_segment);
+ IndexSegment remote = spy(_segment);
+
+ List<Long> streamedValues =
executeStreaming(plannerSkipping(Set.of(local)), List.of(local, remote));
+
+ // Each segment streams every matching value once
+ List<Long> expectedValues = new ArrayList<>();
+ for (long value = 3; value < NUM_DOCS; value += 10) {
+ expectedValues.add(value);
+ expectedValues.add(value);
+ }
+ assertEquals(streamedValues, expectedValues);
+ verifyFetchLifecycle(remote);
+ verifyNoFetch(local);
+ }
+
+ @Test
+ public void testPrefetchWithoutPlannedSegmentsRunsNoFetchLifecycle()
+ throws Exception {
+ IndexSegment first = spy(_segment);
+ IndexSegment second = spy(_segment);
+
+ assertResult(execute(plannerSkipping(Set.of(first, second)),
List.of(first, second), true), 2);
+
+ verifyNoFetch(first);
+ verifyNoFetch(second);
+ }
+
+ @Test
+ public void testDisabledPrefetchRunsNoFetchLifecycle()
+ throws Exception {
+ IndexSegment segment = spy(_segment);
+
+ assertResult(execute(plannerSkipping(Set.of()), List.of(segment), false),
1);
+
+ verifyNoFetch(segment);
+ }
+
+ /// The default planner never returns null, so every segment keeps the fetch
lifecycle as before.
+ @Test
+ public void testDefaultPlannerRunsFetchLifecycleForEverySegment()
+ throws Exception {
+ IndexSegment first = spy(_segment);
+ IndexSegment second = spy(_segment);
+
+ assertResult(execute(new DefaultFetchPlanner(), List.of(first, second),
true), 2);
+
+ for (IndexSegment segment : List.of(first, second)) {
+ verify(segment).prefetch(any());
+ verify(segment).acquire(any());
+ verify(segment, atLeastOnce()).release(any());
+ }
+ }
+
+ /// A planner that plans with the default planner, except for the given
segments, which need no fetch.
+ private static FetchPlanner plannerSkipping(Set<IndexSegment>
segmentsWithoutFetch) {
+ FetchPlanner defaultPlanner = new DefaultFetchPlanner();
+ return new FetchPlanner() {
+ @Override
+ public FetchContext planFetchForPruning(IndexSegment indexSegment,
QueryContext queryContext) {
+ return segmentsWithoutFetch.contains(indexSegment) ? null
+ : defaultPlanner.planFetchForPruning(indexSegment, queryContext);
+ }
+
+ @Override
+ public FetchContext planFetchForProcessing(IndexSegment indexSegment,
QueryContext queryContext) {
+ return segmentsWithoutFetch.contains(indexSegment) ? null
+ : defaultPlanner.planFetchForProcessing(indexSegment,
queryContext);
+ }
+ };
+ }
+
+ private DataTable execute(FetchPlanner fetchPlanner, List<IndexSegment>
segments, boolean enablePrefetch)
+ throws Exception {
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext(QUERY);
+ queryContext.setEndTimeMs(System.currentTimeMillis() + 60_000);
+ queryContext.setEnablePrefetch(enablePrefetch);
+ List<SegmentContext> segmentContexts = new ArrayList<>(segments.size());
+ for (IndexSegment segment : segments) {
+ segmentContexts.add(new SegmentContext(segment));
+ }
+ try (QueryThreadContext ignored = QueryThreadContext.openForSseTest()) {
+ InstanceResponseBlock responseBlock = planMakerWith(fetchPlanner)
+ .makeInstancePlan(segmentContexts, queryContext,
_executorService).execute();
+ assertTrue(responseBlock.getExceptions().isEmpty(), "Query failed: " +
responseBlock.getExceptions());
+ return responseBlock.toDataTable();
+ }
+ }
+
+ /// Runs the streaming selection with prefetch enabled, and returns the
streamed values of `m` in ascending order.
+ private List<Long> executeStreaming(FetchPlanner fetchPlanner,
List<IndexSegment> segments)
+ throws Exception {
+ QueryContext queryContext =
QueryContextConverterUtils.getQueryContext(STREAMING_QUERY);
+ queryContext.setEndTimeMs(System.currentTimeMillis() + 60_000);
+ queryContext.setEnablePrefetch(true);
+ List<SegmentContext> segmentContexts = new ArrayList<>(segments.size());
+ for (IndexSegment segment : segments) {
+ segmentContexts.add(new SegmentContext(segment));
+ }
+ // The streamer is invoked on the response thread, not by the combine
workers
+ List<Long> streamedValues = new ArrayList<>();
+ ResultsBlockStreamer streamer = block -> {
+ for (Object[] row : ((SelectionResultsBlock) block).getRows()) {
+ streamedValues.add(((Number) row[0]).longValue());
+ }
+ };
+ try (QueryThreadContext ignored = QueryThreadContext.openForSseTest()) {
+ InstanceResponseBlock responseBlock = planMakerWith(fetchPlanner)
+ .makeStreamingInstancePlan(segmentContexts, queryContext,
_executorService, streamer).execute();
+ assertTrue(responseBlock.getExceptions().isEmpty(), "Query failed: " +
responseBlock.getExceptions());
+ }
+ Collections.sort(streamedValues);
+ return streamedValues;
+ }
+
+ /// The plan maker takes the registered planner at construction.
+ private static InstancePlanMakerImplV2 planMakerWith(FetchPlanner
fetchPlanner) {
+ try (MockedStatic<FetchPlannerRegistry> registry =
mockStatic(FetchPlannerRegistry.class)) {
+ registry.when(FetchPlannerRegistry::getPlanner).thenReturn(fetchPlanner);
+ return new InstancePlanMakerImplV2();
+ }
+ }
+
+ private static void assertResult(DataTable dataTable, int numSegments) {
+ assertEquals(dataTable.getNumberOfRows(), 1);
+ assertEquals(dataTable.getLong(0, 0), MATCHES_PER_SEGMENT * numSegments);
+ assertEquals(dataTable.getDouble(0, 1), SUM_PER_SEGMENT * numSegments);
+ }
+
+ /// The segment is prefetched, acquired and released with one fetch context
planned for the filter column.
+ private static void verifyFetchLifecycle(IndexSegment segment) {
+ ArgumentCaptor<FetchContext> fetchContext =
ArgumentCaptor.forClass(FetchContext.class);
+ verify(segment).prefetch(fetchContext.capture());
+ verify(segment).acquire(fetchContext.getValue());
+ verify(segment, atLeastOnce()).release(fetchContext.getValue());
+
assertTrue(fetchContext.getValue().getColumnToIndexList().containsKey("d"));
+ }
+
+ private static void verifyNoFetch(IndexSegment segment) {
+ verify(segment, never()).prefetch(any());
+ verify(segment, never()).acquire(any());
+ verify(segment, never()).release(any());
+ }
+}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPrunerTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPrunerTest.java
index d7b31670c16..fba0bde5e1d 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPrunerTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPrunerTest.java
@@ -37,6 +37,9 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.apache.pinot.core.query.prefetch.DefaultFetchPlanner;
+import org.apache.pinot.core.query.prefetch.FetchPlanner;
+import org.apache.pinot.core.query.prefetch.FetchPlannerRegistry;
import org.apache.pinot.core.query.request.context.QueryContext;
import
org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
@@ -66,6 +69,7 @@ import org.apache.pinot.spi.query.RemoteAccessStats;
import org.apache.pinot.spi.utils.ReadMode;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.mockito.ArgumentCaptor;
+import org.mockito.MockedStatic;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
@@ -74,6 +78,7 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@@ -369,6 +374,51 @@ public class BloomFilterSegmentPrunerTest {
}
}
+ @Test
+ public void testPrefetchSkipsSegmentsWithoutPlannedFetch()
+ throws Exception {
+ // The planner returns null for the local segments (no fetch needed): they
are pruned without any fetch call, while
+ // the remote segment keeps the lifecycle
+ IndexSegment localKept = mockIndexSegment(new String[]{"1.0"});
+ IndexSegment localPruned = mockIndexSegment(new String[]{"2.0"});
+ IndexSegment remoteKept = mockIndexSegment(new String[]{"1.0"});
+ Set<IndexSegment> localSegments = Set.of(localKept, localPruned);
+ FetchPlanner defaultPlanner = new DefaultFetchPlanner();
+ FetchPlanner fetchPlanner = new FetchPlanner() {
+ @Override
+ public FetchContext planFetchForPruning(IndexSegment indexSegment,
QueryContext queryContext) {
+ return localSegments.contains(indexSegment) ? null
+ : defaultPlanner.planFetchForPruning(indexSegment, queryContext);
+ }
+
+ @Override
+ public FetchContext planFetchForProcessing(IndexSegment indexSegment,
QueryContext queryContext) {
+ return localSegments.contains(indexSegment) ? null
+ : defaultPlanner.planFetchForProcessing(indexSegment,
queryContext);
+ }
+ };
+ BloomFilterSegmentPruner pruner = new BloomFilterSegmentPruner();
+ // The pruner takes the registered planner at init
+ try (MockedStatic<FetchPlannerRegistry> registry =
mockStatic(FetchPlannerRegistry.class)) {
+ registry.when(FetchPlannerRegistry::getPlanner).thenReturn(fetchPlanner);
+ pruner.init(new PinotConfiguration());
+ }
+
+ assertEquals(pruner.prune(List.of(localKept, localPruned, remoteKept),
prefetchQuery()),
+ List.of(localKept, remoteKept));
+ for (IndexSegment local : localSegments) {
+ verify(local, never()).prefetch(any(FetchContext.class));
+ verify(local, never()).acquire(any(FetchContext.class));
+ verify(local, never()).release(any(FetchContext.class));
+ }
+ assertReleasedPrefetch(remoteKept, 1);
+
+ // With only local segments, pruning is unchanged and still makes no fetch
call
+ assertEquals(pruner.prune(List.of(localKept, localPruned),
prefetchQuery()), List.of(localKept));
+ verify(localKept, never()).prefetch(any(FetchContext.class));
+ verify(localPruned, never()).prefetch(any(FetchContext.class));
+ }
+
private QueryContext prefetchQuery() {
QueryContext query = QueryContextConverterUtils.getQueryContext(
"SELECT COUNT(*) FROM testTable WHERE column = 1.0");
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
index b0ac2d625dd..01540afff9e 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/ImmutableSegmentImpl.java
@@ -341,6 +341,11 @@ public class ImmutableSegmentImpl implements
ImmutableSegment {
return _segmentMetadata.getSchema().getPhysicalColumnNames();
}
+ @Override
+ public boolean requiresFetch() {
+ return _segmentDirectory.requiresFetch();
+ }
+
@Override
public void prefetch(FetchContext fetchContext) {
_segmentDirectory.prefetch(fetchContext);
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
index dcc6605f312..b7a45dcf323 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/IndexSegment.java
@@ -122,6 +122,14 @@ public interface IndexSegment {
/// values from the same segment.
Object getValue(int docId, String column);
+ /// Returns whether this segment serves buffers through [#prefetch],
[#acquire] and [#release]. When it returns false,
+ /// those calls are no-ops, so a fetch planner can skip planning a
[FetchContext] for the segment, letting query
+ /// execution skip the fetch lifecycle altogether. Segments that override
the fetch calls must return true. The value
+ /// must not change over the lifetime of the segment.
+ default boolean requiresFetch() {
+ return false;
+ }
+
/// Hints the segment to begin prefetching buffers for specified columns.
/// Typically, this should be an async call made before operating on the
segment.
/// @param fetchContext context for this segment's fetch
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
index 8a4267f949f..2cda5a5bea3 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectory.java
@@ -110,6 +110,14 @@ public abstract class SegmentDirectory implements
Closeable {
/// @return a set of columns with such index type.
public abstract Set<String> getColumnsWithIndex(IndexType<?, ?, ?> type);
+ /// Returns whether this segment directory serves buffers through
[#prefetch], [#acquire] and [#release]. When it
+ /// returns false, those calls are no-ops, so a fetch planner can skip
planning a [FetchContext] for the segment,
+ /// letting query execution skip the fetch lifecycle altogether. Directories
that override the fetch calls must
+ /// return true. The value must not change over the lifetime of the
directory.
+ public boolean requiresFetch() {
+ return false;
+ }
+
/// This is a hint to the segment directory, to begin prefetching buffers
for given context.
/// Typically, this should be an async call made before operating on the
segment.
/// @param fetchContext context for this segment's fetch
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
