feat: Add array reading support to native_datafusion scan

native/core/src/execution/serde.rs

@@ -111,7 +111,7 @@ pub fn to_arrow_datatype(dt_value: &DataType) -> ArrowDataType {
         {
             DatatypeStruct::List(info) => {
                 let field = Field::new(
-                    "item",
+                    "element",
                     to_arrow_datatype(info.element_type.as_ref().unwrap()),
                     info.contains_null,
                 );

spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala

@@ -42,9 +42,6 @@ object ParquetGenerator {
     DataTypes.createDecimalType(10, 2),
     DataTypes.createDecimalType(36, 18),
     DataTypes.DateType,
-    DataTypes.TimestampType,
-    // TimestampNTZType only in Spark 3.4+
-    // DataTypes.TimestampNTZType,
     DataTypes.StringType,
     DataTypes.BinaryType)
 
@@ -58,6 +55,12 @@ object ParquetGenerator {
     val dataTypes = ListBuffer[DataType]()
     dataTypes.appendAll(primitiveTypes)
 
+    if (options.generateTimestamps) {
+      dataTypes += DataTypes.TimestampType
+      // TimestampNTZType only in Spark 3.4+
+      // dataTypes += DataTypes.TimestampNTZType,
+    }
+
     if (options.generateStruct) {
       dataTypes += StructType(
         primitiveTypes.zipWithIndex.map(x => StructField(s"c${x._2}", x._1, true)))
@@ -212,8 +215,9 @@ object ParquetGenerator {
 }
 
 case class DataGenOptions(
-    allowNull: Boolean,
-    generateNegativeZero: Boolean,
-    generateArray: Boolean,
-    generateStruct: Boolean,
-    generateMap: Boolean)
+    allowNull: Boolean = true,
+    generateNegativeZero: Boolean = true,
+    generateTimestamps: Boolean = true,
+    generateArray: Boolean = false,
+    generateStruct: Boolean = false,
+    generateMap: Boolean = false)

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

@@ -122,8 +122,9 @@ object CometNativeScanExec extends DataTypeSupport {
   }
 
   override def isAdditionallySupported(dt: DataType): Boolean = {
-    // TODO add array and map
+    // TODO add map support
     dt match {
+      case s: ArrayType => isTypeSupported(s.elementType)
       case s: StructType => s.fields.map(_.dataType).forall(isTypeSupported)
       case _ => false
     }

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

@@ -26,7 +26,6 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.functions.{array, col, expr, lit, udf}
-import org.apache.spark.sql.types.StructType
 
 import org.apache.comet.CometSparkSessionExtensions.{isSpark34Plus, isSpark35Plus}
 import org.apache.comet.testing.{DataGenOptions, ParquetGenerator}
@@ -55,17 +54,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp
       val filename = path.toString
       val random = new Random(42)
       withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
-        ParquetGenerator.makeParquetFile(
-          random,
-          spark,
-          filename,
-          100,
-          DataGenOptions(
-            allowNull = true,
-            generateNegativeZero = true,
-            generateArray = false,
-            generateStruct = false,
-            generateMap = false))
+        ParquetGenerator.makeParquetFile(random, spark, filename, 100, DataGenOptions())
       }
       val table = spark.read.parquet(filename)
       table.createOrReplaceTempView("t1")
@@ -79,38 +68,31 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp
     }
   }
 
-  test("array_remove - test all types (convert from Parquet)") {
+  test("array_remove - test arrays (native_datafusion reader)") {
     withTempDir { dir =>
       val path = new Path(dir.toURI.toString, "test.parquet")
       val filename = path.toString
       val random = new Random(42)
       withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
-        val options = DataGenOptions(
-          allowNull = true,
-          generateNegativeZero = true,
-          generateArray = true,
-          generateStruct = true,
-          generateMap = false)
-        ParquetGenerator.makeParquetFile(random, spark, filename, 100, options)
+        ParquetGenerator.makeParquetFile(
+          random,
+          spark,
+          filename,
+          100,
+          DataGenOptions(
+            generateArray = true,
+            // native_datafusion does not support timestamps correctly yet
+            generateTimestamps = false))
       }
-      withSQLConf(
-        CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false",
-        CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true",
-        CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true") {
+      withSQLConf(CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION) {
         val table = spark.read.parquet(filename)
         table.createOrReplaceTempView("t1")
         // test with array of each column
-        for (field <- table.schema.fields) {
-          val fieldName = field.name
+        for (fieldName <- table.schema.fieldNames) {
           sql(s"SELECT array($fieldName, $fieldName) as a, $fieldName as b FROM t1")
             .createOrReplaceTempView("t2")
           val df = sql("SELECT array_remove(a, b) FROM t2")
-          field.dataType match {
-            case _: StructType =>
-            // skip due to https://github.com/apache/datafusion-comet/issues/1314
-            case _ =>
-              checkSparkAnswer(df)
-          }
+          checkSparkAnswerAndOperator(df)
         }
       }
     }