fix: array to array cast

Author: manuzhang

native/spark-expr/src/conversion_funcs/cast.rs

@@ -294,6 +294,7 @@ pub(crate) fn cast_array(
     };
 
     let cast_result = match (&from_type, to_type) {
+        (Null, _) => Ok(cast_with_options(&array, to_type, &native_cast_options)?),
         (Utf8, Boolean) => spark_cast_utf8_to_boolean::<i32>(&array, eval_mode),
         (LargeUtf8, Boolean) => spark_cast_utf8_to_boolean::<i64>(&array, eval_mode),
         (Utf8, Timestamp(_, _)) => {
@@ -366,8 +367,22 @@ pub(crate) fn cast_array(
             cast_options,
         )?),
         (List(_), Utf8) => Ok(cast_array_to_string(array.as_list(), cast_options)?),
-        (List(_), List(_)) if can_cast_types(&from_type, to_type) => {
-            Ok(cast_with_options(&array, to_type, &CAST_OPTIONS)?)
+        (List(from), List(to))
+            if can_cast_types(&from_type, to_type)
+                || (matches!(from.data_type(), Decimal128(_, _))
+                    && matches!(to.data_type(), Boolean)) =>
+        {
+            let list_array = array.as_list::<i32>();
+            Ok(Arc::new(ListArray::new(
+                Arc::clone(to),
+                list_array.offsets().clone(),
+                cast_array(
+                    Arc::clone(list_array.values()),
+                    to.data_type(),
+                    cast_options,
+                )?,
+                list_array.nulls().cloned(),
+            )) as ArrayRef)
         }
         (Map(_, _), Map(_, _)) => Ok(cast_map_to_map(&array, &from_type, to_type, cast_options)?),
         (UInt8 | UInt16 | UInt32 | UInt64, Int8 | Int16 | Int32 | Int64)
@@ -803,7 +818,8 @@ fn cast_binary_formatter(value: &[u8]) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::StringArray;
+    use arrow::array::{BooleanArray, Decimal128Array, ListArray, StringArray};
+    use arrow::buffer::OffsetBuffer;
     use arrow::datatypes::TimestampMicrosecondType;
     use arrow::datatypes::{Field, Fields};
     #[test]
@@ -955,8 +971,6 @@ mod tests {
 
     #[test]
     fn test_cast_i32_array_to_string() {
-        use arrow::array::ListArray;
-        use arrow::buffer::OffsetBuffer;
         let values_array = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(1), None, None]);
         let offsets_buffer = OffsetBuffer::<i32>::new(vec![0, 3, 5, 6, 6].into());
         let item_field = Arc::new(Field::new("item", DataType::Int32, true));
@@ -977,4 +991,42 @@ mod tests {
         assert_eq!(r#"[null]"#, string_array.value(2));
         assert_eq!(r#"[]"#, string_array.value(3));
     }
+
+    #[test]
+    fn test_cast_decimal_array_to_boolean_array() {
+        let values_array =
+            Decimal128Array::from(vec![Some(0), Some(100), None, Some(-100), Some(0)])
+                .with_precision_and_scale(10, 2)
+                .unwrap();
+        let offsets_buffer = OffsetBuffer::<i32>::new(vec![0, 2, 3, 5].into());
+        let item_field = Arc::new(Field::new("item", DataType::Decimal128(10, 2), true));
+        let list_array = Arc::new(ListArray::new(
+            item_field,
+            offsets_buffer,
+            Arc::new(values_array),
+            Some(vec![true, false, true].into()),
+        ));
+
+        let casted_array = cast_array(
+            list_array,
+            &DataType::List(Arc::new(Field::new("item", DataType::Boolean, true))),
+            &SparkCastOptions::new(EvalMode::Legacy, "UTC", false),
+        )
+        .unwrap();
+        let casted_list = casted_array.as_list::<i32>();
+
+        assert_eq!(3, casted_list.len());
+        assert!(!casted_list.is_null(0));
+        assert!(casted_list.is_null(1));
+        assert!(!casted_list.is_null(2));
+
+        let values = casted_list
+            .values()
+            .as_any()
+            .downcast_ref::<BooleanArray>()
+            .expect("Casted list values should be a BooleanArray");
+        let expected =
+            BooleanArray::from(vec![Some(false), Some(true), None, Some(true), Some(false)]);
+        assert_eq!(&expected, values);
+    }
 }

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

@@ -1266,6 +1266,30 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("cast ArrayType to ArrayType") {
+    val types = Seq(
+      BooleanType,
+      StringType,
+      ByteType,
+      IntegerType,
+      LongType,
+      ShortType,
+      DecimalType(10, 2),
+      DecimalType(38, 18))
+    for (fromType <- types) {
+      for (toType <- types) {
+        if (fromType != toType &&
+          !tags
+            .get(s"cast $fromType to $toType")
+            .exists(s => s.contains("org.scalatest.Ignore")) &&
+          Cast.canCast(fromType, toType) &&
+          CometCast.isSupported(fromType, toType, None, CometEvalMode.LEGACY) == Compatible()) {
+          castTest(generateArrays(100, fromType), ArrayType(toType))
+        }
+      }
+    }
+  }
+
   private def generateFloats(): DataFrame = {
     withNulls(gen.generateFloats(dataSize)).toDF("a")
   }
@@ -1294,10 +1318,10 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     withNulls(gen.generateLongs(dataSize)).toDF("a")
   }
 
-  private def generateArrays(rowSize: Int, elementType: DataType): DataFrame = {
+  private def generateArrays(rowNum: Int, elementType: DataType): DataFrame = {
     import scala.collection.JavaConverters._
     val schema = StructType(Seq(StructField("a", ArrayType(elementType), true)))
-    spark.createDataFrame(gen.generateRows(rowSize, schema).asJava, schema)
+    spark.createDataFrame(gen.generateRows(rowNum, schema).asJava, schema)
   }
 
   // https://github.com/apache/datafusion-comet/issues/2038