Jackie-Jiang commented on code in PR #20068:
URL: https://github.com/apache/pinot/pull/20068#discussion_r4406607083
##########
pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java:
##########
@@ -218,35 +219,43 @@ public void setGroupByTrimThreshold(int
groupByTrimThreshold) {
_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.
+ 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) {
applyQueryOptions(queryContext);
boundGroupingSetsBaseAggregation(segmentContexts, queryContext);
- 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));
- }
- }
+ List<PlanNode> planNodes = new ArrayList<>(segmentContexts.size());
+ List<SegmentContext> fetchSegmentContexts = new ArrayList<>();
+ List<FetchContext> fetchContexts = new ArrayList<>();
Review Comment:
Minor allocation cleanup: both fetch lists are now allocated even when
prefetch is disabled, and the enabled path loses the previous `numSegments`
capacity hint. Use empty lists when prefetch is disabled and pre-size the
mutable lists when enabled to avoid repeated growth for queries fetching many
segments. Please apply the same cleanup to `makeStreamingInstancePlan()`. This
is a bounded allocation improvement, not a demonstrated material performance
regression.
##########
pinot-core/src/main/java/org/apache/pinot/core/query/prefetch/FetchPlanner.java:
##########
@@ -27,21 +28,27 @@
/// 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.
+/// A non-null context, even an empty one, keeps the lifecycle.
Review Comment:
Minor: this promises that any non-null context, including an empty one,
preserves the fetch lifecycle, but `BloomFilterSegmentPruner.prefetch()` still
skips empty contexts via `fetchContext != null && !fetchContext.isEmpty()`.
Please scope this guarantee to `planFetchForProcessing()` and document that
pruning skips both null and empty contexts. The existing pruning behavior can
remain unchanged.
##########
pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerPrefetchTest.java:
##########
@@ -0,0 +1,227 @@
+/**
+ * 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.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.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'";
+ 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);
+
+ ArgumentCaptor<FetchContext> fetchContext =
ArgumentCaptor.forClass(FetchContext.class);
+ verify(remote).prefetch(fetchContext.capture());
+ verify(remote).acquire(fetchContext.getValue());
+ verify(remote, atLeastOnce()).release(fetchContext.getValue());
+
assertTrue(fetchContext.getValue().getColumnToIndexList().containsKey("d"));
+ 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();
Review Comment:
Optional coverage improvement: all tests in this class execute a
non-streaming aggregation through `makeInstancePlan()`. Please retain a mixed
local/remote streaming-selection case using `makeStreamingInstancePlan()` and a
`ResultsBlockStreamer`, asserting the streamed rows and selective
prefetch/acquire/release calls. Streaming selection exercises the separate
start/stop and multi-block lifecycle, so this would protect the streaming
change as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
