[CALCITE-2884] Implement JSON_INSERT, JSON_REPLACE, JSON_SET functions

Close apache#3020
JiajunBernoulli · Jan 5, 2023 · 787aa9e · 787aa9e
1 parent 9054682
commit 787aa9e
Show file tree

Hide file tree

Showing 15 changed files with 542 additions and 20 deletions.
diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
@@ -134,10 +134,13 @@
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.FROM_BASE64;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ILIKE;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_DEPTH;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_INSERT;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_KEYS;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_LENGTH;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_PRETTY;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_REMOVE;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_REPLACE;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_SET;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_STORAGE_SIZE;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.JSON_TYPE;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.LEFT;
@@ -644,11 +647,14 @@ Builder populate2() {
       defineMethod(JSON_QUERY, BuiltInMethod.JSON_QUERY.method, NullPolicy.ARG0);
       defineMethod(JSON_TYPE, BuiltInMethod.JSON_TYPE.method, NullPolicy.ARG0);
       defineMethod(JSON_DEPTH, BuiltInMethod.JSON_DEPTH.method, NullPolicy.ARG0);
+      defineMethod(JSON_INSERT, BuiltInMethod.JSON_INSERT.method, NullPolicy.ARG0);
       defineMethod(JSON_KEYS, BuiltInMethod.JSON_KEYS.method, NullPolicy.ARG0);
       defineMethod(JSON_PRETTY, BuiltInMethod.JSON_PRETTY.method, NullPolicy.ARG0);
       defineMethod(JSON_LENGTH, BuiltInMethod.JSON_LENGTH.method, NullPolicy.ARG0);
       defineMethod(JSON_REMOVE, BuiltInMethod.JSON_REMOVE.method, NullPolicy.ARG0);
       defineMethod(JSON_STORAGE_SIZE, BuiltInMethod.JSON_STORAGE_SIZE.method, NullPolicy.ARG0);
+      defineMethod(JSON_REPLACE, BuiltInMethod.JSON_REPLACE.method, NullPolicy.ARG0);
+      defineMethod(JSON_SET, BuiltInMethod.JSON_SET.method, NullPolicy.ARG0);
       defineMethod(JSON_OBJECT, BuiltInMethod.JSON_OBJECT.method, NullPolicy.NONE);
       defineMethod(JSON_ARRAY, BuiltInMethod.JSON_ARRAY.method, NullPolicy.NONE);
       aggMap.put(JSON_OBJECTAGG.with(SqlJsonConstructorNullClause.ABSENT_ON_NULL),

diff --git a/core/src/main/java/org/apache/calcite/runtime/CalciteResource.java b/core/src/main/java/org/apache/calcite/runtime/CalciteResource.java
@@ -950,30 +950,39 @@ ExInstWithCause<CalciteException> failedToAccessField(
   @BaseMessage("While executing SQL [{0}] on JDBC sub-schema")
   ExInst<RuntimeException> exceptionWhilePerformingQueryOnJdbcSubSchema(String sql);
 
-  @BaseMessage("Not a valid input for JSON_TYPE: ''{0}''")
+  @BaseMessage("Invalid input for JSON_TYPE: ''{0}''")
   ExInst<CalciteException> invalidInputForJsonType(String value);
 
-  @BaseMessage("Not a valid input for JSON_DEPTH: ''{0}''")
+  @BaseMessage("Invalid input for JSON_DEPTH: ''{0}''")
   ExInst<CalciteException> invalidInputForJsonDepth(String value);
 
   @BaseMessage("Cannot serialize object to JSON: ''{0}''")
   ExInst<CalciteException> exceptionWhileSerializingToJson(String value);
 
-  @BaseMessage("Not a valid input for JSON_LENGTH: ''{0}''")
+  @BaseMessage("Invalid input for JSON_LENGTH: ''{0}''")
   ExInst<CalciteException> invalidInputForJsonLength(String value);
 
-  @BaseMessage("Not a valid input for JSON_KEYS: ''{0}''")
+  @BaseMessage("Invalid input for JSON_KEYS: ''{0}''")
   ExInst<CalciteException> invalidInputForJsonKeys(String value);
 
   @BaseMessage("Invalid input for JSON_REMOVE: document: ''{0}'', jsonpath expressions: ''{1}''")
   ExInst<CalciteException> invalidInputForJsonRemove(String value, String pathSpecs);
 
-  @BaseMessage("Not a valid input for JSON_STORAGE_SIZE: ''{0}''")
+  @BaseMessage("Invalid input for JSON_STORAGE_SIZE: ''{0}''")
   ExInst<CalciteException> invalidInputForJsonStorageSize(String value);
 
-  @BaseMessage("Not a valid input for REGEXP_REPLACE: ''{0}''")
+  @BaseMessage("Invalid input for REGEXP_REPLACE: ''{0}''")
   ExInst<CalciteException> invalidInputForRegexpReplace(String value);
 
+  @BaseMessage("Invalid input for JSON_INSERT: jsonDoc: ''{0}'', kvs: ''{1}''")
+  ExInst<CalciteException> invalidInputForJsonInsert(String jsonDoc, String kvs);
+
+  @BaseMessage("Invalid input for JSON_REPLACE: jsonDoc: ''{0}'', kvs: ''{1}''")
+  ExInst<CalciteException> invalidInputForJsonReplace(String jsonDoc, String kvs);
+
+  @BaseMessage("Invalid input for JSON_SET: jsonDoc: ''{0}'', kvs: ''{1}''")
+  ExInst<CalciteException> invalidInputForJsonSet(String jsonDoc, String kvs);
+
   @BaseMessage("Illegal xslt specified : ''{0}''")
   ExInst<CalciteException> illegalXslt(String xslt);
 

diff --git a/core/src/main/java/org/apache/calcite/runtime/JsonFunctions.java b/core/src/main/java/org/apache/calcite/runtime/JsonFunctions.java
@@ -77,6 +77,8 @@ public class JsonFunctions {
       new DefaultPrettyPrinter().withObjectIndenter(
           DefaultIndenter.SYSTEM_LINEFEED_INSTANCE.withLinefeed("\n"));
 
+  private static final String JSON_ROOT_PATH = "$";
+
   private JsonFunctions() {
   }
 
@@ -658,6 +660,130 @@ public static Integer jsonStorageSize(JsonValueContext input) {
     }
   }
 
+  private static String jsonModify(JsonValueContext jsonDoc, JsonModifyMode type, Object... kvs) {
+    int step = type == JsonModifyMode.REMOVE ? 1 : 2;
+    assert kvs.length % step == 0;
+    String result = null;
+    DocumentContext ctx = JsonPath.parse(jsonDoc.obj(),
+        Configuration
+            .builder()
+            .options(Option.SUPPRESS_EXCEPTIONS)
+            .jsonProvider(JSON_PATH_JSON_PROVIDER)
+            .mappingProvider(JSON_PATH_MAPPING_PROVIDER)
+            .build());
+
+    for (int i = 0; i < kvs.length; i += step) {
+      String k = (String) kvs[i];
+      Object v = kvs[i + step - 1];
+      if (!(JsonPath.isPathDefinite(k) && k.contains(JSON_ROOT_PATH))) {
+        throw RESOURCE.validationError(k).ex();
+      }
+      switch (type) {
+      case REPLACE:
+        if (k.equals(JSON_ROOT_PATH)) {
+          result = jsonize(v);
+        } else {
+          if (ctx.read(k) != null) {
+            ctx.set(k, v);
+          }
+        }
+        break;
+      case INSERT:
+        if (!k.equals(JSON_ROOT_PATH) && ctx.read(k) == null) {
+          insertToJson(ctx, k, v);
+        }
+        break;
+      case SET:
+        if (k.equals(JSON_ROOT_PATH)) {
+          result = jsonize(v);
+        } else {
+          if (ctx.read(k) != null) {
+            ctx.set(k, v);
+          } else {
+            insertToJson(ctx, k, v);
+          }
+        }
+        break;
+      case REMOVE:
+        if (ctx.read(k) != null) {
+          ctx.delete(k);
+        }
+        break;
+      }
+    }
+
+    return result == null || result.isEmpty() ? ctx.jsonString() : result;
+  }
+
+  private static void insertToJson(DocumentContext ctx, String path, Object value) {
+    final String parentPath;
+    final String key;
+
+    //The following paragraph of logic is mainly used to obtain the parent node path of path and
+    //the key value that should be inserted when the preant Path is map.
+    //eg: path is $.a ,parentPath is $
+    //eg: path is $.a[1] ,parentPath is $.a
+    Integer dotIndex = path.lastIndexOf(".");
+    Integer leftBracketIndex = path.lastIndexOf("[");
+    if (dotIndex.equals(-1) && leftBracketIndex.equals(-1)) {
+      parentPath = path;
+      key = path;
+    } else if (!dotIndex.equals(-1) && leftBracketIndex.equals(-1)) {
+      parentPath = path.substring(0, dotIndex);
+      key = path.substring(dotIndex + 1);
+    } else if (dotIndex.equals(-1) && !leftBracketIndex.equals(-1)) {
+      parentPath = path.substring(0, leftBracketIndex);
+      key = path.substring(leftBracketIndex + 1);
+    } else {
+      int position = dotIndex > leftBracketIndex ? dotIndex : leftBracketIndex;
+      parentPath = path.substring(0, position);
+      key = path.substring(position);
+    }
+
+    Object obj = ctx.read(parentPath);
+    if (obj instanceof Map) {
+      ctx.put(parentPath, key, value.toString());
+    } else if (obj instanceof Collection) {
+      ctx.add(parentPath, value);
+    }
+  }
+
+  public static String jsonReplace(String jsonDoc, Object... kvs) {
+    return jsonReplace(jsonValueExpression(jsonDoc), kvs);
+  }
+
+  public static String jsonReplace(JsonValueContext input, Object... kvs) {
+    try {
+      return jsonModify(input, JsonModifyMode.REPLACE, kvs);
+    } catch (Exception ex) {
+      throw RESOURCE.invalidInputForJsonReplace(input.toString(), jsonize(kvs)).ex();
+    }
+  }
+
+  public static String jsonInsert(String jsonDoc, Object... kvs) {
+    return jsonInsert(jsonValueExpression(jsonDoc), kvs);
+  }
+
+  public static String jsonInsert(JsonValueContext input, Object... kvs) {
+    try {
+      return jsonModify(input, JsonModifyMode.INSERT, kvs);
+    } catch (Exception ex) {
+      throw RESOURCE.invalidInputForJsonInsert(input.toString(), jsonize(kvs)).ex();
+    }
+  }
+
+  public static String jsonSet(String jsonDoc, Object... kvs) {
+    return jsonSet(jsonValueExpression(jsonDoc), kvs);
+  }
+
+  public static String jsonSet(JsonValueContext input, Object... kvs) {
+    try {
+      return jsonModify(input, JsonModifyMode.SET, kvs);
+    } catch (Exception ex) {
+      throw RESOURCE.invalidInputForJsonSet(input.toString(), jsonize(kvs)).ex();
+    }
+  }
+
   public static boolean isJsonValue(String input) {
     try {
       dejsonize(input);
@@ -818,4 +944,14 @@ public enum PathMode {
     UNKNOWN,
     NONE
   }
+
+  /**
+   * Used in the JsonModify function.
+   */
+  public enum JsonModifyMode {
+    REPLACE,
+    INSERT,
+    SET,
+    REMOVE
+  }
 }
diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlJsonModifyFunction.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlJsonModifyFunction.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.sql.fun;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.SqlCall;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperandCountRange;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SqlOperandCountRanges;
+import org.apache.calcite.sql.type.SqlOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlTypeTransforms;
+import org.apache.calcite.sql.type.SqlTypeUtil;
+import org.apache.calcite.sql.validate.SqlValidator;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.util.Locale;
+
+import static org.apache.calcite.util.Static.RESOURCE;
+
+/**
+ * Definition of the <code>JSON_INSERT</code>, <code>JSON_REPLACE</code>
+ * and <code>JSON_SET</code> JSON Modify functions.
+ */
+public class SqlJsonModifyFunction extends SqlFunction {
+  public SqlJsonModifyFunction(String name) {
+    super(name,
+        SqlKind.OTHER_FUNCTION,
+        ReturnTypes.cascade(ReturnTypes.VARCHAR_2000,
+            SqlTypeTransforms.FORCE_NULLABLE),
+        null,
+        OperandTypes.ANY,
+        SqlFunctionCategory.SYSTEM);
+  }
+
+  @Override public SqlOperandCountRange getOperandCountRange() {
+    return SqlOperandCountRanges.from(3);
+  }
+
+  @Override protected void checkOperandCount(SqlValidator validator,
+      @Nullable SqlOperandTypeChecker argType, SqlCall call) {
+    assert (call.operandCount() >= 3) && (call.operandCount() % 2 == 1);
+  }
+
+  @Override public boolean checkOperandTypes(SqlCallBinding callBinding,
+      boolean throwOnFailure) {
+    final int count = callBinding.getOperandCount();
+    for (int i = 1; i < count; i += 2) {
+      RelDataType nameType = callBinding.getOperandType(i);
+      if (!SqlTypeUtil.isCharacter(nameType)) {
+        if (throwOnFailure) {
+          throw callBinding.newError(RESOURCE.expectedCharacter());
+        }
+        return false;
+      }
+      if (nameType.isNullable()) {
+        if (throwOnFailure) {
+          throw callBinding.newError(
+              RESOURCE.argumentMustNotBeNull(
+                  callBinding.operand(i).toString()));
+        }
+        return false;
+      }
+    }
+    return true;
+  }
+
+  @Override public String getSignatureTemplate(int operandsCount) {
+    assert operandsCount % 2 == 1;
+    StringBuilder sb = new StringBuilder();
+    sb.append("{0}(");
+    for (int i = 1; i < operandsCount; i++) {
+      sb.append(String.format(Locale.ROOT, "{%d} ", i + 1));
+    }
+    sb.append("{1})");
+    return sb.toString();
+  }
+}
diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
@@ -305,6 +305,15 @@ private SqlLibraryOperators() {
   @LibraryOperator(libraries = {MYSQL})
   public static final SqlFunction JSON_STORAGE_SIZE = new SqlJsonStorageSizeFunction();
 
+  @LibraryOperator(libraries = {MYSQL})
+  public static final SqlFunction JSON_INSERT = new SqlJsonModifyFunction("JSON_INSERT");
+
+  @LibraryOperator(libraries = {MYSQL})
+  public static final SqlFunction JSON_REPLACE = new SqlJsonModifyFunction("JSON_REPLACE");
+
+  @LibraryOperator(libraries = {MYSQL})
+  public static final SqlFunction JSON_SET = new SqlJsonModifyFunction("JSON_SET");
+
   @LibraryOperator(libraries = {MYSQL, ORACLE})
   public static final SqlFunction REGEXP_REPLACE = new SqlRegexpReplaceFunction();
 

diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlStdOperatorTable.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlStdOperatorTable.java
@@ -1415,6 +1415,15 @@ public class SqlStdOperatorTable extends ReflectiveSqlOperatorTable {
   @Deprecated // to be removed before 2.0
   public static final SqlFunction JSON_STORAGE_SIZE = SqlLibraryOperators.JSON_STORAGE_SIZE;
 
+  @Deprecated // to be removed before 2.0
+  public static final SqlFunction JSON_INSERT = SqlLibraryOperators.JSON_INSERT;
+
+  @Deprecated // to be removed before 2.0
+  public static final SqlFunction JSON_REPLACE = SqlLibraryOperators.JSON_REPLACE;
+
+  @Deprecated // to be removed before 2.0
+  public static final SqlFunction JSON_SET = SqlLibraryOperators.JSON_SET;
+
   public static final SqlJsonArrayAggAggFunction JSON_ARRAYAGG =
       new SqlJsonArrayAggAggFunction(SqlKind.JSON_ARRAYAGG,
           SqlJsonConstructorNullClause.ABSENT_ON_NULL);

diff --git a/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java b/core/src/main/java/org/apache/calcite/util/BuiltInMethod.java
@@ -379,10 +379,13 @@ public enum BuiltInMethod {
   JSON_TYPE(JsonFunctions.class, "jsonType", String.class),
   JSON_DEPTH(JsonFunctions.class, "jsonDepth", String.class),
   JSON_KEYS(JsonFunctions.class, "jsonKeys", String.class),
+  JSON_INSERT(JsonFunctions.class, "jsonInsert", String.class, Object.class),
   JSON_PRETTY(JsonFunctions.class, "jsonPretty", String.class),
   JSON_LENGTH(JsonFunctions.class, "jsonLength", String.class),
+  JSON_REPLACE(JsonFunctions.class, "jsonReplace", String.class, Object.class),
   JSON_REMOVE(JsonFunctions.class, "jsonRemove", String.class),
   JSON_STORAGE_SIZE(JsonFunctions.class, "jsonStorageSize", String.class),
+  JSON_SET(JsonFunctions.class, "jsonSet", String.class, Object.class),
   JSON_OBJECTAGG_ADD(JsonFunctions.class, "jsonObjectAggAdd", Map.class,
       String.class, Object.class, SqlJsonConstructorNullClause.class),
   JSON_ARRAY(JsonFunctions.class, "jsonArray",

diff --git a/core/src/main/resources/org/apache/calcite/runtime/CalciteResource.properties b/core/src/main/resources/org/apache/calcite/runtime/CalciteResource.properties
@@ -310,14 +310,17 @@ NullKeyOfJsonObjectNotAllowed=Null key of JSON object is not allowed
 QueryExecutionTimeoutReached=Timeout of ''{0}'' ms for query execution is reached. Query execution started at ''{1}''
 AmbiguousSortOrderInJsonArrayAggFunc=Including both WITHIN GROUP(...) and inside ORDER BY in a single JSON_ARRAYAGG call is not allowed
 ExceptionWhilePerformingQueryOnJdbcSubSchema = While executing SQL [{0}] on JDBC sub-schema
-InvalidInputForJsonType=Not a valid input for JSON_TYPE: ''{0}''
-InvalidInputForJsonDepth=Not a valid input for JSON_DEPTH: ''{0}''
+InvalidInputForJsonType=Invalid input for JSON_TYPE: ''{0}''
+InvalidInputForJsonDepth=Invalid input for JSON_DEPTH: ''{0}''
 ExceptionWhileSerializingToJson=Cannot serialize object to JSON: ''{0}''
-InvalidInputForJsonLength=Not a valid input for JSON_LENGTH: ''{0}''
-InvalidInputForJsonKeys=Not a valid input for JSON_KEYS: ''{0}''
+InvalidInputForJsonLength=Invalid input for JSON_LENGTH: ''{0}''
+InvalidInputForJsonKeys=Invalid input for JSON_KEYS: ''{0}''
 InvalidInputForJsonRemove=Invalid input for JSON_REMOVE: document: ''{0}'', jsonpath expressions: ''{1}''
-InvalidInputForJsonStorageSize=Not a valid input for JSON_STORAGE_SIZE: ''{0}''
-InvalidInputForRegexpReplace=Not a valid input for REGEXP_REPLACE: ''{0}''
+InvalidInputForJsonStorageSize=Invalid input for JSON_STORAGE_SIZE: ''{0}''
+InvalidInputForRegexpReplace=Invalid input for REGEXP_REPLACE: ''{0}''
+InvalidInputForJsonInsert=Invalid input for JSON_INSERT: jsonDoc: ''{0}'', kvs: ''{1}''
+InvalidInputForJsonReplace=Invalid input for JSON_REPLACE: jsonDoc: ''{0}'', kvs: ''{1}''
+InvalidInputForJsonSet=Invalid input for JSON_SET: jsonDoc: ''{0}'', kvs: ''{1}''
 IllegalXslt=Illegal xslt specified : ''{0}''
 InvalidInputForXmlTransform=Invalid input for XMLTRANSFORM xml: ''{0}''
 InvalidInputForExtractValue=Invalid input for EXTRACTVALUE: xml: ''{0}'', xpath expression: ''{1}''