From bb8bd30a455ceefb8f5a134c2d14555320c4f668 Mon Sep 17 00:00:00 2001 From: YANGDB Date: Mon, 9 Dec 2024 13:03:41 -0800 Subject: [PATCH] update IT tests for `json_delete` Signed-off-by: YANGDB --- .../FlintSparkPPLJsonFunctionITSuite.scala | 47 +++++++++++++++++++ .../expression/function/SerializableUdf.java | 34 +++++--------- .../function/SerializableJsonUdfTest.java | 42 ++++------------- 3 files changed, 68 insertions(+), 55 deletions(-) diff --git a/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala b/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala index fca758101..ed652275d 100644 --- a/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala +++ b/integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala @@ -5,11 +5,16 @@ package org.opensearch.flint.spark.ppl +import java.util + +import org.opensearch.sql.expression.function.SerializableUdf.visit + import org.apache.spark.sql.{AnalysisException, QueryTest, Row} import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar} 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 @@ -385,4 +390,46 @@ class FlintSparkPPLJsonFunctionITSuite null)) assertSameRows(expectedSeq, frame) } + + test("test json_delete() function: one key") { + val frame = sql(s""" + | source = $testTable + | | eval result = json_delete('$validJson1',json_array('age')) | head 1 | fields result + | """.stripMargin) + assertSameRows(Seq(Row("{\"account_number\":1,\"balance\":39225,\"gender\":\"M\"}")), frame) + + val logicalPlan: LogicalPlan = frame.queryExecution.logical + val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test")) + val keysExpression = UnresolvedFunction("array", Seq(Literal("age")), isDistinct = false) + val jsonObjExp = + Literal("{\"account_number\":1,\"balance\":39225,\"age\":32,\"gender\":\"M\"}") + 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: multiple keys") { + val frame = sql(s""" + | source = $testTable + | | eval result = json_delete('$validJson1',json_array('age','gender')) | head 1 | fields result + | """.stripMargin) + assertSameRows(Seq(Row("{\"account_number\":1,\"balance\":39225}")), frame) + + val logicalPlan: LogicalPlan = frame.queryExecution.logical + val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test")) + val keysExpression = + UnresolvedFunction("array", Seq(Literal("age"), Literal("gender")), isDistinct = false) + val jsonObjExp = + Literal("{\"account_number\":1,\"balance\":39225,\"age\":32,\"gender\":\"M\"}") + 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/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 ea7b7d2dc..89888a4e8 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 @@ -6,7 +6,6 @@ package org.opensearch.sql.expression.function; import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.module.scala.DefaultScalaModule; import inet.ipaddr.AddressStringException; import inet.ipaddr.IPAddressString; import inet.ipaddr.IPAddressStringParameters; @@ -16,11 +15,13 @@ import scala.Function2; import scala.Option; import scala.Serializable; +import scala.collection.JavaConverters; +import scala.collection.mutable.WrappedArray; import scala.runtime.AbstractFunction2; +import java.util.Collection; import java.util.List; import java.util.Map; -import java.util.stream.Collectors; import static org.opensearch.sql.ppl.utils.DataTypeTransformer.seq; @@ -40,9 +41,9 @@ abstract class SerializableAbstractFunction2 extends AbstractFunction * @param keysToRemove The list of keys to remove. * @return A new JSON string without the specified keys. */ - Function2, String> jsonDeleteFunction = new SerializableAbstractFunction2<>() { + Function2, String> jsonDeleteFunction = new SerializableAbstractFunction2<>() { @Override - public String apply(String jsonStr, List keysToRemove) { + public String apply(String jsonStr, WrappedArray keysToRemove) { if (jsonStr == null) { return null; } @@ -55,8 +56,9 @@ public String apply(String jsonStr, List keysToRemove) { } } - private void removeKeys(Map map, List keysToRemove) { - for (String key : keysToRemove) { + private void removeKeys(Map map, WrappedArray keysToRemove) { + Collection keys = JavaConverters.asJavaCollection(keysToRemove); + for (String key : keys) { String[] keyParts = key.split("\\."); Map currentMap = map; for (int i = 0; i < keyParts.length - 1; i++) { @@ -64,10 +66,9 @@ private void removeKeys(Map map, List keysToRemove) { if (currentMap.containsKey(currentKey) && currentMap.get(currentKey) instanceof Map) { currentMap = (Map) currentMap.get(currentKey); } else { - return; // Path not found, exit + return; } } - // Remove the final key if it exists currentMap.remove(keyParts[keyParts.length - 1]); } } @@ -107,7 +108,7 @@ public String apply(String jsonStr, List> pathValuePai return objectMapper.writeValueAsString(jsonMap); } catch (Exception e) { - return null; // Return null if parsing fails + return null; } } }; @@ -143,22 +144,11 @@ public String apply(String jsonStr, List>> pathVa return objectMapper.writeValueAsString(jsonMap); } catch (Exception e) { - return null; // Return null if parsing fails + return null; } } }; - - /** - * Check if a key matches the given path expression. - * - * @param key The key to check. - * @param path The path expression (e.g., "a.b"). - * @return True if the key matches, false otherwise. - */ - private static boolean matchesKey(String key, String path) { - return key.equals(path) || key.startsWith(path + "."); - } - + Function2 cidrFunction = new SerializableAbstractFunction2<>() { IPAddressStringParameters valOptions = new IPAddressStringParameters.Builder() diff --git a/ppl-spark-integration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java b/ppl-spark-integration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java index 996569611..8572e01fb 100644 --- a/ppl-spark-integration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java +++ b/ppl-spark-integration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java @@ -5,6 +5,7 @@ package org.opensearch.sql.expression.function; import org.junit.Test; +import scala.collection.mutable.WrappedArray; import java.util.ArrayList; import java.util.Arrays; @@ -12,7 +13,6 @@ import java.util.List; import java.util.Map; -import static java.util.Collections.singletonList; import static org.apache.derby.vti.XmlVTI.asList; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; @@ -26,7 +26,7 @@ public class SerializableJsonUdfTest { public void testJsonDeleteFunctionRemoveSingleKey() { String jsonStr = "{\"key1\":\"value1\",\"key2\":\"value2\",\"key3\":\"value3\"}"; String expectedJson = "{\"key1\":\"value1\",\"key3\":\"value3\"}"; - String result = jsonDeleteFunction.apply(jsonStr, singletonList("key2")); + String result = jsonDeleteFunction.apply(jsonStr, WrappedArray.make(new String[]{"key2"})); assertEquals(expectedJson, result); } @@ -35,7 +35,7 @@ public void testJsonDeleteFunctionRemoveNestedKey() { // Correctly escape double quotes within the JSON string String jsonStr = "{\"key1\":\"value1\",\"key2\":{ \"key3\":\"value3\",\"key4\":\"value4\" }}"; String expectedJson = "{\"key1\":\"value1\",\"key2\":{\"key4\":\"value4\"}}"; - String result = jsonDeleteFunction.apply(jsonStr, singletonList("key2.key3")); + String result = jsonDeleteFunction.apply(jsonStr, WrappedArray.make(new String[]{"key2.key3"})); assertEquals(expectedJson, result); } @@ -43,7 +43,7 @@ public void testJsonDeleteFunctionRemoveNestedKey() { public void testJsonDeleteFunctionRemoveSingleArrayedKey() { String jsonStr = "{\"key1\":\"value1\",\"key2\":\"value2\",\"keyArray\":[\"value1\",\"value2\"]}"; String expectedJson = "{\"key1\":\"value1\",\"key2\":\"value2\"}"; - String result = jsonDeleteFunction.apply(jsonStr, singletonList("keyArray")); + String result = jsonDeleteFunction.apply(jsonStr, WrappedArray.make(new String[]{"keyArray"})); assertEquals(expectedJson, result); } @@ -51,7 +51,7 @@ public void testJsonDeleteFunctionRemoveSingleArrayedKey() { public void testJsonDeleteFunctionRemoveMultipleKeys() { String jsonStr = "{\"key1\":\"value1\",\"key2\":\"value2\",\"key3\":\"value3\"}"; String expectedJson = "{\"key3\":\"value3\"}"; - String result = jsonDeleteFunction.apply(jsonStr, Arrays.asList("key1", "key2")); + String result = jsonDeleteFunction.apply(jsonStr, WrappedArray.make(new String[]{"key1", "key2"})); assertEquals(expectedJson, result); } @@ -59,27 +59,27 @@ public void testJsonDeleteFunctionRemoveMultipleKeys() { public void testJsonDeleteFunctionRemoveMultipleSomeAreNestedKeys() { String jsonStr = "{\"key1\":\"value1\",\"key2\":{ \"key3\":\"value3\",\"key4\":\"value4\" }}"; String expectedJson = "{\"key2\":{\"key3\":\"value3\"}}"; - String result = jsonDeleteFunction.apply(jsonStr, Arrays.asList("key1", "key2.key4")); + String result = jsonDeleteFunction.apply(jsonStr, WrappedArray.make(new String[]{"key1", "key2.key4"})); assertEquals(expectedJson, result); } @Test public void testJsonDeleteFunctionNoKeysRemoved() { String jsonStr = "{\"key1\":\"value1\",\"key2\":\"value2\"}"; - String result = jsonDeleteFunction.apply(jsonStr, Collections.emptyList()); + String result = jsonDeleteFunction.apply(jsonStr, WrappedArray.make(new String[0])); assertEquals(jsonStr, result); } @Test public void testJsonDeleteFunctionNullJson() { - String result = jsonDeleteFunction.apply(null, Collections.singletonList("key1")); + String result = jsonDeleteFunction.apply(null, WrappedArray.make(new String[]{"key1"})); assertNull(result); } @Test public void testJsonDeleteFunctionInvalidJson() { String invalidJson = "invalid_json"; - String result = jsonDeleteFunction.apply(invalidJson, Collections.singletonList("key1")); + String result = jsonDeleteFunction.apply(invalidJson, WrappedArray.make(new String[]{"key1"})); assertNull(result); } @@ -186,67 +186,43 @@ public void testJsonExtendFunctionWithNonArrayPath() { @Test public void testJsonExtendFunctionAddValuesToExistingArray() { - // Initial JSON string String jsonStr = "{\"key1\":\"value1\",\"key2\":[\"value2\"]}"; - - // Path-value pairs to extend List>> pathValuePairs = new ArrayList<>(); pathValuePairs.add( Map.entry("key2", Arrays.asList("value3", "value4"))); - // Expected JSON after extension String expectedJson = "{\"key1\":\"value1\",\"key2\":[\"value2\",\"value3\",\"value4\"]}"; - - // Apply the function String result = jsonExtendFunction.apply(jsonStr, pathValuePairs); - // Assert that the result matches the expected JSON assertEquals(expectedJson, result); } @Test public void testJsonExtendFunctionAddNewArray() { - // Initial JSON string String jsonStr = "{\"key1\":\"value1\"}"; - - // Path-value pairs to add List>> pathValuePairs = new ArrayList<>(); pathValuePairs.add( Map.entry("key2", Arrays.asList("value2", "value3"))); - // Expected JSON after adding new array String expectedJson = "{\"key1\":\"value1\",\"key2\":[\"value2\",\"value3\"]}"; - - // Apply the function String result = jsonExtendFunction.apply(jsonStr, pathValuePairs); - // Assert that the result matches the expected JSON assertEquals(expectedJson, result); } @Test public void testJsonExtendFunctionHandleEmptyValues() { - // Initial JSON string String jsonStr = "{\"key1\":\"value1\",\"key2\":[\"value2\"]}"; - - // Path-value pairs with an empty list of values to add List>> pathValuePairs = new ArrayList<>(); pathValuePairs.add( Map.entry("key2", Collections.emptyList())); - // Expected JSON should remain unchanged String expectedJson = "{\"key1\":\"value1\",\"key2\":[\"value2\"]}"; - - // Apply the function String result = jsonExtendFunction.apply(jsonStr, pathValuePairs); - // Assert that the result matches the expected JSON assertEquals(expectedJson, result); } @Test public void testJsonExtendFunctionHandleNullInput() { - // Apply the function with null input String result = jsonExtendFunction.apply(null, Collections.singletonList( Map.entry("key2", List.of("value2")))); - - // Assert that the result is null assertEquals(null, result); } }