Skip to content

Commit f20b0ee

Browse files
acarbonettoacarbonetto
committed
Consolidate nested function calls
Signed-off-by: Andrew Carbonetto <andrew.carbonetto@improving.com>
1 parent 0937560 commit f20b0ee

File tree

4 files changed

+80
-95
lines changed

4 files changed

+80
-95
lines changed

‎docs/ppl-lang/functions/ppl-json.md‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -206,12 +206,12 @@ Example:
206206

207207
**Description**
208208

209-
`json_set(json_string, array(<path>, <value>, ...))` Inserts or updates one or more values at the corresponding paths in the specified JSON object.
209+
`json_set(json_string, array(path1, value1, path2, value2, ...))` Inserts or updates one or more values at the corresponding paths in the specified JSON object.
210210

211211
**Argument type:**
212212
- \<json_string\> must be a JSON_STRING.
213213
- \<path\> must be a STRING.
214-
- \<value\> can be a JSON_STRING.
214+
- \<value\> must be a JSON_STRING.
215215

216216
**Return type:** JSON_STRING
217217

@@ -232,9 +232,9 @@ Example:
232232

233233
**Description**
234234

235-
`json_delete(json_string, [keys list])` Deletes json elements from a json object based on json specific keys. Return the updated object after keys deletion .
235+
`json_delete(json_string, array(key1, key2, ...))` Deletes json elements from a json object based on json specific keys. Return the updated object after keys deletion .
236236

237-
**Arguments type:** JSON_STRING, List<STRING>
237+
**Arguments type:** JSON_STRING, List<JSON_STRING>
238238

239239
**Return type:** JSON_STRING
240240

@@ -270,7 +270,7 @@ Example:
270270

271271
**Description**
272272

273-
`json_append(json_string, [path_key, value,...])` appends values to end of an array at path_key within the json elements. Return the updated json object after appending.
273+
`json_append(json_string, [path_key1, value1, path_key2, value2, ...])` appends values to end of an array at path_key within the json elements. Return the updated json object after appending.
274274

275275
**Argument type:**
276276
- \<json_string\> must be a JSON_STRING.
@@ -319,7 +319,7 @@ Example:
319319

320320
**Description**
321321

322-
`json_extend(json_string, [path_key, value,...])` extends values to end of an array at path_key within the json elements. Return the updated json object after extending.
322+
`json_extend(json_string, [path_key1, value1, path_key2, value2, ...])` extends values to end of an array at path_key within the json elements. Return the updated json object after extending.
323323

324324
**Argument type:**
325325
- \<json_string\> must be a JSON_STRING.
@@ -370,7 +370,7 @@ Example:
370370

371371
`json_keys(jsonStr)` Returns all the keys of the outermost JSON object as an array.
372372

373-
**Argument type:** STRING
373+
**Argument type:** JSON_STRING
374374

375375
A STRING expression of a valid JSON object format.
376376

‎ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/JsonUtils.java‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import java.util.LinkedHashMap;
1212
import java.util.List;
1313
import java.util.Map;
14+
import java.util.function.BiFunction;
1415

1516
public interface JsonUtils {
1617
ObjectMapper objectMapper = new ObjectMapper();
@@ -25,6 +26,18 @@ static Object parseValue(String value) {
2526
}
2627
}
2728

29+
/**
30+
* update a nested value in a json object.
31+
*/
32+
static void updateNestedValue(List <Object> args){
33+
Object currentObj = args.get(0);
34+
String[] pathParts = (String[]) args.get(1);
35+
int depth = (Integer) args.get(2);
36+
Object valueToUpdate = args.get(3);
37+
38+
updateNestedValue(currentObj, pathParts, depth, valueToUpdate);
39+
}
40+
2841
/**
2942
* update a nested value in a json object.
3043
*
@@ -33,7 +46,7 @@ static Object parseValue(String value) {
3346
* @param depth - current traversal depth
3447
* @param valueToUpdate - value to update
3548
*/
36-
static void updateNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToUpdate) {
49+
private static void updateNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToUpdate) {
3750
if (currentObj == null || depth >= pathParts.length) {
3851
return;
3952
}
@@ -60,6 +73,30 @@ static void updateNestedValue(Object currentObj, String[] pathParts, int depth,
6073
}
6174
}
6275

76+
/**
77+
* append nested value to the json object.
78+
*/
79+
static void appendNestedValue(List <Object> args) {
80+
Object currentObj = args.get(0);
81+
String[] pathParts = (String[]) args.get(1);
82+
int depth = (Integer) args.get(2);
83+
Object valueToUpdate = args.get(3);
84+
85+
appendNestedValue(currentObj, pathParts, depth, valueToUpdate, false);
86+
}
87+
88+
/**
89+
* append nested value to the json object.
90+
*/
91+
static void extendNestedValue(List <Object> args) {
92+
Object currentObj = args.get(0);
93+
String[] pathParts = (String[]) args.get(1);
94+
int depth = (Integer) args.get(2);
95+
Object valueToUpdate = args.get(3);
96+
97+
appendNestedValue(currentObj, pathParts, depth, valueToUpdate, true);
98+
}
99+
63100
/**
64101
* append nested value to the json object.
65102
*
@@ -69,7 +106,7 @@ static void updateNestedValue(Object currentObj, String[] pathParts, int depth,
69106
* @param valueToAppend - value to add list
70107
* @param flattenValue - if value should be flattened first
71108
*/
72-
static void appendNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToAppend, boolean flattenValue) {
109+
private static void appendNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToAppend, boolean flattenValue) {
73110
if (currentObj == null || depth >= pathParts.length) {
74111
return;
75112
}

‎ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java‎

Lines changed: 26 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import inet.ipaddr.AddressStringException;
99
import inet.ipaddr.IPAddressString;
1010
import inet.ipaddr.IPAddressStringParameters;
11+
import java.util.function.Consumer;
1112
import org.apache.spark.sql.catalyst.expressions.Expression;
1213
import org.apache.spark.sql.catalyst.expressions.ScalaUDF;
1314
import org.apache.spark.sql.types.DataTypes;
@@ -34,7 +35,6 @@
3435
import java.util.List;
3536
import java.util.Map;
3637

37-
import static org.opensearch.sql.expression.function.JsonUtils.updateNestedValue;
3838
import static org.opensearch.sql.expression.function.JsonUtils.appendNestedValue;
3939
import static org.opensearch.sql.expression.function.JsonUtils.objectMapper;
4040
import static org.opensearch.sql.expression.function.JsonUtils.parseValue;
@@ -52,9 +52,6 @@ abstract class SerializableAbstractFunction2<T1, T2, R> extends AbstractFunction
5252
abstract class SerializableAbstractFunction3<T1, T2, T3, R> extends AbstractFunction3<T1, T2, T3, R>
5353
implements Serializable {}
5454

55-
abstract class SerializeableUpdateNestedFunction<T1, T2, R> extends AbstractFunction2<T1, T2, R>
56-
implements Serializable {}
57-
5855
/**
5956
* Remove specified keys from a JSON string.
6057
*
@@ -86,15 +83,10 @@ private void removeKeys(Map<String, Object> map, WrappedArray<String> keysToRemo
8683
}
8784
};
8885

89-
/**
90-
* Update the specified key-value pairs in a JSON string. If the key doesn't exist, the key-value is added.
91-
*
92-
* @param jsonStr The input JSON string.
93-
* @param elements A list of key-values pairs where the key is the key/path and value item is the updated value.
94-
* @return A new JSON string with updated values.
95-
*/
96-
Function2<String, WrappedArray<String>, String> jsonSetFunction = new SerializableAbstractFunction2<>() {
97-
public String apply(String jsonStr, WrappedArray<String> elements) {
86+
abstract class NestedFunction2<T1, T2, R> extends AbstractFunction2<T1, T2, R>
87+
implements Serializable {
88+
89+
protected String updateNestedJson(String jsonStr, WrappedArray<String> elements, Consumer<List> updateFieldConsumer) {
9890
if (jsonStr == null) {
9991
return null;
10092
}
@@ -119,7 +111,7 @@ public String apply(String jsonStr, WrappedArray<String> elements) {
119111
String[] pathParts = path.split("\\.");
120112
Object parsedValue = parseValue(iter.next());
121113

122-
updateNestedValue(jsonMap, pathParts, 0, parsedValue);
114+
updateFieldConsumer.accept(List.of(jsonMap, pathParts, 0, parsedValue));
123115
}
124116

125117
// Convert the updated map back to JSON
@@ -128,6 +120,20 @@ public String apply(String jsonStr, WrappedArray<String> elements) {
128120
return null;
129121
}
130122
}
123+
}
124+
125+
/**
126+
* Update the specified key-value pairs in a JSON string. If the key doesn't exist, the key-value is added.
127+
*
128+
* @param jsonStr The input JSON string.
129+
* @param elements A list of key-values pairs where the key is the key/path and value item is the updated value.
130+
* @return A new JSON string with updated values.
131+
*/
132+
Function2<String, WrappedArray<String>, String> jsonSetFunction = new NestedFunction2<>() {
133+
@Override
134+
public String apply(String jsonStr, WrappedArray<String> elements) {
135+
return updateNestedJson(jsonStr, elements, JsonUtils::updateNestedValue);
136+
}
131137
};
132138

133139
/**
@@ -137,40 +143,10 @@ public String apply(String jsonStr, WrappedArray<String> elements) {
137143
* @param elements A list of path-values where the first item is the path and subsequent items are values to append.
138144
* @return The updated JSON string.
139145
*/
140-
Function2<String, WrappedArray<String>, String> jsonAppendFunction = new SerializableAbstractFunction2<>() {
146+
Function2<String, WrappedArray<String>, String> jsonAppendFunction = new NestedFunction2<>() {
147+
@Override
141148
public String apply(String jsonStr, WrappedArray<String> elements) {
142-
if (jsonStr == null) {
143-
return null;
144-
}
145-
try {
146-
List<String> pathValues = JavaConverters.mutableSeqAsJavaList(elements);
147-
// don't update if the list is empty, or the list is not key-value pairs
148-
if (pathValues.isEmpty()) {
149-
return jsonStr;
150-
}
151-
152-
// Parse the JSON string into a Map
153-
Map<String, Object> jsonMap = objectMapper.readValue(jsonStr, Map.class);
154-
155-
// Iterate through the key-value pairs and update the json
156-
var iter = pathValues.iterator();
157-
while (iter.hasNext()) {
158-
String path = iter.next();
159-
if (!iter.hasNext()) {
160-
// no value provided and cannot update anything
161-
break;
162-
}
163-
String[] pathParts = path.split("\\.");
164-
Object parsedValue = parseValue(iter.next());
165-
166-
appendNestedValue(jsonMap, pathParts, 0, parsedValue, false);
167-
}
168-
169-
// Convert the updated map back to JSON
170-
return objectMapper.writeValueAsString(jsonMap);
171-
} catch (Exception e) {
172-
return null;
173-
}
149+
return updateNestedJson(jsonStr, elements, JsonUtils::appendNestedValue);
174150
}
175151
};
176152

@@ -181,40 +157,10 @@ public String apply(String jsonStr, WrappedArray<String> elements) {
181157
* @param elements A list of path-values where the first item is the path and subsequent items are values to append.
182158
* @return The updated JSON string.
183159
*/
184-
Function2<String, WrappedArray<String>, String> jsonExtendFunction = new SerializableAbstractFunction2<>() {
160+
Function2<String, WrappedArray<String>, String> jsonExtendFunction = new NestedFunction2<>() {
161+
@Override
185162
public String apply(String jsonStr, WrappedArray<String> elements) {
186-
if (jsonStr == null) {
187-
return null;
188-
}
189-
try {
190-
List<String> pathValues = JavaConverters.mutableSeqAsJavaList(elements);
191-
// don't update if the list is empty, or the list is not key-value pairs
192-
if (pathValues.isEmpty()) {
193-
return jsonStr;
194-
}
195-
196-
// Parse the JSON string into a Map
197-
Map<String, Object> jsonMap = objectMapper.readValue(jsonStr, Map.class);
198-
199-
// Iterate through the key-value pairs and update the json
200-
var iter = pathValues.iterator();
201-
while (iter.hasNext()) {
202-
String path = iter.next();
203-
if (!iter.hasNext()) {
204-
// no value provided and cannot update anything
205-
break;
206-
}
207-
String[] pathParts = path.split("\\.");
208-
Object parsedValue = parseValue(iter.next());
209-
210-
appendNestedValue(jsonMap, pathParts, 0, parsedValue, true);
211-
}
212-
213-
// Convert the updated map back to JSON
214-
return objectMapper.writeValueAsString(jsonMap);
215-
} catch (Exception e) {
216-
return null;
217-
}
163+
return updateNestedJson(jsonStr, elements, JsonUtils::extendNestedValue);
218164
}
219165
};
220166

‎ppl-spark-integration/src/test/scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanJsonFunctionsTranslatorTestSuite.scala‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -197,14 +197,14 @@ class PPLLogicalPlanJsonFunctionsTranslatorTestSuite
197197
planTransformer.visit(
198198
plan(
199199
pplParser,
200-
"""source=t | eval result = json_set('{"a":[{"b":1},{"c":2}]}', array('a.b', 3, 'a.c', 4))"""),
200+
"""source=t | eval result = json_set('{"a":[{"b":1},{"c":2}]}', array('a.b', '3', 'a.c', '4'))"""),
201201
context)
202202

203203
val table = UnresolvedRelation(Seq("t"))
204204
val keysExpression =
205205
UnresolvedFunction(
206206
"array",
207-
Seq(Literal("a.b"), Literal(3), Literal("a.c"), Literal(4)),
207+
Seq(Literal("a.b"), Literal("3"), Literal("a.c"), Literal("4")),
208208
isDistinct = false)
209209
val jsonObjExp =
210210
Literal("""{"a":[{"b":1},{"c":2}]}""")
@@ -242,14 +242,14 @@ class PPLLogicalPlanJsonFunctionsTranslatorTestSuite
242242
planTransformer.visit(
243243
plan(
244244
pplParser,
245-
"""source=t | eval result = json_append('{"a":[{"b":1},{"c":2}]}', array('a.b','c','d'))"""),
245+
"""source=t | eval result = json_append('{"a":[{"b":1},{"c":2}]}', array('a.b','c','a.d','e'))"""),
246246
context)
247247

248248
val table = UnresolvedRelation(Seq("t"))
249249
val keysExpression =
250250
UnresolvedFunction(
251251
"array",
252-
Seq(Literal("a.b"), Literal("c"), Literal("d")),
252+
Seq(Literal("a.b"), Literal("c"), Literal("a.d"), Literal("e")),
253253
isDistinct = false)
254254
val jsonObjExp =
255255
Literal("""{"a":[{"b":1},{"c":2}]}""")
@@ -266,7 +266,7 @@ class PPLLogicalPlanJsonFunctionsTranslatorTestSuite
266266
planTransformer.visit(
267267
plan(
268268
pplParser,
269-
"""source=t | eval result = json_extend('{"a":[{"b":1},{"c":2}]}', array('a.b',array('c','d')))"""),
269+
"""source=t | eval result = json_extend('{"a":[{"b":1},{"c":2}]}', array('a.b',array('c','d'), 'a.e',array('f','g')))"""),
270270
context)
271271

272272
val table = UnresolvedRelation(Seq("t"))
@@ -275,7 +275,9 @@ class PPLLogicalPlanJsonFunctionsTranslatorTestSuite
275275
"array",
276276
Seq(
277277
Literal("a.b"),
278-
UnresolvedFunction("array", Seq(Literal("c"), Literal("d")), isDistinct = false)),
278+
UnresolvedFunction("array", Seq(Literal("c"), Literal("d")), isDistinct = false),
279+
Literal("a.e"),
280+
UnresolvedFunction("array", Seq(Literal("f"), Literal("g")), isDistinct = false)),
279281
isDistinct = false)
280282
val jsonObjExp =
281283
Literal("""{"a":[{"b":1},{"c":2}]}""")

0 commit comments

Comments
 (0)