Cast pushdown

ssheikin · May 10, 2024 · 4bb9dee · 4bb9dee
1 parent 8bd3fce
commit 4bb9dee
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 14 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteCast.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteCast.java
@@ -13,61 +13,65 @@
  */
 package io.trino.plugin.jdbc.expression;
 
-import com.google.common.collect.ImmutableList;
 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.plugin.jdbc.QueryParameter;
+import io.trino.plugin.jdbc.WriteMapping;
+import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.expression.Call;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.type.Type;
 
-import java.util.List;
 import java.util.Optional;
+import java.util.function.BiFunction;
 
 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.spi.expression.StandardFunctions.CAST_FUNCTION_NAME;
 import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
 
 public class RewriteCast
         implements ConnectorExpressionRule<Call, ParameterizedExpression>
 {
     private static final Capture<ConnectorExpression> VALUE = newCapture();
-    private static final Capture<List<ConnectorExpression>> EXPRESSIONS = newCapture();
+    private final BiFunction<ConnectorSession, Type, WriteMapping> toWriteMapping;
+
+    public RewriteCast(BiFunction<ConnectorSession, Type, WriteMapping> toWriteMapping)
+    {
+        this.toWriteMapping = requireNonNull(toWriteMapping, "toWriteMapping is null");
+    }
 
     @Override
     public Pattern<Call> getPattern()
     {
         return call()
                 .with(functionName().equalTo(CAST_FUNCTION_NAME))
+                .with(argumentCount().equalTo(1))
                 .with(argument(0).matching(expression().capturedAs(VALUE)));
     }
 
     @Override
     public Optional<ParameterizedExpression> rewrite(Call call, Captures captures, RewriteContext<ParameterizedExpression> context)
     {
-        Type targetType = call.getType();
+        Type trinoType = call.getType();
         ConnectorExpression capturedValue = captures.get(VALUE);
 
         Optional<ParameterizedExpression> value = context.defaultRewrite(capturedValue);
         if (value.isEmpty()) {
-            return Optional.empty();
-        }
-
-        ImmutableList.Builder<QueryParameter> parameters = ImmutableList.builder();
-        Optional<ParameterizedExpression> rewritten = context.defaultRewrite(capturedValue);
-        if (rewritten.isEmpty()) {
             // if argument is a call chain that can't be rewritten, then we can't push it down
             return Optional.empty();
         }
 
+        String targetType = toWriteMapping.apply(context.getSession(), trinoType).getDataType();
+
         return Optional.of(new ParameterizedExpression(
-                format("CAST(%s AS %s)", value.get().expression(), targetType.getDisplayName()),
-                parameters.build()));
+                format("CAST(%s AS %s)", value.get().expression(), targetType),
+                value.get().parameters()));
     }
 }
diff --git a/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java b/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java
@@ -308,7 +308,7 @@ public PostgreSqlClient(
         this.connectorExpressionRewriter = JdbcConnectorExpressionRewriterBuilder.newBuilder()
                 .addStandardRules(this::quoted)
                 .add(new RewriteIn())
-                .add(new RewriteCast())
+                .add(new RewriteCast(this::toWriteMapping))
                 .withTypeClass("integer_type", ImmutableSet.of("tinyint", "smallint", "integer", "bigint"))
                 .withTypeClass("numeric_type", ImmutableSet.of("tinyint", "smallint", "integer", "bigint", "decimal", "real", "double"))
                 .withTypeClass("string_type", ImmutableSet.of("char", "varchar"))

diff --git a/...rino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java b/...rino-postgresql/src/test/java/io/trino/plugin/postgresql/TestPostgreSqlConnectorTest.java
@@ -62,6 +62,7 @@
 import static io.trino.sql.planner.assertions.PlanMatchPattern.node;
 import static io.trino.sql.planner.assertions.PlanMatchPattern.tableScan;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_AGGREGATION_PUSHDOWN;
+import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_JOIN_PUSHDOWN;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_JOIN_PUSHDOWN_WITH_FULL_JOIN;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_LIMIT_PUSHDOWN;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_PREDICATE_PUSHDOWN_WITH_VARCHAR_EQUALITY;
@@ -1013,6 +1014,19 @@ public void testInPredicatePushdown()
         }
     }
 
+    @Test
+    public void testJoinWithCastInCriteriaPushdown()
+    {
+        if (!hasBehavior(SUPPORTS_JOIN_PUSHDOWN)) {
+            return;
+        }
+
+        Session session = joinPushdownEnabled(getSession());
+
+        assertThat(query(session, "SELECT c.name, o.orderdate FROM customer c JOIN orders o ON CAST(c.custkey AS varchar(20)) = CAST(o.custkey AS varchar(21))")).isFullyPushedDown();
+        assertThat(query(session, "SELECT c.name, o.orderdate FROM customer c JOIN orders o ON CAST((c.custkey + 123) AS varchar(20)) = CAST(o.custkey AS varchar(21))")).isFullyPushedDown();
+    }
+
     @Override
     protected String errorMessageForInsertIntoNotNullColumn(String columnName)
     {