diff --git a/spark/src/main/java/org/opensearch/sql/spark/functions/response/DefaultSparkSqlFunctionResponseHandle.java b/spark/src/main/java/org/opensearch/sql/spark/functions/response/DefaultSparkSqlFunctionResponseHandle.java index 9d0cd59cf8..8fc417cd80 100644 --- a/spark/src/main/java/org/opensearch/sql/spark/functions/response/DefaultSparkSqlFunctionResponseHandle.java +++ b/spark/src/main/java/org/opensearch/sql/spark/functions/response/DefaultSparkSqlFunctionResponseHandle.java @@ -86,7 +86,7 @@ private static LinkedHashMap extractRow( linkedHashMap.put( column.getName(), new ExprTimestampValue(row.getString(column.getName()))); } else if (type == ExprCoreType.STRING) { - linkedHashMap.put(column.getName(), new ExprStringValue(row.getString(column.getName()))); + linkedHashMap.put(column.getName(), new ExprStringValue(jsonString(row, column.getName()))); } else { throw new RuntimeException("Result contains invalid data type"); } @@ -137,6 +137,10 @@ private ExprCoreType getDataType(String sparkDataType) { } } + private static String jsonString(JSONObject jsonObject, String key) { + return jsonObject.has(key) ? jsonObject.getString(key) : ""; + } + @Override public boolean hasNext() { return responseIterator.hasNext(); diff --git a/spark/src/test/java/org/opensearch/sql/spark/functions/SparkSqlFunctionTableScanOperatorTest.java b/spark/src/test/java/org/opensearch/sql/spark/functions/SparkSqlFunctionTableScanOperatorTest.java index 586f0ef2d8..188cd695a3 100644 --- a/spark/src/test/java/org/opensearch/sql/spark/functions/SparkSqlFunctionTableScanOperatorTest.java +++ b/spark/src/test/java/org/opensearch/sql/spark/functions/SparkSqlFunctionTableScanOperatorTest.java @@ -10,6 +10,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.when; +import static org.opensearch.sql.data.model.ExprValueUtils.stringValue; import static org.opensearch.sql.spark.constants.TestConstants.QUERY; import static org.opensearch.sql.spark.utils.TestUtils.getJson; @@ -173,4 +174,60 @@ void testQuerySchema() { ExecutionEngine.Schema expectedSchema = new ExecutionEngine.Schema(columns); assertEquals(expectedSchema, sparkSqlFunctionTableScanOperator.schema()); } + + /** https://github.com/opensearch-project/sql/issues/2210. */ + @Test + @SneakyThrows + void issue2210() { + SparkQueryRequest sparkQueryRequest = new SparkQueryRequest(); + sparkQueryRequest.setSql(QUERY); + + SparkSqlFunctionTableScanOperator sparkSqlFunctionTableScanOperator = + new SparkSqlFunctionTableScanOperator(sparkClient, sparkQueryRequest); + + when(sparkClient.sql(any())).thenReturn(new JSONObject(getJson("issue2210.json"))); + sparkSqlFunctionTableScanOperator.open(); + assertTrue(sparkSqlFunctionTableScanOperator.hasNext()); + assertEquals( + new ExprTupleValue( + new LinkedHashMap<>() { + { + put("col_name", stringValue("day")); + put("data_type", stringValue("int")); + put("comment", stringValue("")); + } + }), + sparkSqlFunctionTableScanOperator.next()); + assertEquals( + new ExprTupleValue( + new LinkedHashMap<>() { + { + put("col_name", stringValue("# Partition Information")); + put("data_type", stringValue("")); + put("comment", stringValue("")); + } + }), + sparkSqlFunctionTableScanOperator.next()); + assertEquals( + new ExprTupleValue( + new LinkedHashMap<>() { + { + put("col_name", stringValue("# col_name")); + put("data_type", stringValue("data_type")); + put("comment", stringValue("comment")); + } + }), + sparkSqlFunctionTableScanOperator.next()); + assertEquals( + new ExprTupleValue( + new LinkedHashMap<>() { + { + put("col_name", stringValue("day")); + put("data_type", stringValue("int")); + put("comment", stringValue("")); + } + }), + sparkSqlFunctionTableScanOperator.next()); + Assertions.assertFalse(sparkSqlFunctionTableScanOperator.hasNext()); + } } diff --git a/spark/src/test/resources/issue2210.json b/spark/src/test/resources/issue2210.json new file mode 100644 index 0000000000..dec24efdc2 --- /dev/null +++ b/spark/src/test/resources/issue2210.json @@ -0,0 +1,17 @@ +{ + "data": { + "result": [ + "{'col_name':'day','data_type':'int'}", + "{'col_name':'# Partition Information','data_type':'','comment':''}", + "{'col_name':'# col_name','data_type':'data_type','comment':'comment'}", + "{'col_name':'day','data_type':'int'}" + ], + "schema": [ + "{'column_name':'col_name','data_type':'string'}", + "{'column_name':'data_type','data_type':'string'}", + "{'column_name':'comment','data_type':'string'}" + ], + "stepId": "s-123456789", + "applicationId": "application-abc" + } +}