From 9874bde6ce11f7f279de009c9b0ddad647d64400 Mon Sep 17 00:00:00 2001 From: Nikolay Laptev Date: Tue, 18 Jan 2022 07:58:54 -0800 Subject: [PATCH] Simplify cardinality on map keys and values functions A new optimizer rule is added to simplify expressions like `cardinality(map_keys(m))` into `cardinality((m))`. Same for `map_values` function. --- .../presto/sql/planner/PlanOptimizers.java | 2 + .../rule/SimplifyCardinalityMap.java | 25 ++++++ .../rule/SimplifyCardinalityMapRewriter.java | 84 +++++++++++++++++++ .../rule/TestSimplifyCardinalityMap.java | 72 ++++++++++++++++ 4 files changed, 183 insertions(+) create mode 100644 presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMap.java create mode 100644 presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMapRewriter.java create mode 100644 presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestSimplifyCardinalityMap.java diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java index 46ee924b6068c..5e331c4200747 100644 --- a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java +++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java @@ -101,6 +101,7 @@ import com.facebook.presto.sql.planner.iterative.rule.RewriteFilterWithExternalFunctionToProject; import com.facebook.presto.sql.planner.iterative.rule.RewriteSpatialPartitioningAggregation; import com.facebook.presto.sql.planner.iterative.rule.RuntimeReorderJoinSides; +import com.facebook.presto.sql.planner.iterative.rule.SimplifyCardinalityMap; import com.facebook.presto.sql.planner.iterative.rule.SimplifyCountOverConstant; import com.facebook.presto.sql.planner.iterative.rule.SimplifyExpressions; import com.facebook.presto.sql.planner.iterative.rule.SimplifyRowExpressions; @@ -288,6 +289,7 @@ public PlanOptimizers( .addAll(new DesugarAtTimeZone(metadata, sqlParser).rules()) .addAll(new DesugarCurrentUser().rules()) .addAll(new DesugarTryExpression().rules()) + .addAll(new SimplifyCardinalityMap().rules()) .addAll(new DesugarRowSubscript(metadata, sqlParser).rules()) .build()), new IterativeOptimizer( diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMap.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMap.java new file mode 100644 index 0000000000000..4985dc6d2bbff --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMap.java @@ -0,0 +1,25 @@ +/* + * Licensed 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 com.facebook.presto.sql.planner.iterative.rule; + +import static com.facebook.presto.sql.planner.iterative.rule.SimplifyCardinalityMapRewriter.rewrite; + +public class SimplifyCardinalityMap + extends ExpressionRewriteRuleSet +{ + public SimplifyCardinalityMap() + { + super((expression, context) -> rewrite(expression)); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMapRewriter.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMapRewriter.java new file mode 100644 index 0000000000000..96f32eca78b45 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/SimplifyCardinalityMapRewriter.java @@ -0,0 +1,84 @@ +/* + * Licensed 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 com.facebook.presto.sql.planner.iterative.rule; +import com.facebook.presto.sql.tree.Expression; +import com.facebook.presto.sql.tree.ExpressionRewriter; +import com.facebook.presto.sql.tree.ExpressionTreeRewriter; +import com.facebook.presto.sql.tree.FunctionCall; +import com.facebook.presto.sql.tree.QualifiedName; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; + +import java.util.Set; + +/** + * Transforms: + * 
+ * - Cardinality(Map_Values(map))
+ *     - X
+ *
+ * Into: + * 
+ * - Cardinality(map)
+ *     - X
+ *
+ */ +public class SimplifyCardinalityMapRewriter +{ + private static final Set MAP_FUNCTIONS = ImmutableSet.of(QualifiedName.of("map_values"), QualifiedName.of("map_keys")); + + private SimplifyCardinalityMapRewriter() {} + + public static Expression rewrite(Expression expression) + { + return ExpressionTreeRewriter.rewriteWith(new Visitor(), expression); + } + + private static class Visitor + extends ExpressionRewriter + { + @Override + public Expression rewriteFunctionCall(FunctionCall node, Void context, ExpressionTreeRewriter treeRewriter) + { + ImmutableList.Builder rewrittenArguments = ImmutableList.builder(); + + if (node.getName().equals(QualifiedName.of("cardinality"))) { + for (Expression argument : node.getArguments()) { + if (argument instanceof FunctionCall) { + FunctionCall functionCall = (FunctionCall) argument; + if (MAP_FUNCTIONS.contains(functionCall.getName()) && functionCall.getArguments().size() == 1) { + rewrittenArguments.add(treeRewriter.rewrite(functionCall.getArguments().get(0), context)); + continue; + } + } + rewrittenArguments.add(treeRewriter.rewrite(argument, context)); + } + return newFunctionIfRewritten(node, rewrittenArguments); + } + for (Expression argument : node.getArguments()) { + rewrittenArguments.add(treeRewriter.rewrite(argument, context)); + } + return newFunctionIfRewritten(node, rewrittenArguments); + } + + private Expression newFunctionIfRewritten(FunctionCall node, ImmutableList.Builder rewrittenArguments) + { + if (!node.getArguments().equals(rewrittenArguments.build())) { + return new FunctionCall(node.getName(), rewrittenArguments.build()); + } + return node; + } + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestSimplifyCardinalityMap.java b/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestSimplifyCardinalityMap.java new file mode 100644 index 0000000000000..9ea13c01dd712 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestSimplifyCardinalityMap.java @@ -0,0 +1,72 @@ +/* + * Licensed 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 com.facebook.presto.sql.planner.iterative.rule; + +import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest; +import com.facebook.presto.sql.planner.iterative.rule.test.PlanBuilder; +import org.testng.annotations.Test; + +import static com.facebook.presto.sql.planner.iterative.rule.SimplifyCardinalityMapRewriter.rewrite; +import static org.testng.Assert.assertEquals; + +public class TestSimplifyCardinalityMap + extends BaseRuleTest +{ + @Test + public void testRewriteMapValuesCardinality() + { + assertRewritten("cardinality(map_values(m))", "cardinality(m)"); + } + + @Test + public void testRewriteMapValuesMixedCasesCardinality() + { + assertRewritten("CaRDinality(map_values(m))", "cardinaLITY(m)"); + } + + @Test + public void testNoRewriteMapValuesCardinality() + { + assertRewritten("cardinality(map(ARRAY[1,3], ARRAY[2,4]))", "cardinality(map(ARRAY[1,3], ARRAY[2,4]))"); + } + + @Test + public void testNestedRewriteMapValuesCardinality() + { + assertRewritten( + "cardinality(map(ARRAY[cardinality(map_values(m_1)),3], ARRAY[2,cardinality(map_values(m_2))]))", + "cardinality(map(ARRAY[cardinality(m_1),3], ARRAY[2,cardinality(m_2)]))"); + } + + @Test + public void testNestedRewriteMapKeysCardinality() + { + assertRewritten( + "cardinality(map(ARRAY[cardinality(map_keys(m_1)),3], ARRAY[2,cardinality(map_keys(m_2))]))", + "cardinality(map(ARRAY[cardinality(m_1),3], ARRAY[2,cardinality(m_2)]))"); + } + + @Test + public void testAnotherNestedRewriteMapValuesCardinality() + { + assertRewritten( + "cardinality(map(ARRAY[cardinality(map_values(map(ARRAY[1,3], ARRAY[2,4]))),3], ARRAY[2,cardinality(map_values(m_2))]))", + "cardinality(map(ARRAY[cardinality(map(ARRAY[1,3], ARRAY[2,4])),3], ARRAY[2,cardinality(m_2)]))"); + } + + private static void assertRewritten(String from, String to) + { + assertEquals(rewrite(PlanBuilder.expression(from)), PlanBuilder.expression(to)); + } +}