From 8f6b3b4c72419d2aed97b4be91dfa79c6611bd6f Mon Sep 17 00:00:00 2001 From: Gonzalo Ortiz <1913993+gortiz@users.noreply.github.com> Date: Fri, 22 May 2026 14:20:32 +0200 Subject: [PATCH] Make PlannerContext implement Calcite Context for rule-time access PlannerContext now implements org.apache.calcite.plan.Context and passes itself as the context to both the opt planner (LogicalPlanner backed by optProgram) and the trait planner (LogicalPlanner backed by traitProgram). This allows Calcite rules contributed via the new RuleSetCustomizer SPI (see #18387) to read per-query options at match time without reflection: PlannerContext ctx = call.getPlanner().getContext().unwrap(PlannerContext.class); String runtime = ctx.getOptions().get("workerRuntime"); unwrap() also delegates to _envConfig, so any rule that currently unwraps QueryEnvironment.Config continues to work unchanged. A package-private constructor and a public forTesting() factory are added so unit tests can construct a PlannerContext without going through QueryEnvironment (which requires a full broker setup). --- .../pinot/query/context/PlannerContext.java | 65 +++++++++++++-- .../query/context/PlannerContextTest.java | 79 +++++++++++++++++++ 2 files changed, 136 insertions(+), 8 deletions(-) create mode 100644 pinot-query-planner/src/test/java/org/apache/pinot/query/context/PlannerContextTest.java 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); + } +}