This is an automated email from the ASF dual-hosted git repository.
gortiz 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 f77730e57c7 Make PlannerContext implement Calcite Context for
rule-time access (#18567)
f77730e57c7 is described below
commit f77730e57c735127501a16292d3dcdc90f40b5b9
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Mon May 25 15:05:20 2026 +0200
Make PlannerContext implement Calcite Context for rule-time access (#18567)
---
.../apache/pinot/query/context/PlannerContext.java | 65 +++++++++++++++---
.../pinot/query/context/PlannerContextTest.java | 79 ++++++++++++++++++++++
2 files changed, 136 insertions(+), 8 deletions(-)
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 6e5abd8f660..5bfdfc7d248 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 @@ import org.apache.pinot.query.validate.Validator;
/**
- * 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.
+ * <p>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.
+ *
+ * <p>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 class PlannerContext implements AutoCloseable {
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<String, String> 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<String, String> 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 class PlannerContext implements AutoCloseable {
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> C unwrap(Class<C> 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 00000000000..ab965e53e70
--- /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<String, String> 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);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
