From d17e337aa4ce1bcd66cf11f0d5a7d3313fc74471 Mon Sep 17 00:00:00 2001 From: Andrew Carbonetto Date: Thu, 13 Feb 2025 13:49:31 -0800 Subject: [PATCH] Refactor nested traverse Signed-off-by: Andrew Carbonetto --- .../sql/expression/function/JsonUtils.java | 141 +++++++----------- .../expression/function/SerializableUdf.java | 56 +------ 2 files changed, 65 insertions(+), 132 deletions(-) diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/JsonUtils.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/JsonUtils.java index a9c99fd51..7f151eb00 100644 --- a/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/JsonUtils.java +++ b/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/JsonUtils.java @@ -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(); @@ -26,27 +25,13 @@ static Object parseValue(String value) { } } - /** - * update a nested value in a json object. - */ - static void updateNestedValue(List 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 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; } @@ -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 list = (List) 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 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 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 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 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 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 list = (List) existingValue; + list.add(value); } + } - if (currentObj instanceof Map) { - Map currentMap = (Map) currentObj; - String currentKey = pathParts[depth]; + static void extendObjectValue(Map 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 existingList = (List) existingValue; - if (flattenValue && valueToAppend instanceof List) { - existingList.addAll((List) valueToAppend); - } else { - existingList.add(valueToAppend); - } - } + if (existingValue instanceof List) { + List existingList = (List) 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 list = (List) currentObj; - for (Object item : list) { - if (item instanceof Map) { - appendNestedValue(item, pathParts, depth, valueToAppend, flattenValue); - } + existingList.add(value); } } } diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java index 9afe7d3ff..a9de800e7 100644 --- a/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java +++ b/ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java @@ -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; @@ -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 { @@ -83,45 +80,6 @@ private void removeKeys(Map map, WrappedArray keysToRemo } }; - abstract class NestedFunction2 extends AbstractFunction2 - implements Serializable { - - protected String updateNestedJson(String jsonStr, WrappedArray elements, Consumer updateFieldConsumer) { - if (jsonStr == null) { - return null; - } - try { - List 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 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. * @@ -129,10 +87,10 @@ protected String updateNestedJson(String jsonStr, WrappedArray elements, * @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> jsonSetFunction = new NestedFunction2<>() { + Function2, String> jsonSetFunction = new SerializableAbstractFunction2<>() { @Override public String apply(String jsonStr, WrappedArray elements) { - return updateNestedJson(jsonStr, elements, JsonUtils::updateNestedValue); + return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), (obj, key, value) -> obj.put(key, value)); } }; @@ -143,10 +101,10 @@ public String apply(String jsonStr, WrappedArray 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> jsonAppendFunction = new NestedFunction2<>() { + Function2, String> jsonAppendFunction = new SerializableAbstractFunction2<>() { @Override public String apply(String jsonStr, WrappedArray elements) { - return updateNestedJson(jsonStr, elements, JsonUtils::appendNestedValue); + return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), JsonUtils::appendObjectValue); } }; @@ -157,10 +115,10 @@ public String apply(String jsonStr, WrappedArray 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> jsonExtendFunction = new NestedFunction2<>() { + Function2, String> jsonExtendFunction = new SerializableAbstractFunction2<>() { @Override public String apply(String jsonStr, WrappedArray elements) { - return updateNestedJson(jsonStr, elements, JsonUtils::extendNestedValue); + return updateNestedJson(jsonStr, JavaConverters.mutableSeqAsJavaList(elements), JsonUtils::extendObjectValue); } };