diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java index d5f79b8946c..f82acdb14a4 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java @@ -246,6 +246,7 @@ public enum BuiltinFunctionName { JSON_ARRAY(FunctionName.of("json_array")), JSON_ARRAY_LENGTH(FunctionName.of("json_array_length")), JSON_EXTRACT(FunctionName.of("json_extract")), + JSON_EXTRACT_ALL(FunctionName.of("json_extract_all"), true), JSON_KEYS(FunctionName.of("json_keys")), JSON_SET(FunctionName.of("json_set")), JSON_DELETE(FunctionName.of("json_delete")), diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java index 1097a9a5b0d..4692aac996e 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java @@ -52,6 +52,7 @@ import org.opensearch.sql.expression.function.jsonUDF.JsonArrayLengthFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonDeleteFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonExtendFunctionImpl; +import org.opensearch.sql.expression.function.jsonUDF.JsonExtractAllFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonExtractFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonFunctionImpl; import org.opensearch.sql.expression.function.jsonUDF.JsonKeysFunctionImpl; @@ -111,6 +112,8 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable { new JsonArrayLengthFunctionImpl().toUDF("JSON_ARRAY_LENGTH"); public static final SqlOperator JSON_EXTRACT = new JsonExtractFunctionImpl().toUDF("JSON_EXTRACT"); + public static final SqlOperator JSON_EXTRACT_ALL = + new JsonExtractAllFunctionImpl().toUDF("JSON_EXTRACT_ALL"); public static final SqlOperator JSON_KEYS = new JsonKeysFunctionImpl().toUDF("JSON_KEYS"); public static final SqlOperator JSON_SET = new JsonSetFunctionImpl().toUDF("JSON_SET"); public static final SqlOperator JSON_DELETE = new JsonDeleteFunctionImpl().toUDF("JSON_DELETE"); diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java index eeeed029b69..9245a719c66 100644 --- a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java +++ b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java @@ -99,6 +99,7 @@ import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_DELETE; import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_EXTEND; import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_EXTRACT; +import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_EXTRACT_ALL; import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_KEYS; import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_OBJECT; import static org.opensearch.sql.expression.function.BuiltinFunctionName.JSON_SET; @@ -867,6 +868,7 @@ void populate() { registerOperator(JSON_DELETE, PPLBuiltinOperators.JSON_DELETE); registerOperator(JSON_APPEND, PPLBuiltinOperators.JSON_APPEND); registerOperator(JSON_EXTEND, PPLBuiltinOperators.JSON_EXTEND); + registerOperator(JSON_EXTRACT_ALL, PPLBuiltinOperators.JSON_EXTRACT_ALL); // internal // Register operators with a different type checker diff --git a/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImpl.java b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImpl.java new file mode 100644 index 00000000000..1f91c87bb77 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImpl.java @@ -0,0 +1,218 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.jsonUDF; + +import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY; + +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import java.io.IOException; +import java.util.Collection; +import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Stack; +import org.apache.calcite.adapter.enumerable.NotNullImplementor; +import org.apache.calcite.adapter.enumerable.NullPolicy; +import org.apache.calcite.adapter.enumerable.RexImpTable; +import org.apache.calcite.adapter.enumerable.RexToLixTranslator; +import org.apache.calcite.linq4j.tree.Expression; +import org.apache.calcite.linq4j.tree.Types; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.schema.impl.ScalarFunctionImpl; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlReturnTypeInference; +import org.apache.calcite.sql.type.SqlTypeFamily; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.sql.expression.function.ImplementorUDF; +import org.opensearch.sql.expression.function.UDFOperandMetadata; + +/** + * UDF which extract all the fields from JSON to a MAP. Items are collected from input JSON and + * stored with the key of their path in the JSON. This UDF is designed to be used for `spath` + * command without path param. See {@ref JsonExtractAllFunctionImplTest} for the detailed spec. + */ +public class JsonExtractAllFunctionImpl extends ImplementorUDF { + private static final String ARRAY_SUFFIX = "{}"; + private static final JsonFactory JSON_FACTORY = new JsonFactory(); + + public JsonExtractAllFunctionImpl() { + super(new JsonExtractAllImplementor(), NullPolicy.ANY); + } + + @Override + public SqlReturnTypeInference getReturnTypeInference() { + return ReturnTypes.explicit( + TYPE_FACTORY.createMapType( + TYPE_FACTORY.createSqlType(SqlTypeName.VARCHAR), + TYPE_FACTORY.createSqlType(SqlTypeName.ANY), + true)); + } + + @Override + public UDFOperandMetadata getOperandMetadata() { + return UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.STRING)); + } + + public static class JsonExtractAllImplementor implements NotNullImplementor { + @Override + public Expression implement( + RexToLixTranslator translator, RexCall call, List translatedOperands) { + ScalarFunctionImpl function = + (ScalarFunctionImpl) + ScalarFunctionImpl.create( + Types.lookupMethod(JsonExtractAllFunctionImpl.class, "eval", Object[].class)); + return function.getImplementor().implement(translator, call, RexImpTable.NullAs.NULL); + } + } + + public static Object eval(Object... args) { + if (args.length < 1) { + return null; + } + + String jsonStr = (String) args[0]; + if (jsonStr == null || jsonStr.trim().isEmpty()) { + return null; + } + + return parseJson(jsonStr); + } + + private static Map parseJson(String jsonStr) { + Map resultMap = new HashMap<>(); + Stack pathStack = new Stack<>(); + + try (JsonParser parser = JSON_FACTORY.createParser(jsonStr)) { + JsonToken token; + + while ((token = parser.nextToken()) != null) { + switch (token) { + case START_OBJECT: + break; + + case END_OBJECT: + if (!pathStack.isEmpty() && !isInArray(pathStack)) { + pathStack.pop(); + } + break; + + case START_ARRAY: + pathStack.push(ARRAY_SUFFIX); + break; + + case END_ARRAY: + pathStack.pop(); + if (!pathStack.isEmpty() && !isInArray(pathStack)) { + pathStack.pop(); + } + break; + + case FIELD_NAME: + String fieldName = parser.currentName(); + pathStack.push(fieldName); + break; + + case VALUE_STRING: + case VALUE_NUMBER_INT: + case VALUE_NUMBER_FLOAT: + case VALUE_TRUE: + case VALUE_FALSE: + case VALUE_NULL: + if (pathStack.isEmpty()) { + // ignore top level value + return null; + } + + appendValue(resultMap, buildPath(pathStack), extractValue(parser, token)); + + if (!isInArray(pathStack)) { + pathStack.pop(); + } + break; + default: + // Skip other tokens + break; + } + } + } catch (IOException e) { + // ignore exception, and current result will be returned + } + return resultMap; + } + + @SuppressWarnings("unchecked") + private static void appendValue(Map resultMap, String path, Object value) { + Object existingValue = resultMap.get(path); + if (existingValue == null) { + resultMap.put(path, value); + } else if (existingValue instanceof List) { + ((List) existingValue).add(value); + } else { + resultMap.put(path, list(existingValue, value)); + } + } + + private static List list(Object... args) { + List result = new LinkedList<>(); + for (Object arg : args) { + result.add(arg); + } + return result; + } + + private static boolean isInArray(List path) { + return path.size() >= 1 && path.getLast().equals(ARRAY_SUFFIX); + } + + private static Object extractValue(JsonParser parser, JsonToken token) throws IOException { + switch (token) { + case VALUE_STRING: + return parser.getValueAsString(); + case VALUE_NUMBER_INT: + return getIntValue(parser); + case VALUE_NUMBER_FLOAT: + return parser.getDoubleValue(); + case VALUE_TRUE: + return true; + case VALUE_FALSE: + return false; + case VALUE_NULL: + return null; + default: + return parser.getValueAsString(); + } + } + + private static Object getIntValue(JsonParser parser) throws IOException { + if (parser.getNumberType() == JsonParser.NumberType.INT) { + return parser.getIntValue(); + } else if (parser.getNumberType() == JsonParser.NumberType.LONG) { + return parser.getLongValue(); + } else { + // store as double in case of BIG_INTEGER (exceed LONG precision) + return parser.getBigIntegerValue().doubleValue(); + } + } + + private static String buildPath(Collection pathStack) { + StringBuilder builder = new StringBuilder(); + for (String path : pathStack) { + if (ARRAY_SUFFIX.equals(path)) { + builder.append(ARRAY_SUFFIX); + } else if (!path.isEmpty()) { + if (!builder.isEmpty()) { + builder.append("."); + } + builder.append(path); + } + } + return builder.toString(); + } +} diff --git a/core/src/test/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImplTest.java b/core/src/test/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImplTest.java new file mode 100644 index 00000000000..5a010a17422 --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImplTest.java @@ -0,0 +1,360 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.expression.function.jsonUDF; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +public class JsonExtractAllFunctionImplTest { + + private final JsonExtractAllFunctionImpl function = new JsonExtractAllFunctionImpl(); + + @SuppressWarnings("unchecked") + private Map assertValidMapResult(Object result) { + assertNotNull(result); + assertTrue(result instanceof Map); + return (Map) result; + } + + @SuppressWarnings("unchecked") + private List assertListValue(Map map, String key) { + Object value = map.get(key); + assertNotNull(value); + assertTrue(value instanceof List); + return (List) value; + } + + private void assertListEquals(List actual, Object... expected) { + assertEquals(expected.length, actual.size()); + for (int i = 0; i < expected.length; i++) { + assertEquals(expected[i], actual.get(i)); + } + } + + private void assertMapListValue(Map map, String key, Object... expectedValues) { + List list = assertListValue(map, key); + assertListEquals(list, expectedValues); + } + + private void assertMapValue(Map map, String key, Object expectedValue) { + assertEquals(expectedValue, map.get(key)); + } + + private Map eval(String json) { + Object result = JsonExtractAllFunctionImpl.eval(json); + return assertValidMapResult(result); + } + + @Test + public void testReturnTypeInference() { + assertNotNull(function.getReturnTypeInference(), "Return type inference should not be null"); + } + + @Test + public void testOperandMetadata() { + assertNotNull(function.getOperandMetadata(), "Operand metadata should not be null"); + } + + @Test + public void testFunctionConstructor() { + JsonExtractAllFunctionImpl testFunction = new JsonExtractAllFunctionImpl(); + + assertNotNull(testFunction, "Function should be properly initialized"); + } + + @Test + public void testNoArguments() { + Object result = JsonExtractAllFunctionImpl.eval(); + + assertNull(result); + } + + @Test + public void testNullInput() { + Object result = JsonExtractAllFunctionImpl.eval((String) null); + + assertNull(result); + } + + @Test + public void testEmptyString() { + Object result = JsonExtractAllFunctionImpl.eval(""); + + assertNull(result); + } + + @Test + public void testWhitespaceString() { + Object result = JsonExtractAllFunctionImpl.eval(" "); + + assertNull(result); + } + + @Test + public void testEmptyJsonObject() { + Map map = eval("{}"); + + assertTrue(map.isEmpty()); + } + + @Test + public void testSimpleJsonObject() throws Exception { + Map map = eval("{\"name\": \"John\", \"age\": 30}"); + + assertEquals("John", map.get("name")); + assertEquals(30, map.get("age")); + assertEquals(2, map.size()); + } + + @Test + public void testInvalidJsonReturnResults() { + Map map = eval("{\"name\": \"John\", \"age\":}"); + + assertEquals("John", map.get("name")); + assertEquals(1, map.size()); + } + + @Test + public void testNonObjectJsonArray() { + Map map = eval("[1, 2, 3]"); + + assertMapListValue(map, "{}", 1, 2, 3); + assertEquals(1, map.size()); + } + + @Test + public void testTopLevelArrayOfObjects() { + Map map = eval("[{\"age\": 1}, {\"age\": 2}]"); + + assertMapListValue(map, "{}.age", 1, 2); + assertEquals(1, map.size()); + } + + @Test + public void testTopLevelArrayOfComplexObjects() { + Map map = + eval("[{\"name\": \"John\", \"age\": 30}, {\"name\": \"Jane\", \"age\": 25}]"); + + assertMapListValue(map, "{}.name", "John", "Jane"); + assertMapListValue(map, "{}.age", 30, 25); + assertEquals(2, map.size()); + } + + @Test + public void testNonObjectJsonPrimitive() { + Object result = JsonExtractAllFunctionImpl.eval("\"just a string\""); + + assertNull(result); + } + + @Test + public void testNonObjectJsonNumber() { + Object result = JsonExtractAllFunctionImpl.eval("42"); + + assertNull(result); + } + + @Test + public void testSingleLevelNesting() { + Map map = eval("{\"user\": {\"name\": \"John\"}, \"system\": \"linux\"}"); + + assertEquals("John", map.get("user.name")); + assertEquals("linux", map.get("system")); + assertEquals(2, map.size()); + } + + @Test + public void testMultiLevelNesting() { + Map map = eval("{\"a\": {\"b\": {\"c\": \"value\"}}}"); + + assertEquals("value", map.get("a.b.c")); + assertEquals(1, map.size()); + } + + @Test + public void testMixedNestedAndFlat() { + Map map = + eval("{\"name\": \"John\", \"address\": {\"city\": \"NYC\", \"zip\": \"10001\"}}"); + + assertEquals("John", map.get("name")); + assertEquals("NYC", map.get("address.city")); + assertEquals("10001", map.get("address.zip")); + assertEquals(3, map.size()); + } + + @Test + public void testDeeplyNestedStructure() { + Map map = + eval("{\"level1\": {\"level2\": {\"level3\": {\"level4\": {\"level5\": \"deep\"}}}}}"); + + assertEquals("deep", map.get("level1.level2.level3.level4.level5")); + assertEquals(1, map.size()); + } + + @Test + public void testSimpleArray() { + Map map = eval("{\"tags\": [\"a\", \"b\", \"c\"]}"); + + assertMapListValue(map, "tags{}", "a", "b", "c"); + assertEquals(1, map.size()); + } + + @Test + public void testArrayOfObjects() { + Map map = eval("{\"users\": [{\"name\": \"John\"}, {\"name\": \"Jane\"}]}"); + + assertMapListValue(map, "users{}.name", "John", "Jane"); + assertEquals(1, map.size()); + } + + @Test + public void testNestedArray() { + Map map = eval("{\"data\": {\"items\": [1, 2, 3]}}"); + + assertMapListValue(map, "data.items{}", 1, 2, 3); + assertEquals(1, map.size()); + } + + @Test + public void testNested() { + Map map = + eval( + "{\"data\": {\"items\": [[1, 2, {\"hello\": 3}], 4], \"other\": 5}, \"another\": [6," + + " [7, 8], 9]}"); + + assertMapListValue(map, "data.items{}{}", 1, 2); + assertMapValue(map, "data.items{}{}.hello", 3); + assertMapValue(map, "data.items{}", 4); + assertMapValue(map, "data.other", 5); + assertMapListValue(map, "another{}", 6, 9); + assertMapListValue(map, "another{}{}", 7, 8); + assertEquals(6, map.size()); + } + + @Test + public void testEmptyArray() { + Map map = eval("{\"empty\": []}"); + + Object emptyValue = map.get("empty{}"); + assertNull(emptyValue); + } + + @Test + public void testStringValues() { + Map map = eval("{\"text\": \"hello world\", \"empty\": \"\"}"); + + assertMapValue(map, "text", "hello world"); + assertMapValue(map, "empty", ""); + assertEquals(2, map.size()); + } + + @Test + public void testNumericValues() { + Map map = + eval( + "{\"int\": 42, \"long\": 9223372036854775807, \"hugeNumber\": 9223372036854775808," + + " \"double\": 3.14159}"); + + assertEquals(4, map.size()); + assertEquals(42, map.get("int")); + assertEquals(9223372036854775807L, map.get("long")); + assertEquals(9223372036854775808.0, map.get("hugeNumber")); + assertEquals(3.14159, map.get("double")); + } + + @Test + public void testBooleanValues() { + Map map = eval("{\"isTrue\": true, \"isFalse\": false}"); + + assertEquals(true, map.get("isTrue")); + assertEquals(false, map.get("isFalse")); + assertEquals(2, map.size()); + } + + @Test + public void testNullValues() { + Map map = eval("{\"nullValue\": null, \"notNull\": \"value\"}"); + + assertNull(map.get("nullValue")); + assertEquals("value", map.get("notNull")); + assertEquals(2, map.size()); + } + + @Test + public void testMixedTypesInArray() { + Map map = eval("{\"mixed\": [\"string\", 42, true, null, 3.14]}"); + + List mixed = (List) assertListValue(map, "mixed{}"); + assertEquals(5, mixed.size()); + assertEquals("string", mixed.get(0)); + assertEquals(42, mixed.get(1)); + assertEquals(true, mixed.get(2)); + assertNull(mixed.get(3)); + assertEquals(3.14, mixed.get(4)); + assertEquals(1, map.size()); + } + + @Test + public void testSpecialCharactersInKeys() { + Map map = + eval( + "{\"key.with.dots\": \"value1\", \"key-with-dashes\": \"value2\"," + + " \"key_with_underscores\": \"value3\"}"); + + assertEquals("value1", map.get("key.with.dots")); + assertEquals("value2", map.get("key-with-dashes")); + assertEquals("value3", map.get("key_with_underscores")); + assertEquals(3, map.size()); + } + + @Test + public void testUnicodeCharacters() { + Map map = eval("{\"unicode\": \"こんにちは\", \"emoji\": \"🚀\", \"🚀\": 1}"); + + assertEquals("こんにちは", map.get("unicode")); + assertEquals("🚀", map.get("emoji")); + assertEquals(1, map.get("🚀")); + assertEquals(3, map.size()); + } + + @Test + public void testComplexNestedStructure() { + Map map = + eval( + "{\"user\": {\"profile\": {\"name\": \"John\", \"contacts\": [{\"type\": \"email\"," + + " \"value\": \"john@example.com\"}, {\"type\": \"phone\", \"value\":" + + " \"123-456-7890\"}]}, \"preferences\": {\"theme\": \"dark\", \"notifications\":" + + " true}}}"); + + assertEquals("John", map.get("user.profile.name")); + assertMapListValue(map, "user.profile.contacts{}.type", "email", "phone"); + assertMapListValue(map, "user.profile.contacts{}.value", "john@example.com", "123-456-7890"); + assertEquals("dark", map.get("user.preferences.theme")); + assertEquals(true, map.get("user.preferences.notifications")); + assertEquals(5, map.size()); + } + + @Test + public void testLargeJsonObject() { + StringBuilder jsonBuilder = new StringBuilder("{"); + for (int i = 0; i < 100; i++) { + if (i > 0) jsonBuilder.append(","); + jsonBuilder.append("\"field").append(i).append("\": ").append(i); + } + jsonBuilder.append("}"); + + Map map = eval(jsonBuilder.toString()); + assertEquals(100, map.size()); + assertEquals(0, map.get("field0")); + assertEquals(99, map.get("field99")); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/JsonExtractAllFunctionIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/JsonExtractAllFunctionIT.java new file mode 100644 index 00000000000..b3819aa55ec --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/standalone/JsonExtractAllFunctionIT.java @@ -0,0 +1,372 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.standalone; + +import java.io.IOException; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.calcite.plan.Contexts; +import org.apache.calcite.plan.RelTraitDef; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.schema.SchemaPlus; +import org.apache.calcite.sql.parser.SqlParser; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.tools.Frameworks; +import org.apache.calcite.tools.Programs; +import org.apache.calcite.tools.RelBuilder; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.calcite.CalcitePlanContext; +import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelRunners; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.executor.QueryType; +import org.opensearch.sql.expression.function.BuiltinFunctionName; +import org.opensearch.sql.expression.function.PPLFuncImpTable; + +/** Integration test for internal function JSON_EXTRACT_ALL in Calcite PPL. */ +public class JsonExtractAllFunctionIT extends CalcitePPLIntegTestCase { + + private static final String RESULT_FIELD = "result"; + private static final String ID_FIELD = "id"; + TestContext context; + + @Override + public void init() throws IOException { + super.init(); + context = createTestContext(); + enableCalcite(); + } + + @Test + public void testJsonExtractAllWithNullInput() throws Exception { + RelDataType stringType = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR); + RexNode nullJson = context.rexBuilder.makeNullLiteral(stringType); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, nullJson); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + assertNull(resultSet.getObject(1)); + }); + } + + @Test + public void testJsonExtractAllWithSimpleObject() throws Exception { + String jsonString = "{\"name\": \"John\", \"age\": 30}"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(jsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + assertEquals("John", map.get("name")); + assertEquals(30, map.get("age")); + assertEquals(2, map.size()); + }); + } + + private Map getMap(ResultSet resultSet, int columnIndex) throws SQLException { + Object result = resultSet.getObject(columnIndex); + assertNotNull(result); + assertTrue(result instanceof Map); + + @SuppressWarnings("unchecked") + Map map = (Map) result; + System.out.println( + "map: " + + map.entrySet().stream() + .map(entry -> entry.getKey() + "=" + entry.getValue()) + .collect(Collectors.joining(", "))); + return map; + } + + @Test + public void testJsonExtractAllWithNestedObject() throws Exception { + String jsonString = "{\"user\": {\"name\": \"John\", \"age\": 30}, \"active\": true}"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(jsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + assertEquals("John", map.get("user.name")); + assertEquals(30, map.get("user.age")); + assertEquals(true, map.get("active")); + assertEquals(3, map.size()); + }); + } + + @Test + public void testJsonExtractAllWithArray() throws Exception { + String jsonString = "{\"tags\": [\"java\", \"sql\", \"opensearch\"]}"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(jsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + List tags = getList(map, "tags{}"); + + assertEquals(3, tags.size()); + assertEquals("java", tags.get(0)); + assertEquals("sql", tags.get(1)); + assertEquals("opensearch", tags.get(2)); + }); + } + + @Test + public void testJsonExtractAllWithArrayOfObjects() throws Exception { + String jsonString = "{\"users\": [{\"name\": \"John\"}, {\"name\": \"Jane\"}]}"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(jsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + List names = getList(map, "users{}.name"); + assertEquals(2, names.size()); + assertEquals("John", names.get(0)); + assertEquals("Jane", names.get(1)); + assertEquals(1, map.size()); // Only flattened key should exist + }); + } + + @Test + public void testJsonExtractAllWithTopLevelArray() throws Exception { + String jsonString = "[{\"id\": 1}, {\"id\": 2}]"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(jsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + List ids = getList(map, "{}.id"); + assertEquals(2, ids.size()); + assertEquals(1, ids.get(0)); + assertEquals(2, ids.get(1)); + assertEquals(1, map.size()); + }); + } + + @SuppressWarnings("unchecked") + private List getList(Map map, String key) { + Object value = map.get(key); + assertNotNull(value); + assertTrue(value instanceof List); + + return (List) value; + } + + @Test + public void testJsonExtractAllWithEmptyObject() throws Exception { + String jsonString = "{}"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(jsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + assertTrue(map.isEmpty()); + }); + } + + @Test + public void testJsonExtractAllWithInvalidJson() throws Exception { + String invalidJsonString = "{\"name\": \"John\", \"age\":}"; + RexNode jsonLiteral = context.rexBuilder.makeLiteral(invalidJsonString); + + RexNode jsonExtractAllCall = + PPLFuncImpTable.INSTANCE.resolve( + context.rexBuilder, BuiltinFunctionName.JSON_EXTRACT_ALL, jsonLiteral); + + RelNode relNode = + context + .relBuilder + .values(new String[] {ID_FIELD}, 1) + .project(context.relBuilder.alias(jsonExtractAllCall, RESULT_FIELD)) + .build(); + + executeRelNodeAndVerify( + context.planContext, + relNode, + resultSet -> { + assertTrue(resultSet.next()); + verifyColumns(resultSet, RESULT_FIELD); + + Map map = getMap(resultSet, 1); + assertEquals("John", map.get("name")); + assertEquals(1, map.size()); + }); + } + + @FunctionalInterface + private interface ResultVerifier { + void verify(ResultSet resultSet) throws SQLException; + } + + private static class TestContext { + final CalcitePlanContext planContext; + final RelBuilder relBuilder; + final RexBuilder rexBuilder; + + TestContext(CalcitePlanContext planContext, RelBuilder relBuilder, RexBuilder rexBuilder) { + this.planContext = planContext; + this.relBuilder = relBuilder; + this.rexBuilder = rexBuilder; + } + } + + private TestContext createTestContext() { + CalcitePlanContext planContext = createCalcitePlanContext(); + return new TestContext(planContext, planContext.relBuilder, planContext.rexBuilder); + } + + private void executeRelNodeAndVerify( + CalcitePlanContext planContext, RelNode relNode, ResultVerifier verifier) + throws SQLException { + try (PreparedStatement statement = OpenSearchRelRunners.run(planContext, relNode)) { + ResultSet resultSet = statement.executeQuery(); + verifier.verify(resultSet); + } + } + + private void verifyColumns(ResultSet resultSet, String... expectedColumnNames) + throws SQLException { + assertEquals(expectedColumnNames.length, resultSet.getMetaData().getColumnCount()); + + for (int i = 0; i < expectedColumnNames.length; i++) { + String expectedName = expectedColumnNames[i]; + String actualName = resultSet.getMetaData().getColumnName(i + 1); + assertEquals(expectedName, actualName); + } + } + + private CalcitePlanContext createCalcitePlanContext() { + // Create a Frameworks.ConfigBuilder similar to CalcitePPLAbstractTest + final SchemaPlus rootSchema = Frameworks.createRootSchema(true); + Frameworks.ConfigBuilder config = + Frameworks.newConfigBuilder() + .parserConfig(SqlParser.Config.DEFAULT) + .defaultSchema(rootSchema) + .traitDefs((List) null) + .programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2)); + + config.context(Contexts.of(RelBuilder.Config.DEFAULT)); + + Settings settings = getSettings(); + return CalcitePlanContext.create( + config.build(), settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), QueryType.PPL); + } +}