Skip to content

Commit e6b80e5

Browse files
yuancuyuancu
committed
Make getValiadtor thread-safe & improve type checking for JsonSet
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 229ddb7 commit e6b80e5

4 files changed

Lines changed: 82 additions & 33 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java‎

Lines changed: 35 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ public class CalcitePlanContext {
9090
@Setter private static SqlOperatorTableProvider operatorTableProvider;
9191

9292
/** Cached SqlValidator instance (lazy initialized). */
93-
private SqlValidator validator;
93+
private volatile SqlValidator validator;
9494

9595
private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType queryType) {
9696
this.config = config;
@@ -128,36 +128,41 @@ private CalcitePlanContext(CalcitePlanContext parent) {
128128
*/
129129
public SqlValidator getValidator() {
130130
if (validator == null) {
131-
final CalciteServerStatement statement;
132-
try {
133-
statement = connection.createStatement().unwrap(CalciteServerStatement.class);
134-
} catch (SQLException e) {
135-
throw new RuntimeException(e);
131+
synchronized (this) {
132+
// Double-Checked Locking for thread-safety
133+
if (validator == null) {
134+
final CalciteServerStatement statement;
135+
try {
136+
statement = connection.createStatement().unwrap(CalciteServerStatement.class);
137+
} catch (SQLException e) {
138+
throw new RuntimeException(e);
139+
}
140+
if (operatorTableProvider == null) {
141+
throw new IllegalStateException(
142+
"SqlOperatorTableProvider must be set before creating CalcitePlanContext");
143+
}
144+
SqlValidator.Config validatorConfig =
145+
SqlValidator.Config.DEFAULT
146+
.withTypeCoercionRules(PplTypeCoercionRule.instance())
147+
.withTypeCoercionFactory(PplTypeCoercion::create)
148+
// Use lenient conformance for PPL compatibility
149+
.withConformance(OpenSearchSparkSqlDialect.DEFAULT.getConformance())
150+
// Use Spark SQL's NULL collation (NULLs sorted LOW/FIRST)
151+
.withDefaultNullCollation(NullCollation.LOW)
152+
// This ensures that coerced arguments are replaced with cast version in sql
153+
// select list because coercion is performed during select list expansion during
154+
// sql validation. Affects 4356.yml
155+
// See SqlValidatorImpl#validateSelectList and AggConverter#translateAgg
156+
.withIdentifierExpansion(true);
157+
validator =
158+
PplValidator.create(
159+
statement,
160+
config,
161+
operatorTableProvider.getOperatorTable(),
162+
TYPE_FACTORY,
163+
validatorConfig);
164+
}
136165
}
137-
if (operatorTableProvider == null) {
138-
throw new IllegalStateException(
139-
"SqlOperatorTableProvider must be set before creating CalcitePlanContext");
140-
}
141-
SqlValidator.Config validatorConfig =
142-
SqlValidator.Config.DEFAULT
143-
.withTypeCoercionRules(PplTypeCoercionRule.instance())
144-
.withTypeCoercionFactory(PplTypeCoercion::create)
145-
// Use lenient conformance for PPL compatibility
146-
.withConformance(OpenSearchSparkSqlDialect.DEFAULT.getConformance())
147-
// Use Spark SQL's NULL collation (NULLs sorted LOW/FIRST)
148-
.withDefaultNullCollation(NullCollation.LOW)
149-
// This ensures that coerced arguments are replaced with cast version in sql select
150-
// list because coercion is performed during select list expansion during sql
151-
// validation. Affects 4356.yml
152-
// See SqlValidatorImpl#validateSelectList and AggConverter#translateAgg
153-
.withIdentifierExpansion(true);
154-
validator =
155-
PplValidator.create(
156-
statement,
157-
config,
158-
operatorTableProvider.getOperatorTable(),
159-
TYPE_FACTORY,
160-
validatorConfig);
161166
}
162167
return validator;
163168
}

‎core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/TransformFunctionImpl.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,9 @@ public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFail
135135
}
136136
return IntStream.range(0, 2)
137137
.allMatch(
138-
i -> checkSingleOperandType(callBinding, callBinding.operand(i), i, false));
138+
i ->
139+
checkSingleOperandType(
140+
callBinding, callBinding.operand(i), i, throwOnFailure));
139141
}
140142

141143
@Override

‎core/src/main/java/org/opensearch/sql/expression/function/UDFOperandMetadata.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,9 @@ public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFail
8989

9090
@Override
9191
public SqlOperandCountRange getOperandCountRange() {
92+
if (allowedParamTypes == null || allowedParamTypes.isEmpty()) {
93+
return SqlOperandCountRanges.between(0, 0);
94+
}
9295
int max = Integer.MIN_VALUE;
9396
int min = Integer.MAX_VALUE;
9497
for (List<ExprType> paramTypes : allowedParamTypes) {

‎core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonSetFunctionImpl.java‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package org.opensearch.sql.expression.function.jsonUDF;
77

8+
import static org.apache.calcite.util.Static.RESOURCE;
89
import static org.opensearch.sql.calcite.utils.PPLReturnTypes.STRING_FORCE_NULLABLE;
910
import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.*;
1011

@@ -18,11 +19,18 @@
1819
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
1920
import org.apache.calcite.linq4j.tree.Expression;
2021
import org.apache.calcite.linq4j.tree.Types;
22+
import org.apache.calcite.rel.type.RelDataType;
2123
import org.apache.calcite.rex.RexCall;
2224
import org.apache.calcite.runtime.JsonFunctions;
2325
import org.apache.calcite.schema.impl.ScalarFunctionImpl;
24-
import org.apache.calcite.sql.type.OperandTypes;
26+
import org.apache.calcite.sql.SqlCallBinding;
27+
import org.apache.calcite.sql.SqlOperandCountRange;
28+
import org.apache.calcite.sql.SqlOperator;
29+
import org.apache.calcite.sql.fun.SqlJsonModifyFunction;
30+
import org.apache.calcite.sql.type.SqlOperandCountRanges;
31+
import org.apache.calcite.sql.type.SqlOperandTypeChecker;
2532
import org.apache.calcite.sql.type.SqlReturnTypeInference;
33+
import org.opensearch.sql.calcite.utils.OpenSearchTypeUtil;
2634
import org.opensearch.sql.expression.function.ImplementorUDF;
2735
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2836

@@ -38,7 +46,38 @@ public SqlReturnTypeInference getReturnTypeInference() {
3846

3947
@Override
4048
public UDFOperandMetadata getOperandMetadata() {
41-
return UDFOperandMetadata.wrap(OperandTypes.ONE_OR_MORE);
49+
return UDFOperandMetadata.wrap(
50+
new SqlOperandTypeChecker() {
51+
/**
52+
* Copied from {@link SqlJsonModifyFunction#checkOperandTypes(SqlCallBinding, boolean)}
53+
* (Calcite 1.41)
54+
*/
55+
@Override
56+
public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure) {
57+
final int count = callBinding.getOperandCount();
58+
for (int i = 1; i < count; i += 2) {
59+
RelDataType nameType = callBinding.getOperandType(i);
60+
if (!OpenSearchTypeUtil.isCharacter(nameType)) {
61+
if (throwOnFailure) {
62+
throw callBinding.newError(RESOURCE.expectedCharacter());
63+
}
64+
return false;
65+
}
66+
}
67+
return true;
68+
}
69+
70+
@Override
71+
public SqlOperandCountRange getOperandCountRange() {
72+
return SqlOperandCountRanges.from(3);
73+
}
74+
75+
@Override
76+
public String getAllowedSignatures(SqlOperator op, String opName) {
77+
return "(json_string: STRING, path1: STRING, value1: ANY, path2: STRING, value2: ANY"
78+
+ " ...)";
79+
}
80+
});
4281
}
4382

4483
public static class JsonSetImplementor implements NotNullImplementor {

0 commit comments

Comments
 (0)