Add table handle to ProjectFunctionRewriter.rewrite

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/projection/ProjectFunctionRewriter.java

@@ -19,6 +19,7 @@
 import io.trino.plugin.base.projection.ProjectFunctionRule.RewriteContext;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.expression.ConnectorExpression;
 
 import java.util.Iterator;
@@ -39,8 +40,9 @@ public ProjectFunctionRewriter(ConnectorExpressionRewriter<ExpressionResult> con
         this.rules = ImmutableSet.copyOf(requireNonNull(rules, "rules is null"));
     }
 
-    public Optional<ProjectionResult> rewrite(ConnectorSession session, ConnectorExpression projectionExpression, Map<String, ColumnHandle> assignments)
+    public Optional<ProjectionResult> rewrite(ConnectorSession session, ConnectorTableHandle handle, ConnectorExpression projectionExpression, Map<String, ColumnHandle> assignments)
     {
+        requireNonNull(handle, "handle is null");
         requireNonNull(projectionExpression, "projectionExpression is null");
         requireNonNull(assignments, "assignments is null");
 
@@ -70,7 +72,7 @@ public Optional<ExpressionResult> rewriteExpression(ConnectorExpression expressi
                 Iterator<Match> matches = rule.getPattern().match(projectionExpression, context).iterator();
                 while (matches.hasNext()) {
                     Match match = matches.next();
-                    Optional<ProjectionResult> rewritten = rule.rewrite(projectionExpression, match.captures(), context);
+                    Optional<ProjectionResult> rewritten = rule.rewrite(handle, projectionExpression, match.captures(), context);
                     if (rewritten.isPresent()) {
                         return rewritten;
                     }

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/projection/ProjectFunctionRule.java

@@ -17,6 +17,7 @@
 import io.trino.matching.Pattern;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.expression.ConnectorExpression;
 
 import java.util.Map;
@@ -34,7 +35,7 @@ default boolean isEnabled(ConnectorSession session)
         return true;
     }
 
-    Optional<ProjectionResult> rewrite(ConnectorExpression projectionExpression, Captures captures, RewriteContext<ExpressionResult> context);
+    Optional<ProjectionResult> rewrite(ConnectorTableHandle handle, ConnectorExpression projectionExpression, Captures captures, RewriteContext<ExpressionResult> context);
 
     interface RewriteContext<ExpressionResult>
     {

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java

@@ -240,9 +240,9 @@ public Optional<ParameterizedExpression> convertPredicate(ConnectorSession sessi
     }
 
     @Override
-    public Optional<JdbcExpression> convertProjection(ConnectorSession session, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
+    public Optional<JdbcExpression> convertProjection(ConnectorSession session, JdbcTableHandle handle, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
     {
-        return delegate.convertProjection(session, expression, assignments);
+        return delegate.convertProjection(session, handle, expression, assignments);
     }
 
     @Override

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultJdbcMetadata.java

@@ -387,7 +387,7 @@ private Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProject
         Set<ConnectorExpression> translatedExpressions = new HashSet<>();
 
         for (ConnectorExpression projection : ImmutableSet.copyOf(projections)) {
-            RewrittenExpression rewrittenExpression = rewriteExpression(session, nextSyntheticColumnId, projection, assignments, translatedExpressions);
+            RewrittenExpression rewrittenExpression = rewriteExpression(session, handle, nextSyntheticColumnId, projection, assignments, translatedExpressions);
             nextSyntheticColumnId = rewrittenExpression.nextSyntheticColumnId();
             newVariablesBuilder.putAll(rewrittenExpression.syntheticVariables());
             columnExpressionsBuilder.putAll(rewrittenExpression.columnExpressions());
@@ -430,6 +430,7 @@ private Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProject
 
     private RewrittenExpression rewriteExpression(
             ConnectorSession session,
+            JdbcTableHandle handle,
             int nextSyntheticColumnId,
             ConnectorExpression projection,
             Map<String, ColumnHandle> assignments,
@@ -453,7 +454,7 @@ private RewrittenExpression rewriteExpression(
                     ImmutableSet.of((JdbcColumnHandle) assignments.get(variable.getName())),
                     ImmutableList.of(new Assignment(variable.getName(), assignments.get(variable.getName()), variable.getType())));
         }
-        Optional<JdbcExpression> convertedExpression = jdbcClient.convertProjection(session, projection, assignments);
+        Optional<JdbcExpression> convertedExpression = jdbcClient.convertProjection(session, handle, projection, assignments);
         if (convertedExpression.isPresent()) {
             String columnName = SYNTHETIC_COLUMN_NAME_PREFIX + nextSyntheticColumnId;
             JdbcColumnHandle newColumn = JdbcColumnHandle.builder()
@@ -481,6 +482,7 @@ private RewrittenExpression rewriteExpression(
         for (ConnectorExpression child : projection.getChildren()) {
             RewrittenExpression rewrittenExpression = rewriteExpression(
                     session,
+                    handle,
                     nextSyntheticColumnId,
                     child,
                     assignments,

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java

@@ -163,9 +163,9 @@ public Optional<ParameterizedExpression> convertPredicate(ConnectorSession sessi
     }
 
     @Override
-    public Optional<JdbcExpression> convertProjection(ConnectorSession session, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
+    public Optional<JdbcExpression> convertProjection(ConnectorSession session, JdbcTableHandle handle, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
     {
-        return delegate().convertProjection(session, expression, assignments);
+        return delegate().convertProjection(session, handle, expression, assignments);
     }
 
     @Override

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java

@@ -100,7 +100,7 @@ default Optional<ParameterizedExpression> convertPredicate(ConnectorSession sess
         return Optional.empty();
     }
 
-    default Optional<JdbcExpression> convertProjection(ConnectorSession session, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
+    default Optional<JdbcExpression> convertProjection(ConnectorSession session, JdbcTableHandle handle, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
     {
         return Optional.empty();
     }

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java

@@ -185,9 +185,9 @@ public Optional<ParameterizedExpression> convertPredicate(ConnectorSession sessi
     }
 
     @Override
-    public Optional<JdbcExpression> convertProjection(ConnectorSession session, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
+    public Optional<JdbcExpression> convertProjection(ConnectorSession session, JdbcTableHandle handle, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
     {
-        return stats.getConvertProjection().wrap(() -> delegate().convertProjection(session, expression, assignments));
+        return stats.getConvertProjection().wrap(() -> delegate().convertProjection(session, handle, expression, assignments));
     }
 
     @Override

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

@@ -840,9 +840,9 @@ public Optional<ParameterizedExpression> convertPredicate(ConnectorSession sessi
     }
 
     @Override
-    public Optional<JdbcExpression> convertProjection(ConnectorSession session, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
+    public Optional<JdbcExpression> convertProjection(ConnectorSession session, JdbcTableHandle handle, ConnectorExpression expression, Map<String, ColumnHandle> assignments)
     {
-        return projectFunctionRewriter.rewrite(session, expression, assignments);
+        return projectFunctionRewriter.rewrite(session, handle, expression, assignments);
     }
 
     private static Optional<JdbcTypeHandle> toTypeHandle(DecimalType decimalType)

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/rule/RewriteDotProductFunction.java

@@ -22,6 +22,7 @@
 import io.trino.plugin.jdbc.JdbcTypeHandle;
 import io.trino.plugin.jdbc.QueryParameter;
 import io.trino.plugin.jdbc.expression.ParameterizedExpression;
+import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.expression.Call;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.expression.FunctionName;
@@ -60,7 +61,7 @@ public Pattern<? extends ConnectorExpression> getPattern()
     }
 
     @Override
-    public Optional<JdbcExpression> rewrite(ConnectorExpression projectionExpression, Captures captures, RewriteContext<ParameterizedExpression> context)
+    public Optional<JdbcExpression> rewrite(ConnectorTableHandle handle, ConnectorExpression projectionExpression, Captures captures, RewriteContext<ParameterizedExpression> context)
     {
         ConnectorExpression call = captures.get(CALL);
 

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/rule/RewriteStringReverseFunction.java

@@ -22,6 +22,7 @@
 import io.trino.plugin.jdbc.JdbcExpression;
 import io.trino.plugin.jdbc.JdbcTypeHandle;
 import io.trino.plugin.jdbc.expression.ParameterizedExpression;
+import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.expression.Call;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.expression.FunctionName;
@@ -57,7 +58,7 @@ public Pattern<? extends ConnectorExpression> getPattern()
     }
 
     @Override
-    public Optional<JdbcExpression> rewrite(ConnectorExpression projectionExpression, Captures captures, RewriteContext<ParameterizedExpression> context)
+    public Optional<JdbcExpression> rewrite(ConnectorTableHandle handle, ConnectorExpression projectionExpression, Captures captures, RewriteContext<ParameterizedExpression> context)
     {
         Variable argument = captures.get(ARGUMENT);
         JdbcTypeHandle typeHandle = ((JdbcColumnHandle) context.getAssignment(argument.getName())).getJdbcTypeHandle();

plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/rule/RewriteVectorDistanceFunction.java

@@ -24,6 +24,7 @@
 import io.trino.plugin.jdbc.QueryParameter;
 import io.trino.plugin.jdbc.expression.ParameterizedExpression;
 import io.trino.spi.block.Block;
+import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.expression.Call;
 import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.expression.Constant;
@@ -76,7 +77,7 @@ public Pattern<? extends ConnectorExpression> getPattern()
     }
 
     @Override
-    public Optional<JdbcExpression> rewrite(ConnectorExpression projectionExpression, Captures captures, RewriteContext<ParameterizedExpression> context)
+    public Optional<JdbcExpression> rewrite(ConnectorTableHandle handle, ConnectorExpression projectionExpression, Captures captures, RewriteContext<ParameterizedExpression> context)
     {
         Optional<ParameterizedExpression> leftExpression = rewrite(captures.get(LEFT_ARGUMENT), context);
         if (leftExpression.isEmpty()) {