Skip to content

Commit f5e5c48

Browse files
mohsakakasiafi
authored and
xin-zhang2
committed
Rewrite exclude_columns table function into projection
Changes adapted from trino/PR#25493 Author: kasiafi Co-authored-by: kasiafi <[email protected]>
1 parent f228923 commit f5e5c48

File tree

7 files changed

+167
-7
lines changed

7 files changed

+167
-7
lines changed

‎presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@
137137
import com.facebook.presto.sql.planner.iterative.rule.RewriteCaseExpressionPredicate;
138138
import com.facebook.presto.sql.planner.iterative.rule.RewriteCaseToMap;
139139
import com.facebook.presto.sql.planner.iterative.rule.RewriteConstantArrayContainsToInExpression;
140+
import com.facebook.presto.sql.planner.iterative.rule.RewriteExcludeColumnsFunctionToProjection;
140141
import com.facebook.presto.sql.planner.iterative.rule.RewriteFilterWithExternalFunctionToProject;
141142
import com.facebook.presto.sql.planner.iterative.rule.RewriteSpatialPartitioningAggregation;
142143
import com.facebook.presto.sql.planner.iterative.rule.RewriteTableFunctionToTableScan;
@@ -888,7 +889,7 @@ public PlanOptimizers(
888889
ruleStats,
889890
statsCalculator,
890891
costCalculator,
891-
ImmutableSet.of(new RewriteTableFunctionToTableScan(metadata))));
892+
ImmutableSet.of(new RewriteTableFunctionToTableScan(metadata), new RewriteExcludeColumnsFunctionToProjection())));
892893

893894
if (!noExchange) {
894895
builder.add(new ReplicateSemiJoinInDelete()); // Must run before AddExchanges

‎presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlannerUtils.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,9 @@ public static PlanNode addOverrideProjection(PlanNode source, PlanNodeIdAllocato
205205
|| source.getOutputVariables().stream().distinct().count() != source.getOutputVariables().size()) {
206206
return source;
207207
}
208+
if (source instanceof ProjectNode && ((ProjectNode) source).getAssignments().getMap().equals(variableMap)) {
209+
return source;
210+
}
208211
Assignments.Builder assignmentsBuilder = Assignments.builder();
209212
assignmentsBuilder.putAll(source.getOutputVariables().stream().collect(toImmutableMap(identity(), x -> variableMap.containsKey(x) ? variableMap.get(x) : x)));
210213
return new ProjectNode(source.getSourceLocation(), planNodeIdAllocator.getNextId(), source, assignmentsBuilder.build(), LOCAL);

‎presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteExcludeColumnsFunctionToProjection.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.presto.sql.planner.iterative.rule;
15+
16+
import com.facebook.presto.matching.Captures;
17+
import com.facebook.presto.matching.Pattern;
18+
import com.facebook.presto.operator.table.ExcludeColumns.ExcludeColumnsFunctionHandle;
19+
import com.facebook.presto.spi.plan.Assignments;
20+
import com.facebook.presto.spi.plan.ProjectNode;
21+
import com.facebook.presto.spi.relation.VariableReferenceExpression;
22+
import com.facebook.presto.sql.planner.iterative.Rule;
23+
import com.facebook.presto.sql.planner.plan.TableFunctionProcessorNode;
24+
25+
import java.util.List;
26+
import java.util.NoSuchElementException;
27+
28+
import static com.facebook.presto.sql.planner.plan.Patterns.tableFunctionProcessor;
29+
import static com.google.common.base.Preconditions.checkState;
30+
import static com.google.common.collect.Iterators.getOnlyElement;
31+
32+
public class RewriteExcludeColumnsFunctionToProjection
33+
implements Rule<TableFunctionProcessorNode>
34+
{
35+
private static final Pattern<TableFunctionProcessorNode> PATTERN = tableFunctionProcessor();
36+
37+
@Override
38+
public Pattern<TableFunctionProcessorNode> getPattern()
39+
{
40+
return PATTERN;
41+
}
42+
43+
@Override
44+
public Result apply(TableFunctionProcessorNode node, Captures captures, Context context)
45+
{
46+
if (!(node.getHandle().getFunctionHandle() instanceof ExcludeColumnsFunctionHandle)) {
47+
return Result.empty();
48+
}
49+
50+
List<VariableReferenceExpression> inputSymbols = getOnlyElement(node.getRequiredVariables().iterator());
51+
List<VariableReferenceExpression> outputSymbols = node.getOutputVariables();
52+
53+
checkState(inputSymbols.size() == outputSymbols.size(), "inputSymbols size differs from outputSymbols size");
54+
Assignments.Builder assignments = Assignments.builder();
55+
for (int i = 0; i < outputSymbols.size(); i++) {
56+
assignments.put(outputSymbols.get(i), inputSymbols.get(i));
57+
}
58+
59+
return Result.ofPlanNode(new ProjectNode(
60+
node.getId(),
61+
node.getSource().orElseThrow(NoSuchElementException::new),
62+
assignments.build()));
63+
}
64+
}

‎presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/QueryCardinalityUtil.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.facebook.presto.spi.plan.PlanNode;
2020
import com.facebook.presto.spi.plan.ProjectNode;
2121
import com.facebook.presto.spi.plan.ValuesNode;
22+
import com.facebook.presto.spi.plan.WindowNode;
2223
import com.facebook.presto.sql.planner.iterative.GroupReference;
2324
import com.facebook.presto.sql.planner.iterative.Lookup;
2425
import com.facebook.presto.sql.planner.plan.EnforceSingleRowNode;
@@ -102,6 +103,12 @@ public Range<Long> visitEnforceSingleRow(EnforceSingleRowNode node, Void context
102103
return Range.singleton(1L);
103104
}
104105

106+
@Override
107+
public Range<Long> visitWindow(WindowNode node, Void context)
108+
{
109+
return node.getSource().accept(this, null);
110+
}
111+
105112
@Override
106113
public Range<Long> visitAggregation(AggregationNode node, Void context)
107114
{

‎presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1676,6 +1676,16 @@ public void testOffsetWithLimit()
16761676
.withAlias("row_num", new RowNumberSymbolMatcher()))))));
16771677
}
16781678

1679+
@Test
1680+
public void testRewriteExcludeColumnsFunctionToProjection()
1681+
{
1682+
assertPlan("SELECT *\n" +
1683+
"FROM TABLE(system.builtin.exclude_columns(\n" +
1684+
" INPUT => TABLE(orders),\n" +
1685+
" COLUMNS => DESCRIPTOR(comment)))\n",
1686+
output(tableScan("orders")));
1687+
}
1688+
16791689
private Session noJoinReordering()
16801690
{
16811691
return Session.builder(this.getQueryRunner().getDefaultSession())

‎presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestTableFunctionInvocation.java‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,6 @@
5555
import static com.facebook.presto.sql.planner.assertions.TableFunctionMatcher.DescriptorArgumentValue.descriptorArgument;
5656
import static com.facebook.presto.sql.planner.assertions.TableFunctionMatcher.DescriptorArgumentValue.nullDescriptor;
5757
import static com.facebook.presto.sql.planner.assertions.TableFunctionMatcher.TableArgumentValue.Builder.tableArgument;
58-
import static com.facebook.presto.sql.tree.BooleanLiteral.TRUE_LITERAL;
5958

6059
public class TestTableFunctionInvocation
6160
extends BasePlanTest
@@ -205,22 +204,22 @@ public void testPruneTableFunctionColumns()
205204
ImmutableList.of(ImmutableList.of("a", "b")))
206205
.requiredSymbols(ImmutableList.of(ImmutableList.of("a")))
207206
.specification(specification(ImmutableList.of(), ImmutableList.of(), ImmutableMap.of())),
208-
values(ImmutableList.of("a", "b"), ImmutableList.of(ImmutableList.of(new LongLiteral("1"), TRUE_LITERAL))))));
207+
project(ImmutableMap.of("a", expression("INTEGER'1'"), "b", expression("BOOLEAN'true'")), values(1)))));
209208

210209
// no table function outputs are referenced. All pass-through symbols are pruned from the TableFunctionProcessorNode. The unused symbol "b" is pruned from the source values node.
211210
assertPlan("SELECT 'constant' c FROM TABLE(test.system.pass_through_function(input => TABLE(SELECT 1, true) t(a, b)))",
212211
strictOutput(
213212
ImmutableList.of("c"),
214213
strictProject(
215-
ImmutableMap.of("c", expression("'constant'")),
214+
ImmutableMap.of("c", expression("VARCHAR'constant'")),
216215
tableFunctionProcessor(
217216
builder -> builder
218217
.name("pass_through_function")
219218
.properOutputs(ImmutableList.of("x"))
220219
.passThroughSymbols(ImmutableList.of(ImmutableList.of()))
221220
.requiredSymbols(ImmutableList.of(ImmutableList.of("a")))
222221
.specification(specification(ImmutableList.of(), ImmutableList.of(), ImmutableMap.of())),
223-
values(ImmutableList.of("a"), ImmutableList.of(ImmutableList.of(new LongLiteral("1"))))))));
222+
project(ImmutableMap.of("a", expression("INTEGER'1'")), values(1))))));
224223
}
225224

226225
@Test
@@ -265,7 +264,9 @@ public void testRemoveRedundantTableFunction()
265264
project(
266265
rowNumber(
267266
builder -> builder.partitionBy(ImmutableList.of()),
268-
values(ImmutableList.of("c"), ImmutableList.of(ImmutableList.of(new LongLiteral("2")))))
269-
.withAlias("input_2_row_number", new RowNumberSymbolMatcher()))))));
267+
project(
268+
ImmutableMap.of("c", expression("INTEGER'2'")),
269+
values(1))
270+
).withAlias("input_2_row_number", new RowNumberSymbolMatcher()))))));
270271
}
271272
}

‎presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestRewriteExcludeColumnsFunctionToProjection.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.presto.sql.planner.iterative.rule;
15+
16+
import com.facebook.presto.operator.table.ExcludeColumns.ExcludeColumnsFunctionHandle;
17+
import com.facebook.presto.spi.relation.VariableReferenceExpression;
18+
import com.facebook.presto.sql.planner.assertions.PlanMatchPattern;
19+
import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
20+
import com.google.common.collect.ImmutableList;
21+
import com.google.common.collect.ImmutableMap;
22+
import org.testng.annotations.Test;
23+
24+
import static com.facebook.presto.common.type.BigintType.BIGINT;
25+
import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
26+
import static com.facebook.presto.common.type.SmallintType.SMALLINT;
27+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.expression;
28+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values;
29+
30+
public class TestRewriteExcludeColumnsFunctionToProjection
31+
extends BaseRuleTest
32+
{
33+
@Test
34+
public void rewriteExcludeColumnsFunction()
35+
{
36+
tester().assertThat(new RewriteExcludeColumnsFunctionToProjection())
37+
.on(p -> {
38+
VariableReferenceExpression a = p.variable("a", BOOLEAN);
39+
VariableReferenceExpression b = p.variable("b", BIGINT);
40+
VariableReferenceExpression c = p.variable("c", SMALLINT);
41+
VariableReferenceExpression x = p.variable("x", BIGINT);
42+
VariableReferenceExpression y = p.variable("y", SMALLINT);
43+
return p.tableFunctionProcessor(
44+
builder -> builder
45+
.name("exclude_columns")
46+
.properOutputs(x, y)
47+
.pruneWhenEmpty()
48+
.requiredSymbols(ImmutableList.of(ImmutableList.of(b, c)))
49+
.connectorHandle(new ExcludeColumnsFunctionHandle())
50+
.source(p.values(a, b, c)));
51+
})
52+
.matches(PlanMatchPattern.strictProject(
53+
ImmutableMap.of(
54+
"x", expression("b"),
55+
"y", expression("c")),
56+
values("a", "b", "c")));
57+
}
58+
59+
@Test
60+
public void doNotRewriteOtherFunction()
61+
{
62+
tester().assertThat(new RewriteExcludeColumnsFunctionToProjection())
63+
.on(p -> {
64+
VariableReferenceExpression a = p.variable("a", BOOLEAN);
65+
VariableReferenceExpression b = p.variable("b", BIGINT);
66+
VariableReferenceExpression c = p.variable("c", SMALLINT);
67+
return p.tableFunctionProcessor(
68+
builder -> builder
69+
.name("testing_function")
70+
.requiredSymbols(ImmutableList.of(ImmutableList.of(b, c)))
71+
.source(p.values(a, b, c)));
72+
}).doesNotFire();
73+
}
74+
}

0 commit comments

Comments
 (0)