addressing PR comments

Signed-off-by: Kenrick Yap <[email protected]>
opensearch-project · Feb 12, 2025 · 0d1cc28 · 0d1cc28
1 parent 80f44e2
commit 0d1cc28
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 38 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/expression/json/JsonFunctions.java b/core/src/main/java/org/opensearch/sql/expression/json/JsonFunctions.java
@@ -5,6 +5,7 @@
 
 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;
@@ -40,6 +41,7 @@ private DefaultFunctionResolver jsonFunction() {
   private DefaultFunctionResolver jsonExtract() {
     return define(
         BuiltinFunctionName.JSON_EXTRACT.getName(),
-        impl(JsonUtils::extractJson, UNDEFINED, STRING, STRING));
+        impl(JsonUtils::extractJsonPaths, UNDEFINED, STRING, ARRAY),
+        impl(JsonUtils::extractJsonPath, UNDEFINED, STRING, STRING));
   }
 }
diff --git a/core/src/main/java/org/opensearch/sql/utils/JsonUtils.java b/core/src/main/java/org/opensearch/sql/utils/JsonUtils.java
@@ -14,6 +14,7 @@
 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;
@@ -88,29 +89,48 @@ public static ExprValue castJson(ExprValue json) {
    * @param path JSON path (e.g. "$.hello")
    * @return ExprValue of value at given path of json string.
    */
-  public static ExprValue extractJson(ExprValue json, ExprValue path) {
+  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();
 
-    if (jsonString.isEmpty()) {
+    return extractJson(jsonString, jsonPath);
+  }
+
+  public static ExprValue extractJsonPaths(ExprValue json, ExprValue paths) {
+    List<ExprValue> pathList = paths.collectionValue();
+    List<ExprValue> resultList = new ArrayList<>();
+
+    for (ExprValue path : pathList) {
+      resultList.add(extractJsonPath(json, path));
+    }
+
+    return new ExprCollectionValue(resultList);
+  }
+
+  private static ExprValue extractJson(String json, String path) {
+    if (json.isEmpty() || json.equals("null")) {
       return LITERAL_NULL;
     }
 
     try {
-      Object results = JsonPath.parse(jsonString).read(jsonPath);
+      Object results = JsonPath.parse(json).read(path);
       return ExprValueUtils.fromObjectValue(results);
-    } catch (PathNotFoundException e) {
+    } catch (PathNotFoundException ignored) {
       return LITERAL_NULL;
-    } catch (InvalidPathException e) {
+    } catch (InvalidPathException invalidPathException) {
       final String errorFormat = "JSON path '%s' is not valid. Error details: %s";
-      throw new SemanticCheckException(String.format(errorFormat, path, e.getMessage()), e);
-    } catch (InvalidJsonException e) {
+      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, e.getMessage()), e);
+      throw new SemanticCheckException(
+          String.format(errorFormat, json, invalidJsonException.getMessage()),
+          invalidJsonException);
     }
   }
 

diff --git a/core/src/test/java/org/opensearch/sql/expression/json/JsonFunctionsTest.java b/core/src/test/java/org/opensearch/sql/expression/json/JsonFunctionsTest.java
@@ -197,11 +197,6 @@ void json_extract_search() {
     execute_extract_json(expected, "{\"a\":1}", "$.a");
   }
 
-  @Test
-  void json_extract_search_arrays_out_of_bound() {
-    execute_extract_json(LITERAL_NULL, "{\"a\":[1,2,3]}", "$.a[4]");
-  }
-
   @Test
   void json_extract_search_arrays() {
     String jsonArray = "{\"a\":[1,2.3,\"abc\",true,null,{\"c\":{\"d\":1}},[1,2,3]]}";
@@ -252,17 +247,20 @@ void json_extract_returns_null() {
 
     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(new ExprStringValue("$.a")))
-            .valueOf());
+        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(new ExprStringValue("$.a")))
-            .valueOf());
+        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
@@ -271,13 +269,9 @@ void json_extract_throws_SemanticCheckException() {
     SemanticCheckException invalidPathError =
         assertThrows(
             SemanticCheckException.class,
-            () ->
-                DSL.jsonExtract(
-                        DSL.literal(new ExprStringValue("{\"a\":1}")),
-                        DSL.literal(new ExprStringValue("$a")))
-                    .valueOf());
+            () -> 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"
+        "JSON path '$a' is not valid. Error details: Illegal character at position 1 expected"
             + " '.' or '['",
         invalidPathError.getMessage());
 
@@ -287,14 +281,13 @@ void json_extract_throws_SemanticCheckException() {
             SemanticCheckException.class,
             () ->
                 DSL.jsonExtract(
-                        DSL.literal(new ExprStringValue("{\"invalid\":\"json\", \"string\"}")),
-                        DSL.literal(new ExprStringValue("$.a")))
+                        DSL.literal("{\"invalid\":\"json\", \"string\"}"), DSL.literal("$.a"))
                     .valueOf());
     assertTrue(
         invalidJsonError
             .getMessage()
             .startsWith(
-                "JSON string '\"{\"invalid\":\"json\", \"string\"}\"' is not valid. Error"
+                "JSON string '{\"invalid\":\"json\", \"string\"}' is not valid. Error"
                     + " details:"));
   }
 
@@ -303,18 +296,21 @@ void json_extract_throws_ExpressionEvaluationException() {
     // null path
     assertThrows(
         ExpressionEvaluationException.class,
-        () ->
-            DSL.jsonExtract(
-                    DSL.literal(new ExprStringValue("{\"a\":1}")), DSL.literal(LITERAL_NULL))
-                .valueOf());
+        () -> DSL.jsonExtract(DSL.literal("{\"a\":1}"), DSL.literal(LITERAL_NULL)).valueOf());
 
     // missing path
     assertThrows(
         ExpressionEvaluationException.class,
-        () ->
-            DSL.jsonExtract(
-                    DSL.literal(new ExprStringValue("{\"a\":1}")), DSL.literal(LITERAL_MISSING))
-                .valueOf());
+        () -> 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]}";
+
+    execute_extract_json(LITERAL_NULL, objectJson, "($.foo, $bar2)");
   }
 
   private static void execute_extract_json(ExprValue expected, String json, String path) {

diff --git a/docs/user/ppl/functions/json.rst b/docs/user/ppl/functions/json.rst
@@ -75,14 +75,14 @@ Argument type: STRING, STRING
 
 Return type: STRING/BOOLEAN/DOUBLE/INTEGER/NULL/STRUCT/ARRAY
 
-- Returns a JSON array if `path` points to multiple results (e.g. $.a[*]) or if the `path` points to 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.
 - Throws ExpressionEvaluationException if `path` is missing.
 
 Example::
 
-    > source=json_test | where json_valid(json_string) | eval json_extract=json_extract(json_string, '$.b') | fields test_name, json_string, json_extract
+    os> source=json_test | where json_valid(json_string) | eval json_extract=json_extract(json_string, '$.b') | fields test_name, json_string, json_extract
     fetched rows / total rows = 6/6
     +---------------------+-------------------------------------+-------------------+
     | test_name           | json_string                         | json_extract      |

diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java
@@ -390,6 +390,17 @@ 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()