functionality implemented

Signed-off-by: Kenrick Yap <[email protected]>

Author: kenrickyap

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 Aggregator avg(Expression... expressions) {
     return aggregate(BuiltinFunctionName.AVG, expressions);
   }

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")),
 
   /** NULL Test. */
   IS_NULL(FunctionName.of("is null")),

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,10 @@ private DefaultFunctionResolver jsonFunction() {
         BuiltinFunctionName.JSON.getName(),
         impl(nullMissingHandling(JsonUtils::castJson), UNDEFINED, STRING));
   }
+
+  private DefaultFunctionResolver jsonExtract() {
+    return define(
+        BuiltinFunctionName.JSON_EXTRACT.getName(),
+        impl(nullMissingHandling(JsonUtils::extractJson), UNDEFINED, STRING, STRING));
+  }
 }

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

@@ -7,10 +7,22 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.jayway.jsonpath.Configuration;
+import com.jayway.jsonpath.DocumentContext;
+import com.jayway.jsonpath.InvalidJsonException;
+import com.jayway.jsonpath.JsonPath;
+
+import java.io.ObjectInputFilter;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.jayway.jsonpath.Option;
+import com.jayway.jsonpath.PathNotFoundException;
+import com.jayway.jsonpath.spi.json.JacksonJsonProvider;
+import com.jayway.jsonpath.spi.mapper.JacksonMappingProvider;
 import lombok.experimental.UtilityClass;
 import org.opensearch.sql.data.model.ExprCollectionValue;
 import org.opensearch.sql.data.model.ExprDoubleValue;
@@ -46,16 +58,59 @@ public static ExprValue isValidJson(ExprValue jsonExprValue) {
 
   /** Converts a JSON encoded string to an Expression object. */
   public static ExprValue castJson(ExprValue json) {
+    JsonNode jsonNode = jsonToNode(json);
+    return processJsonNode(jsonNode);
+  }
+
+  public static ExprValue extractJson(ExprValue json, ExprValue path) {
+    String jsonString = json.stringValue();
+    String jsonPath = path.stringValue();
+
+    try {
+      Configuration config = Configuration.builder()
+              .options(Option.AS_PATH_LIST)
+              .build();
+      List<String> resultPaths = JsonPath.using(config).parse(jsonString).read(jsonPath);
+
+      List<ExprValue> elements = new LinkedList<>();
+
+      for (String resultPath : resultPaths) {
+        Object result = JsonPath.parse(jsonString).read(resultPath);
+        String resultJsonString = new ObjectMapper().writeValueAsString(result);
+        try {
+          elements.add(processJsonNode(jsonStringToNode(resultJsonString)));
+        } catch (SemanticCheckException e) {
+          elements.add(new ExprStringValue(resultJsonString));
+        }
+      }
+
+      if (elements.size() == 1) {
+        return elements.get(0);
+      } else {
+        return new ExprCollectionValue(elements);
+      }
+    } catch (PathNotFoundException e) {
+      return LITERAL_NULL;
+    } catch (InvalidJsonException | JsonProcessingException e) {
+      final String errorFormat = "JSON string '%s' is not valid. Error details: %s";
+      throw new SemanticCheckException(String.format(errorFormat, json, e.getMessage()), e);
+    }
+  }
+
+  private static JsonNode jsonToNode(ExprValue json) {
+    return jsonStringToNode(json.stringValue());
+  }
+
+  private static JsonNode jsonStringToNode(String jsonString) {
     ObjectMapper objectMapper = new ObjectMapper();
     JsonNode jsonNode;
     try {
-      jsonNode = objectMapper.readTree(json.stringValue());
+      jsonNode = objectMapper.readTree(jsonString);
     } catch (JsonProcessingException e) {
       final String errorFormat = "JSON string '%s' is not valid. Error details: %s";
-      throw new SemanticCheckException(String.format(errorFormat, json, e.getMessage()), e);
+      throw new SemanticCheckException(String.format(errorFormat, jsonString, e.getMessage()), e);
     }
-
-    return processJsonNode(jsonNode);
+    return jsonNode;
   }
 
   private static ExprValue processJsonNode(JsonNode jsonNode) {

ppl/src/main/antlr/OpenSearchPPLLexer.g4

@@ -335,6 +335,7 @@ CIDRMATCH:                          'CIDRMATCH';
 // JSON FUNCTIONS
 JSON_VALID:                         'JSON_VALID';
 JSON:                               'JSON';
+JSON_EXTRACT:                       'JSON_EXTRACT';
 
 // FLOWCONTROL FUNCTIONS
 IFNULL:                             'IFNULL';

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -704,6 +704,7 @@ positionFunctionName
 
 jsonFunctionName
    : JSON
+   | JSON_EXTRACT
    ;
 
 // operators