Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion vortex-datafusion/src/convert/exprs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ pub trait ExpressionConvertor: Send + Sync {
}
}

/// The default [`ExpressionConvertor`].
/// The default [`ExpressionConvertor`] implementation.
#[derive(Default)]
pub struct DefaultExpressionConvertor {}

Expand Down
19 changes: 16 additions & 3 deletions vortex-datafusion/src/convert/mod.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,32 @@
// 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;
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<D: ?Sized>: Sized {
pub trait FromDataFusion<D: ?Sized>: 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<D> {
/// First-party trait for implementing fallible conversions from Vortex to DataFusion types.
pub trait TryToDataFusion<D> {
/// Try to convert this Vortex type from the input DataFusion type.
fn try_to_df(&self) -> VortexResult<D>;
}
29 changes: 26 additions & 3 deletions vortex-datafusion/src/convert/scalars.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,10 +119,16 @@ impl TryToDataFusion<ScalarValue> for Scalar {
.cloned()
.map(|b| Vec::<u8>::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();
Expand Down Expand Up @@ -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}");
}
}
3 changes: 1 addition & 2 deletions vortex-datafusion/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions vortex-datafusion/src/persistent/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
2 changes: 1 addition & 1 deletion vortex-datafusion/src/persistent/reader.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down
2 changes: 2 additions & 0 deletions vortex-datafusion/src/persistent/sink.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,15 @@ 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,
session: VortexSession,
}

impl VortexSink {
/// Creates a new [`VortexSink`] instance.
pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self {
Self {
config,
Expand Down
5 changes: 5 additions & 0 deletions vortex-datafusion/src/persistent/source.rs
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,11 @@ impl VortexSource {
self
}

/// Returns the predicate this source is going to push down
pub fn predicate(&self) -> Option<&Arc<dyn PhysicalExpr>> {
self.vortex_predicate.as_ref()
}

fn create_vortex_opener(
&self,
object_store: Arc<dyn ObjectStore>,
Expand Down
Loading