Feature: PPL - json_set function

async-query-core/src/main/antlr/OpenSearchPPLLexer.g4

@@ -384,10 +384,10 @@ TO_JSON_STRING:                     'TO_JSON_STRING';
 JSON_EXTRACT:                       'JSON_EXTRACT';
 JSON_KEYS:                          'JSON_KEYS';
 JSON_VALID:                         'JSON_VALID';
+JSON_SET:                           'JSON_SET';
 //JSON_APPEND:                        'JSON_APPEND';
 //JSON_DELETE:                        'JSON_DELETE';
 //JSON_EXTEND:                        'JSON_EXTEND';
-//JSON_SET:                           'JSON_SET';
 //JSON_ARRAY_ALL_MATCH:               'JSON_ARRAY_ALL_MATCH';
 //JSON_ARRAY_ANY_MATCH:               'JSON_ARRAY_ANY_MATCH';
 //JSON_ARRAY_FILTER:                  'JSON_ARRAY_FILTER';

async-query-core/src/main/antlr/OpenSearchPPLParser.g4

@@ -872,10 +872,10 @@ jsonFunctionName
    | JSON_EXTRACT
    | JSON_KEYS
    | JSON_VALID
+   | JSON_SET
 //   | JSON_APPEND
 //   | JSON_DELETE
 //   | JSON_EXTEND
-//   | JSON_SET
 //   | JSON_ARRAY_ALL_MATCH
 //   | JSON_ARRAY_ANY_MATCH
 //   | JSON_ARRAY_FILTER

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,14 @@ 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 jsonSet(Expression... expressions) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.JSON_SET, 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,8 @@ public enum BuiltinFunctionName {
   /** Json Functions. */
   JSON_VALID(FunctionName.of("json_valid")),
   JSON(FunctionName.of("json")),
+  JSON_EXTRACT(FunctionName.of("json_extract")),
+  JSON_SET(FunctionName.of("json_set")),
 
   /** GEOSPATIAL Functions. */
   GEOIP(FunctionName.of("geoip")),

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

@@ -5,8 +5,21 @@
 
 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.BYTE;
+import static org.opensearch.sql.data.type.ExprCoreType.DATE;
+import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE;
+import static org.opensearch.sql.data.type.ExprCoreType.FLOAT;
+import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
+import static org.opensearch.sql.data.type.ExprCoreType.INTERVAL;
+import static org.opensearch.sql.data.type.ExprCoreType.IP;
+import static org.opensearch.sql.data.type.ExprCoreType.LONG;
+import static org.opensearch.sql.data.type.ExprCoreType.SHORT;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+import static org.opensearch.sql.data.type.ExprCoreType.STRUCT;
+import static org.opensearch.sql.data.type.ExprCoreType.TIME;
+import static org.opensearch.sql.data.type.ExprCoreType.TIMESTAMP;
 import static org.opensearch.sql.data.type.ExprCoreType.UNDEFINED;
 import static org.opensearch.sql.expression.function.FunctionDSL.define;
 import static org.opensearch.sql.expression.function.FunctionDSL.impl;
@@ -23,6 +36,8 @@ public class JsonFunctions {
   public void register(BuiltinFunctionRepository repository) {
     repository.register(jsonValid());
     repository.register(jsonFunction());
+    repository.register(jsonExtract());
+    repository.register(jsonSet());
   }
 
   private DefaultFunctionResolver jsonValid() {
@@ -35,4 +50,32 @@ 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));
+  }
+
+  private DefaultFunctionResolver jsonSet() {
+    return define(
+        BuiltinFunctionName.JSON_SET.getName(),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, BYTE),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, SHORT),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, INTEGER),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, LONG),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, FLOAT),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, DOUBLE),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, STRING),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, BOOLEAN),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, DATE),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, TIME),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, TIMESTAMP),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, INTERVAL),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, IP),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, STRUCT),
+        impl(nullMissingHandling(JsonUtils::setJson), UNDEFINED, STRING, STRING, ARRAY));
+  }
 }

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

@@ -12,6 +12,14 @@
 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.InvalidPathException;
+import com.jayway.jsonpath.JsonPath;
+import com.jayway.jsonpath.Option;
+import com.jayway.jsonpath.PathNotFoundException;
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
@@ -79,6 +87,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.isNull() || json.isMissing()) {
+        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:
@@ -109,4 +169,60 @@ private static ExprValue processJsonNode(JsonNode jsonNode) {
         return LITERAL_NULL;
     }
   }
+
+  /**
+   * Perform an upsert operation against the incoming jsonString value with provided jsonPath and
+   * value.
+   *
+   * @param json jsonObject in String format.
+   * @param path upsert reference in the form of JsonPath.
+   * @param valueToInsert value to be added
+   * @return JsonString after the upsert operation.
+   */
+  public static ExprValue setJson(ExprValue json, ExprValue path, ExprValue valueToInsert) {
+
+    String jsonString = json.stringValue();
+    String jsonPathString = path.stringValue();
+    Object valueToInsertObj = valueToInsert.value();
+    Configuration conf =
+        Configuration.defaultConfiguration().addOptions(Option.SUPPRESS_EXCEPTIONS);
+    try {
+      JsonPath jsonPath = JsonPath.compile(jsonPathString);
+      DocumentContext docContext = JsonPath.using(conf).parse(jsonString);
+      Object readResult = docContext.read(jsonPath);
+      if (readResult == null) {
+        recursiveCreate(docContext, jsonPathString, valueToInsertObj);
+      } else {
+        docContext.set(jsonPathString, valueToInsertObj);
+      }
+      return new ExprStringValue(docContext.jsonString());
+
+    } catch (InvalidPathException e) {
+      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) {
+      final String errorFormat = "JSON object '%s' is not valid. Error details: %s";
+      throw new SemanticCheckException(String.format(errorFormat, json, e.getMessage()), e);
+    }
+  }
+
+  /**
+   * Helper method to handle recursive scenario.
+   *
+   * @param docContext incoming Json in Java object form.
+   * @param path path in String to perform insertion.
+   * @param value value to be inserted with given path.
+   */
+  private static DocumentContext recursiveCreate(
+      DocumentContext docContext, String path, Object value) {
+    final int pos = path.lastIndexOf('.');
+    final String parent = path.substring(0, pos);
+    final String current = path.substring(pos + 1);
+    // Attempt to read the current path as it is, trigger the recursive in case of deep insert.
+    if (docContext.read(parent) == null) {
+      recursiveCreate(docContext, parent, new LinkedHashMap<>());
+    }
+    return docContext.put(parent, current, value);
+  }
 }