Skip to content

Commit

Permalink
update IT tests for json_delete
Browse files Browse the repository at this point in the history 
Signed-off-by: YANGDB <[email protected]>
  • Loading branch information
@YANG-DB
YANG-DB committed Dec 9, 2024
1 parent bf6e607 commit bb8bd30
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 55 deletions.
47 changes: 47 additions & 0 deletions ...c/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}

}
34 changes: 12 additions & 22 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 @@ -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;
Expand All @@ -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;

Expand All @@ -40,9 +41,9 @@ abstract class SerializableAbstractFunction2<T1, T2, R> extends AbstractFunction
* @param keysToRemove The list of keys to remove.
* @return A new JSON string without the specified keys.
*/
Function2<String, List<String>, String> jsonDeleteFunction = new SerializableAbstractFunction2<>() {
Function2<String, WrappedArray<String>, String> jsonDeleteFunction = new SerializableAbstractFunction2<>() {
@Override
public String apply(String jsonStr, List<String> keysToRemove) {
public String apply(String jsonStr, WrappedArray<String> keysToRemove) {
if (jsonStr == null) {
return null;
}
Expand All @@ -55,19 +56,19 @@ public String apply(String jsonStr, List<String> keysToRemove) {
}
}

private void removeKeys(Map<String, Object> map, List<String> keysToRemove) {
for (String key : keysToRemove) {
private void removeKeys(Map<String, Object> map, WrappedArray<String> keysToRemove) {
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; // Path not found, exit
return;
}
}
// Remove the final key if it exists
currentMap.remove(keyParts[keyParts.length - 1]);
}
}
Expand Down Expand Up @@ -107,7 +108,7 @@ public String apply(String jsonStr, List<Map.Entry<String, String>> pathValuePai

return objectMapper.writeValueAsString(jsonMap);
} catch (Exception e) {
return null; // Return null if parsing fails
return null;
}
}
};
Expand Down Expand Up @@ -143,22 +144,11 @@ public String apply(String jsonStr, List<Map.Entry<String, List<String>>> 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<String, String, Boolean> cidrFunction = new SerializableAbstractFunction2<>() {

IPAddressStringParameters valOptions = new IPAddressStringParameters.Builder()
Expand Down
42 changes: 9 additions & 33 deletions ...gration/src/test/java/org/opensearch/sql/expression/function/SerializableJsonUdfTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,14 @@
package org.opensearch.sql.expression.function;

import org.junit.Test;
import scala.collection.mutable.WrappedArray;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
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;
Expand All @@ -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);
}

Expand All @@ -35,51 +35,51 @@ 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);
}

@Test
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);
}

@Test
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);
}

@Test
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);
}

Expand Down Expand Up @@ -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<Map.Entry<String, List<String>>> 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<Map.Entry<String, List<String>>> 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<Map.Entry<String, List<String>>> 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);
}
}

0 comments on commit bb8bd30

Please sign in to comment.