Jackie-Jiang commented on code in PR #13943:
URL: https://github.com/apache/pinot/pull/13943#discussion_r1803542898


##########
pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java:
##########
@@ -100,26 +111,64 @@ public class QueryEnvironment {
   private final FrameworkConfig _config;
   private final CalciteCatalogReader _catalogReader;
   private final HepProgram _optProgram;
-  private final HepProgram _traitProgram;
-
-  // Pinot extensions
-  private final TableCache _tableCache;
-  private final WorkerManager _workerManager;
+  private final Config _envConfig;
 
-  public QueryEnvironment(String database, TableCache tableCache, @Nullable 
WorkerManager workerManager) {
-    PinotCatalog catalog = new PinotCatalog(tableCache, database);
+  public QueryEnvironment(Config config) {
+    _envConfig = config;
+    String database = config.getDatabase();
+    PinotCatalog catalog = new PinotCatalog(config.getTableCache(), database);
     CalciteSchema rootSchema = CalciteSchema.createRootSchema(false, false, 
database, catalog);
     _config = 
Frameworks.newConfigBuilder().traitDefs().operatorTable(PinotOperatorTable.instance())
         
.defaultSchema(rootSchema.plus()).sqlToRelConverterConfig(PinotRuleUtils.PINOT_SQL_TO_REL_CONFIG).build();
     _catalogReader = new CalciteCatalogReader(rootSchema, List.of(database), 
_typeFactory, CONNECTION_CONFIG);
     _optProgram = getOptProgram();
-    _traitProgram = getTraitProgram();
-    _tableCache = tableCache;
-    _workerManager = workerManager;
   }
 
-  private PlannerContext getPlannerContext() {
-    return new PlannerContext(_config, _catalogReader, _typeFactory, 
_optProgram, _traitProgram);
+  public QueryEnvironment(String database, TableCache tableCache, @Nullable 
WorkerManager workerManager) {
+    this(configBuilder()
+        .database(database)
+        .tableCache(tableCache)
+        .workerManager(workerManager)
+        .build());
+  }
+
+  /**
+   * Returns a planner context that can be used to either parse, explain or 
execute a query.
+   */
+  private PlannerContext getPlannerContext(SqlNodeAndOptions 
sqlNodeAndOptions) {
+    WorkerManager workerManager = getWorkerManager(sqlNodeAndOptions);
+    HepProgram traitProgram = getTraitProgram(workerManager);
+    return new PlannerContext(_config, _catalogReader, _typeFactory, 
_optProgram, traitProgram);
+  }
+
+  @Nullable
+  private WorkerManager getWorkerManager(SqlNodeAndOptions sqlNodeAndOptions) {
+    String useImplicitColocatedOptionValue = sqlNodeAndOptions.getOptions()
+        
.get(CommonConstants.Broker.Request.QueryOptionKey.IMPLICIT_COLOCATE_JOIN);
+    WorkerManager workerManager = _envConfig.getWorkerManager();
+
+    if (useImplicitColocatedOptionValue == null) {
+      return _envConfig.useImplicitColocatedByDefault() ? workerManager : null;
+    }
+    switch (useImplicitColocatedOptionValue.toLowerCase()) {
+      case "true":
+        Objects.requireNonNull(workerManager, "WorkerManager is required for 
implicit colocated join");

Review Comment:
   If I understand the intention correctly, we want to fail the planning if 
`partitionMetadataManager` is unavailable instead of `workerManager`. 
`workerManager` can be `null` in certain query path independent of the value of 
`partitionMetadataManager`. We need to ensure query with option never hits that 
code path



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org

Reply via email to