Make PlannerContext implement Calcite Context for rule-time option access

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.
+ * <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 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<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 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> 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();

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);
+  }
+}