Normalize OR expressions referencing the same symbol as IN

core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/NormalizeOrExpressionRewriter.java

@@ -0,0 +1,93 @@
+/*
+ * 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 io.trino.sql.planner.iterative.rule;
+
+import com.google.common.collect.ImmutableList;
+import io.trino.sql.tree.ComparisonExpression;
+import io.trino.sql.tree.Expression;
+import io.trino.sql.tree.ExpressionRewriter;
+import io.trino.sql.tree.ExpressionTreeRewriter;
+import io.trino.sql.tree.InListExpression;
+import io.trino.sql.tree.InPredicate;
+import io.trino.sql.tree.LogicalExpression;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static io.trino.sql.ExpressionUtils.and;
+import static io.trino.sql.ExpressionUtils.or;
+import static io.trino.sql.tree.ComparisonExpression.Operator.EQUAL;
+import static io.trino.sql.tree.LogicalExpression.Operator.AND;
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
+
+public final class NormalizeOrExpressionRewriter
+{
+    public static Expression normalizeOrExpression(Expression expression)
+    {
+        return ExpressionTreeRewriter.rewriteWith(new Visitor(), expression);
+    }
+
+    private NormalizeOrExpressionRewriter() {}
+
+    private static class Visitor
+            extends ExpressionRewriter<Void>
+    {
+        @Override
+        public Expression rewriteLogicalExpression(LogicalExpression node, Void context, ExpressionTreeRewriter<Void> treeRewriter)
+        {
+            List<Expression> terms = node.getTerms().stream()
+                    .map(expression -> treeRewriter.rewrite(expression, context))
+                    .collect(toImmutableList());
+
+            if (node.getOperator() == AND) {
+                return and(terms);
+            }
+
+            List<InPredicate> comparisons = terms.stream()
+                    .filter(NormalizeOrExpressionRewriter::isEqualityComparisonExpression)
+                    .map(ComparisonExpression.class::cast)
+                    .collect(groupingBy(
+                            ComparisonExpression::getLeft,
+                            LinkedHashMap::new,
+                            mapping(ComparisonExpression::getRight, Collectors.toList())))
+                    .entrySet().stream()
+                    .filter(entry -> entry.getValue().size() > 1)
+                    .map(entry -> new InPredicate(entry.getKey(), new InListExpression(entry.getValue())))
+                    .collect(Collectors.toList());
+
+            Set<Expression> expressionToSkip = comparisons.stream()
+                    .map(InPredicate::getValue)
+                    .collect(toImmutableSet());
+
+            List<Expression> others = terms.stream()
+                    .filter(expression -> !isEqualityComparisonExpression(expression) || !expressionToSkip.contains(((ComparisonExpression) expression).getLeft()))
+                    .collect(Collectors.toList());
+
+            return or(ImmutableList.<Expression>builder()
+                    .addAll(others)
+                    .addAll(comparisons)
+                    .build());
+        }
+    }
+
+    private static boolean isEqualityComparisonExpression(Expression expression)
+    {
+        return expression instanceof ComparisonExpression && ((ComparisonExpression) expression).getOperator() == EQUAL;
+    }
+}

core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/SimplifyExpressions.java

@@ -31,6 +31,7 @@
 import java.util.Set;
 
 import static io.trino.sql.planner.iterative.rule.ExtractCommonPredicatesExpressionRewriter.extractCommonPredicates;
+import static io.trino.sql.planner.iterative.rule.NormalizeOrExpressionRewriter.normalizeOrExpression;
 import static io.trino.sql.planner.iterative.rule.PushDownNegationsExpressionRewriter.pushDownNegations;
 import static java.util.Objects.requireNonNull;
 
@@ -47,6 +48,7 @@ public static Expression rewrite(Expression expression, Session session, SymbolA
         Map<NodeRef<Expression>, Type> expressionTypes = typeAnalyzer.getTypes(session, symbolAllocator.getTypes(), expression);
         expression = pushDownNegations(plannerContext.getMetadata(), expression, expressionTypes);
         expression = extractCommonPredicates(plannerContext.getMetadata(), expression);
+        expression = normalizeOrExpression(expression);
         expressionTypes = typeAnalyzer.getTypes(session, symbolAllocator.getTypes(), expression);
         ExpressionInterpreter interpreter = new ExpressionInterpreter(expression, plannerContext, session, expressionTypes);
         Object optimized = interpreter.optimize(NoOpSymbolResolver.INSTANCE);

core/trino-main/src/test/java/io/trino/sql/planner/iterative/rule/TestSimplifyExpressions.java

@@ -311,6 +311,16 @@ public void testPushesDownNegationsNumericTypes()
         assertSimplifiesNumericTypes("NOT (1 > R2)", "NOT (1 > R2)");
     }
 
+    @Test
+    public void testRewriteOrExpression()
+    {
+        assertSimplifiesNumericTypes("I1 = 1 OR I1 = 2 ", "I1 IN (1, 2)");
+        // TODO: Implement rule for Merging IN expression
+        assertSimplifiesNumericTypes("I1 = 1 OR I1 = 2 OR I1 IN (3, 4)", "I1 IN (3, 4) OR I1 IN (1, 2)");
+        assertSimplifiesNumericTypes("I1 = 1 OR I1 = 2 OR I1 = I2", "I1 IN (1, 2, I2)");
+        assertSimplifiesNumericTypes("I1 = 1 OR I1 = 2 OR I2 = 3 OR I2 = 4", "I1 IN (1, 2) OR I2 IN (3, 4)");
+    }
+
     private static void assertSimplifiesNumericTypes(String expression, String expected)
     {
         ParsingOptions parsingOptions = new ParsingOptions();