Skip to content

Commit

Permalink
Refactor nested traverse
Browse files Browse the repository at this point in the history 
Signed-off-by: Andrew Carbonetto <[email protected]>
  • Loading branch information
@acarbonetto
acarbonetto committed Feb 13, 2025
1 parent ab531cd commit d17e337
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 132 deletions.
141 changes: 58 additions & 83 deletions ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/JsonUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;

public interface JsonUtils {
ObjectMapper objectMapper = new ObjectMapper();
Expand All @@ -26,27 +25,13 @@ static Object parseValue(String value) {
}
}

/**
* update a nested value in a json object.
*/
static void updateNestedValue(List <Object> args){
Object currentObj = args.get(0);
String[] pathParts = (String[]) args.get(1);
int depth = (Integer) args.get(2);
Object valueToUpdate = args.get(3);

updateNestedValue(currentObj, pathParts, depth, valueToUpdate);
@FunctionalInterface
interface UpdateConsumer {
void apply(Map<String, Object> obj, String key, Object value);
}

/**
* update a nested value in a json object.
*
* @param currentObj - json object to traverse
* @param pathParts - key at path to update
* @param depth - current traversal depth
* @param valueToUpdate - value to update
*/
private static void updateNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToUpdate) {
private static void traverseNestedObject(Object currentObj, String[] pathParts, int depth,
Object valueToUpdate, UpdateConsumer updateObjectFunction) {
if (currentObj == null || depth >= pathParts.length) {
return;
}
Expand All @@ -56,90 +41,80 @@ private static void updateNestedValue(Object currentObj, String[] pathParts, int
String currentKey = pathParts[depth];

if (depth == pathParts.length - 1) {
currentMap.put(currentKey, valueToUpdate);
updateObjectFunction.apply(currentMap, currentKey, valueToUpdate);
} else {
// Continue traversing
currentMap.computeIfAbsent(currentKey, k -> new LinkedHashMap<>()); // Create map if not present
updateNestedValue(currentMap.get(currentKey), pathParts, depth + 1, valueToUpdate);
currentMap.computeIfAbsent(currentKey,
k -> new LinkedHashMap<>()); // Create map if not present
traverseNestedObject(currentMap.get(currentKey), pathParts, depth + 1,
valueToUpdate, updateObjectFunction);
}
} else if (currentObj instanceof List) {
// If the current object is a list, process each map in the list
List<Object> list = (List<Object>) currentObj;
for (Object item : list) {
if (item instanceof Map) {
updateNestedValue(item, pathParts, depth, valueToUpdate);
traverseNestedObject(item, pathParts, depth, valueToUpdate, updateObjectFunction);
}
}
}
}

/**
* append nested value to the json object.
*/
static void appendNestedValue(List <Object> args) {
Object currentObj = args.get(0);
String[] pathParts = (String[]) args.get(1);
int depth = (Integer) args.get(2);
Object valueToUpdate = args.get(3);

appendNestedValue(currentObj, pathParts, depth, valueToUpdate, false);
}
static String updateNestedJson(String jsonStr, List<String> pathValues, UpdateConsumer updateFieldConsumer) {
if (jsonStr == null) {
return null;
}
// don't update if the list is empty, or the list is not key-value pairs
if (pathValues.isEmpty()) {
return jsonStr;
}
try {
// Parse the JSON string into a Map
Map<String, Object> jsonMap = objectMapper.readValue(jsonStr, Map.class);

// Iterate through the key-value pairs and update the json
var iter = pathValues.iterator();
while (iter.hasNext()) {
String path = iter.next();
if (!iter.hasNext()) {
// no value provided and cannot update anything
break;
}
String[] pathParts = path.split("\\.");
Object parsedValue = parseValue(iter.next());

/**
* append nested value to the json object.
*/
static void extendNestedValue(List <Object> args) {
Object currentObj = args.get(0);
String[] pathParts = (String[]) args.get(1);
int depth = (Integer) args.get(2);
Object valueToUpdate = args.get(3);
traverseNestedObject(jsonMap, pathParts, 0, parsedValue, updateFieldConsumer);
}

appendNestedValue(currentObj, pathParts, depth, valueToUpdate, true);
// Convert the updated map back to JSON
return objectMapper.writeValueAsString(jsonMap);
} catch (Exception e) {
return null;
}
}

/**
* append nested value to the json object.
*
* @param currentObj - json object to traverse
* @param pathParts - key at path to update
* @param depth - current traversal depth
* @param valueToAppend - value to add list
* @param flattenValue - if value should be flattened first
*/
private static void appendNestedValue(Object currentObj, String[] pathParts, int depth, Object valueToAppend, boolean flattenValue) {
if (currentObj == null || depth >= pathParts.length) {
return;
static void appendObjectValue(Map<String, Object> obj, String key, Object value) {
// If it's the last key, append to the array
obj.computeIfAbsent(key, k -> new ArrayList<>()); // Create list if not present
Object existingValue = obj.get(key);

if (existingValue instanceof List) {
List<Object> list = (List<Object>) existingValue;
list.add(value);
}
}

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

if (depth == pathParts.length - 1) {
// If it's the last key, append to the array
currentMap.computeIfAbsent(currentKey, k -> new ArrayList<>()); // Create list if not present
Object existingValue = currentMap.get(currentKey);

if (existingValue instanceof List) {
List<Object> existingList = (List<Object>) existingValue;
if (flattenValue && valueToAppend instanceof List) {
existingList.addAll((List) valueToAppend);
} else {
existingList.add(valueToAppend);
}
}
if (existingValue instanceof List) {
List<Object> existingList = (List<Object>) existingValue;
if (value instanceof List) {
existingList.addAll((List) value);
} else {
// Continue traversing
currentMap.computeIfAbsent(currentKey, k -> new LinkedHashMap<>()); // Create map if not present
appendNestedValue(currentMap.get(currentKey), pathParts, depth + 1, valueToAppend, flattenValue);
}
} else if (currentObj instanceof List) {
// If the current object is a list, process each map in the list
List<Object> list = (List<Object>) currentObj;
for (Object item : list) {
if (item instanceof Map) {
appendNestedValue(item, pathParts, depth, valueToAppend, flattenValue);
}
existingList.add(value);
}
}
}
Expand Down
56 changes: 7 additions & 49 deletions ...ark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import inet.ipaddr.AddressStringException;
import inet.ipaddr.IPAddressString;
import inet.ipaddr.IPAddressStringParameters;
import java.util.function.Consumer;
import org.apache.spark.sql.catalyst.expressions.Expression;
import org.apache.spark.sql.catalyst.expressions.ScalaUDF;
import org.apache.spark.sql.types.DataTypes;
Expand All @@ -32,13 +31,11 @@
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.Collection;
import java.util.List;
import java.util.Map;

import static org.opensearch.sql.expression.function.JsonUtils.appendNestedValue;
import static org.opensearch.sql.expression.function.JsonUtils.objectMapper;
import static org.opensearch.sql.expression.function.JsonUtils.parseValue;
import static org.opensearch.sql.expression.function.JsonUtils.removeNestedKey;
import static org.opensearch.sql.expression.function.JsonUtils.updateNestedJson;
import static org.opensearch.sql.ppl.utils.DataTypeTransformer.seq;

public interface SerializableUdf {
Expand Down Expand Up @@ -83,56 +80,17 @@ private void removeKeys(Map<String, Object> map, WrappedArray<String> keysToRemo
}
};

abstract class NestedFunction2<T1, T2, R> extends AbstractFunction2<T1, T2, R>
implements Serializable {

protected String updateNestedJson(String jsonStr, WrappedArray<String> elements, Consumer<List> updateFieldConsumer) {
if (jsonStr == null) {
return null;
}
try {
List<String> pathValues = JavaConverters.mutableSeqAsJavaList(elements);
// don't update if the list is empty, or the list is not key-value pairs
if (pathValues.isEmpty()) {
return jsonStr;
}

// Parse the JSON string into a Map
Map<String, Object> jsonMap = objectMapper.readValue(jsonStr, Map.class);

// Iterate through the key-value pairs and update the json
var iter = pathValues.iterator();
while (iter.hasNext()) {
String path = iter.next();
if (!iter.hasNext()) {
// no value provided and cannot update anything
break;
}
String[] pathParts = path.split("\\.");
Object parsedValue = parseValue(iter.next());

updateFieldConsumer.accept(List.of(jsonMap, pathParts, 0, parsedValue));
}

// Convert the updated map back to JSON
return objectMapper.writeValueAsString(jsonMap);
} catch (Exception e) {
return null;
}
}
}

/**
* Update the specified key-value pairs in a JSON string. If the key doesn't exist, the key-value is added.
*
* @param jsonStr The input JSON string.
* @param elements A list of key-values pairs where the key is the key/path and value item is the updated value.
* @return A new JSON string with updated values.
*/
Function2<String, WrappedArray<String>, String> jsonSetFunction = new NestedFunction2<>() {
Function2<String, WrappedArray<String>, String> jsonSetFunction = new SerializableAbstractFunction2<>() {
@Override
public String apply(String jsonStr, WrappedArray<String> elements) {
return updateNestedJson(jsonStr, elements, JsonUtils::updateNestedValue);
return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), (obj, key, value) -> obj.put(key, value));
}
};

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

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

Expand Down

0 comments on commit d17e337

Please sign in to comment.