update IT tests for json_delete

Signed-off-by: YANGDB <[email protected]>
opensearch-project · Dec 9, 2024 · 9dbc89e · 9dbc89e
1 parent bb8bd30
commit 9dbc89e
Show file tree

Hide file tree

Showing 3 changed files with 110 additions and 9 deletions.
diff --git a/...c/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala b/...c/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala
@@ -14,7 +14,6 @@ import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFu
 import org.apache.spark.sql.catalyst.expressions.{Alias, EqualTo, Literal, Not}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.streaming.StreamTest
-import org.apache.spark.sql.types.StringType
 
 class FlintSparkPPLJsonFunctionITSuite
     extends QueryTest
@@ -432,4 +431,75 @@ class FlintSparkPPLJsonFunctionITSuite
     comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
   }
 
+  test("test json_delete() function: nested key") {
+    val frame = sql(s"""
+                       | source = $testTable
+                       | | eval result = json_delete('$validJson2',json_array('f2.f3')) | head 1 | fields result
+                       | """.stripMargin)
+    assertSameRows(Seq(Row("{\"f1\":\"abc\",\"f2\":{\"f4\":\"b\"}}")), frame)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val keysExpression =
+      UnresolvedFunction("array", Seq(Literal("f2.f3")), isDistinct = false)
+    val jsonObjExp =
+      Literal("{\"f1\":\"abc\",\"f2\":{\"f3\":\"a\",\"f4\":\"b\"}}")
+    val jsonFunc =
+      Alias(visit("json_delete", util.List.of(jsonObjExp, keysExpression)), "result")()
+    val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
+    val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
+    val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+
+  test("test json_delete() function: multi depth keys ") {
+    val frame = sql(s"""
+                       | source = $testTable
+                       | | eval result = json_delete('$validJson5',json_array('teacher', 'student.rank')) | head 1 | fields result
+                       | """.stripMargin)
+    assertSameRows(Seq(Row("{\"student\":[{\"name\":\"Bob\"},{\"name\":\"Charlie\"}]}")), frame)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val keysExpression =
+      UnresolvedFunction(
+        "array",
+        Seq(Literal("teacher"), Literal("student.rank")),
+        isDistinct = false)
+    val jsonObjExp =
+      Literal(
+        "{\"teacher\":\"Alice\",\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
+    val jsonFunc =
+      Alias(visit("json_delete", util.List.of(jsonObjExp, keysExpression)), "result")()
+    val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
+    val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
+    val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+
+  test("test json_delete() function: key not found") {
+    val frame = sql(s"""
+                       | source = $testTable
+                       | | eval result = json_delete('$validJson5',json_array('none')) | head 1 | fields result
+                       | """.stripMargin)
+    assertSameRows(
+      Seq(Row(
+        "{\"teacher\":\"Alice\",\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")),
+      frame)
+
+    val logicalPlan: LogicalPlan = frame.queryExecution.logical
+    val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
+    val keysExpression =
+      UnresolvedFunction("array", Seq(Literal("none")), isDistinct = false)
+    val jsonObjExp =
+      Literal(
+        "{\"teacher\":\"Alice\",\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
+    val jsonFunc =
+      Alias(visit("json_delete", util.List.of(jsonObjExp, keysExpression)), "result")()
+    val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
+    val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
+    val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
+    comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
+  }
+
 }
diff --git a/...ark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java b/...ark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java
@@ -60,16 +60,39 @@ private void removeKeys(Map<String, Object> map, WrappedArray<String> keysToRemo
             Collection<String> keys = JavaConverters.asJavaCollection(keysToRemove);
             for (String key : keys) {
                 String[] keyParts = key.split("\\.");
-                Map<String, Object> currentMap = map;
-                for (int i = 0; i < keyParts.length - 1; i++) {
-                    String currentKey = keyParts[i];
-                    if (currentMap.containsKey(currentKey) && currentMap.get(currentKey) instanceof Map) {
-                        currentMap = (Map<String, Object>) currentMap.get(currentKey);
-                    } else {
-                        return; 
+                removeNestedKey(map, keyParts, 0);
+            }
+        }
+
+        private void removeNestedKey(Object currentObj, String[] keyParts, int depth) {
+            if (currentObj == null || depth >= keyParts.length) {
+                return;
+            }
+
+            if (currentObj instanceof Map) {
+                Map<String, Object> currentMap = (Map<String, Object>) currentObj;
+                String currentKey = keyParts[depth];
+
+                if (depth == keyParts.length - 1) {
+                    // If it's the last key, remove it from the map
+                    currentMap.remove(currentKey);
+                } else {
+                    // If not the last key, continue traversing
+                    if (currentMap.containsKey(currentKey)) {
+                        Object nextObj = currentMap.get(currentKey);
+
+                        if (nextObj instanceof List) {
+                            // If the value is a list, process each item in the list
+                            List<Object> list = (List<Object>) nextObj;
+                            for (int i = 0; i < list.size(); i++) {
+                                removeNestedKey(list.get(i), keyParts, depth + 1);
+                            }
+                        } else {
+                            // Continue traversing if it's a map
+                            removeNestedKey(nextObj, keyParts, depth + 1);
+                        }
                     }
                 }
-                currentMap.remove(keyParts[keyParts.length - 1]);
             }
         }
     };

diff --git a/...gration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java b/...gration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java
@@ -63,6 +63,14 @@ public void testJsonDeleteFunctionRemoveMultipleSomeAreNestedKeys() {
         assertEquals(expectedJson, result);
     }
 
+    @Test
+    public void testJsonDeleteFunctionRemoveMultipleKeysNestedArrayKeys() {
+        String jsonStr = "{\"key1\":\"value1\",\"key2\":[{ \"a\":\"valueA\",\"key3\":\"value3\"}, {\"a\":\"valueA\",\"key4\":\"value4\"}]}";
+        String expectedJson = "{\"key2\":[{\"key3\":\"value3\"},{\"key4\":\"value4\"}]}";
+        String result = jsonDeleteFunction.apply(jsonStr,  WrappedArray.make(new String[]{"key1", "key2.a"}));
+        assertEquals(expectedJson, result);
+    }
+
     @Test
     public void testJsonDeleteFunctionNoKeysRemoved() {
         String jsonStr = "{\"key1\":\"value1\",\"key2\":\"value2\"}";