diff --git a/Cargo.lock b/Cargo.lock index 6bf6d046c1b..693082a77f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9623,6 +9623,7 @@ name = "vortex-datafusion" version = "0.1.0" dependencies = [ "anyhow", + "arrow-array", "arrow-schema", "async-trait", "datafusion 54.0.0", diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 2215bd4db8e..4aaefec35ff 100644 --- a/vortex-datafusion/Cargo.toml +++ b/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 } diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index 95c088742b3..06a2d0be6b9 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/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 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 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 { + let scalar = scalar.as_struct(); + let struct_fields = scalar.struct_fields(); + let (fields, arrays): (Vec, 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::>>()? + .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::>(); + 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(()) + } } diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 26d9d86ce18..ecfaa72d01e 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/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 = struct_dtype .names() .iter()