From ffbbfdd83c0214a606c07512eb6f2cf962996587 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 16 Jun 2026 11:39:19 -0400 Subject: [PATCH 1/3] struct scalar conversation in vortex-datafusion Signed-off-by: Nemo Yu --- vortex-datafusion/Cargo.toml | 1 + vortex-datafusion/src/convert/scalars.rs | 113 ++++++++++++++++++++++- vortex-datafusion/src/convert/schema.rs | 15 ++- 3 files changed, 126 insertions(+), 3 deletions(-) diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 2215bd4db8e..c2e1673c635 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -34,6 +34,7 @@ object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "fs"] } tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } +arrow-array = { workspace = true } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-utils = { workspace = true, features = ["dashmap"] } diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index 95c088742b3..f2b49f3c7a7 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,7 @@ use vortex::dtype::i256; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; @@ -113,7 +120,42 @@ 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_fields, _) => { + let scalar = self.as_struct(); + 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}"))?; + let field = Field::new(name.as_ref(), array.data_type().clone(), nullable); + Ok((field, 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}"))? + }; + ScalarValue::Struct(Arc::new(struct_array)) + } DType::Union(..) => todo!("union scalar conversion"), DType::Variant(_) => vortex_bail!("Variant scalars aren't supported with DF"), DType::Extension(ext) => { @@ -288,6 +330,28 @@ impl FromDataFusion for Scalar { } } ScalarValue::Dictionary(_, v) => Scalar::from_df(v.as_ref()), + ScalarValue::Struct(array) => { + let nullable = array.is_null(0); + let dtype = DType::from_arrow((array.data_type(), nullable.into())); + if nullable { + Scalar::null(dtype) + } else { + let children = array + .columns() + .iter() + .map(|column| { + Scalar::from_df( + &ScalarValue::try_from_array(&**column, 0).unwrap_or_else(|e| { + unimplemented!( + "Can't convert struct field to a Vortex scalar: {e}" + ) + }), + ) + }) + .collect::>(); + Scalar::struct_(dtype, children) + } + } _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"), } } @@ -301,8 +365,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 +757,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..172bff84791 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -8,6 +8,7 @@ use datafusion_common::Result as DFResult; use datafusion_common::exec_datafusion_err; use vortex::array::arrow::ArrowSession; use vortex::dtype::DType; +use vortex::dtype::StructFields; /// Calculate the physical Arrow schema for a Vortex file given its DType and the expected logical schema. /// @@ -85,9 +86,9 @@ 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 { + if let Some(struct_dtype) = struct_fields(dtype) { let physical_fields: Vec = struct_dtype .names() .iter() @@ -204,6 +205,16 @@ fn calculate_physical_field_type( }) } +/// The struct fields of `dtype` if it is a struct, or of an extension type whose storage is +/// (eventually) a struct -- e.g. the native geo Point over `Struct`. +fn struct_fields(dtype: &DType) -> Option<&StructFields> { + match dtype { + DType::Struct(fields, _) => Some(fields), + DType::Extension(ext) => struct_fields(ext.storage_dtype()), + _ => None, + } +} + #[cfg(test)] mod tests { use std::sync::Arc; From a16cfc5f34537116236e0c08eef85598226e839a Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 16 Jun 2026 13:20:03 -0400 Subject: [PATCH 2/3] address review comments Signed-off-by: Nemo Yu --- Cargo.lock | 1 + vortex-datafusion/src/convert/scalars.rs | 123 ++++++++++++----------- vortex-datafusion/src/convert/schema.rs | 18 ++-- 3 files changed, 72 insertions(+), 70 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 57be369d7c3..c73c9c121eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9607,6 +9607,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/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index f2b49f3c7a7..06a2d0be6b9 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -21,6 +21,7 @@ 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; @@ -120,42 +121,7 @@ impl TryToDataFusion for Scalar { ), DType::List(..) => todo!("list scalar conversion"), DType::FixedSizeList(..) => todo!("fixed-size list scalar conversion"), - DType::Struct(struct_fields, _) => { - let scalar = self.as_struct(); - 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}"))?; - let field = Field::new(name.as_ref(), array.data_type().clone(), nullable); - Ok((field, 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}"))? - }; - ScalarValue::Struct(Arc::new(struct_array)) - } + 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) => { @@ -330,33 +296,74 @@ impl FromDataFusion for Scalar { } } ScalarValue::Dictionary(_, v) => Scalar::from_df(v.as_ref()), - ScalarValue::Struct(array) => { - let nullable = array.is_null(0); - let dtype = DType::from_arrow((array.data_type(), nullable.into())); - if nullable { - Scalar::null(dtype) - } else { - let children = array - .columns() - .iter() - .map(|column| { - Scalar::from_df( - &ScalarValue::try_from_array(&**column, 0).unwrap_or_else(|e| { - unimplemented!( - "Can't convert struct field to a Vortex scalar: {e}" - ) - }), - ) - }) - .collect::>(); - Scalar::struct_(dtype, children) - } - } + 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; diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 172bff84791..ecfaa72d01e 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -8,7 +8,6 @@ use datafusion_common::Result as DFResult; use datafusion_common::exec_datafusion_err; use vortex::array::arrow::ArrowSession; use vortex::dtype::DType; -use vortex::dtype::StructFields; /// Calculate the physical Arrow schema for a Vortex file given its DType and the expected logical schema. /// @@ -88,7 +87,12 @@ fn calculate_physical_field_type( // For struct types, recursively check each field. DataType::Struct(logical_fields) => { - if let Some(struct_dtype) = struct_fields(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() @@ -205,16 +209,6 @@ fn calculate_physical_field_type( }) } -/// The struct fields of `dtype` if it is a struct, or of an extension type whose storage is -/// (eventually) a struct -- e.g. the native geo Point over `Struct`. -fn struct_fields(dtype: &DType) -> Option<&StructFields> { - match dtype { - DType::Struct(fields, _) => Some(fields), - DType::Extension(ext) => struct_fields(ext.storage_dtype()), - _ => None, - } -} - #[cfg(test)] mod tests { use std::sync::Arc; From f52a2c4b16b07944349e892d6f31524173125f73 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 16 Jun 2026 13:32:24 -0400 Subject: [PATCH 3/3] fix toml order Signed-off-by: Nemo Yu --- vortex-datafusion/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index c2e1673c635..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 } @@ -34,7 +35,6 @@ object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "fs"] } tokio-stream = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } -arrow-array = { workspace = true } vortex = { workspace = true, features = ["object_store", "tokio", "files"] } vortex-utils = { workspace = true, features = ["dashmap"] }