J-HowHuang commented on code in PR #20068: URL: https://github.com/apache/pinot/pull/20068#discussion_r4408595052
########## 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: added one test for streaming selection, and verify the prefetch/acquire/release cycles -- 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]
