Skip to content

Commit d17e337

Browse files
acarbonettoacarbonetto
committed
Refactor nested traverse
Signed-off-by: Andrew Carbonetto <andrew.carbonetto@improving.com>
1 parent ab531cd commit d17e337

File tree

2 files changed

+65
-132
lines changed

2 files changed

+65
-132
lines changed

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

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

1615
public interface JsonUtils {
1716
ObjectMapper objectMapper = new ObjectMapper();
@@ -26,27 +25,13 @@ static Object parseValue(String value) {
2625
}
2726
}
2827

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);
28+
@FunctionalInterface
29+
interface UpdateConsumer {
30+
void apply(Map<String, Object> obj, String key, Object value);
3931
}
4032

41-
/**
42-
* update a nested value in a json object.
43-
*
44-
* @param currentObj - json object to traverse
45-
* @param pathParts - key at path to update
46-
* @param depth - current traversal depth
47-
* @param valueToUpdate - value to update
48-
*/
49-
private static void updateNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToUpdate) {
33+
private static void traverseNestedObject(Object currentObj, String[] pathParts, int depth,
34+
Object valueToUpdate, UpdateConsumer updateObjectFunction) {
5035
if (currentObj == null || depth >= pathParts.length) {
5136
return;
5237
}
@@ -56,90 +41,80 @@ private static void updateNestedValue(Object currentObj, String[] pathParts, int
5641
String currentKey = pathParts[depth];
5742

5843
if (depth == pathParts.length - 1) {
59-
currentMap.put(currentKey, valueToUpdate);
44+
updateObjectFunction.apply(currentMap, currentKey, valueToUpdate);
6045
} else {
6146
// Continue traversing
62-
currentMap.computeIfAbsent(currentKey, k -> new LinkedHashMap<>()); // Create map if not present
63-
updateNestedValue(currentMap.get(currentKey), pathParts, depth + 1, valueToUpdate);
47+
currentMap.computeIfAbsent(currentKey,
48+
k -> new LinkedHashMap<>()); // Create map if not present
49+
traverseNestedObject(currentMap.get(currentKey), pathParts, depth + 1,
50+
valueToUpdate, updateObjectFunction);
6451
}
6552
} else if (currentObj instanceof List) {
6653
// If the current object is a list, process each map in the list
6754
List<Object> list = (List<Object>) currentObj;
6855
for (Object item : list) {
6956
if (item instanceof Map) {
70-
updateNestedValue(item, pathParts, depth, valueToUpdate);
57+
traverseNestedObject(item, pathParts, depth, valueToUpdate, updateObjectFunction);
7158
}
7259
}
7360
}
7461
}
7562

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-
}
63+
static String updateNestedJson(String jsonStr, List<String> pathValues, UpdateConsumer updateFieldConsumer) {
64+
if (jsonStr == null) {
65+
return null;
66+
}
67+
// don't update if the list is empty, or the list is not key-value pairs
68+
if (pathValues.isEmpty()) {
69+
return jsonStr;
70+
}
71+
try {
72+
// Parse the JSON string into a Map
73+
Map<String, Object> jsonMap = objectMapper.readValue(jsonStr, Map.class);
74+
75+
// Iterate through the key-value pairs and update the json
76+
var iter = pathValues.iterator();
77+
while (iter.hasNext()) {
78+
String path = iter.next();
79+
if (!iter.hasNext()) {
80+
// no value provided and cannot update anything
81+
break;
82+
}
83+
String[] pathParts = path.split("\\.");
84+
Object parsedValue = parseValue(iter.next());
8785

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);
86+
traverseNestedObject(jsonMap, pathParts, 0, parsedValue, updateFieldConsumer);
87+
}
9688

97-
appendNestedValue(currentObj, pathParts, depth, valueToUpdate, true);
89+
// Convert the updated map back to JSON
90+
return objectMapper.writeValueAsString(jsonMap);
91+
} catch (Exception e) {
92+
return null;
93+
}
9894
}
9995

100-
/**
101-
* append nested value to the json object.
102-
*
103-
* @param currentObj - json object to traverse
104-
* @param pathParts - key at path to update
105-
* @param depth - current traversal depth
106-
* @param valueToAppend - value to add list
107-
* @param flattenValue - if value should be flattened first
108-
*/
109-
private static void appendNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToAppend, boolean flattenValue) {
110-
if (currentObj == null || depth >= pathParts.length) {
111-
return;
96+
static void appendObjectValue(Map<String, Object> obj, String key, Object value) {
97+
// If it's the last key, append to the array
98+
obj.computeIfAbsent(key, k -> new ArrayList<>()); // Create list if not present
99+
Object existingValue = obj.get(key);
100+
101+
if (existingValue instanceof List) {
102+
List<Object> list = (List<Object>) existingValue;
103+
list.add(value);
112104
}
105+
}
113106

114-
if (currentObj instanceof Map) {
115-
Map<String, Object> currentMap = (Map<String, Object>) currentObj;
116-
String currentKey = pathParts[depth];
107+
static void extendObjectValue(Map<String, Object> obj, String key, Object value) {
108+
// If it's the last key, append to the array
109+
obj.computeIfAbsent(key, k -> new ArrayList<>()); // Create list if not present
110+
Object existingValue = obj.get(key);
117111

118-
if (depth == pathParts.length - 1) {
119-
// If it's the last key, append to the array
120-
currentMap.computeIfAbsent(currentKey, k -> new ArrayList<>()); // Create list if not present
121-
Object existingValue = currentMap.get(currentKey);
122-
123-
if (existingValue instanceof List) {
124-
List<Object> existingList = (List<Object>) existingValue;
125-
if (flattenValue && valueToAppend instanceof List) {
126-
existingList.addAll((List) valueToAppend);
127-
} else {
128-
existingList.add(valueToAppend);
129-
}
130-
}
112+
if (existingValue instanceof List) {
113+
List<Object> existingList = (List<Object>) existingValue;
114+
if (value instanceof List) {
115+
existingList.addAll((List) value);
131116
} else {
132-
// Continue traversing
133-
currentMap.computeIfAbsent(currentKey, k -> new LinkedHashMap<>()); // Create map if not present
134-
appendNestedValue(currentMap.get(currentKey), pathParts, depth + 1, valueToAppend, flattenValue);
135-
}
136-
} else if (currentObj instanceof List) {
137-
// If the current object is a list, process each map in the list
138-
List<Object> list = (List<Object>) currentObj;
139-
for (Object item : list) {
140-
if (item instanceof Map) {
141-
appendNestedValue(item, pathParts, depth, valueToAppend, flattenValue);
142-
}
117+
existingList.add(value);
143118
}
144119
}
145120
}

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

Lines changed: 7 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import inet.ipaddr.AddressStringException;
99
import inet.ipaddr.IPAddressString;
1010
import inet.ipaddr.IPAddressStringParameters;
11-
import java.util.function.Consumer;
1211
import org.apache.spark.sql.catalyst.expressions.Expression;
1312
import org.apache.spark.sql.catalyst.expressions.ScalaUDF;
1413
import org.apache.spark.sql.types.DataTypes;
@@ -32,13 +31,11 @@
3231
import java.time.ZoneId;
3332
import java.time.ZonedDateTime;
3433
import java.util.Collection;
35-
import java.util.List;
3634
import java.util.Map;
3735

38-
import static org.opensearch.sql.expression.function.JsonUtils.appendNestedValue;
3936
import static org.opensearch.sql.expression.function.JsonUtils.objectMapper;
40-
import static org.opensearch.sql.expression.function.JsonUtils.parseValue;
4137
import static org.opensearch.sql.expression.function.JsonUtils.removeNestedKey;
38+
import static org.opensearch.sql.expression.function.JsonUtils.updateNestedJson;
4239
import static org.opensearch.sql.ppl.utils.DataTypeTransformer.seq;
4340

4441
public interface SerializableUdf {
@@ -83,56 +80,17 @@ private void removeKeys(Map<String, Object> map, WrappedArray<String> keysToRemo
8380
}
8481
};
8582

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) {
90-
if (jsonStr == null) {
91-
return null;
92-
}
93-
try {
94-
List<String> pathValues = JavaConverters.mutableSeqAsJavaList(elements);
95-
// don't update if the list is empty, or the list is not key-value pairs
96-
if (pathValues.isEmpty()) {
97-
return jsonStr;
98-
}
99-
100-
// Parse the JSON string into a Map
101-
Map<String, Object> jsonMap = objectMapper.readValue(jsonStr, Map.class);
102-
103-
// Iterate through the key-value pairs and update the json
104-
var iter = pathValues.iterator();
105-
while (iter.hasNext()) {
106-
String path = iter.next();
107-
if (!iter.hasNext()) {
108-
// no value provided and cannot update anything
109-
break;
110-
}
111-
String[] pathParts = path.split("\\.");
112-
Object parsedValue = parseValue(iter.next());
113-
114-
updateFieldConsumer.accept(List.of(jsonMap, pathParts, 0, parsedValue));
115-
}
116-
117-
// Convert the updated map back to JSON
118-
return objectMapper.writeValueAsString(jsonMap);
119-
} catch (Exception e) {
120-
return null;
121-
}
122-
}
123-
}
124-
12583
/**
12684
* Update the specified key-value pairs in a JSON string. If the key doesn't exist, the key-value is added.
12785
*
12886
* @param jsonStr The input JSON string.
12987
* @param elements A list of key-values pairs where the key is the key/path and value item is the updated value.
13088
* @return A new JSON string with updated values.
13189
*/
132-
Function2<String, WrappedArray<String>, String> jsonSetFunction = new NestedFunction2<>() {
90+
Function2<String, WrappedArray<String>, String> jsonSetFunction = new SerializableAbstractFunction2<>() {
13391
@Override
13492
public String apply(String jsonStr, WrappedArray<String> elements) {
135-
return updateNestedJson(jsonStr, elements, JsonUtils::updateNestedValue);
93+
return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), (obj, key, value) -> obj.put(key, value));
13694
}
13795
};
13896

@@ -143,10 +101,10 @@ public String apply(String jsonStr, WrappedArray<String> elements) {
143101
* @param elements A list of path-values where the first item is the path and subsequent items are values to append.
144102
* @return The updated JSON string.
145103
*/
146-
Function2<String, WrappedArray<String>, String> jsonAppendFunction = new NestedFunction2<>() {
104+
Function2<String, WrappedArray<String>, String> jsonAppendFunction = new SerializableAbstractFunction2<>() {
147105
@Override
148106
public String apply(String jsonStr, WrappedArray<String> elements) {
149-
return updateNestedJson(jsonStr, elements, JsonUtils::appendNestedValue);
107+
return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), JsonUtils::appendObjectValue);
150108
}
151109
};
152110

@@ -157,10 +115,10 @@ public String apply(String jsonStr, WrappedArray<String> elements) {
157115
* @param elements A list of path-values where the first item is the path and subsequent items are values to append.
158116
* @return The updated JSON string.
159117
*/
160-
Function2<String, WrappedArray<String>, String> jsonExtendFunction = new NestedFunction2<>() {
118+
Function2<String, WrappedArray<String>, String> jsonExtendFunction = new SerializableAbstractFunction2<>() {
161119
@Override
162120
public String apply(String jsonStr, WrappedArray<String> elements) {
163-
return updateNestedJson(jsonStr, elements, JsonUtils::extendNestedValue);
121+
return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), JsonUtils::extendObjectValue);
164122
}
165123
};
166124

0 commit comments

Comments
 (0)