Skip to content

Commit

Permalink
[CALCITE-6365] Support for RETURNING clause of JSON_QUERY
Browse files Browse the repository at this point in the history
  • Loading branch information
@dawidwys
dawidwys committed May 10, 2024
1 parent f854ef5 commit 07be4b2
Show file tree
Hide file tree
Showing 12 changed files with 308 additions and 39 deletions.
7 changes: 6 additions & 1 deletion core/src/main/codegen/templates/Parser.jj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6706,7 +6706,7 @@ SqlNode JsonQueryWrapperBehavior() :

SqlCall JsonQueryFunctionCall() :
{
final SqlNode[] args = new SqlNode[5];
final SqlNode[] args = new SqlNode[6];
SqlNode e;
List<SqlNode> commonSyntax;
final Span span;
Expand All @@ -6718,6 +6718,11 @@ SqlCall JsonQueryFunctionCall() :
args[0] = commonSyntax.get(0);
args[1] = commonSyntax.get(1);
}
[
e = JsonReturningClause() {
args[5] = e;
}
]
[
e = JsonQueryWrapperBehavior() <WRAPPER> {
args[2] = e;
Expand Down
30 changes: 29 additions & 1 deletion core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -945,7 +945,7 @@ Builder populate2() {
BuiltInMethod.JSON_EXISTS3.method);
map.put(JSON_VALUE,
new JsonValueImplementor(BuiltInMethod.JSON_VALUE.method));
defineReflective(JSON_QUERY, BuiltInMethod.JSON_QUERY.method);
map.put(JSON_QUERY, new JsonQueryImplementor(BuiltInMethod.JSON_QUERY.method));
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);
Expand Down Expand Up @@ -2903,6 +2903,34 @@ private static class JsonValueImplementor extends MethodImplementor {
}
}

/**
* Implementor for JSON_QUERY function. Passes the jsonize flag depending on the output type.
*/
private static class JsonQueryImplementor extends MethodImplementor {
JsonQueryImplementor(Method method) {
super(method, NullPolicy.ARG0, false);
}

@Override Expression implementSafe(RexToLixTranslator translator,
RexCall call, List<Expression> argValueList) {
final List<Expression> newOperands = new ArrayList<>(argValueList);

final Expression jsonize;
if (SqlTypeUtil.inCharFamily(call.getType())) {
jsonize = TRUE_EXPR;
} else {
jsonize = FALSE_EXPR;
}
newOperands.add(jsonize);

List<Expression> argValueList0 =
EnumUtils.fromInternal(method.getParameterTypes(), newOperands);
final Expression target =
Expressions.new_(method.getDeclaringClass());
return Expressions.call(target, method, argValueList0);
}
}

/** Implementor for binary operators. */
private static class BinaryImplementor extends AbstractRexCallImplementor {
/** Types that can be arguments to comparison operators such as
Expand Down
3 changes: 3 additions & 0 deletions core/src/main/java/org/apache/calcite/runtime/CalciteResource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -993,6 +993,9 @@ ExInstWithCause<CalciteException> failedToAccessField(
@BaseMessage("Illegal error behavior ''{0}'' specified in JSON_VALUE function")
ExInst<CalciteException> illegalErrorBehaviorInJsonQueryFunc(String errorBehavior);

@BaseMessage("EMPTY_OBJECT is illegal for given return type")
ExInst<CalciteException> illegalEmptyObjectInJsonQueryFunc();

@BaseMessage("Null key of JSON object is not allowed")
ExInst<CalciteException> nullKeyOfJsonObjectNotAllowed();

Expand Down
53 changes: 37 additions & 16 deletions core/src/main/java/org/apache/calcite/runtime/JsonFunctions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -310,30 +310,35 @@ public JsonPathContext jsonApiCommonSyntaxWithCache(String input,
}
}

public @Nullable String jsonQuery(String input,
public @Nullable Object jsonQuery(
String input,
String pathSpec,
SqlJsonQueryWrapperBehavior wrapperBehavior,
SqlJsonQueryEmptyOrErrorBehavior emptyBehavior,
SqlJsonQueryEmptyOrErrorBehavior errorBehavior) {
SqlJsonQueryEmptyOrErrorBehavior errorBehavior,
boolean jsonize) {
return jsonQuery(
jsonApiCommonSyntaxWithCache(input, pathSpec),
wrapperBehavior, emptyBehavior, errorBehavior);
wrapperBehavior, emptyBehavior, errorBehavior, jsonize);
}

public @Nullable String jsonQuery(JsonValueContext input,
public @Nullable Object jsonQuery(JsonValueContext input,
String pathSpec,
SqlJsonQueryWrapperBehavior wrapperBehavior,
SqlJsonQueryEmptyOrErrorBehavior emptyBehavior,
SqlJsonQueryEmptyOrErrorBehavior errorBehavior) {
SqlJsonQueryEmptyOrErrorBehavior errorBehavior,
boolean jsonize) {
return jsonQuery(
jsonApiCommonSyntax(input, pathSpec),
wrapperBehavior, emptyBehavior, errorBehavior);
wrapperBehavior, emptyBehavior, errorBehavior, jsonize);
}

public @Nullable String jsonQuery(JsonPathContext context,
public @Nullable Object jsonQuery(
JsonPathContext context,
SqlJsonQueryWrapperBehavior wrapperBehavior,
SqlJsonQueryEmptyOrErrorBehavior emptyBehavior,
SqlJsonQueryEmptyOrErrorBehavior errorBehavior) {
SqlJsonQueryEmptyOrErrorBehavior errorBehavior,
boolean jsonize) {
final Exception exc;
if (context.hasException()) {
exc = context.exc;
Expand Down Expand Up @@ -369,9 +374,9 @@ && isScalarObject(value)) {
case NULL:
return null;
case EMPTY_ARRAY:
return "[]";
return jsonQueryEmptyArray(jsonize);
case EMPTY_OBJECT:
return "{}";
return jsonQueryEmptyObject(jsonize);
default:
throw RESOURCE.illegalEmptyBehaviorInJsonQueryFunc(
emptyBehavior.toString()).ex();
Expand All @@ -381,10 +386,14 @@ && isScalarObject(value)) {
RESOURCE.arrayOrObjectValueRequiredInStrictModeOfJsonQueryFunc(
value.toString()).ex();
} else {
try {
return jsonize(value);
} catch (Exception e) {
exc = e;
if (jsonize) {
try {
return jsonize(value);
} catch (Exception e) {
exc = e;
}
} else {
return value;
}
}
}
Expand All @@ -394,14 +403,26 @@ && isScalarObject(value)) {
case NULL:
return null;
case EMPTY_ARRAY:
return "[]";
return jsonQueryEmptyArray(jsonize);
case EMPTY_OBJECT:
return "{}";
return jsonQueryEmptyObject(jsonize);
default:
throw RESOURCE.illegalErrorBehaviorInJsonQueryFunc(
errorBehavior.toString()).ex();
}
}

private static Object jsonQueryEmptyArray(boolean jsonize) {
return jsonize ? "[]" : Collections.emptyList();
}

private static String jsonQueryEmptyObject(boolean jsonize) {
if (jsonize) {
return "{}";
} else {
throw RESOURCE.illegalEmptyObjectInJsonQueryFunc().ex();
}
}
}

public static String jsonObject(SqlJsonConstructorNullClause nullClause,
Expand Down
95 changes: 87 additions & 8 deletions core/src/main/java/org/apache/calcite/sql/fun/SqlJsonQueryFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
*/
package org.apache.calcite.sql.fun;

import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.sql.SqlBasicCall;
import org.apache.calcite.sql.SqlCall;
import org.apache.calcite.sql.SqlFunction;
import org.apache.calcite.sql.SqlFunctionCategory;
Expand All @@ -24,15 +27,25 @@
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.SqlLiteral;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlOperatorBinding;
import org.apache.calcite.sql.SqlWriter;
import org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.calcite.sql.type.OperandTypes;
import org.apache.calcite.sql.type.ReturnTypes;
import org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.sql.type.SqlTypeTransforms;
import org.apache.calcite.sql.type.SqlTypeUtil;

import com.google.common.collect.ImmutableList;

import org.checkerframework.checker.nullness.qual.Nullable;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;

import static java.util.Objects.requireNonNull;

/**
Expand All @@ -41,15 +54,45 @@
public class SqlJsonQueryFunction extends SqlFunction {
public SqlJsonQueryFunction() {
super("JSON_QUERY", SqlKind.OTHER_FUNCTION,
ReturnTypes.VARCHAR_2000.andThen(SqlTypeTransforms.FORCE_NULLABLE),
ReturnTypes.cascade(
opBinding ->
explicitTypeSpec(opBinding)
.map(t -> deriveExplicitType(opBinding, t))
.orElseGet(() -> getDefaultType(opBinding)),
SqlTypeTransforms.FORCE_NULLABLE),
null,
OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER,
SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.ANY),
OperandTypes.family(
ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER,
SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.ANY, SqlTypeFamily.ANY),
i -> i >= 5),
SqlFunctionCategory.SYSTEM);
}

/** Returns VARCHAR(2000) as default. */
private static RelDataType getDefaultType(SqlOperatorBinding opBinding) {
final RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
final RelDataType baseType = typeFactory.createSqlType(SqlTypeName.VARCHAR, 2000);
return typeFactory.createTypeWithNullability(baseType, true);
}

private static RelDataType deriveExplicitType(SqlOperatorBinding opBinding, RelDataType type) {
if (SqlTypeName.ARRAY == type.getSqlTypeName()) {
RelDataType elementType = Objects.requireNonNull(type.getComponentType());
RelDataType nullableElementType = deriveExplicitType(opBinding, elementType);
return SqlTypeUtil.createArrayType(
opBinding.getTypeFactory(),
nullableElementType,
true);
}
return opBinding.getTypeFactory().createTypeWithNullability(type, true);
}

@Override public @Nullable String getSignatureTemplate(int operandsCount) {
return "{0}({1} {2} {3} WRAPPER {4} ON EMPTY {5} ON ERROR)";
if (operandsCount == 6) {
return "{0}({1} {2} RETURNING {6} {3} WRAPPER {4} ON EMPTY {5} ON ERROR)";
} else {
return "{0}({1} {2} {3} WRAPPER {4} ON EMPTY {5} ON ERROR)";
}
}

@Override public void unparse(SqlWriter writer, SqlCall call, int leftPrec,
Expand All @@ -58,6 +101,10 @@ public SqlJsonQueryFunction() {
call.operand(0).unparse(writer, 0, 0);
writer.sep(",", true);
call.operand(1).unparse(writer, 0, 0);
if (call.operandCount() == 6) {
writer.keyword("RETURNING");
call.operand(5).unparse(writer, 0, 0);
}
final SqlJsonQueryWrapperBehavior wrapperBehavior =
getEnumValue(call.operand(2));
switch (wrapperBehavior) {
Expand All @@ -83,16 +130,32 @@ public SqlJsonQueryFunction() {

@Override public SqlCall createCall(@Nullable SqlLiteral functionQualifier,
SqlParserPos pos, @Nullable SqlNode... operands) {
final List<SqlNode> args = new ArrayList<>();
args.add(Objects.requireNonNull(operands[0]));
args.add(Objects.requireNonNull(operands[1]));

if (operands[2] == null) {
operands[2] = SqlLiteral.createSymbol(SqlJsonQueryWrapperBehavior.WITHOUT_ARRAY, pos);
args.add(SqlLiteral.createSymbol(SqlJsonQueryWrapperBehavior.WITHOUT_ARRAY, pos));
} else {
args.add(operands[2]);
}
if (operands[3] == null) {
operands[3] = SqlLiteral.createSymbol(SqlJsonQueryEmptyOrErrorBehavior.NULL, pos);
args.add(SqlLiteral.createSymbol(SqlJsonQueryEmptyOrErrorBehavior.NULL, pos));
} else {
args.add(operands[3]);
}
if (operands[4] == null) {
operands[4] = SqlLiteral.createSymbol(SqlJsonQueryEmptyOrErrorBehavior.NULL, pos);
args.add(SqlLiteral.createSymbol(SqlJsonQueryEmptyOrErrorBehavior.NULL, pos));
} else {
args.add(operands[4]);
}
return super.createCall(functionQualifier, pos, operands);

if (operands.length >= 6 && operands[5] != null) {
args.add(operands[5]);
}

pos = pos.plusAll(operands);
return new SqlBasicCall(this, args, pos, functionQualifier);
}

private static void unparseEmptyOrErrorBehavior(SqlWriter writer,
Expand All @@ -119,4 +182,20 @@ private static void unparseEmptyOrErrorBehavior(SqlWriter writer,
private static <E extends Enum<E>> E getEnumValue(SqlNode operand) {
return (E) requireNonNull(((SqlLiteral) operand).getValue(), "operand.value");
}

public static boolean hasExplicitTypeSpec(List<SqlNode> operands) {
return operands.size() >= 6;
}

public static List<SqlNode> removeTypeSpecOperands(SqlCall call) {
return call.getOperandList().subList(0, 5);
}

/** Returns the optional explicit returning type specification. * */
private static Optional<RelDataType> explicitTypeSpec(SqlOperatorBinding opBinding) {
if (opBinding.getOperandCount() >= 6) {
return Optional.of(opBinding.getOperandType(5));
}
return Optional.empty();
}
}
Loading

0 comments on commit 07be4b2

Please sign in to comment.