From 734821a6053e0ea7e47101ff98d8de9c04267a02 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 18 Jun 2026 10:22:50 +0100 Subject: [PATCH] Touch up some of the interfaces and docs in vortex-datafusion Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs.rs | 2 +- vortex-datafusion/src/convert/mod.rs | 19 +++++++++++--- vortex-datafusion/src/convert/scalars.rs | 29 +++++++++++++++++++--- vortex-datafusion/src/lib.rs | 3 +-- vortex-datafusion/src/persistent/mod.rs | 1 + vortex-datafusion/src/persistent/reader.rs | 2 +- vortex-datafusion/src/persistent/sink.rs | 2 ++ vortex-datafusion/src/persistent/source.rs | 5 ++++ 8 files changed, 53 insertions(+), 10 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index b6d741fc5ec..c14576c5301 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -106,7 +106,7 @@ pub trait ExpressionConvertor: Send + Sync { } } -/// The default [`ExpressionConvertor`]. +/// The default [`ExpressionConvertor`] implementation. #[derive(Default)] pub struct DefaultExpressionConvertor {} diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 050987522b7..6a6fb8aef08 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -1,6 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Utilities and interface to convert DataFusion types to Vortex types. +//! +//! Currently includes: +//! [`ExpressionConvertor`] - Controls the rewrite of DataFusion expressions to Vortex expressions, and whether they can +//! be pushed into the underlying scan. A default implementation is provided. +//! [`FromDataFusion`] - Converts a DataFusion type into a Vortex type infallible. +//! [TryToDataFusion] - Fallibly converts a Vortex type to a DataFusion type. + use vortex::error::VortexResult; pub(crate) mod exprs; @@ -8,12 +16,17 @@ mod scalars; pub(crate) mod schema; pub(crate) mod stats; +pub use exprs::DefaultExpressionConvertor; +pub use exprs::ExpressionConvertor; + /// First-party trait for implementing conversion from DataFusion types to Vortex types. -pub(crate) trait FromDataFusion: Sized { +pub trait FromDataFusion: Sized { + /// Convert to this Vortex type from the input DataFusion type. fn from_df(df: &D) -> Self; } -/// First-party trait for implementing conversion from Vortex to DataFusion types. -pub(crate) trait TryToDataFusion { +/// First-party trait for implementing fallible conversions from Vortex to DataFusion types. +pub trait TryToDataFusion { + /// Try to convert this Vortex type from the input DataFusion type. fn try_to_df(&self) -> VortexResult; } diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index 06a2d0be6b9..86c1309df57 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -119,10 +119,16 @@ impl TryToDataFusion for Scalar { .cloned() .map(|b| Vec::::from(b.into_inner())), ), - DType::List(..) => todo!("list scalar conversion"), - DType::FixedSizeList(..) => todo!("fixed-size list scalar conversion"), + dtype @ DType::List(..) => vortex_bail!( + "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" + ), + dtype @ DType::FixedSizeList(..) => vortex_bail!( + "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" + ), DType::Struct(..) => struct_to_df(self)?, - DType::Union(..) => todo!("union scalar conversion"), + dtype @ DType::Union(..) => vortex_bail!( + "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" + ), DType::Variant(_) => vortex_bail!("Variant scalars aren't supported with DF"), DType::Extension(ext) => { let storage_scalar = self.as_extension().to_storage_scalar(); @@ -809,4 +815,21 @@ mod tests { assert!(Scalar::from_df(&df).is_null()); Ok(()) } + + #[rstest] + #[case::list(Scalar::null(DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + Nullability::Nullable + )))] + #[case::fixed_size_list(Scalar::null(DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), + 2, + Nullability::Nullable + )))] + #[case::union(Scalar::null(DType::Union(Nullability::Nullable)))] + fn unsupported_vortex_scalars_return_errors(#[case] scalar: Scalar) { + let err = scalar.try_to_df().unwrap_err(); + + assert!(err.to_string().contains("unsupported scalar type"), "{err}"); + } } diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index d3115c56b09..64992849383 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -88,14 +88,13 @@ use std::fmt::Debug; use datafusion_common::stats::Precision as DFPrecision; use vortex::expr::stats::Precision; -mod convert; +pub mod convert; mod persistent; pub mod v2; #[cfg(test)] mod tests; -pub use convert::exprs::ExpressionConvertor; pub use persistent::*; /// Extension trait to convert our [`Precision`] to DataFusion's diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs index 17565937a0f..53fefacafa4 100644 --- a/vortex-datafusion/src/persistent/mod.rs +++ b/vortex-datafusion/src/persistent/mod.rs @@ -36,6 +36,7 @@ pub use access_plan::VortexAccessPlan; pub use format::VortexFormat; pub use format::VortexFormatFactory; pub use format::VortexTableOptions; +pub use sink::VortexSink; pub use source::VortexSource; #[cfg(test)] diff --git a/vortex-datafusion/src/persistent/reader.rs b/vortex-datafusion/src/persistent/reader.rs index 357b7531dce..66afb43d75a 100644 --- a/vortex-datafusion/src/persistent/reader.rs +++ b/vortex-datafusion/src/persistent/reader.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Factory for creating [`VortexReadAt`] instances from [`PartitionedFile`]s. +//! Factory for creating [`VortexReadAt`] instances for [`PartitionedFile`]s. use std::fmt::Debug; use std::sync::Arc; diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 9f5e8ad184d..9c02d72c120 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -32,6 +32,7 @@ use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::session::VortexSession; +/// Implements [`DataSink`] for writing Vortex files. pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, @@ -39,6 +40,7 @@ pub struct VortexSink { } impl VortexSink { + /// Creates a new [`VortexSink`] instance. pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { Self { config, diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 086c75adda7..896f8e39829 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -306,6 +306,11 @@ impl VortexSource { self } + /// Returns the predicate this source is going to push down + pub fn predicate(&self) -> Option<&Arc> { + self.vortex_predicate.as_ref() + } + fn create_vortex_opener( &self, object_store: Arc,