Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
package org.apache.pinot.calcite.rel.rules;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.Map;
import javax.annotation.Nullable;
import org.apache.calcite.plan.RelOptRuleCall;
Expand Down Expand Up @@ -75,7 +76,7 @@ public void onMatch(RelOptRuleCall call) {
@Nullable
RelHint explicitHint = getTableOptionHint(tableScan);
TableOptions tableOptions = calculateTableOptions(explicitHint, implicitTableOptions, tableScan);
RelNode newRel = withNewTableOptions(tableScan, tableOptions);
RelNode newRel = withNewTableOptions(tableScan, tableOptions, explicitHint);
call.transformTo(newRel);
}

Expand All @@ -102,23 +103,31 @@ private static RelHint getTableOptionHint(TableScan tableScan) {
}

/**
* Returns a new node which is a copy of the given table scan with the new table options hint.
* Returns a new node which is a copy of the given table scan with a rebuilt table options hint: options from the
* explicit hint that are not modeled by {@link TableOptions} (e.g. is_replicated) are carried over, and the merged
* (inferred + explicitly overridden) partition options overwrite the explicitly supplied ones.
*/
private static RelNode withNewTableOptions(TableScan tableScan, TableOptions tableOptions) {
private static RelNode withNewTableOptions(TableScan tableScan, TableOptions tableOptions,
@Nullable RelHint explicitHint) {
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()));

// Seed with the explicitly supplied options so that the ones not rebuilt from the given table options (e.g.
// is_replicated) are carried over, then overwrite with the merged (inferred + explicitly overridden) partition
// options.
Map<String, String> kvOptions =
explicitHint != null ? new LinkedHashMap<>(explicitHint.kvOptions) : new LinkedHashMap<>();
kvOptions.put(PinotHintOptions.TableHintOptions.PARTITION_KEY, tableOptions.getPartitionKey());
kvOptions.put(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION, tableOptions.getPartitionFunction());
kvOptions.put(PinotHintOptions.TableHintOptions.PARTITION_SIZE, String.valueOf(tableOptions.getPartitionSize()));
if (tableOptions.getPartitionParallelism() != null) {
builder.hintOption(PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM,
kvOptions.put(PinotHintOptions.TableHintOptions.PARTITION_PARALLELISM,
String.valueOf(tableOptions.getPartitionParallelism()));
}

RelHint.Builder builder = RelHint.builder(PinotHintOptions.TABLE_HINT_OPTIONS);
kvOptions.forEach(builder::hintOption);
newHints.add(builder.build());

return tableScan.withHints(newHints);
Expand Down Expand Up @@ -152,7 +161,7 @@ private static TableOptions calculateTableOptions(
*/
private static ImmutableTableOptions overridePartitionKey(ImmutableTableOptions base, TableScan tableScan,
Map<String, String> kvOptions) {
String partitionKey = kvOptions.get(kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol

String partitionKey = kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_KEY);
if (partitionKey == null || partitionKey.equals(base.getPartitionKey())) {
return base;
}
Expand All @@ -165,7 +174,7 @@ private static ImmutableTableOptions overridePartitionKey(ImmutableTableOptions
*/
private static ImmutableTableOptions overridePartitionFunction(ImmutableTableOptions base,
TableScan tableScan, Map<String, String> kvOptions) {
String partitionFunction = kvOptions.get(kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol + 1

String partitionFunction = kvOptions.get(PinotHintOptions.TableHintOptions.PARTITION_FUNCTION);
if (partitionFunction == null || partitionFunction.equals(base.getPartitionFunction())) {
return base;
}
Expand Down
32 changes: 32 additions & 0 deletions pinot-query-planner/src/test/resources/queries/JoinPlans.json
Original file line number Diff line number Diff line change
Expand Up @@ -834,6 +834,38 @@
"\n"
]
},
{
"description": "Simple local replicated join with partition hint inference enabled (is_replicated should be preserved when the table options hint is rebuilt with the inferred partition options)",
"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 (is_replicated should be preserved 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'",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading