chore: Enable more complex type tests (apache#1753)

Author: andygrove

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -2465,17 +2465,14 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         val df = spark.read.parquet(dir.toString())
 
         checkSparkAnswerAndOperator(df.select("nested1.id"))
-
         checkSparkAnswerAndOperator(df.select("nested1.id", "nested1.nested2.id"))
-
-        // unsupported cast from Int64 to Struct([Field { name: "id", data_type: Int64, ...
-        // checkSparkAnswerAndOperator(df.select("nested1.nested2.id"))
+        checkSparkAnswerAndOperator(df.select("nested1.nested2.id"))
       }
     }
   }
 
-  // TODO this is not using DataFusion's ParquetExec for some reason
-  ignore("get_struct_field with DataFusion ParquetExec - read entire struct") {
+  test("get_struct_field with DataFusion ParquetExec - read entire struct") {
+    assume(usingDataSourceExec(conf))
     withTempPath { dir =>
       // create input file with Comet disabled
       withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
@@ -2501,13 +2498,19 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
           CometConf.COMET_EXPLAIN_FALLBACK_ENABLED.key -> "true") {
 
           val df = spark.read.parquet(dir.toString())
-          checkSparkAnswerAndOperator(df.select("nested1"))
+          if (v1List.isEmpty) {
+            checkSparkAnswer(df.select("nested1"))
+          } else {
+            checkSparkAnswerAndOperator(df.select("nested1"))
+          }
         }
       }
     }
   }
 
-  ignore("read map[int, int] from parquet") {
+  test("read map[int, int] from parquet") {
+    assume(usingDataSourceExec(conf))
+
     withTempPath { dir =>
 // create input file with Comet disabled
       withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
@@ -2523,15 +2526,63 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       Seq("", "parquet").foreach { v1List =>
         withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1List) {
           val df = spark.read.parquet(dir.toString())
-          checkSparkAnswerAndOperator(df.select("map1"))
-          checkSparkAnswerAndOperator(df.select(map_keys(col("map1"))))
-          checkSparkAnswerAndOperator(df.select(map_values(col("map1"))))
+          if (v1List.isEmpty) {
+            checkSparkAnswer(df.select("map1"))
+          } else {
+            checkSparkAnswerAndOperator(df.select("map1"))
+          }
+          // we fall back to Spark for map_keys and map_values
+          checkSparkAnswer(df.select(map_keys(col("map1"))))
+          checkSparkAnswer(df.select(map_values(col("map1"))))
         }
       }
     }
   }
 
-  ignore("read array[int] from parquet") {
+  // repro for https://github.com/apache/datafusion-comet/issues/1754
+  ignore("read map[struct, struct] from parquet") {
+    assume(usingDataSourceExec(conf))
+
+    withTempPath { dir =>
+      // create input file with Comet disabled
+      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+        val df = spark
+          .range(5)
+          .withColumn("id2", col("id"))
+          .withColumn("id3", col("id"))
+          // Spark does not allow null as a key but does allow null as a
+          // value, and the entire map be null
+          .select(
+            when(
+              col("id") > 1,
+              map(
+                struct(col("id"), col("id2"), col("id3")),
+                when(col("id") > 2, struct(col("id"), col("id2"), col("id3"))))).alias("map1"))
+        df.write.parquet(dir.toString())
+      }
+
+      Seq("", "parquet").foreach { v1List =>
+        withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1List) {
+          val df = spark.read.parquet(dir.toString())
+          df.createOrReplaceTempView("tbl")
+          if (v1List.isEmpty) {
+            checkSparkAnswer(df.select("map1"))
+          } else {
+            checkSparkAnswerAndOperator(df.select("map1"))
+          }
+          // we fall back to Spark for map_keys and map_values
+          checkSparkAnswer(df.select(map_keys(col("map1"))))
+          checkSparkAnswer(df.select(map_values(col("map1"))))
+          checkSparkAnswer(spark.sql("SELECT map_keys(map1).id2 FROM tbl"))
+          checkSparkAnswer(spark.sql("SELECT map_values(map1).id2 FROM tbl"))
+        }
+      }
+    }
+  }
+
+  test("read array[int] from parquet") {
+    assume(usingDataSourceExec(conf))
+
     withTempPath { dir =>
 // create input file with Comet disabled
       withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
@@ -2546,8 +2597,13 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       Seq("", "parquet").foreach { v1List =>
         withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1List) {
           val df = spark.read.parquet(dir.toString())
-          checkSparkAnswerAndOperator(df.select("array1"))
-          checkSparkAnswerAndOperator(df.select(element_at(col("array1"), lit(1))))
+          if (v1List.isEmpty) {
+            checkSparkAnswer(df.select("array1"))
+            checkSparkAnswer(df.select(element_at(col("array1"), lit(1))))
+          } else {
+            checkSparkAnswerAndOperator(df.select("array1"))
+            checkSparkAnswerAndOperator(df.select(element_at(col("array1"), lit(1))))
+          }
         }
       }
     }