gortiz commented on code in PR #13943:
URL: https://github.com/apache/pinot/pull/13943#discussion_r1802916505
##########
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:
Sure. WorkerManager will be null or not depending on how QueryEnvironment is
initialized.
But `inferPartitionHint` (the new name of `useImplicitColocatedOptionValue`)
is decided by the query option. What I'm doing here is to fail if WorkerManager
was initialized to null AND the user is explicitly asking to
`inferPartitionHint`. I think that is better than the alternative, which is to
do not apply infer partition hint when explicitly asked because the worker
manager wasn't enabled at startup time.
--
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
