From 97c4621d8dc58e740b8d3b45dc839b49f2233b1c Mon Sep 17 00:00:00 2001
From: Gonzalo Ortiz <gortiz@users.noreply.github.com>
Date: Tue, 1 Oct 2024 13:06:22 +0200
Subject: [PATCH] colocated-without-hints: Improve useImplicitColocated usage

---
 .../rel/rules/PinotImplicitTableHintRule.java | 16 +++++++-
 .../apache/pinot/query/QueryEnvironment.java  | 38 +++++++++++--------
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRule.java
index 14ed49be3eef..65d1367e0026 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRule.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRule.java
@@ -252,11 +252,25 @@ public interface PartitionTableFinder {
     @Nullable
     TablePartitionInfo getTablePartitionInfo(String tableNameWithType);
 
+    default boolean isEnabled() {
+      return true;
+    }
+
     /**
      * A partition table finder that always returns null, meaning that the table partition info is not found.
      */
     static PartitionTableFinder disabled() {
-      return (table) -> null;
+      return new PartitionTableFinder() {
+        @Override
+        public @Nullable TablePartitionInfo getTablePartitionInfo(String tableNameWithType) {
+          return null;
+        }
+
+        @Override
+        public boolean isEnabled() {
+          return false;
+        }
+      };
     }
   }
 
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
index 104b72efdb97..5d6dde9b435c 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -136,32 +136,39 @@ public QueryEnvironment(String database, TableCache tableCache, @Nullable Worker
    * Returns a planner context that can be used to either parse, explain or execute a query.
    */
   private PlannerContext getPlannerContext(SqlNodeAndOptions sqlNodeAndOptions) {
-    boolean useImplicitColocated;
-    PinotImplicitTableHintRule.PartitionTableFinder ptf;
+    PinotImplicitTableHintRule.PartitionTableFinder ptf = getPartitionTableFinder(sqlNodeAndOptions);
+    HepProgram traitProgram = getTraitProgram(ptf);
+    return new PlannerContext(_config, _catalogReader, _typeFactory, _optProgram, traitProgram);
+  }
+
+  private PinotImplicitTableHintRule.PartitionTableFinder getPartitionTableFinder(SqlNodeAndOptions sqlNodeAndOptions) {
     String useImplicitColocatedOptionValue = sqlNodeAndOptions.getOptions()
         .get(CommonConstants.Broker.Request.QueryOptionKey.IMPLICIT_COLOCATE_JOIN);
+    WorkerManager workerManager = _envConfig.getWorkerManager();
     if (Boolean.parseBoolean(useImplicitColocatedOptionValue)) {
-      useImplicitColocated = true;
-      Objects.requireNonNull(_envConfig.getWorkerManager(), "WorkerManager is required for implicit colocated join");
-      ptf = _envConfig.getWorkerManager()::getTablePartitionInfo;
-    } else {
-      useImplicitColocated = _envConfig.useImplicitColocatedByDefault();
-      WorkerManager workerManager = _envConfig.getWorkerManager();
+      Objects.requireNonNull(workerManager, "WorkerManager is required for implicit colocated join");
+      return workerManager::getTablePartitionInfo;
+    } else if (useImplicitColocatedOptionValue == null) {
+      boolean useImplicitColocated = _envConfig.useImplicitColocatedByDefault();
       if (useImplicitColocated && workerManager != null) {
-        ptf = workerManager::getTablePartitionInfo;
+        return workerManager::getTablePartitionInfo;
       } else {
-        ptf = PinotImplicitTableHintRule.PartitionTableFinder.disabled();
+        return PinotImplicitTableHintRule.PartitionTableFinder.disabled();
       }
+    } else if ("false".equalsIgnoreCase(useImplicitColocatedOptionValue)) {
+      return PinotImplicitTableHintRule.PartitionTableFinder.disabled();
+    } else {
+      throw new RuntimeException("Invalid value for query option '"
+          + CommonConstants.Broker.Request.QueryOptionKey.IMPLICIT_COLOCATE_JOIN + "': "
+          + useImplicitColocatedOptionValue);
     }
-    HepProgram traitProgram = getTraitProgram(ptf, useImplicitColocated);
-    return new PlannerContext(_config, _catalogReader, _typeFactory, _optProgram, traitProgram);
   }
 
   /**
    * Returns the planner context that should be used only for parsing queries.
    */
   private PlannerContext getParsingPlannerContext() {
-    HepProgram traitProgram = getTraitProgram(PinotImplicitTableHintRule.PartitionTableFinder.disabled(), false);
+    HepProgram traitProgram = getTraitProgram(PinotImplicitTableHintRule.PartitionTableFinder.disabled());
     return new PlannerContext(_config, _catalogReader, _typeFactory, _optProgram, traitProgram);
   }
 
@@ -428,8 +435,7 @@ private static HepProgram getOptProgram() {
     return hepProgramBuilder.build();
   }
 
-  private static HepProgram getTraitProgram(
-      PinotImplicitTableHintRule.PartitionTableFinder tablePartitionTableFinder, boolean useImplicitColocated) {
+  private static HepProgram getTraitProgram(PinotImplicitTableHintRule.PartitionTableFinder tablePartitionTableFinder) {
     HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
 
     // Set the match order as BOTTOM_UP.
@@ -442,7 +448,7 @@ private static HepProgram getTraitProgram(
     }
 
     // apply RelDistribution trait to all nodes
-    if (useImplicitColocated) {
+    if (tablePartitionTableFinder.isEnabled()) {
       hepProgramBuilder.addRuleInstance(PinotImplicitTableHintRule.withPartitionTableFinder(tablePartitionTableFinder));
     }
     hepProgramBuilder.addRuleInstance(PinotRelDistributionTraitRule.INSTANCE);