Skip to content

Commit eed5aa9

Browse files
xinyualxinyual
committed
align with splunk
Signed-off-by: xinyual <xinyual@amazon.com>
1 parent 7690d78 commit eed5aa9

14 files changed

Lines changed: 277 additions & 310 deletions

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

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

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

8+
import static org.apache.calcite.sql.SqlJsonConstructorNullClause.NULL_ON_NULL;
89
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
910
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.getLegacyTypeName;
1011
import static org.opensearch.sql.expression.function.BuiltinFunctionName.*;
@@ -16,6 +17,7 @@
1617
import java.util.List;
1718
import java.util.Map;
1819
import java.util.Optional;
20+
import java.util.stream.Stream;
1921
import org.apache.calcite.rel.type.RelDataType;
2022
import org.apache.calcite.rex.RexBuilder;
2123
import org.apache.calcite.rex.RexNode;
@@ -82,6 +84,23 @@ default List<RelDataType> getParams() {
8284
}
8385
}
8486

87+
public interface FunctionImpAny extends FunctionImp {
88+
RexNode resolve(RexBuilder builder, List<RexNode> additonalArgs, RexNode... args);
89+
90+
@Override
91+
default RexNode resolve(RexBuilder builder, RexNode... args) {
92+
if (args.length != 2) {
93+
throw new IllegalArgumentException("This function requires exactly 2 arguments");
94+
}
95+
return resolve(builder, args[0], args[1]);
96+
}
97+
98+
@Override
99+
default List<RelDataType> getParams() {
100+
return null;
101+
}
102+
}
103+
85104
/** The singleton instance. */
86105
public static final PPLFuncImpTable INSTANCE;
87106

@@ -305,14 +324,26 @@ void populate() {
305324
registerOperator(ARRAY, PPLBuiltinOperators.ARRAY);
306325

307326
// Register Json function
327+
register(
328+
JSON_ARRAY,
329+
((builder, args) ->
330+
builder.makeCall(
331+
SqlStdOperatorTable.JSON_ARRAY,
332+
Stream.concat(Stream.of(builder.makeFlag(NULL_ON_NULL)), Arrays.stream(args))
333+
.toArray(RexNode[]::new))));
334+
register(
335+
JSON_OBJECT,
336+
((builder, args) ->
337+
builder.makeCall(
338+
SqlStdOperatorTable.JSON_OBJECT,
339+
Stream.concat(Stream.of(builder.makeFlag(NULL_ON_NULL)), Arrays.stream(args))
340+
.toArray(RexNode[]::new))));
308341
registerOperator(JSON, PPLBuiltinOperators.JSON);
309-
registerOperator(JSON_OBJECT, PPLBuiltinOperators.JSON_OBJECT);
310-
registerOperator(JSON_ARRAY, PPLBuiltinOperators.JSON_ARRAY);
311342
registerOperator(TO_JSON_STRING, PPLBuiltinOperators.TO_JSON_STRING);
312343
registerOperator(JSON_ARRAY_LENGTH, PPLBuiltinOperators.JSON_ARRAY_LENGTH);
313344
registerOperator(JSON_EXTRACT, PPLBuiltinOperators.JSON_EXTRACT);
314345
registerOperator(JSON_KEYS, PPLBuiltinOperators.JSON_KEYS);
315-
registerOperator(JSON_VALID, PPLBuiltinOperators.JSON_VALID);
346+
registerOperator(JSON_VALID, SqlStdOperatorTable.IS_JSON_VALUE);
316347
registerOperator(JSON_SET, PPLBuiltinOperators.JSON_SET);
317348
registerOperator(JSON_DELETE, PPLBuiltinOperators.JSON_DELETE);
318349
registerOperator(JSON_APPEND, PPLBuiltinOperators.JSON_APPEND);

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import com.jayway.jsonpath.PathNotFoundException;
1818
import com.jayway.jsonpath.spi.json.JacksonJsonNodeJsonProvider;
1919
import com.jayway.jsonpath.spi.mapper.JacksonMappingProvider;
20+
import java.util.ArrayList;
2021
import java.util.Arrays;
2122
import java.util.List;
2223
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
@@ -26,6 +27,7 @@
2627
import org.apache.calcite.linq4j.tree.Expression;
2728
import org.apache.calcite.linq4j.tree.Types;
2829
import org.apache.calcite.rex.RexCall;
30+
import org.apache.calcite.runtime.JsonFunctions;
2931
import org.apache.calcite.schema.impl.ScalarFunctionImpl;
3032
import org.apache.calcite.sql.type.SqlReturnTypeInference;
3133
import org.opensearch.sql.expression.function.ImplementorUDF;
@@ -59,7 +61,19 @@ public static Object eval(Object... args) throws JsonProcessingException {
5961
throw new RuntimeException(
6062
"Json append function needs corresponding path and values, but current get: " + keys);
6163
}
62-
return jsonAppendIfArray(jsonStr, keys, false);
64+
JsonNode root = verifyInput(args[0]);
65+
List<Object> expands = new ArrayList<>();
66+
for (int i = 0; i < keys.size(); i += 2) {
67+
List<String> expandedPaths = expandJsonPath(root, convertToJsonPath(keys.get(i).toString()));
68+
for (String expandedPath : expandedPaths) {
69+
expands.add(
70+
expandedPath
71+
+ ".meaninglessKey"); // We add meaningless Key since calcite json_insert can only
72+
// insert when the path point to null
73+
expands.add(keys.get(i + 1));
74+
}
75+
}
76+
return JsonFunctions.jsonInsert(jsonStr, expands.toArray());
6377
}
6478

6579
public static String jsonAppendIfArray(Object json, List<Object> pathValueMap, boolean isExtend) {

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

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,19 @@
55

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

8-
import static org.apache.calcite.sql.type.SqlTypeUtil.createArrayType;
8+
import static org.opensearch.sql.calcite.utils.PPLReturnTypes.STRING_FORCE_NULLABLE;
99

10-
import java.util.Arrays;
1110
import java.util.List;
1211
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
1312
import org.apache.calcite.adapter.enumerable.NullPolicy;
1413
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
1514
import org.apache.calcite.linq4j.tree.Expression;
1615
import org.apache.calcite.linq4j.tree.Expressions;
1716
import org.apache.calcite.linq4j.tree.Types;
18-
import org.apache.calcite.rel.type.RelDataTypeFactory;
1917
import org.apache.calcite.rex.RexCall;
18+
import org.apache.calcite.runtime.JsonFunctions;
19+
import org.apache.calcite.sql.SqlJsonConstructorNullClause;
2020
import org.apache.calcite.sql.type.SqlReturnTypeInference;
21-
import org.apache.calcite.sql.type.SqlTypeName;
2221
import org.opensearch.sql.expression.function.ImplementorUDF;
2322

2423
/**
@@ -37,13 +36,7 @@ public JsonArrayFunctionImpl() {
3736

3837
@Override
3938
public SqlReturnTypeInference getReturnTypeInference() {
40-
return sqlOperatorBinding -> {
41-
RelDataTypeFactory typeFactory = sqlOperatorBinding.getTypeFactory();
42-
return createArrayType(
43-
typeFactory,
44-
typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.ANY), true),
45-
true);
46-
};
39+
return STRING_FORCE_NULLABLE;
4740
}
4841

4942
public static class JsonArrayImplementor implements NotNullImplementor {
@@ -57,6 +50,6 @@ public Expression implement(
5750
}
5851

5952
public static Object eval(Object... args) {
60-
return Arrays.asList(args);
53+
return JsonFunctions.jsonArray(SqlJsonConstructorNullClause.NULL_ON_NULL, args);
6154
}
6255
}

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

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,11 @@
55

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

8+
import static org.apache.calcite.runtime.JsonFunctions.jsonRemove;
89
import static org.opensearch.sql.calcite.utils.PPLReturnTypes.STRING_FORCE_NULLABLE;
910
import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.*;
1011

1112
import com.fasterxml.jackson.core.JsonProcessingException;
12-
import com.fasterxml.jackson.databind.JsonNode;
13-
import com.jayway.jsonpath.DocumentContext;
14-
import com.jayway.jsonpath.JsonPath;
15-
import com.jayway.jsonpath.PathNotFoundException;
1613
import java.util.Arrays;
1714
import java.util.List;
1815
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
@@ -50,24 +47,11 @@ public Expression implement(
5047

5148
public static Object eval(Object... args) throws JsonProcessingException {
5249
List<Object> jsonPaths = Arrays.asList(args).subList(1, args.length);
53-
JsonNode root = verifyInput(args[0]);
54-
DocumentContext ctx;
55-
if (args[0] instanceof String) {
56-
ctx = JsonPath.parse(args[0].toString());
57-
} else {
58-
ctx = JsonPath.parse(args[0]);
59-
}
60-
for (Object originalPath : jsonPaths) {
61-
String jsonPath = convertToJsonPath(originalPath.toString());
62-
try {
63-
Object matches = ctx.read(jsonPath); // verify whether it's a valid path
64-
} catch (PathNotFoundException e) {
65-
continue;
66-
}
67-
// Resolve path tokens
68-
PathTokenizer tokenizer = new PathTokenizer(jsonPath);
69-
root = deletePath(root, tokenizer);
70-
}
71-
return root.toString();
50+
String[] pathSpecs =
51+
jsonPaths.stream()
52+
.map(Object::toString)
53+
.map(JsonUtils::convertToJsonPath)
54+
.toArray(String[]::new);
55+
return jsonRemove(args[0].toString(), pathSpecs);
7256
}
7357
}

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

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@
66
package org.opensearch.sql.expression.function.jsonUDF;
77

88
import static org.opensearch.sql.calcite.utils.PPLReturnTypes.STRING_FORCE_NULLABLE;
9-
import static org.opensearch.sql.expression.function.jsonUDF.JsonAppendFunctionImpl.jsonAppendIfArray;
9+
import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.*;
1010

1111
import com.fasterxml.jackson.core.JsonProcessingException;
12+
import com.fasterxml.jackson.databind.JsonNode;
13+
import java.util.ArrayList;
1214
import java.util.Arrays;
1315
import java.util.List;
1416
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
@@ -18,6 +20,7 @@
1820
import org.apache.calcite.linq4j.tree.Expression;
1921
import org.apache.calcite.linq4j.tree.Types;
2022
import org.apache.calcite.rex.RexCall;
23+
import org.apache.calcite.runtime.JsonFunctions;
2124
import org.apache.calcite.schema.impl.ScalarFunctionImpl;
2225
import org.apache.calcite.sql.type.SqlReturnTypeInference;
2326
import org.opensearch.sql.expression.function.ImplementorUDF;
@@ -49,8 +52,36 @@ public static Object eval(Object... args) throws JsonProcessingException {
4952
List<Object> keys = Arrays.asList(args).subList(1, args.length);
5053
if (keys.size() % 2 != 0) {
5154
throw new RuntimeException(
52-
"Json extend function needs corresponding path and values, but current get: " + keys);
55+
"Json append function needs corresponding path and values, but current get: " + keys);
5356
}
54-
return jsonAppendIfArray(jsonStr, keys, true);
57+
JsonNode root = verifyInput(args[0]);
58+
List<Object> expands = new ArrayList<>();
59+
for (int i = 0; i < keys.size(); i += 2) {
60+
List<String> expandedPaths = expandJsonPath(root, convertToJsonPath(keys.get(i).toString()));
61+
for (String expandedPath : expandedPaths) {
62+
Object value = keys.get(i + 1);
63+
if (value instanceof List<?> targetValues) {
64+
for (Object targetValue : targetValues) {
65+
expands.add(expandedPath + ".meaninglessKey");
66+
expands.add(targetValue);
67+
}
68+
} else if (value instanceof String stringValue) {
69+
try {
70+
List<Object> targetValues = gson.fromJson(stringValue, List.class);
71+
for (Object targetValue : targetValues) {
72+
expands.add(expandedPath + ".meaninglessKey");
73+
expands.add(targetValue);
74+
}
75+
} catch (Exception e) {
76+
expands.add(expandedPath + ".meaninglessKey");
77+
expands.add(value);
78+
}
79+
} else {
80+
expands.add(expandedPath + ".meaninglessKey");
81+
expands.add(value);
82+
}
83+
}
84+
}
85+
return JsonFunctions.jsonInsert(jsonStr, expands.toArray());
5586
}
5687
}

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

Lines changed: 60 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,28 @@
55

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

8-
import static org.apache.calcite.sql.type.SqlTypeUtil.createArrayType;
9-
import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.convertToJsonPath;
10-
import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.gson;
8+
import static org.apache.calcite.sql.SqlJsonQueryEmptyOrErrorBehavior.NULL;
9+
import static org.apache.calcite.sql.SqlJsonQueryWrapperBehavior.WITHOUT_ARRAY;
10+
import static org.opensearch.sql.calcite.utils.PPLReturnTypes.STRING_FORCE_NULLABLE;
11+
import static org.opensearch.sql.expression.function.jsonUDF.JsonUtils.*;
1112

12-
import com.jayway.jsonpath.JsonPath;
13+
import com.fasterxml.jackson.databind.JsonNode;
1314
import java.util.ArrayList;
1415
import java.util.Arrays;
16+
import java.util.Collection;
1517
import java.util.List;
18+
import java.util.Map;
1619
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
1720
import org.apache.calcite.adapter.enumerable.NullPolicy;
1821
import org.apache.calcite.adapter.enumerable.RexImpTable;
1922
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
2023
import org.apache.calcite.linq4j.tree.Expression;
2124
import org.apache.calcite.linq4j.tree.Types;
22-
import org.apache.calcite.rel.type.RelDataType;
23-
import org.apache.calcite.rel.type.RelDataTypeFactory;
2425
import org.apache.calcite.rex.RexCall;
26+
import org.apache.calcite.runtime.JsonFunctions;
2527
import org.apache.calcite.schema.impl.ScalarFunctionImpl;
28+
import org.apache.calcite.sql.SqlJsonValueEmptyOrErrorBehavior;
2629
import org.apache.calcite.sql.type.SqlReturnTypeInference;
27-
import org.apache.calcite.sql.type.SqlTypeName;
2830
import org.opensearch.sql.expression.function.ImplementorUDF;
2931

3032
public class JsonExtractFunctionImpl extends ImplementorUDF {
@@ -34,17 +36,7 @@ public JsonExtractFunctionImpl() {
3436

3537
@Override
3638
public SqlReturnTypeInference getReturnTypeInference() {
37-
return sqlOperatorBinding -> {
38-
RelDataTypeFactory typeFactory = sqlOperatorBinding.getTypeFactory();
39-
RelDataType varcharType =
40-
typeFactory.createTypeWithNullability(
41-
typeFactory.createSqlType(SqlTypeName.VARCHAR), true);
42-
if (sqlOperatorBinding.collectOperandTypes().size() > 2) {
43-
return createArrayType(typeFactory, varcharType, true);
44-
} else {
45-
return varcharType;
46-
}
47-
};
39+
return STRING_FORCE_NULLABLE;
4840
}
4941

5042
public static class JsonExtractImplementor implements NotNullImplementor {
@@ -63,27 +55,58 @@ public static Object eval(Object... args) {
6355
if (args.length < 2) {
6456
return null;
6557
}
66-
Object value = args[0];
58+
String jsonStr = (String) args[0];
59+
List<Object> jsonPaths = Arrays.asList(args).subList(1, args.length);
60+
/*
61+
JsonNode root = verifyInput(args[0]);
62+
List<Object> expands = new ArrayList<>();
63+
List<String> results = new ArrayList<>();
64+
for (Object key : keys) {
65+
expands.addAll(expandJsonPath(root, convertToJsonPath(key.toString())));
66+
}
67+
*/
68+
JsonNode root = verifyInput(args[0]);
69+
List<String> expands = new ArrayList<>();
6770
List<Object> results = new ArrayList<>();
68-
List<Object> paths = Arrays.asList(args).subList(1, args.length);
69-
for (Object path : paths) {
70-
String jsonPath = convertToJsonPath(path.toString());
71-
try {
72-
Object result;
73-
if (value instanceof String) {
74-
result = JsonPath.read((String) value, jsonPath);
75-
} else {
76-
result = JsonPath.read(value, jsonPath);
77-
}
78-
result = result != null ? gson.toJson(result) : null;
79-
results.add(result);
80-
} catch (Exception e) {
81-
results.add(null);
82-
}
71+
for (Object key : jsonPaths) {
72+
expands.addAll(expandJsonPath(root, convertToJsonPath(key.toString())));
73+
}
74+
List<String> modeExpands = new ArrayList<>();
75+
for (String expand : expands) {
76+
modeExpands.add(" lax " + expand);
77+
}
78+
JsonFunctions.StatefulFunction a = new JsonFunctions.StatefulFunction();
79+
for (String pathSpec : modeExpands) {
80+
Object queryResult = a.jsonQuery(jsonStr, pathSpec, WITHOUT_ARRAY, NULL, NULL, false);
81+
Object valueResult =
82+
a.jsonValue(
83+
jsonStr,
84+
pathSpec,
85+
SqlJsonValueEmptyOrErrorBehavior.NULL,
86+
null,
87+
SqlJsonValueEmptyOrErrorBehavior.NULL,
88+
null);
89+
results.add(queryResult != null ? queryResult : valueResult);
8390
}
84-
if (paths.size() > 1) {
85-
return results;
91+
if (expands.size() == 1) {
92+
return doJsonize(results.getFirst());
93+
}
94+
return doJsonize(results);
95+
}
96+
97+
private static boolean isScalarObject(Object obj) {
98+
if (obj instanceof Collection) {
99+
return false;
100+
} else {
101+
return !(obj instanceof Map);
102+
}
103+
}
104+
105+
private static String doJsonize(Object candidate) {
106+
if (isScalarObject(candidate)) {
107+
return candidate.toString();
108+
} else {
109+
return JsonFunctions.jsonize(candidate);
86110
}
87-
return results.get(0);
88111
}
89112
}

0 commit comments

Comments
 (0)