colocated-without-hints

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

@@ -128,7 +128,16 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       Long timeoutMsFromQueryOption = QueryOptionsUtils.getTimeoutMs(queryOptions);
       queryTimeoutMs = timeoutMsFromQueryOption != null ? timeoutMsFromQueryOption : _brokerTimeoutMs;
       database = DatabaseUtils.extractDatabaseFromQueryRequest(queryOptions, httpHeaders);
-      QueryEnvironment queryEnvironment = new QueryEnvironment(database, _tableCache, _workerManager);
+      boolean inferPartitionHint = _config.getProperty(CommonConstants.Broker.CONFIG_OF_INFER_PARTITION_HINT,
+          CommonConstants.Broker.DEFAULT_INFER_PARTITION_HINT);
+      //@formatter:off
+      QueryEnvironment queryEnvironment = new QueryEnvironment(QueryEnvironment.configBuilder()
+          .database(database)
+          .tableCache(_tableCache)
+          .workerManager(_workerManager)
+          .defaultInferPartitionHint(inferPartitionHint)
+          .build());
+      //@formatter:on
       switch (sqlNodeAndOptions.getSqlNode().getKind()) {
         case EXPLAIN:
           boolean askServers = QueryOptionsUtils.isExplainAskingServers(queryOptions)

pinot-query-planner/pom.xml

@@ -49,6 +49,10 @@
       <groupId>org.codehaus.janino</groupId>
       <artifactId>commons-compiler</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.immutables</groupId>
+      <artifactId>value-annotations</artifactId>
+    </dependency>
 
     <dependency>
       <groupId>org.testng</groupId>

pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java

@@ -87,9 +87,28 @@ public static class JoinHintOptions {
   }
 
   public static class TableHintOptions {
+    /**
+     * Indicates how many partitions the table must be partitioned by.
+     * This must be equal to the partition count of the table in
+     * {@code tableIndexConfig.segmentPartitionConfig.columnPartitionMap}.
+     */
     public static final String PARTITION_KEY = "partition_key";
+    /**
+     * The function to use to partition the table.
+     * This must be equal to {@code functionName} in {@code tableIndexConfig.segmentPartitionConfig.columnPartitionMap}.
+     */
     public static final String PARTITION_FUNCTION = "partition_function";
+    /**
+     * The size of each partition.
+     * This must be equal to {@code numPartition} in {@code tableIndexConfig.segmentPartitionConfig.columnPartitionMap}.
+     */
     public static final String PARTITION_SIZE = "partition_size";
+    /**
+     * The number of workers per partition.
+     *
+     * How many threads to use in the following stage after partition is joined.
+     * When partition info is set, each partition is processed as a separate query in the leaf stage.
+     */
     public static final String PARTITION_PARALLELISM = "partition_parallelism";
   }
 }

pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintStrategyTable.java

@@ -19,9 +19,11 @@
 package org.apache.pinot.calcite.rel.hint;
 
 import java.util.List;
+import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import org.apache.calcite.rel.hint.HintPredicates;
 import org.apache.calcite.rel.hint.HintStrategyTable;
+import org.apache.calcite.rel.hint.Hintable;
 import org.apache.calcite.rel.hint.RelHint;
 import org.apache.pinot.spi.utils.BooleanUtils;
 
@@ -38,6 +40,29 @@ private PinotHintStrategyTable() {
           .hintStrategy(PinotHintOptions.JOIN_HINT_OPTIONS, HintPredicates.JOIN)
           .hintStrategy(PinotHintOptions.TABLE_HINT_OPTIONS, HintPredicates.TABLE_SCAN).build();
 
+
+  /**
+   * Get the first hint that has the specified name.
+   */
+  @Nullable
+  public static RelHint getHint(Hintable hintable, String hintName) {
+    return hintable.getHints().stream()
+        .filter(relHint -> relHint.hintName.equals(hintName))
+        .findFirst()
+        .orElse(null);
+  }
+
+  /**
+   * Get the first hint that satisfies the predicate.
+   */
+  @Nullable
+  public static RelHint getHint(Hintable hintable, Predicate<RelHint> predicate) {
+    return hintable.getHints().stream()
+        .filter(predicate)
+        .findFirst()
+        .orElse(null);
+  }
+
   /**
    * Check if a hint-able {@link org.apache.calcite.rel.RelNode} contains a specific {@link RelHint} by name.
    *

pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRule.java

@@ -0,0 +1,223 @@
+/**
+ * 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.calcite.rel.rules;
+
+import java.util.ArrayList;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.hint.RelHint;
+import org.apache.calcite.rel.logical.LogicalTableScan;
+import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
+import org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable;
+import org.apache.pinot.query.planner.logical.RelToPlanNodeConverter;
+import org.apache.pinot.query.routing.WorkerManager;
+import org.immutables.value.Value;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@Value.Enclosing
+public class PinotImplicitTableHintRule extends RelRule<RelRule.Config> {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(PinotImplicitTableHintRule.class);
+  private final WorkerManager _workerManager;
+
+  private PinotImplicitTableHintRule(Config config) {
+    super(config);
+    _workerManager = config.getWorkerManager();
+  }
+
+  public static PinotImplicitTableHintRule withWorkerManager(WorkerManager workerManager) {
+    return new PinotImplicitTableHintRule(ImmutablePinotImplicitTableHintRule.Config.builder()
+        .operandSupplier(b0 -> b0.operand(LogicalTableScan.class).anyInputs())
+        .workerManager(workerManager)
+        .build()
+    );
+  }
+
+  @Override
+  public boolean matches(RelOptRuleCall call) {
+    LogicalTableScan tableScan = call.rel(0);
+
+    RelHint explicitHint = getTableOptionHint(tableScan);
+
+    if (explicitHint == null) {
+      return true;
+    }
+    // we don't want to apply this rule if the explicit hint is complete
+    Map<String, String> kvOptions = explicitHint.kvOptions;
+    return kvOptions.containsKey(PinotHintOptions.TableHintOptions.PARTITION_KEY)
+        && kvOptions.containsKey(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION)
+        && kvOptions.containsKey(PinotHintOptions.TableHintOptions.PARTITION_SIZE);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    LogicalTableScan tableScan = call.rel(0);
+
+    String tableName = RelToPlanNodeConverter.getTableNameFromTableScan(tableScan);
+    @Nullable
+    TableOptions implicitTableOptions = _workerManager.inferTableOptions(tableName);
+    if (implicitTableOptions == null) {
+      return;
+    }
+
+    @Nullable
+    RelHint explicitHint = getTableOptionHint(tableScan);
+    TableOptions tableOptions = calculateTableOptions(explicitHint, implicitTableOptions, tableScan);
+    RelNode newRel = withNewTableOptions(tableScan, tableOptions);
+    call.transformTo(newRel);
+  }
+
+  /**
+   * Get the table option hint from the table scan, if any.
+   */
+  @Nullable
+  private static RelHint getTableOptionHint(LogicalTableScan tableScan) {
+    return PinotHintStrategyTable.getHint(tableScan, PinotHintOptions.TABLE_HINT_OPTIONS);
+  }
+
+  /**
+   * Returns a new node which is a copy of the given table scan with the new table options hint.
+   */
+  private static RelNode withNewTableOptions(LogicalTableScan tableScan, TableOptions tableOptions) {
+    ArrayList<RelHint> newHints = new ArrayList<>(tableScan.getHints());
+
+    newHints.removeIf(relHint -> relHint.hintName.equals(PinotHintOptions.TABLE_HINT_OPTIONS));
+
+    RelHint.Builder builder = RelHint.builder(PinotHintOptions.TABLE_HINT_OPTIONS)
+        .hintOption(PinotHintOptions.TableHintOptions.PARTITION_KEY, tableOptions.getPartitionKey())
+        .hintOption(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, tableOptions.getPartitionFunction())
+        .hintOption(PinotHintOptions.TableHintOptions.PARTITION_SIZE, String.valueOf(tableOptions.getPartitionSize()));
+
+    if (tableOptions.getPartitionParallelism() != null) {
+      builder.hintOption(PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM,
+          String.valueOf(tableOptions.getPartitionParallelism()));
+    }
+
+    newHints.add(builder.build());
+
+    return tableScan.withHints(newHints);
+  }
+
+  /**
+   * Creates a new table options hint based on the given table partition info and the explicit hint, if any.
+   *
+   * Any explicit hint will override the implicit hint obtained from the table partition info.
+   */
+  private static TableOptions calculateTableOptions(
+      @Nullable RelHint relHint, TableOptions implicitTableOptions, LogicalTableScan tableScan) {
+    if (relHint == null) {
+      return implicitTableOptions;
+    }
+
+    // there is a hint, check fill default data and obtain the partition parallelism if supplied
+    Map<String, String> kvOptions = relHint.kvOptions;
+
+    ImmutableTableOptions newTableOptions = ImmutableTableOptions.copyOf(implicitTableOptions);
+    newTableOptions = overridePartitionKey(newTableOptions, tableScan, kvOptions);
+    newTableOptions = overridePartitionFunction(newTableOptions, tableScan, kvOptions);
+    newTableOptions = overridePartitionSize(newTableOptions, tableScan, kvOptions);
+    newTableOptions = overridePartitionParallelism(newTableOptions, tableScan, kvOptions);
+
+    return newTableOptions;
+  }
+
+  /**
+   * Returns a table options hint with the partition key overridden by the hint, if any.
+   */
+  private static ImmutableTableOptions overridePartitionKey(ImmutableTableOptions base, LogicalTableScan tableScan,
+      Map<String, String> kvOptions) {
+    String partitionKey = kvOptions.get(kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY));
+    if (partitionKey == null || partitionKey.equals(base.getPartitionKey())) {
+      return base;
+    }
+    LOGGER.debug("Override implicit table hint for {} with explicit partition key: {}", tableScan, partitionKey);
+    return base.withPartitionKey(partitionKey);
+  }
+
+  /**
+   * Returns a table options hint with the partition function overridden by the hint, if any.
+   */
+  private static ImmutableTableOptions overridePartitionFunction(ImmutableTableOptions base, LogicalTableScan tableScan,
+      Map<String, String> kvOptions) {
+    String partitionFunction = kvOptions.get(kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION));
+    if (partitionFunction == null || partitionFunction.equals(base.getPartitionFunction())) {
+      return base;
+    }
+    LOGGER.debug("Override implicit table hint for {} with explicit partition function: {}", tableScan,
+        partitionFunction);
+    return base.withPartitionFunction(partitionFunction);
+  }
+
+  /**
+   * Returns a table options hint with the partition parallelism overridden by the hint, if any.
+   */
+  private static ImmutableTableOptions overridePartitionParallelism(ImmutableTableOptions base,
+      LogicalTableScan tableScan, Map<String, String> kvOptions) {
+    String partitionParallelismStr = kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM);
+    if (partitionParallelismStr == null) {
+      return base;
+    }
+    try {
+      int partitionParallelism = Integer.parseInt(partitionParallelismStr);
+      LOGGER.debug("Override implicit table hint for {} with explicit partition parallelism: {}", tableScan,
+          partitionParallelism);
+      return base.withPartitionParallelism(partitionParallelism);
+    } catch (NumberFormatException e) {
+      throw new IllegalArgumentException(
+          "Invalid partition parallelism: " + partitionParallelismStr + " for table: " + tableScan);
+    }
+  }
+
+  /**
+   * Returns a table options hint with the partition size overridden by the hint, if any.
+   */
+  private static ImmutableTableOptions overridePartitionSize(ImmutableTableOptions base, LogicalTableScan tableScan,
+      Map<String, String> kvOptions) {
+    String partitionSizeStr = kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_SIZE);
+    if (partitionSizeStr == null) {
+      return base;
+    }
+    try {
+      int explicitPartitionSize = Integer.parseInt(partitionSizeStr);
+      if (explicitPartitionSize == base.getPartitionSize()) {
+        return base;
+      }
+      LOGGER.debug("Override implicit table hint for {} with explicit partition size: {}", tableScan,
+          explicitPartitionSize);
+      return base.withPartitionSize(explicitPartitionSize);
+    } catch (NumberFormatException e) {
+      throw new IllegalArgumentException("Invalid partition size: " + partitionSizeStr + " for table: " + tableScan);
+    }
+  }
+
+  @Value.Immutable
+  public interface Config extends RelRule.Config {
+    @Nullable
+    WorkerManager getWorkerManager();
+
+    @Override
+    default PinotImplicitTableHintRule toRule() {
+      return new PinotImplicitTableHintRule(this);
+    }
+  }
+}

pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/TableOptions.java

@@ -0,0 +1,38 @@
+/**
+ * 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.calcite.rel.rules;
+
+import javax.annotation.Nullable;
+import org.immutables.value.Value;
+
+
+/**
+ * An internal interface used to generate the table options hint.
+ */
+@Value.Immutable
+public interface TableOptions {
+  String getPartitionKey();
+
+  String getPartitionFunction();
+
+  int getPartitionSize();
+
+  @Nullable
+  Integer getPartitionParallelism();
+}