PPL: Add json_extract function

core/build.gradle

@@ -54,6 +54,7 @@ dependencies {
     api "com.fasterxml.jackson.core:jackson-core:${versions.jackson}"
     api "com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}"
     api "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}"
+    api group: 'com.jayway.jsonpath', name: 'json-path', version: '2.9.0'
     api group: 'com.google.code.gson', name: 'gson', version: '2.8.9'
     api group: 'com.tdunning', name: 't-digest', version: '3.3'
     api project(':common')

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -687,6 +687,10 @@ public static FunctionExpression jsonValid(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.JSON_VALID, expressions);
   }
 
+  public static FunctionExpression jsonExtract(Expression... expressions) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.JSON_EXTRACT, expressions);
+  }
+
   public static FunctionExpression stringToJson(Expression value) {
     return compile(FunctionProperties.None, BuiltinFunctionName.JSON, value);
   }

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -207,6 +207,7 @@ public enum BuiltinFunctionName {
   /** Json Functions. */
   JSON_VALID(FunctionName.of("json_valid")),
   JSON(FunctionName.of("json")),
+  JSON_EXTRACT(FunctionName.of("json_extract")),
 
   /** GEOSPATIAL Functions. */
   GEOIP(FunctionName.of("geoip")),

core/src/main/java/org/opensearch/sql/expression/json/JsonFunctions.java

@@ -23,6 +23,7 @@ public class JsonFunctions {
   public void register(BuiltinFunctionRepository repository) {
     repository.register(jsonValid());
     repository.register(jsonFunction());
+    repository.register(jsonExtract());
   }
 
   private DefaultFunctionResolver jsonValid() {
@@ -35,4 +36,12 @@ private DefaultFunctionResolver jsonFunction() {
         BuiltinFunctionName.JSON.getName(),
         impl(nullMissingHandling(JsonUtils::castJson), UNDEFINED, STRING));
   }
+
+  private DefaultFunctionResolver jsonExtract() {
+    return define(
+        BuiltinFunctionName.JSON_EXTRACT.getName(),
+        impl(JsonUtils::extractJson, UNDEFINED, STRING, STRING),
+        impl(JsonUtils::extractJson, UNDEFINED, STRING, STRING, STRING),
+        impl(JsonUtils::extractJson, UNDEFINED, STRING, STRING, STRING, STRING));
+  }
 }

core/src/main/java/org/opensearch/sql/utils/JsonUtils.java

@@ -6,12 +6,18 @@
 package org.opensearch.sql.utils;
 
 import static org.opensearch.sql.data.model.ExprValueUtils.LITERAL_FALSE;
+import static org.opensearch.sql.data.model.ExprValueUtils.LITERAL_MISSING;
 import static org.opensearch.sql.data.model.ExprValueUtils.LITERAL_NULL;
 import static org.opensearch.sql.data.model.ExprValueUtils.LITERAL_TRUE;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.jayway.jsonpath.InvalidJsonException;
+import com.jayway.jsonpath.InvalidPathException;
+import com.jayway.jsonpath.JsonPath;
+import com.jayway.jsonpath.PathNotFoundException;
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
@@ -79,6 +85,58 @@ public static ExprValue castJson(ExprValue json) {
     return processJsonNode(jsonNode);
   }
 
+  /**
+   * Extract value of JSON string at given JSON path.
+   *
+   * @param json JSON string (e.g. "{\"hello\": \"world\"}").
+   * @param paths list of JSON path (e.g. "$.hello")
+   * @return ExprValue of value at given path of json string.
+   */
+  public static ExprValue extractJson(ExprValue json, ExprValue... paths) {
+    List<ExprValue> resultList = new ArrayList<>();
+
+    for (ExprValue path : paths) {
+      System.out.println("Processing path: " + path);
+      if (json == LITERAL_NULL || json == LITERAL_MISSING) {
+        return json;
+      }
+
+      String jsonString = json.stringValue();
+      String jsonPath = path.stringValue();
+
+      resultList.add(extractJsonPath(jsonString, jsonPath));
+    }
+
+    if (resultList.size() == 1) {
+      return resultList.getFirst();
+    } else {
+      return new ExprCollectionValue(resultList);
+    }
+  }
+
+  private static ExprValue extractJsonPath(String json, String path) {
+    if (json.isEmpty() || json.equals("null")) {
+      return LITERAL_NULL;
+    }
+
+    try {
+      Object results = JsonPath.parse(json).read(path);
+      return ExprValueUtils.fromObjectValue(results);
+    } catch (PathNotFoundException ignored) {
+      return LITERAL_NULL;
+    } catch (InvalidPathException invalidPathException) {
+      final String errorFormat = "JSON path '%s' is not valid. Error details: %s";
+      throw new SemanticCheckException(
+          String.format(errorFormat, path, invalidPathException.getMessage()),
+          invalidPathException);
+    } catch (InvalidJsonException invalidJsonException) {
+      final String errorFormat = "JSON string '%s' is not valid. Error details: %s";
+      throw new SemanticCheckException(
+          String.format(errorFormat, json, invalidJsonException.getMessage()),
+          invalidJsonException);
+    }
+  }
+
   private static ExprValue processJsonNode(JsonNode jsonNode) {
     switch (jsonNode.getNodeType()) {
       case ARRAY:

core/src/test/java/org/opensearch/sql/expression/json/JsonFunctionsTest.java

@@ -22,6 +22,7 @@
 import org.opensearch.sql.data.model.ExprBooleanValue;
 import org.opensearch.sql.data.model.ExprCollectionValue;
 import org.opensearch.sql.data.model.ExprDoubleValue;
+import org.opensearch.sql.data.model.ExprFloatValue;
 import org.opensearch.sql.data.model.ExprIntegerValue;
 import org.opensearch.sql.data.model.ExprLongValue;
 import org.opensearch.sql.data.model.ExprNullValue;
@@ -32,6 +33,7 @@
 import org.opensearch.sql.exception.ExpressionEvaluationException;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
+import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.LiteralExpression;
 
@@ -216,5 +218,153 @@ void json_returnsSemanticCheckException() {
                     SemanticCheckException.class,
                     () -> DSL.castJson(expr).valueOf(),
                     "Expected to throw SemanticCheckException when calling castJson with " + expr));
+
+    // invalid type
+    assertThrows(
+        SemanticCheckException.class, () -> DSL.castJson(DSL.literal("invalid")).valueOf());
+
+    // missing bracket
+    assertThrows(SemanticCheckException.class, () -> DSL.castJson(DSL.literal("{{[}}")).valueOf());
+
+    // missing quote
+    assertThrows(
+        SemanticCheckException.class, () -> DSL.castJson(DSL.literal("\"missing quote")).valueOf());
+  }
+
+  @Test
+  void json_extract_search() {
+    ExprValue expected = new ExprIntegerValue(1);
+    execute_extract_json(expected, "{\"a\":1}", "$.a");
+  }
+
+  @Test
+  void json_extract_search_arrays() {
+    String jsonArray = "{\"a\":[1,2.3,\"abc\",true,null,{\"c\":{\"d\":1}},[1,2,3]]}";
+    List<ExprValue> expectedExprValues =
+        List.of(
+            new ExprIntegerValue(1),
+            new ExprFloatValue(2.3),
+            new ExprStringValue("abc"),
+            LITERAL_TRUE,
+            LITERAL_NULL,
+            ExprTupleValue.fromExprValueMap(
+                Map.of("c", ExprTupleValue.fromExprValueMap(Map.of("d", new ExprIntegerValue(1))))),
+            new ExprCollectionValue(
+                List.of(
+                    new ExprIntegerValue(1), new ExprIntegerValue(2), new ExprIntegerValue(3))));
+
+    // extract specific index from JSON list
+    for (int i = 0; i < expectedExprValues.size(); i++) {
+      String path = String.format("$.a[%d]", i);
+      execute_extract_json(expectedExprValues.get(i), jsonArray, path);
+    }
+
+    // extract nested object
+    ExprValue nestedExpected =
+        ExprTupleValue.fromExprValueMap(Map.of("d", new ExprIntegerValue(1)));
+    execute_extract_json(nestedExpected, jsonArray, "$.a[5].c");
+
+    // extract * from JSON list
+    ExprValue starExpected = new ExprCollectionValue(expectedExprValues);
+    execute_extract_json(starExpected, jsonArray, "$.a[*]");
+  }
+
+  @Test
+  void json_extract_returns_null() {
+    List<String> jsonStrings =
+        List.of(
+            "{\"a\":\"1\",\"b\":\"2\"}",
+            "{\"a\":1,\"b\":{\"c\":2,\"d\":3}}",
+            "{\"arr1\": [1,2,3], \"arr2\": [4,5,6]}",
+            "[1, 2, 3, 4]",
+            "[{\"a\":1,\"b\":2}, {\"c\":3,\"d\":2}]",
+            "\"abc\"",
+            "1234",
+            "12.34",
+            "true",
+            "false",
+            "");
+
+    jsonStrings.forEach(str -> execute_extract_json(LITERAL_NULL, str, "$.a.path_not_found_key"));
+
+    // null string literal
+    assertEquals(LITERAL_NULL, DSL.jsonExtract(DSL.literal("null"), DSL.literal("$.a")).valueOf());
+
+    // null json
+    assertEquals(
+        LITERAL_NULL, DSL.jsonExtract(DSL.literal(LITERAL_NULL), DSL.literal("$.a")).valueOf());
+
+    // missing json
+    assertEquals(
+        LITERAL_MISSING,
+        DSL.jsonExtract(DSL.literal(LITERAL_MISSING), DSL.literal("$.a")).valueOf());
+
+    // array out of bounds
+    execute_extract_json(LITERAL_NULL, "{\"a\":[1,2,3]}", "$.a[4]");
+  }
+
+  @Test
+  void json_extract_throws_SemanticCheckException() {
+    // invalid path
+    SemanticCheckException invalidPathError =
+        assertThrows(
+            SemanticCheckException.class,
+            () -> DSL.jsonExtract(DSL.literal("{\"a\":1}"), DSL.literal("$a")).valueOf());
+    assertEquals(
+        "JSON path '$a' is not valid. Error details: Illegal character at position 1 expected"
+            + " '.' or '['",
+        invalidPathError.getMessage());
+
+    // invalid json
+    SemanticCheckException invalidJsonError =
+        assertThrows(
+            SemanticCheckException.class,
+            () ->
+                DSL.jsonExtract(
+                        DSL.literal("{\"invalid\":\"json\", \"string\"}"), DSL.literal("$.a"))
+                    .valueOf());
+    assertTrue(
+        invalidJsonError
+            .getMessage()
+            .startsWith(
+                "JSON string '{\"invalid\":\"json\", \"string\"}' is not valid. Error"
+                    + " details:"));
+  }
+
+  @Test
+  void json_extract_throws_ExpressionEvaluationException() {
+    // null path
+    assertThrows(
+        ExpressionEvaluationException.class,
+        () -> DSL.jsonExtract(DSL.literal("{\"a\":1}"), DSL.literal(LITERAL_NULL)).valueOf());
+
+    // missing path
+    assertThrows(
+        ExpressionEvaluationException.class,
+        () -> DSL.jsonExtract(DSL.literal("{\"a\":1}"), DSL.literal(LITERAL_MISSING)).valueOf());
+  }
+
+  @Test
+  void json_extract_search_list_of_paths() {
+    final String objectJson =
+        "{\"foo\": \"foo\", \"fuzz\": true, \"bar\": 1234, \"bar2\": 12.34, \"baz\": null, "
+            + "\"obj\": {\"internal\": \"value\"}, \"arr\": [\"string\", true, null]}";
+
+    ExprValue expected =
+        new ExprCollectionValue(
+            List.of(new ExprStringValue("foo"), new ExprFloatValue(12.34), LITERAL_NULL));
+    Expression pathExpr1 = DSL.literal(ExprValueUtils.stringValue("$.foo"));
+    Expression pathExpr2 = DSL.literal(ExprValueUtils.stringValue("$.bar2"));
+    Expression pathExpr3 = DSL.literal(ExprValueUtils.stringValue("$.potato"));
+    Expression jsonExpr = DSL.literal(ExprValueUtils.stringValue(objectJson));
+    ExprValue actual = DSL.jsonExtract(jsonExpr, pathExpr1, pathExpr2, pathExpr3).valueOf();
+    assertEquals(expected, actual);
+  }
+
+  private static void execute_extract_json(ExprValue expected, String json, String path) {
+    Expression pathExpr = DSL.literal(ExprValueUtils.stringValue(path));
+    Expression jsonExpr = DSL.literal(ExprValueUtils.stringValue(json));
+    ExprValue actual = DSL.jsonExtract(jsonExpr, pathExpr).valueOf();
+    assertEquals(expected, actual);
   }
 }