diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java index 6e5abd8f660d..5bfdfc7d2483 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/context/PlannerContext.java @@ -18,13 +18,15 @@ */ package org.apache.pinot.query.context; +import com.google.common.annotations.VisibleForTesting; import java.util.Collections; import java.util.HashMap; import java.util.Map; import javax.annotation.Nullable; -import org.apache.calcite.plan.Contexts; +import org.apache.calcite.plan.Context; import org.apache.calcite.plan.RelOptPlanner; import org.apache.calcite.plan.hep.HepProgram; +import org.apache.calcite.plan.hep.HepProgramBuilder; import org.apache.calcite.prepare.PlannerImpl; import org.apache.calcite.prepare.Prepare; import org.apache.calcite.rel.RelDistributionTraitDef; @@ -38,12 +40,16 @@ /** - * PlannerContext is an object that holds all contextual information during planning phase. + * Holds all per-query contextual information used during the planning phase. * - * TODO: currently we don't support option or query rewrite. - * It is used to hold per query context for query planning, which cannot be shared across queries. + *

This class implements {@link Context} so that Calcite rules can retrieve it directly from the + * planner: {@code call.getPlanner().getContext().unwrap(PlannerContext.class)}. Both the opt planner + * and the trait planner expose this instance as their context. + * + *

Callers may also unwrap {@link QueryEnvironment.Config} to access broker-wide defaults: + * {@code call.getPlanner().getContext().unwrap(QueryEnvironment.Config.class)}. */ -public class PlannerContext implements AutoCloseable { +public class PlannerContext implements AutoCloseable, Context { private final PlannerImpl _planner; private final SqlValidator _validator; @@ -63,16 +69,42 @@ public PlannerContext(FrameworkConfig config, Prepare.CatalogReader catalogReade SqlExplainFormat sqlExplainFormat, @Nullable PhysicalPlannerContext physicalPlannerContext) { _planner = new PlannerImpl(config); _validator = new Validator(config.getOperatorTable(), catalogReader, typeFactory); - _relOptPlanner = new LogicalPlanner(optProgram, Contexts.EMPTY_CONTEXT, config.getTraitDefs()); - _relTraitPlanner = new LogicalPlanner(traitProgram, Contexts.of(envConfig), - Collections.singletonList(RelDistributionTraitDef.INSTANCE)); _options = options; _envConfig = envConfig; + _relOptPlanner = new LogicalPlanner(optProgram, this, config.getTraitDefs()); + _relTraitPlanner = new LogicalPlanner(traitProgram, this, + Collections.singletonList(RelDistributionTraitDef.INSTANCE)); _plannerOutput = new HashMap<>(); _sqlExplainFormat = sqlExplainFormat; _physicalPlannerContext = physicalPlannerContext; } + /** + * Test factory: creates a minimal {@link PlannerContext} without going through + * {@link org.apache.pinot.query.QueryEnvironment}, suitable for unit tests. + */ + @VisibleForTesting + public static PlannerContext forTesting(Map options, QueryEnvironment.Config envConfig) { + return new PlannerContext(options, envConfig); + } + + /** + * Minimal constructor for use in unit tests. Creates no-op planners backed by an empty HEP program. + */ + @VisibleForTesting + PlannerContext(Map options, QueryEnvironment.Config envConfig) { + _planner = null; + _validator = null; + _options = options; + _envConfig = envConfig; + HepProgram emptyProgram = new HepProgramBuilder().build(); + _relOptPlanner = new LogicalPlanner(emptyProgram, this); + _relTraitPlanner = new LogicalPlanner(emptyProgram, this); + _plannerOutput = new HashMap<>(); + _sqlExplainFormat = null; + _physicalPlannerContext = null; + } + public PlannerImpl getPlanner() { return _planner; } @@ -97,6 +129,23 @@ public QueryEnvironment.Config getEnvConfig() { return _envConfig; } + /** + * Unwraps this context. Returns {@code this} when asked for {@link PlannerContext} or + * {@link Context}, and delegates to {@link #_envConfig} when asked for + * {@link QueryEnvironment.Config} so that existing rules remain compatible. + */ + @Override + @Nullable + public C unwrap(Class clazz) { + if (clazz.isInstance(this)) { + return clazz.cast(this); + } + if (clazz.isInstance(_envConfig)) { + return clazz.cast(_envConfig); + } + return null; + } + @Override public void close() { _planner.close(); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/context/PlannerContextTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/context/PlannerContextTest.java new file mode 100644 index 000000000000..ab965e53e708 --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/context/PlannerContextTest.java @@ -0,0 +1,79 @@ +/** + * 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.query.context; + +import java.util.Map; +import org.apache.calcite.plan.Context; +import org.apache.pinot.query.QueryEnvironment; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; + + +public class PlannerContextTest { + + @Test + public void testUnwrapReturnsSelf() { + QueryEnvironment.Config config = mock(QueryEnvironment.Config.class); + PlannerContext ctx = PlannerContext.forTesting(Map.of(), config); + + assertSame(ctx.unwrap(PlannerContext.class), ctx); + assertSame(ctx.unwrap(Context.class), ctx); + } + + @Test + public void testUnwrapReturnsEnvConfig() { + QueryEnvironment.Config config = mock(QueryEnvironment.Config.class); + PlannerContext ctx = PlannerContext.forTesting(Map.of(), config); + + assertSame(ctx.unwrap(QueryEnvironment.Config.class), config); + } + + @Test + public void testUnwrapReturnsNullForUnknownType() { + QueryEnvironment.Config config = mock(QueryEnvironment.Config.class); + PlannerContext ctx = PlannerContext.forTesting(Map.of(), config); + + assertNull(ctx.unwrap(String.class)); + } + + @Test + public void testPlannersExposeContextInstance() { + QueryEnvironment.Config config = mock(QueryEnvironment.Config.class); + PlannerContext ctx = PlannerContext.forTesting(Map.of("k", "v"), config); + + // Both planners must expose this PlannerContext via their context, so rules can + // call call.getPlanner().getContext().unwrap(PlannerContext.class) to read options. + assertSame(ctx.getRelOptPlanner().getContext().unwrap(PlannerContext.class), ctx); + assertSame(ctx.getRelTraitPlanner().getContext().unwrap(PlannerContext.class), ctx); + } + + @Test + public void testOptionsAreAccessibleThroughUnwrap() { + QueryEnvironment.Config config = mock(QueryEnvironment.Config.class); + Map options = Map.of("workerRuntime", "datafusion"); + PlannerContext ctx = PlannerContext.forTesting(options, config); + + PlannerContext unwrapped = ctx.getRelOptPlanner().getContext().unwrap(PlannerContext.class); + assertSame(unwrapped, ctx); + assertSame(unwrapped.getOptions(), options); + } +}