Fix output schema issue and IT

Signed-off-by: Chen Dai <daichen@amazon.com>
opensearch-project · Jun 13, 2023 · f3ef196 · f3ef196
1 parent cbc8a33
commit f3ef196
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/...egration/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingIndex.scala b/...egration/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingIndex.scala
@@ -71,13 +71,13 @@ class FlintSparkSkippingIndex(
   }
 
   private def getSchema: String = {
-    val indexFieldTypes = indexedColumns.map { indexCol =>
-      val columnName = indexCol.columnName
-      // Data type INT from catalog is not recognized by Spark DataType.fromJson()
-      val columnType = if (indexCol.columnType == "int") "integer" else indexCol.columnType
-      val sparkType = DataType.fromJson("\"" + columnType + "\"")
-      StructField(columnName, sparkType, nullable = false)
-    }
+    val indexFieldTypes =
+      indexedColumns.flatMap(_.outputSchema()).map { case (colName, colType) =>
+        // Data type INT from catalog is not recognized by Spark DataType.fromJson()
+        val columnType = if (colType == "int") "integer" else colType
+        val sparkType = DataType.fromJson("\"" + columnType + "\"")
+        StructField(colName, sparkType, nullable = false)
+      }
 
     val allFieldTypes =
       indexFieldTypes :+ StructField(FILE_PATH_COLUMN, StringType, nullable = false)

diff --git a/...t/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexSuite.scala b/...t/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexSuite.scala
@@ -139,7 +139,10 @@ class FlintSparkSkippingIndexSuite
         |     "address": {
         |       "type": "keyword"
         |     },
-        |     "age": {
+        |     "MinMax_age_0": {
+        |       "type": "integer"
+        |     },
+        |     "MinMax_age_1" : {
         |       "type": "integer"
         |     },
         |     "file_path": {