Bug Fix, handle DESC TABLE response (#2211)

Signed-off-by: Peng Huo <[email protected]>
opensearch-project · Oct 4, 2023 · 9df968a · 9df968a
1 parent 49ea48c
commit 9df968a
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 1 deletion.
diff --git a/...va/org/opensearch/sql/spark/functions/response/DefaultSparkSqlFunctionResponseHandle.java b/...va/org/opensearch/sql/spark/functions/response/DefaultSparkSqlFunctionResponseHandle.java
@@ -86,7 +86,7 @@ private static LinkedHashMap<String, ExprValue> 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/...c/test/java/org/opensearch/sql/spark/functions/SparkSqlFunctionTableScanOperatorTest.java b/...c/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
@@ -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"
+  }
+}