OR expression pushdown into PostgreSQL

This includes OR pushdown framework for JDBC connectors, implemented
initially for PostgreSQL.
Along with this, there comes support for comparison expressions to
support a minimal viable usage example for the OR pushdown.

plugin/trino-base-jdbc/pom.xml

@@ -177,6 +177,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-parser</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-spi</artifactId>

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteComparison.java

@@ -0,0 +1,140 @@
+/*
+ * 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.plugin.jdbc.expression;
+
+import com.google.common.collect.ImmutableSet;
+import io.trino.matching.Capture;
+import io.trino.matching.Captures;
+import io.trino.matching.Pattern;
+import io.trino.plugin.base.expression.ConnectorExpressionRule;
+import io.trino.spi.expression.Call;
+import io.trino.spi.expression.ConnectorExpression;
+import io.trino.spi.expression.FunctionName;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Stream;
+
+import static com.google.common.base.Verify.verify;
+import static com.google.common.base.Verify.verifyNotNull;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static io.trino.matching.Capture.newCapture;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argument;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argumentCount;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.expression;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionName;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type;
+import static io.trino.spi.expression.StandardFunctions.EQUAL_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.expression.StandardFunctions.GREATER_THAN_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.expression.StandardFunctions.GREATER_THAN_OR_EQUAL_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.expression.StandardFunctions.IS_DISTINCT_FROM_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.expression.StandardFunctions.LESS_THAN_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.expression.StandardFunctions.LESS_THAN_OR_EQUAL_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.expression.StandardFunctions.NOT_EQUAL_OPERATOR_FUNCTION_NAME;
+import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+import static java.util.function.Function.identity;
+
+public class RewriteComparison
+        implements ConnectorExpressionRule<Call, String>
+{
+    private static final Capture<ConnectorExpression> LEFT = newCapture();
+    private static final Capture<ConnectorExpression> RIGHT = newCapture();
+
+    public enum ComparisonOperator
+    {
+        EQUAL(EQUAL_OPERATOR_FUNCTION_NAME, "="),
+        NOT_EQUAL(NOT_EQUAL_OPERATOR_FUNCTION_NAME, "<>"),
+        LESS_THAN(LESS_THAN_OPERATOR_FUNCTION_NAME, "<"),
+        LESS_THAN_OR_EQUAL(LESS_THAN_OR_EQUAL_OPERATOR_FUNCTION_NAME, "<="),
+        GREATER_THAN(GREATER_THAN_OPERATOR_FUNCTION_NAME, ">"),
+        GREATER_THAN_OR_EQUAL(GREATER_THAN_OR_EQUAL_OPERATOR_FUNCTION_NAME, ">="),
+        IS_DISTINCT_FROM(IS_DISTINCT_FROM_OPERATOR_FUNCTION_NAME, "IS DISTINCT FROM"),
+        /**/;
+
+        private final FunctionName functionName;
+        private final String operator;
+
+        private static final Map<FunctionName, ComparisonOperator> OPERATOR_BY_FUNCTION_NAME = Stream.of(values())
+                .collect(toImmutableMap(ComparisonOperator::getFunctionName, identity()));
+
+        ComparisonOperator(FunctionName functionName, String operator)
+        {
+            this.functionName = requireNonNull(functionName, "functionName is null");
+            this.operator = requireNonNull(operator, "operator is null");
+        }
+
+        private FunctionName getFunctionName()
+        {
+            return functionName;
+        }
+
+        private String getOperator()
+        {
+            return operator;
+        }
+
+        private static ComparisonOperator forFunctionName(FunctionName functionName)
+        {
+            return verifyNotNull(OPERATOR_BY_FUNCTION_NAME.get(functionName), "Function name not recognized: %s", functionName);
+        }
+    }
+
+    private final Pattern<Call> pattern;
+
+    public RewriteComparison(ComparisonOperator... enabledOperators)
+    {
+        this(ImmutableSet.copyOf(enabledOperators));
+    }
+
+    public RewriteComparison(Set<ComparisonOperator> enabledOperators)
+    {
+        Set<FunctionName> functionNames = enabledOperators.stream()
+                .map(ComparisonOperator::getFunctionName)
+                .collect(toImmutableSet());
+
+        pattern = call()
+                .with(functionName().matching(functionNames::contains))
+                .with(type().equalTo(BOOLEAN))
+                .with(argumentCount().equalTo(2))
+                .with(argument(0).matching(expression().capturedAs(LEFT)))
+                .with(argument(1).matching(expression().capturedAs(RIGHT)));
+    }
+
+    @Override
+    public Pattern<Call> getPattern()
+    {
+        return pattern;
+    }
+
+    @Override
+    public Optional<String> rewrite(Call call, Captures captures, RewriteContext<String> context)
+    {
+        Optional<String> left = context.defaultRewrite(captures.get(LEFT));
+        if (left.isEmpty()) {
+            return Optional.empty();
+        }
+        Optional<String> right = context.defaultRewrite(captures.get(RIGHT));
+        if (right.isEmpty()) {
+            return Optional.empty();
+        }
+        verify(call.getFunctionName().getCatalogSchema().isEmpty()); // filtered out by the pattern
+        ComparisonOperator operator = ComparisonOperator.forFunctionName(call.getFunctionName());
+        return Optional.of(format("(%s) %s (%s)", left.get(), operator.getOperator(), right.get()));
+    }
+}

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteExactNumericConstant.java

@@ -0,0 +1,64 @@
+/*
+ * 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.plugin.jdbc.expression;
+
+import io.trino.matching.Captures;
+import io.trino.matching.Pattern;
+import io.trino.plugin.base.expression.ConnectorExpressionRule;
+import io.trino.spi.expression.Constant;
+import io.trino.spi.type.DecimalType;
+import io.trino.spi.type.Decimals;
+import io.trino.spi.type.Int128;
+import io.trino.spi.type.Type;
+
+import java.util.Optional;
+
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.constant;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type;
+import static io.trino.spi.type.BigintType.BIGINT;
+import static io.trino.spi.type.IntegerType.INTEGER;
+import static io.trino.spi.type.SmallintType.SMALLINT;
+import static io.trino.spi.type.TinyintType.TINYINT;
+
+public class RewriteExactNumericConstant
+        implements ConnectorExpressionRule<Constant, String>
+{
+    private static final Pattern<Constant> PATTERN = constant().with(type().matching(type ->
+            type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT || type instanceof DecimalType));
+
+    @Override
+    public Pattern<Constant> getPattern()
+    {
+        return PATTERN;
+    }
+
+    @Override
+    public Optional<String> rewrite(Constant constant, Captures captures, RewriteContext<String> context)
+    {
+        Type type = constant.getType();
+        if (type == TINYINT || type == SMALLINT || type == INTEGER || type == BIGINT) {
+            return Optional.of(Long.toString((long) constant.getValue()));
+        }
+
+        if (type instanceof DecimalType) {
+            DecimalType decimalType = (DecimalType) type;
+            if (decimalType.isShort()) {
+                return Optional.of(Decimals.toString((long) constant.getValue(), decimalType.getScale()));
+            }
+            return Optional.of(Decimals.toString((Int128) constant.getValue(), decimalType.getScale()));
+        }
+
+        throw new UnsupportedOperationException("Unsupported type: " + type);
+    }
+}

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteOr.java

@@ -0,0 +1,65 @@
+/*
+ * 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.plugin.jdbc.expression;
+
+import io.trino.matching.Captures;
+import io.trino.matching.Pattern;
+import io.trino.plugin.base.expression.ConnectorExpressionRule;
+import io.trino.spi.expression.Call;
+import io.trino.spi.expression.ConnectorExpression;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static com.google.common.base.Verify.verify;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionName;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type;
+import static io.trino.spi.expression.StandardFunctions.OR_FUNCTION_NAME;
+import static io.trino.spi.type.BooleanType.BOOLEAN;
+
+public class RewriteOr
+        implements ConnectorExpressionRule<Call, String>
+{
+    private static final Pattern<Call> PATTERN = call()
+            .with(functionName().equalTo(OR_FUNCTION_NAME))
+            .with(type().equalTo(BOOLEAN));
+
+    @Override
+    public Pattern<Call> getPattern()
+    {
+        return PATTERN;
+    }
+
+    @Override
+    public Optional<String> rewrite(Call call, Captures captures, RewriteContext<String> context)
+    {
+        List<ConnectorExpression> arguments = call.getArguments();
+        verify(!arguments.isEmpty(), "no arguments");
+        List<String> terms = new ArrayList<>(arguments.size());
+        for (ConnectorExpression argument : arguments) {
+            verify(argument.getType() == BOOLEAN, "Unexpected type of OR argument: %s", argument.getType());
+            Optional<String> rewritten = context.defaultRewrite(argument);
+            if (rewritten.isEmpty()) {
+                return Optional.empty();
+            }
+            terms.add(rewritten.get());
+        }
+
+        return Optional.of(terms.stream()
+                .collect(Collectors.joining(") OR (", "(", ")")));
+    }
+}

plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/expression/TestRewriteComparison.java

@@ -0,0 +1,36 @@
+/*
+ * 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.plugin.jdbc.expression;
+
+import io.trino.sql.tree.ComparisonExpression;
+import org.testng.annotations.Test;
+
+import java.util.stream.Stream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TestRewriteComparison
+{
+    @Test
+    public void testOperatorEnumsInSync()
+    {
+        assertThat(
+                Stream.of(RewriteComparison.ComparisonOperator.values())
+                        .map(Enum::name))
+                .containsExactlyInAnyOrder(
+                        Stream.of(ComparisonExpression.Operator.values())
+                                .map(Enum::name)
+                                .toArray(String[]::new));
+    }
+}

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java

@@ -62,8 +62,11 @@
 import io.trino.plugin.jdbc.aggregation.ImplementSum;
 import io.trino.plugin.jdbc.aggregation.ImplementVariancePop;
 import io.trino.plugin.jdbc.aggregation.ImplementVarianceSamp;
+import io.trino.plugin.jdbc.expression.RewriteComparison;
+import io.trino.plugin.jdbc.expression.RewriteExactNumericConstant;
 import io.trino.plugin.jdbc.expression.RewriteLike;
 import io.trino.plugin.jdbc.expression.RewriteLikeWithEscape;
+import io.trino.plugin.jdbc.expression.RewriteOr;
 import io.trino.plugin.jdbc.expression.RewriteVarcharConstant;
 import io.trino.plugin.jdbc.expression.RewriteVariable;
 import io.trino.plugin.jdbc.mapping.IdentifierMapping;
@@ -308,6 +311,10 @@ public PostgreSqlClient(
         connectorExpressionRewriter = new ConnectorExpressionRewriter<>(this::quoted, ImmutableSet.of(
                 new RewriteVariable(),
                 new RewriteVarcharConstant(),
+                new RewriteExactNumericConstant(),
+                new RewriteOr(),
+                // TODO allow all comparison operators for numeric types
+                new RewriteComparison(RewriteComparison.ComparisonOperator.EQUAL, RewriteComparison.ComparisonOperator.NOT_EQUAL),
                 new RewriteLike(),
                 new RewriteLikeWithEscape()));
     }