fix[vortex-array]: propagate preferred type conversion to nested fields (#4390)

asubiotto · web-flow · commit 8b8c9ee3eef1 · 2025-08-28T10:33:58.000+01:00
When arrow_type is None, vortex to arrow conversion should respect
conversion to the encoding's preferred arrow type for all nested fields.
Previous to this commit, this preference was lost for any fields within
structs or lists.

For example, a struct with a utf8 dictionary field would lose the
dictionary encoding.

Signed-off-by: Alfonso Subiotto Marques &lt;alfonso.subiotto@polarsignals.com&gt;
diff --git a/docs/guides/python-integrations.md b/docs/guides/python-integrations.md
@@ -98,7 +98,7 @@ If you have a struct array, use {meth}`.Array.to_arrow_table` to construct an Ar
 >>> struct_arr.to_arrow_table()
 pyarrow.Table
 age: int64
-name: string_view
+name: string
 ----
 age: [[25,31,33,57]]
 name: [["Joseph","Narendra","Angela","Mikhail"]]
diff --git a/vortex-array/src/arrow/compute/to_arrow/canonical.rs b/vortex-array/src/arrow/compute/to_arrow/canonical.rs
@@ -39,14 +39,17 @@ pub(super) struct ToArrowCanonical;
 impl Kernel for ToArrowCanonical {
     #[allow(clippy::cognitive_complexity)]
     fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
-        let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;
+        let ToArrowArgs {
+            array,
+            arrow_type: arrow_type_opt,
+        } = ToArrowArgs::try_from(args)?;
         if !array.is_canonical() {
             // Not handled by this kernel
             return Ok(None);
         }
 
         // Figure out the target Arrow type, or use the canonical type
-        let arrow_type = arrow_type
+        let arrow_type = arrow_type_opt
             .cloned()
             .map(Ok)
             .unwrap_or_else(|| array.dtype().to_arrow_dtype())?;
@@ -136,11 +139,13 @@ impl Kernel for ToArrowCanonical {
                 to_arrow_decimal256(array)
             }
             (Canonical::Struct(array), DataType::Struct(fields)) => {
-                to_arrow_struct(array, fields.as_ref())
+                to_arrow_struct(array, fields.as_ref(), arrow_type_opt.is_none())
+            }
+            (Canonical::List(array), DataType::List(field)) => {
+                to_arrow_list::<i32>(array, field, arrow_type_opt.is_none())
             }
-            (Canonical::List(array), DataType::List(field)) => to_arrow_list::<i32>(array, field),
             (Canonical::List(array), DataType::LargeList(field)) => {
-                to_arrow_list::<i64>(array, field)
+                to_arrow_list::<i64>(array, field, arrow_type_opt.is_none())
             }
             (Canonical::VarBinView(array), DataType::BinaryView) if array.dtype().is_binary() => {
                 to_arrow_varbinview::<BinaryViewType>(array)
@@ -277,7 +282,11 @@ fn to_arrow_decimal256(array: DecimalArray) -> VortexResult<ArrowArrayRef> {
     ))
 }
 
-fn to_arrow_struct(array: StructArray, fields: &[FieldRef]) -> VortexResult<ArrowArrayRef> {
+fn to_arrow_struct(
+    array: StructArray,
+    fields: &[FieldRef],
+    to_preferred: bool,
+) -> VortexResult<ArrowArrayRef> {
     if array.fields().len() != fields.len() {
         vortex_bail!(
             "StructArray has {} fields, but target Arrow type has {} fields",
@@ -301,9 +310,12 @@ fn to_arrow_struct(array: StructArray, fields: &[FieldRef]) -> VortexResult<Arro
                 );
             }
 
-            arr.clone()
-                .into_arrow(field.data_type())
-                .map_err(|err| err.with_context(format!("Failed to canonicalize field {field}")))
+            let result = if to_preferred {
+                arr.clone().into_arrow_preferred()
+            } else {
+                arr.clone().into_arrow(field.data_type())
+            };
+            result.map_err(|err| err.with_context(format!("Failed to canonicalize field {field}")))
         })
         .collect::<VortexResult<Vec<_>>>()?;
 
@@ -341,14 +353,19 @@ fn to_arrow_struct(array: StructArray, fields: &[FieldRef]) -> VortexResult<Arro
 fn to_arrow_list<O: NativePType + OffsetSizeTrait>(
     array: ListArray,
     element: &FieldRef,
+    to_preferred: bool,
 ) -> VortexResult<ArrowArrayRef> {
     // First we cast the offsets into the correct width.
     let offsets_dtype = DType::Primitive(O::PTYPE, array.dtype().nullability());
     let arrow_offsets = cast(array.offsets(), &offsets_dtype)
         .map_err(|err| err.with_context(format!("Failed to cast offsets to {offsets_dtype}")))?
         .to_primitive()?;
 
-    let values = array.elements().clone().into_arrow(element.data_type())?;
+    let values = if to_preferred {
+        array.elements().clone().into_arrow_preferred()?
+    } else {
+        array.elements().clone().into_arrow(element.data_type())?
+    };
     let nulls = array.validity_mask()?.to_null_buffer();
 
     Ok(Arc::new(GenericListArray::new(
diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs
@@ -328,7 +328,8 @@ mod tests {
         .unwrap()
         .into_array();
         let actual = actual.into_arrow_preferred().unwrap();
-        let expected = expected.into_arrow_preferred().unwrap();
+        let expected_arrow_dtype = expected.dtype().to_arrow_dtype().unwrap();
+        let expected = expected.into_arrow(&expected_arrow_dtype).unwrap();
         assert_eq!(actual.data_type(), expected.data_type());
         assert_eq!(&actual, &expected);
     }
@@ -340,7 +341,7 @@ mod tests {
         vec![true, false, true], // Expected: nulls excluded, all dict values match
     )]
     #[case::all_false_case(
-        vec![Some("x"), None, Some("x")], // Dict values: ["x"] 
+        vec![Some("x"), None, Some("x")], // Dict values: ["x"]
         "", // Filter for empty string
         vec![false, false, false], // Expected: all false, no dict values match
     )]
diff --git a/vortex-python/python/vortex/arrays.py b/vortex-python/python/vortex/arrays.py
@@ -80,7 +80,7 @@ def _Array_to_arrow_table(self: _arrays.Array) -> pyarrow.Table:
     >>> array.to_arrow_table()
     pyarrow.Table
     age: int64
-    name: string_view
+    name: string
     ----
     age: [[25,31,33,57]]
     name: [["Joseph","Narendra","Angela","Mikhail"]]
