added multi path use case

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

Author: kenrickyap

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

@@ -5,7 +5,6 @@
 
 package org.opensearch.sql.expression.json;
 
-import static org.opensearch.sql.data.type.ExprCoreType.ARRAY;
 import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 import static org.opensearch.sql.data.type.ExprCoreType.UNDEFINED;
@@ -41,7 +40,8 @@ private DefaultFunctionResolver jsonFunction() {
   private DefaultFunctionResolver jsonExtract() {
     return define(
         BuiltinFunctionName.JSON_EXTRACT.getName(),
-        impl(JsonUtils::extractJsonPaths, UNDEFINED, STRING, ARRAY),
-        impl(JsonUtils::extractJsonPath, UNDEFINED, STRING, STRING));
+        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

@@ -86,32 +86,32 @@ public static ExprValue castJson(ExprValue json) {
    * Extract value of JSON string at given JSON path.
    *
    * @param json JSON string (e.g. "{\"hello\": \"world\"}").
-   * @param path JSON path (e.g. "$.hello")
+   * @param paths list of JSON path (e.g. "$.hello")
    * @return ExprValue of value at given path of json string.
    */
-  public static ExprValue extractJsonPath(ExprValue json, ExprValue path) {
-    if (json == LITERAL_NULL || json == LITERAL_MISSING) {
-      return json;
-    }
-
-    String jsonString = json.stringValue();
-    String jsonPath = path.stringValue();
+  public static ExprValue extractJson(ExprValue json, ExprValue... paths) {
+    List<ExprValue> resultList = new ArrayList<>();
 
-    return extractJson(jsonString, jsonPath);
-  }
+    for (ExprValue path : paths) {
+      System.out.println("Processing path: " + path);
+      if (json == LITERAL_NULL || json == LITERAL_MISSING) {
+        return json;
+      }
 
-  public static ExprValue extractJsonPaths(ExprValue json, ExprValue paths) {
-    List<ExprValue> pathList = paths.collectionValue();
-    List<ExprValue> resultList = new ArrayList<>();
+      String jsonString = json.stringValue();
+      String jsonPath = path.stringValue();
 
-    for (ExprValue path : pathList) {
-      resultList.add(extractJsonPath(json, path));
+      resultList.add(extractJsonPath(jsonString, jsonPath));
     }
 
-    return new ExprCollectionValue(resultList);
+    if (resultList.size() == 1) {
+      return resultList.getFirst();
+    } else {
+      return new ExprCollectionValue(resultList);
+    }
   }
 
-  private static ExprValue extractJson(String json, String path) {
+  private static ExprValue extractJsonPath(String json, String path) {
     if (json.isEmpty() || json.equals("null")) {
       return LITERAL_NULL;
     }

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

@@ -310,7 +310,15 @@ void json_extract_search_list_of_paths() {
         "{\"foo\": \"foo\", \"fuzz\": true, \"bar\": 1234, \"bar2\": 12.34, \"baz\": null, "
             + "\"obj\": {\"internal\": \"value\"}, \"arr\": [\"string\", true, null]}";
 
-    execute_extract_json(LITERAL_NULL, objectJson, "($.foo, $bar2)");
+    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) {

docs/user/ppl/functions/json.rst

@@ -69,12 +69,13 @@ ____________
 Description
 >>>>>>>>>>>
 
-Usage: `json_extract(doc, path)` Extracts a JSON value from a json document based on the path specified.
+Usage: `json_extract(doc, path[, path])` Extracts a JSON value from a json document based on the path specified.
 
 Argument type: STRING, STRING
 
 Return type: STRING/BOOLEAN/DOUBLE/INTEGER/NULL/STRUCT/ARRAY
 
+- Up to 3 paths can be provided, and results of each `path` with be returned in an ARRAY.
 - Returns an ARRAY if `path` points to multiple results (e.g. $.a[*]) or if the `path` points to an array.
 - Return null if `path` is not valid, or if JSON `doc` is MISSING or NULL.
 - Throws SemanticCheckException if `doc` or `path` is malformed.
@@ -95,18 +96,26 @@ Example::
     | json empty string   |                                     | null              |
     +---------------------+-------------------------------------+-------------------+
 
-    > source=json_test | where test_name="json nested list" | eval json_extract=json_extract('{"a":[{"b":1},{"b":2}]}', '$.b[1].c')
+    os> source=json_test | where test_name="json nested list" | eval json_extract=json_extract('{"a":[{"b":1},{"b":2}]}', '$.b[1].c')
     fetched rows / total rows = 1/1
     +---------------------+-------------------------------------+--------------+
     | test_name           | json_string                         | json_extract |
     |---------------------|-------------------------------------|--------------|
     | json nested list    | {"a":"1","b":[{"c":"2"},{"c":"3"}]} | 3            |
     +---------------------+-------------------------------------+--------------+
 
-    > source=json_test | where test_name="json nested list" | eval json_extract=json_extract('{"a":[{"b":1},{"b":2}]}', '$.b[*].c')
+    os> source=json_test | where test_name="json nested list" | eval json_extract=json_extract('{"a":[{"b":1},{"b":2}]}', '$.b[*].c')
     fetched rows / total rows = 1/1
     +---------------------+-------------------------------------+--------------+
     | test_name           | json_string                         | json_extract |
     |---------------------|-------------------------------------|--------------|
     | json nested list    | {"a":"1","b":[{"c":"2"},{"c":"3"}]} | [2,3]        |
     +---------------------+-------------------------------------+--------------+
+
+    os> source=json_test | where test_name="json nested list" | eval json_extract=json_extract('{"a":[{"b":1},{"b":2}]}', '$.a', '$.b[*].c')
+    fetched rows / total rows = 1/1
+    +---------------------+-------------------------------------+--------------+
+    | test_name           | json_string                         | json_extract |
+    |---------------------|-------------------------------------|--------------|
+    | json nested list    | {"a":"1","b":[{"c":"2"},{"c":"3"}]} | [1,[2,3]]    |
+    +---------------------+-------------------------------------+--------------+

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java

@@ -390,17 +390,6 @@ public UnresolvedExpression visitSpanClause(SpanClauseContext ctx) {
     return new Span(visit(ctx.fieldExpression()), visit(ctx.value), SpanUnit.of(unit));
   }
 
-  @Override
-  public UnresolvedExpression visitJsonExtract(
-      OpenSearchPPLParser.JsonExtractFunctionCallContext ctx) {}
-
-  @Override
-  public UnresolvedExpression visitJsonPathString(OpenSearchPPLParser.JsonPathStringContext ctx) {}
-
-  @Override
-  public List<UnresolvedExpression> visitJsonPathList(
-      OpenSearchPPLParser.JsonPathListContext ctx) {}
-
   private QualifiedName visitIdentifiers(List<? extends ParserRuleContext> ctx) {
     return new QualifiedName(
         ctx.stream()