From 79721ec40884e66cbbbd2aa4f301b7c03a0aed5e Mon Sep 17 00:00:00 2001
From: Alex Huang <huangweijun1001@gmail.com>
Date: Sun, 4 Feb 2024 22:11:30 +0800
Subject: [PATCH] Minor: support cast values to fixedsizelist (#5340)

* support cast utf8 to fixedsizelist

* fix ci

* fix chore

* Apply suggestions from code review

Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>

* fix compile

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
---
 arrow-cast/src/cast.rs | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
index bd35096e0645..e321ae1b1b2f 100644
--- a/arrow-cast/src/cast.rs
+++ b/arrow-cast/src/cast.rs
@@ -141,6 +141,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
         }
         (_, List(list_to)) => can_cast_types(from_type, list_to.data_type()),
         (_, LargeList(list_to)) => can_cast_types(from_type, list_to.data_type()),
+        (_, FixedSizeList(list_to,size)) if *size == 1 => {
+            can_cast_types(from_type, list_to.data_type())},
         // cast one decimal type to another decimal type
         (Decimal128(_, _), Decimal128(_, _)) => true,
         (Decimal256(_, _), Decimal256(_, _)) => true,
@@ -802,6 +804,9 @@ pub fn cast_with_options(
         }
         (_, List(ref to)) => cast_values_to_list::<i32>(array, to, cast_options),
         (_, LargeList(ref to)) => cast_values_to_list::<i64>(array, to, cast_options),
+        (_, FixedSizeList(ref to, size)) if *size == 1 => {
+            cast_values_to_fixed_size_list(array, to, *size, cast_options)
+        }
         (Decimal128(_, s1), Decimal128(p2, s2)) => {
             cast_decimal_to_decimal_same_type::<Decimal128Type>(
                 array.as_primitive(),
@@ -3040,6 +3045,18 @@ fn cast_values_to_list<O: OffsetSizeTrait>(
     Ok(Arc::new(list))
 }
 
+/// Helper function that takes a primitive array and casts to a fixed size list array.
+fn cast_values_to_fixed_size_list(
+    array: &dyn Array,
+    to: &FieldRef,
+    size: i32,
+    cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+    let values = cast_with_options(array, to.data_type(), cast_options)?;
+    let list = FixedSizeListArray::new(to.clone(), size, values, None);
+    Ok(Arc::new(list))
+}
+
 /// A specified helper to cast from `GenericBinaryArray` to `GenericStringArray` when they have same
 /// offset size so re-encoding offset is unnecessary.
 fn cast_binary_to_string<O: OffsetSizeTrait>(
@@ -7609,6 +7626,30 @@ mod tests {
         assert_eq!(expected.values(), actual.values());
     }
 
+    #[test]
+    fn test_cast_utf8_to_list() {
+        // DataType::List
+        let array = Arc::new(StringArray::from(vec!["5"])) as ArrayRef;
+        let field = Arc::new(Field::new("", DataType::Int32, false));
+        let list_array = cast(&array, &DataType::List(field.clone())).unwrap();
+        let actual = list_array.as_list_opt::<i32>().unwrap();
+        let expect = ListArray::from_iter_primitive::<Int32Type, _, _>([Some([Some(5)])]);
+        assert_eq!(&expect.value(0), &actual.value(0));
+
+        // DataType::LargeList
+        let list_array = cast(&array, &DataType::LargeList(field.clone())).unwrap();
+        let actual = list_array.as_list_opt::<i64>().unwrap();
+        let expect = LargeListArray::from_iter_primitive::<Int32Type, _, _>([Some([Some(5)])]);
+        assert_eq!(&expect.value(0), &actual.value(0));
+
+        // DataType::FixedSizeList
+        let list_array = cast(&array, &DataType::FixedSizeList(field.clone(), 1)).unwrap();
+        let actual = list_array.as_fixed_size_list_opt().unwrap();
+        let expect =
+            FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>([Some([Some(5)])], 1);
+        assert_eq!(&expect.value(0), &actual.value(0));
+    }
+
     #[test]
     fn test_cast_list_containers() {
         // large-list to list