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 b1c9157bb286..ffc465910d6c 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 @@ -57,8 +57,11 @@ public static PinotImplicitTableHintRule withWorkerManager(WorkerManager workerM public boolean matches(RelOptRuleCall call) { TableScan tableScan = call.rel(0); - // we don't want to apply this rule if the explicit hint is complete - return !isHintComplete(getTableOptionHint(tableScan)); + // we don't want to apply this rule if the explicit hint is complete, or if the table is hinted as replicated + // across all workers (each worker scans all the segments, so partition options are irrelevant) + @Nullable + RelHint explicitHint = getTableOptionHint(tableScan); + return !isHintComplete(explicitHint) && !isReplicated(explicitHint); } @Override @@ -93,6 +96,16 @@ private boolean isHintComplete(@Nullable RelHint hint) { && kvOptions.containsKey(PinotHintOptions.TableHintOptions.PARTITION_SIZE); } + /** + * Determines whether the provided hint declares the table as replicated across all workers. + */ + private static boolean isReplicated(@Nullable RelHint hint) { + if (hint == null || hint.kvOptions == null) { + return false; + } + return Boolean.parseBoolean(hint.kvOptions.get(PinotHintOptions.TableHintOptions.IS_REPLICATED)); + } + /** * Get the table option hint from the table scan, if any. */ @@ -119,6 +132,11 @@ private static RelNode withNewTableOptions(TableScan tableScan, TableOptions tab String.valueOf(tableOptions.getPartitionParallelism())); } + if (tableOptions.isReplicated() != null) { + builder.hintOption(PinotHintOptions.TableHintOptions.IS_REPLICATED, + String.valueOf(tableOptions.isReplicated())); + } + newHints.add(builder.build()); return tableScan.withHints(newHints); @@ -143,16 +161,31 @@ private static TableOptions calculateTableOptions( newTableOptions = overridePartitionFunction(newTableOptions, tableScan, kvOptions); newTableOptions = overridePartitionSize(newTableOptions, tableScan, kvOptions); newTableOptions = overridePartitionParallelism(newTableOptions, tableScan, kvOptions); + newTableOptions = carryOverIsReplicated(newTableOptions, kvOptions); return newTableOptions; } + /** + * Returns a table options hint with the replicated flag carried over from the explicit hint, if any. Note that this + * rule does not match when the table is hinted as replicated across all workers, so this only ever carries over an + * explicit {@code is_replicated='false'}. + */ + private static ImmutableTableOptions carryOverIsReplicated(ImmutableTableOptions base, + Map kvOptions) { + String isReplicated = kvOptions.get(PinotHintOptions.TableHintOptions.IS_REPLICATED); + if (isReplicated == null) { + return base; + } + return base.withIsReplicated(Boolean.parseBoolean(isReplicated)); + } + /** * Returns a table options hint with the partition key overridden by the hint, if any. */ private static ImmutableTableOptions overridePartitionKey(ImmutableTableOptions base, TableScan tableScan, Map kvOptions) { - String partitionKey = kvOptions.get(kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY)); + String partitionKey = kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY); if (partitionKey == null || partitionKey.equals(base.getPartitionKey())) { return base; } @@ -165,7 +198,7 @@ private static ImmutableTableOptions overridePartitionKey(ImmutableTableOptions */ private static ImmutableTableOptions overridePartitionFunction(ImmutableTableOptions base, TableScan tableScan, Map kvOptions) { - String partitionFunction = kvOptions.get(kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION)); + String partitionFunction = kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION); if (partitionFunction == null || partitionFunction.equals(base.getPartitionFunction())) { return base; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/TableOptions.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/TableOptions.java index dfcbaa8ae3e1..158440245518 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/TableOptions.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/TableOptions.java @@ -35,4 +35,11 @@ public interface TableOptions { @Nullable Integer getPartitionParallelism(); + + /** + * Whether the table is replicated across all workers. This is never inferred; it is only ever populated from an + * explicitly supplied table options hint. + */ + @Nullable + Boolean isReplicated(); } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRuleTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRuleTest.java new file mode 100644 index 000000000000..4cc3dcddb97e --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotImplicitTableHintRuleTest.java @@ -0,0 +1,125 @@ +/** + * 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.lang.reflect.Field; +import java.lang.reflect.Modifier; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.calcite.rel.hint.PinotHintOptions; +import org.apache.pinot.query.QueryEnvironmentTestBase; +import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; +import org.apache.pinot.query.planner.physical.DispatchableSubPlan; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.TableScanNode; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + + +/** + * Tests for {@link PinotImplicitTableHintRule}, verifying the table options hint that ends up on the table scan after + * partition hint inference. + */ +public class PinotImplicitTableHintRuleTest extends QueryEnvironmentTestBase { + + @Test + public void testInferenceSkippedForReplicatedTableScan() { + // Table 'b' is hinted as replicated, so partition hint inference should leave its hint untouched (no partition + // options injected) even though 'b' is a partitioned table. Table 'a' has no explicit hint, so inference should + // add the partition options to its scan. + DispatchableSubPlan dispatchableSubPlan = _queryEnvironment.planQuery( + "SET inferPartitionHint=true; SELECT /*+ joinOptions(left_distribution_type='local', " + + "right_distribution_type='local') */ a.col1, b.col2 FROM a JOIN b " + + "/*+ tableOptions(is_replicated = 'true') */ ON a.col1 = b.col1"); + + Map replicatedScanOptions = getTableOptions(dispatchableSubPlan, "b"); + assertEquals(replicatedScanOptions, Map.of(PinotHintOptions.TableHintOptions.IS_REPLICATED, "true")); + + Map inferredScanOptions = getTableOptions(dispatchableSubPlan, "a"); + assertEquals(inferredScanOptions, Map.of( + PinotHintOptions.TableHintOptions.PARTITION_KEY, "col2", + PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, "Hashcode", + PinotHintOptions.TableHintOptions.PARTITION_SIZE, String.valueOf(PARTITION_COUNT))); + } + + @Test + public void testExplicitIsReplicatedFalseCarriedOverDuringInference() { + // An explicit is_replicated='false' should not skip inference, and should be carried over into the rebuilt hint. + DispatchableSubPlan dispatchableSubPlan = _queryEnvironment.planQuery( + "SET inferPartitionHint=true; SELECT * FROM a /*+ tableOptions(is_replicated = 'false') */ LIMIT 10"); + + Map tableOptions = getTableOptions(dispatchableSubPlan, "a"); + assertEquals(tableOptions, Map.of( + PinotHintOptions.TableHintOptions.PARTITION_KEY, "col2", + PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, "Hashcode", + PinotHintOptions.TableHintOptions.PARTITION_SIZE, String.valueOf(PARTITION_COUNT), + PinotHintOptions.TableHintOptions.IS_REPLICATED, "false")); + } + + @Test + public void testTableOptionsModelsAllTableHintOptions() + throws IllegalAccessException { + // PinotImplicitTableHintRule rebuilds the table options hint entirely from TableOptions, so every option defined + // in PinotHintOptions.TableHintOptions must be modeled by TableOptions and emitted in withNewTableOptions — + // otherwise an explicitly supplied value for the unmodeled option would be silently dropped when the hint is + // rebuilt with the inferred partition options. + Set definedOptions = new HashSet<>(); + for (Field field : PinotHintOptions.TableHintOptions.class.getDeclaredFields()) { + if (Modifier.isStatic(field.getModifiers()) && field.getType() == String.class) { + definedOptions.add((String) field.get(null)); + } + } + Set modeledOptions = Set.of(PinotHintOptions.TableHintOptions.PARTITION_KEY, + PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, PinotHintOptions.TableHintOptions.PARTITION_SIZE, + PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM, PinotHintOptions.TableHintOptions.IS_REPLICATED); + assertEquals(definedOptions, modeledOptions, + "A new table hint option must be modeled in TableOptions and emitted in " + + "PinotImplicitTableHintRule#withNewTableOptions (then added to this test), or it will be silently " + + "dropped from explicit hints when partition hint inference is enabled"); + } + + private static Map getTableOptions(DispatchableSubPlan dispatchableSubPlan, String tableName) { + for (DispatchablePlanFragment fragment : dispatchableSubPlan.getQueryStageMap().values()) { + TableScanNode tableScanNode = findTableScan(fragment.getPlanFragment().getFragmentRoot(), tableName); + if (tableScanNode != null) { + Map tableOptions = + tableScanNode.getNodeHint().getHintOptions().get(PinotHintOptions.TABLE_HINT_OPTIONS); + assertNotNull(tableOptions, "No table options hint found on scan of table: " + tableName); + return tableOptions; + } + } + throw new AssertionError("No table scan found for table: " + tableName); + } + + private static TableScanNode findTableScan(PlanNode node, String tableName) { + if (node instanceof TableScanNode && ((TableScanNode) node).getTableName().equals(tableName)) { + return (TableScanNode) node; + } + for (PlanNode input : node.getInputs()) { + TableScanNode found = findTableScan(input, tableName); + if (found != null) { + return found; + } + } + return null; + } +} diff --git a/pinot-query-planner/src/test/resources/queries/JoinPlans.json b/pinot-query-planner/src/test/resources/queries/JoinPlans.json index a5fde2dc6282..ffa6cdc35775 100644 --- a/pinot-query-planner/src/test/resources/queries/JoinPlans.json +++ b/pinot-query-planner/src/test/resources/queries/JoinPlans.json @@ -834,6 +834,38 @@ "\n" ] }, + { + "description": "Simple local replicated join with partition hint inference enabled (inference is skipped for tables hinted as replicated, leaving the explicit tableOptions hint untouched)", + "sql": "SET inferPartitionHint=true; EXPLAIN PLAN FOR SELECT /*+ joinOptions(left_distribution_type = 'local', right_distribution_type = 'local') */ a.col1, b.col2 FROM a JOIN b /*+ tableOptions(is_replicated = 'true') */ ON a.col1 = b.col1", + "output": [ + "Execution Plan", + "\nLogicalProject(col1=[$0], col2=[$2])", + "\n LogicalJoin(condition=[=($0, $1)], joinType=[inner])", + "\n PinotLogicalExchange(distribution=[single])", + "\n LogicalProject(col1=[$0])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n PinotLogicalExchange(distribution=[single])", + "\n LogicalProject(col1=[$0], col2=[$1])", + "\n PinotLogicalTableScan(table=[[default, b]])", + "\n" + ] + }, + { + "description": "Local replicated join with partition hint inference enabled (inference is skipped for tables hinted as replicated even when the hint also carries an explicit partition parallelism, and planning should succeed)", + "sql": "SET inferPartitionHint=true; EXPLAIN PLAN FOR SELECT /*+ joinOptions(left_distribution_type = 'local', right_distribution_type = 'local') */ a.col1, b.col2 FROM a JOIN b /*+ tableOptions(is_replicated = 'true', partition_parallelism = '2') */ ON a.col1 = b.col1", + "output": [ + "Execution Plan", + "\nLogicalProject(col1=[$0], col2=[$2])", + "\n LogicalJoin(condition=[=($0, $1)], joinType=[inner])", + "\n PinotLogicalExchange(distribution=[single])", + "\n LogicalProject(col1=[$0])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n PinotLogicalExchange(distribution=[single])", + "\n LogicalProject(col1=[$0], col2=[$1])", + "\n PinotLogicalTableScan(table=[[default, b]])", + "\n" + ] + }, { "description": "Broadcast join with filter on both left and right table", "sql": "EXPLAIN PLAN FOR SELECT /*+ joinOptions(left_distribution_type = 'local', right_distribution_type = 'local') */ a.col1, b.col2 FROM a JOIN b /*+ tableOptions(is_replicated = 'true') */ ON a.col1 = b.col1 WHERE a.col2 = 'foo' AND b.col2 = 'bar'", diff --git a/pinot-query-planner/src/test/resources/queries/PinotHintablePlans.json b/pinot-query-planner/src/test/resources/queries/PinotHintablePlans.json index 2a95cb1af7d7..412dd54caa75 100644 --- a/pinot-query-planner/src/test/resources/queries/PinotHintablePlans.json +++ b/pinot-query-planner/src/test/resources/queries/PinotHintablePlans.json @@ -21,6 +21,28 @@ "sql": "EXPLAIN PLAN FOR SELECT * FROM a /*+ tableOptions(partition_function='murmur', partition_key='col2', partition_size='4') */ LIMIT 10", "expectedException": ".*Partition function mismatch \\(hint: murmur, table: Hashcode\\).*" }, + { + "description": "explicit partition key in a partial hint should override the inferred partition key (wrong key should throw exception)", + "sql": "SET inferPartitionHint=true; EXPLAIN PLAN FOR SELECT * FROM a /*+ tableOptions(partition_key='col1') */ LIMIT 10", + "expectedException": ".*Partition key: col1 does not match partition column: col2.*" + }, + { + "description": "explicit partition function in a partial hint should override the inferred partition function (wrong function should throw exception)", + "sql": "SET inferPartitionHint=true; EXPLAIN PLAN FOR SELECT * FROM a /*+ tableOptions(partition_function='murmur') */ LIMIT 10", + "expectedException": ".*Partition function mismatch \\(hint: murmur, table: Hashcode\\).*" + }, + { + "description": "explicit partition key in a partial hint matching the inferred partition key should plan successfully", + "sql": "SET inferPartitionHint=true; EXPLAIN PLAN FOR SELECT * FROM a /*+ tableOptions(partition_key='col2') */ LIMIT 10", + "output": [ + "Execution Plan", + "\nLogicalSort(offset=[0], fetch=[10])", + "\n PinotLogicalSortExchange(distribution=[hash], collation=[[]], isSortOnSender=[false], isSortOnReceiver=[false])", + "\n LogicalSort(fetch=[10])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, { "description": "Inner join with group by", "sql": "EXPLAIN PLAN FOR SELECT /*+ aggOptions(is_partitioned_by_group_by_keys='true') */ a.col1, AVG(b.col3) FROM a JOIN b ON a.col1 = b.col2 WHERE a.col3 >= 0 AND a.col2 = 'a' AND b.col3 < 0 GROUP BY a.col1",