Support translation of JsonPathType to ConnectorExpression

trinodb · Jan 10, 2024 · d6dd9f9 · d6dd9f9
1 parent 875fe1b
commit d6dd9f9
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/ConnectorExpressionTranslator.java b/core/trino-main/src/main/java/io/trino/sql/planner/ConnectorExpressionTranslator.java
@@ -21,6 +21,7 @@
 import io.trino.Session;
 import io.trino.connector.system.GlobalSystemConnector;
 import io.trino.metadata.ResolvedFunction;
+import io.trino.operator.scalar.JsonPath;
 import io.trino.plugin.base.expression.ConnectorExpressions;
 import io.trino.security.AllowAllAccessControl;
 import io.trino.spi.connector.CatalogSchemaName;
@@ -63,6 +64,7 @@
 import io.trino.sql.tree.SubscriptExpression;
 import io.trino.sql.tree.SymbolReference;
 import io.trino.type.JoniRegexp;
+import io.trino.type.JsonPathType;
 import io.trino.type.LikePattern;
 import io.trino.type.Re2JRegexp;
 import io.trino.type.Re2JRegexpType;
@@ -333,7 +335,8 @@ private Optional<Expression> translateCall(String functionName, ResolvedFunction
                     return Optional.empty();
                 }
                 Expression expression = translated.get();
-                if ((formalType == JONI_REGEXP || formalType instanceof Re2JRegexpType) && argumentType instanceof VarcharType) {
+                if ((formalType == JONI_REGEXP || formalType instanceof Re2JRegexpType || formalType instanceof JsonPathType)
+                        && argumentType instanceof VarcharType) {
                     // These types are not used in connector expressions, so require special handling when translating back to expressions.
                     expression = new Cast(expression, toSqlType(formalType));
                 }
@@ -814,6 +817,10 @@ private ConnectorExpression constantFor(Expression node)
                 Slice pattern = Slices.utf8Slice(((Re2JRegexp) value).pattern());
                 return new Constant(pattern, createVarcharType(countCodePoints(pattern)));
             }
+            if (type instanceof JsonPathType) {
+                Slice pattern = Slices.utf8Slice(((JsonPath) value).pattern());
+                return new Constant(pattern, createVarcharType(countCodePoints(pattern)));
+            }
             return new Constant(value, type);
         }
 

diff --git a/core/trino-main/src/test/java/io/trino/sql/planner/TestConnectorExpressionTranslator.java b/core/trino-main/src/test/java/io/trino/sql/planner/TestConnectorExpressionTranslator.java
@@ -88,6 +88,7 @@
 import static io.trino.sql.planner.TypeAnalyzer.createTestingTypeAnalyzer;
 import static io.trino.testing.TransactionBuilder.transaction;
 import static io.trino.type.JoniRegexpType.JONI_REGEXP;
+import static io.trino.type.JsonPathType.JSON_PATH;
 import static io.trino.type.LikeFunctions.likePattern;
 import static io.trino.type.LikePatternType.LIKE_PATTERN;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -455,6 +456,24 @@ public void testTranslateRegularExpression()
                 });
     }
 
+    @Test
+    void testTranslateJsonPath()
+    {
+        // JSON path type is considered implementation detail of the engine and is not exposed to connectors
+        // within ConnectorExpression. Instead, it is replaced with a varchar pattern.
+        assertTranslationRoundTrips(
+                BuiltinFunctionCallBuilder.resolve(PLANNER_CONTEXT.getMetadata())
+                        .setName("json_extract_scalar")
+                        .addArgument(VARCHAR_TYPE, new SymbolReference("varchar_symbol_1"))
+                        .addArgument(JSON_PATH, new Cast(new StringLiteral("$.path"), toSqlType(JSON_PATH)))
+                        .build(),
+                new Call(
+                        VARCHAR_TYPE,
+                        new FunctionName("json_extract_scalar"),
+                        List.of(new Variable("varchar_symbol_1", VARCHAR_TYPE),
+                                new Constant(utf8Slice("$.path"), createVarcharType(6)))));
+    }
+
     @Test
     public void testTranslateIn()
     {