From d6dd9f947fdc7969a64e5ee66ddaca16592763fe Mon Sep 17 00:00:00 2001
From: Assaf Bern <assaf.bern@starburstdata.com>
Date: Sun, 24 Dec 2023 13:54:12 +0200
Subject: [PATCH] Support translation of JsonPathType to ConnectorExpression

---
 .../ConnectorExpressionTranslator.java        |  9 ++++++++-
 .../TestConnectorExpressionTranslator.java    | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+), 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
index e9d28509754c..31cb589eaf07 100644
--- 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
index 0db62980c9c5..278a67a4fb85 100644
--- 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()
     {