feat(vortex-datafusion): struct scalar conversion + extension-over-struct scan

vortex-datafusion/Cargo.toml

@@ -14,6 +14,7 @@ rust-version = { workspace = true }
 version = { workspace = true }
 
 [dependencies]
+arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
 async-trait = { workspace = true }
 datafusion-catalog = { workspace = true }

vortex-datafusion/src/convert/scalars.rs

@@ -1,6 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
+use arrow_array::Array;
+use arrow_array::StructArray;
+use arrow_schema::Field;
+use arrow_schema::Fields;
 use datafusion_common::ScalarValue;
 use vortex::buffer::ByteBuffer;
 use vortex::dtype::DType;
@@ -14,6 +20,8 @@ use vortex::dtype::i256;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::error::vortex_bail;
+use vortex::error::vortex_err;
+use vortex::error::vortex_panic;
 use vortex::extension::datetime::AnyTemporal;
 use vortex::extension::datetime::TemporalMetadata;
 use vortex::extension::datetime::TimeUnit;
@@ -113,7 +121,7 @@ impl TryToDataFusion<ScalarValue> for Scalar {
             ),
             DType::List(..) => todo!("list scalar conversion"),
             DType::FixedSizeList(..) => todo!("fixed-size list scalar conversion"),
-            DType::Struct(..) => todo!("struct scalar conversion"),
+            DType::Struct(..) => struct_to_df(self)?,
             DType::Union(..) => todo!("union scalar conversion"),
             DType::Variant(_) => vortex_bail!("Variant scalars aren't supported with DF"),
             DType::Extension(ext) => {
@@ -288,11 +296,74 @@ impl FromDataFusion<ScalarValue> for Scalar {
                 }
             }
             ScalarValue::Dictionary(_, v) => Scalar::from_df(v.as_ref()),
+            ScalarValue::Struct(array) => struct_from_df(array),
             _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"),
         }
     }
 }
 
+/// Converts a Vortex struct scalar to a DataFusion `ScalarValue::Struct`.
+fn struct_to_df(scalar: &Scalar) -> VortexResult<ScalarValue> {
+    let scalar = scalar.as_struct();
+    let struct_fields = scalar.struct_fields();
+    let (fields, arrays): (Vec<Field>, Vec<_>) = struct_fields
+        .names()
+        .iter()
+        .zip(struct_fields.fields())
+        .enumerate()
+        .map(|(idx, (name, field_dtype))| {
+            let nullable = field_dtype.is_nullable();
+            let child = if scalar.is_null() {
+                Scalar::null(field_dtype)
+            } else {
+                scalar
+                    .field_by_idx(idx)
+                    .ok_or_else(|| vortex_err!("missing struct field {name}"))?
+            };
+            let array = child
+                .try_to_df()?
+                .to_array()
+                .map_err(|e| vortex_err!("failed to build struct field array: {e}"))?;
+            Ok((
+                Field::new(name.as_ref(), array.data_type().clone(), nullable),
+                array,
+            ))
+        })
+        .collect::<VortexResult<Vec<_>>>()?
+        .into_iter()
+        .unzip();
+
+    let fields = Fields::from(fields);
+    let struct_array = if scalar.is_null() {
+        StructArray::new_null(fields, 1)
+    } else {
+        StructArray::try_new(fields, arrays, None)
+            .map_err(|e| vortex_err!("failed to build struct scalar array: {e}"))?
+    };
+    Ok(ScalarValue::Struct(Arc::new(struct_array)))
+}
+
+/// Converts a DataFusion `ScalarValue::Struct` (a one-row struct array) to a Vortex struct scalar.
+fn struct_from_df(array: &StructArray) -> Scalar {
+    let dtype = DType::from_arrow((array.data_type(), Nullability::Nullable));
+    if array.is_null(0) {
+        Scalar::null(dtype)
+    } else {
+        let children = array
+            .columns()
+            .iter()
+            .map(|column| {
+                Scalar::from_df(
+                    &ScalarValue::try_from_array(column.as_ref(), 0).unwrap_or_else(|e| {
+                        vortex_panic!("cannot convert struct field to a Vortex scalar: {e}")
+                    }),
+                )
+            })
+            .collect::<Vec<_>>();
+        Scalar::struct_(dtype, children)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use datafusion_common::ScalarValue;
@@ -301,8 +372,10 @@ mod tests {
     use vortex::buffer::ByteBuffer;
     use vortex::dtype::DType;
     use vortex::dtype::DecimalDType;
+    use vortex::dtype::FieldNames;
     use vortex::dtype::Nullability;
     use vortex::dtype::PType;
+    use vortex::dtype::StructFields;
     use vortex::dtype::i256;
     use vortex::scalar::DecimalValue;
     use vortex::scalar::Scalar;
@@ -691,4 +764,49 @@ mod tests {
             .into();
         assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]);
     }
+
+    #[test]
+    fn struct_scalar_round_trips() -> VortexResult<()> {
+        let dtype = DType::Struct(
+            StructFields::new(
+                FieldNames::from(["x", "y"]),
+                vec![
+                    DType::Primitive(PType::F64, Nullability::NonNullable),
+                    DType::Primitive(PType::F64, Nullability::NonNullable),
+                ],
+            ),
+            Nullability::NonNullable,
+        );
+        let original = Scalar::struct_(
+            dtype,
+            vec![Scalar::from(-111.7610f64), Scalar::from(34.8697f64)],
+        );
+
+        let df = original.try_to_df()?;
+        assert!(matches!(df, ScalarValue::Struct(_)));
+
+        // Back through `from_df` and out again yields the identical DataFusion struct value.
+        let back = Scalar::from_df(&df);
+        assert_eq!(back.try_to_df()?, df);
+        Ok(())
+    }
+
+    #[test]
+    fn null_struct_scalar_round_trips() -> VortexResult<()> {
+        let dtype = DType::Struct(
+            StructFields::new(
+                FieldNames::from(["x", "y"]),
+                vec![
+                    DType::Primitive(PType::F64, Nullability::Nullable),
+                    DType::Primitive(PType::F64, Nullability::Nullable),
+                ],
+            ),
+            Nullability::Nullable,
+        );
+
+        let df = Scalar::null(dtype).try_to_df()?;
+        assert!(matches!(df, ScalarValue::Struct(_)));
+        assert!(Scalar::from_df(&df).is_null());
+        Ok(())
+    }
 }

vortex-datafusion/src/convert/schema.rs

@@ -85,9 +85,14 @@ fn calculate_physical_field_type(
         // RunEndEncoded loses its encoding
         DataType::RunEndEncoded(..) => logical_type.clone(),
 
-        // For struct types, recursively check each field
+        // For struct types, recursively check each field.
         DataType::Struct(logical_fields) => {
-            if let DType::Struct(struct_dtype, _) = dtype {
+            // Walk through any extension layers to reach the underlying struct fields.
+            let mut inner = dtype;
+            while let DType::Extension(ext) = inner {
+                inner = ext.storage_dtype();
+            }
+            if let DType::Struct(struct_dtype, _) = inner {
                 let physical_fields: Vec<Field> = struct_dtype
                     .names()
                     .iter()