From 406fbe05b5ef751aa16765076f7fbf3fb716b9ac Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 5 Aug 2026 15:28:41 +0800 Subject: [PATCH 1/2] update Signed-off-by: Huaijin --- vortex-datafusion/src/convert/exprs.rs | 2 + vortex-datafusion/src/convert/mod.rs | 1 + vortex-datafusion/src/persistent/format.rs | 180 +++++++++++++++++++-- 3 files changed, 173 insertions(+), 10 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 847a020b887..8142c54f560 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -52,7 +52,9 @@ use crate::convert::FromDataFusion; /// Result of splitting a projection into Vortex expressions and leftover DataFusion projections. pub struct ProcessedProjection { + /// Projection evaluated by the Vortex scan. pub scan_projection: Expression, + /// Projection evaluated by DataFusion after the Vortex scan. pub leftover_projection: ProjectionExprs, } diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 6a6fb8aef08..72ee2438220 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -18,6 +18,7 @@ pub(crate) mod stats; pub use exprs::DefaultExpressionConvertor; pub use exprs::ExpressionConvertor; +pub use exprs::ProcessedProjection; /// First-party trait for implementing conversion from DataFusion types to Vortex types. pub trait FromDataFusion: Sized { diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 779b74b212e..97188fe9fdb 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -68,6 +68,7 @@ use super::cache::CachedVortexMetadata; use super::sink::VortexSink; use super::source::VortexSource; use crate::PrecisionExt as _; +use crate::convert::ExpressionConvertor; use crate::convert::TryToDataFusion; use crate::convert::stats::is_constant_to_distinct_count; @@ -123,12 +124,17 @@ const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usi pub struct VortexFormat { session: VortexSession, opts: VortexTableOptions, + expression_convertor: Option>, } impl Debug for VortexFormat { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("VortexFormat") .field("opts", &self.opts) + .field( + "has_expression_convertor", + &self.expression_convertor.is_some(), + ) .finish() } } @@ -270,10 +276,23 @@ impl ConfigExtension for VortexTableOptions { /// ``` /// /// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html -#[derive(Debug)] pub struct VortexFormatFactory { session: VortexSession, options: Option, + expression_convertor: Option>, +} + +impl Debug for VortexFormatFactory { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VortexFormatFactory") + .field("session", &self.session) + .field("options", &self.options) + .field( + "has_expression_convertor", + &self.expression_convertor.is_some(), + ) + .finish() + } } impl GetExt for VortexFormatFactory { @@ -297,6 +316,7 @@ impl VortexFormatFactory { Self { session: VortexSession::default(), options: None, + expression_convertor: None, } } @@ -311,6 +331,7 @@ impl VortexFormatFactory { Self { session, options: Some(options), + expression_convertor: None, } } @@ -337,6 +358,15 @@ impl VortexFormatFactory { self.options = Some(options); self } + + /// Sets the [`ExpressionConvertor`] used by formats and sources created by this factory. + pub fn with_expression_convertor( + mut self, + expression_convertor: Arc, + ) -> Self { + self.expression_convertor = Some(expression_convertor); + self + } } impl FileFormatFactory for VortexFormatFactory { @@ -374,14 +404,19 @@ impl FileFormatFactory for VortexFormatFactory { } } - Ok(Arc::new(VortexFormat::new_with_options( - self.session.clone(), - opts, - ))) + let mut format = VortexFormat::new_with_options(self.session.clone(), opts); + if let Some(expression_convertor) = &self.expression_convertor { + format = format.with_expression_convertor(Arc::clone(expression_convertor)); + } + Ok(Arc::new(format)) } fn default(&self) -> Arc { - Arc::new(VortexFormat::new(self.session.clone())) + let mut format = VortexFormat::new(self.session.clone()); + if let Some(expression_convertor) = &self.expression_convertor { + format = format.with_expression_convertor(Arc::clone(expression_convertor)); + } + Arc::new(format) } } @@ -398,7 +433,11 @@ impl VortexFormat { /// Creates a format with explicit [`VortexTableOptions`]. pub fn new_with_options(session: VortexSession, opts: VortexTableOptions) -> Self { - Self { session, opts } + Self { + session, + opts, + expression_convertor: None, + } } /// Returns the format-specific configuration that will be copied into the @@ -406,6 +445,15 @@ impl VortexFormat { pub fn options(&self) -> &VortexTableOptions { &self.opts } + + /// Sets the [`ExpressionConvertor`] used by every [`VortexSource`] created by this format. + pub fn with_expression_convertor( + mut self, + expression_convertor: Arc, + ) -> Self { + self.expression_convertor = Some(expression_convertor); + self + } } #[async_trait] @@ -702,9 +750,12 @@ impl FileFormat for VortexFormat { } fn file_source(&self, table_schema: TableSchema) -> Arc { - Arc::new( - VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()), - ) as _ + let mut source = + VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()); + if let Some(expression_convertor) = &self.expression_convertor { + source = source.with_expression_convertor(Arc::clone(expression_convertor)); + } + Arc::new(source) as _ } } @@ -730,9 +781,87 @@ fn scalar_stat_to_df( #[cfg(test)] mod tests { + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + + use arrow_schema::DataType; + use arrow_schema::Field; + use datafusion_common::ScalarValue; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions as df_expr; + use datafusion_physical_expr::projection::ProjectionExprs; + use datafusion_physical_plan::filter_pushdown::PushedDown; + use vortex::expr::Expression; use super::*; use crate::common_tests::TestSessionContext; + use crate::convert::DefaultExpressionConvertor; + use crate::convert::ProcessedProjection; + + struct RejectingExpressionConvertor { + inner: DefaultExpressionConvertor, + called: Arc, + } + + impl RejectingExpressionConvertor { + fn new(session: VortexSession, called: Arc) -> Self { + Self { + inner: DefaultExpressionConvertor::new(session), + called, + } + } + } + + impl ExpressionConvertor for RejectingExpressionConvertor { + fn can_be_pushed_down(&self, _expr: &Arc, _schema: &Schema) -> bool { + self.called.store(true, Ordering::Relaxed); + false + } + + fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { + self.inner.convert(expr) + } + + fn split_projection( + &self, + source_projection: ProjectionExprs, + input_schema: &Schema, + output_schema: &Schema, + ) -> DFResult { + self.inner + .split_projection(source_projection, input_schema, output_schema) + } + } + + fn expression_convertor_test_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn expression_convertor_test_filter() -> Arc { + let column = Arc::new(df_expr::Column::new("a", 0)) as Arc; + let literal = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))) as Arc; + Arc::new(df_expr::BinaryExpr::new(column, Operator::Gt, literal)) + } + + fn assert_uses_expression_convertor( + format: &dyn FileFormat, + called: &AtomicBool, + ) -> anyhow::Result<()> { + let source = format.file_source(TableSchema::from_file_schema( + expression_convertor_test_schema(), + )); + let result = source.try_pushdown_filters( + vec![expression_convertor_test_filter()], + &ConfigOptions::new(), + )?; + + assert!(called.load(Ordering::Relaxed)); + assert!(matches!(result.filters.as_slice(), [PushedDown::No])); + Ok(()) + } #[tokio::test] async fn create_table() -> anyhow::Result<()> { @@ -804,4 +933,35 @@ mod tests { assert_eq!(source.options().scan_concurrency, opts.scan_concurrency); Ok(()) } + + #[test] + fn format_plumbs_expression_convertor() -> anyhow::Result<()> { + let session = VortexSession::default(); + let called = Arc::new(AtomicBool::new(false)); + let convertor = Arc::new(RejectingExpressionConvertor::new( + session.clone(), + Arc::clone(&called), + )); + let format = VortexFormat::new(session).with_expression_convertor(convertor); + + assert_uses_expression_convertor(&format, &called) + } + + #[test] + fn factory_plumbs_expression_convertor() -> anyhow::Result<()> { + let called = Arc::new(AtomicBool::new(false)); + let convertor = Arc::new(RejectingExpressionConvertor::new( + VortexSession::default(), + Arc::clone(&called), + )); + let factory = VortexFormatFactory::new().with_expression_convertor(convertor); + let ctx = TestSessionContext::default(); + + let format = factory.create(&ctx.session.state(), &Default::default())?; + assert_uses_expression_convertor(format.as_ref(), &called)?; + + called.store(false, Ordering::Relaxed); + let format = FileFormatFactory::default(&factory); + assert_uses_expression_convertor(format.as_ref(), &called) + } } From 3e6edca9d69149cf1ec823b5c72f3ab65291f1f0 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 5 Aug 2026 16:31:43 +0800 Subject: [PATCH 2/2] update Signed-off-by: Huaijin --- vortex-datafusion/src/convert/exprs.rs | 41 +++++++ vortex-datafusion/src/lib.rs | 8 +- vortex-datafusion/src/persistent/format.rs | 123 +++++++++++++++++---- 3 files changed, 149 insertions(+), 23 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs index 8142c54f560..f9c4c65a460 100644 --- a/vortex-datafusion/src/convert/exprs.rs +++ b/vortex-datafusion/src/convert/exprs.rs @@ -72,6 +72,47 @@ pub(crate) fn make_vortex_predicate( } /// Trait for converting DataFusion expressions to Vortex ones. +/// +/// # Implementing a custom convertor +/// +/// ``` +/// use std::sync::Arc; +/// +/// use arrow_schema::Schema; +/// use datafusion_common::Result as DFResult; +/// use datafusion_physical_expr::PhysicalExpr; +/// use datafusion_physical_expr::projection::ProjectionExprs; +/// use vortex::expr::Expression; +/// use vortex_datafusion::convert::DefaultExpressionConvertor; +/// use vortex_datafusion::convert::ExpressionConvertor; +/// use vortex_datafusion::convert::ProcessedProjection; +/// +/// struct CustomExpressionConvertor(DefaultExpressionConvertor); +/// +/// impl ExpressionConvertor for CustomExpressionConvertor { +/// fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool { +/// self.0.can_be_pushed_down(expr, schema) +/// } +/// +/// fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { +/// self.0.convert(expr) +/// } +/// +/// fn split_projection( +/// &self, +/// source_projection: ProjectionExprs, +/// input_schema: &Schema, +/// output_schema: &Schema, +/// ) -> DFResult { +/// self.0 +/// .split_projection(source_projection, input_schema, output_schema) +/// } +/// } +/// +/// let _convertor: Arc = Arc::new(CustomExpressionConvertor( +/// DefaultExpressionConvertor::default(), +/// )); +/// ``` pub trait ExpressionConvertor: Send + Sync { /// Can an expression be pushed down given a specific schema fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool; diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs index fd430403e55..5de551f59ad 100644 --- a/vortex-datafusion/src/lib.rs +++ b/vortex-datafusion/src/lib.rs @@ -166,12 +166,18 @@ mod common_tests { impl TestSessionContext { /// Create a new test session context with the given projection pushdown setting. pub fn new(projection_pushdown: bool) -> Self { - let store = Arc::new(InMemory::new()); let opts = VortexTableOptions { projection_pushdown, ..Default::default() }; let factory = Arc::new(VortexFormatFactory::new().with_options(opts)); + + Self::new_with_factory(factory) + } + + /// Create a new test session context with the given Vortex format factory. + pub fn new_with_factory(factory: Arc) -> Self { + let store = Arc::new(InMemory::new()); let mut session_state_builder = SessionStateBuilder::new() .with_default_features() .with_table_factory( diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 97188fe9fdb..4e59821dbfa 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -784,6 +784,7 @@ mod tests { use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; + use arrow_array::Int32Array; use arrow_schema::DataType; use arrow_schema::Field; use datafusion_common::ScalarValue; @@ -800,27 +801,56 @@ mod tests { use crate::convert::DefaultExpressionConvertor; use crate::convert::ProcessedProjection; - struct RejectingExpressionConvertor { + #[derive(Clone, Copy)] + enum PushdownMode { + Reject, + Delegate, + } + + #[derive(Default)] + struct ExpressionConvertorCalls { + can_be_pushed_down: AtomicBool, + convert: AtomicBool, + } + + impl ExpressionConvertorCalls { + fn reset(&self) { + self.can_be_pushed_down.store(false, Ordering::Relaxed); + self.convert.store(false, Ordering::Relaxed); + } + } + + struct TestExpressionConvertor { inner: DefaultExpressionConvertor, - called: Arc, + pushdown_mode: PushdownMode, + calls: Arc, } - impl RejectingExpressionConvertor { - fn new(session: VortexSession, called: Arc) -> Self { + impl TestExpressionConvertor { + fn new( + session: VortexSession, + pushdown_mode: PushdownMode, + calls: Arc, + ) -> Self { Self { inner: DefaultExpressionConvertor::new(session), - called, + pushdown_mode, + calls, } } } - impl ExpressionConvertor for RejectingExpressionConvertor { - fn can_be_pushed_down(&self, _expr: &Arc, _schema: &Schema) -> bool { - self.called.store(true, Ordering::Relaxed); - false + impl ExpressionConvertor for TestExpressionConvertor { + fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool { + self.calls.can_be_pushed_down.store(true, Ordering::Relaxed); + match self.pushdown_mode { + PushdownMode::Reject => false, + PushdownMode::Delegate => self.inner.can_be_pushed_down(expr, schema), + } } fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { + self.calls.convert.store(true, Ordering::Relaxed); self.inner.convert(expr) } @@ -846,9 +876,9 @@ mod tests { Arc::new(df_expr::BinaryExpr::new(column, Operator::Gt, literal)) } - fn assert_uses_expression_convertor( + fn assert_rejects_pushdown_with_expression_convertor( format: &dyn FileFormat, - called: &AtomicBool, + calls: &ExpressionConvertorCalls, ) -> anyhow::Result<()> { let source = format.file_source(TableSchema::from_file_schema( expression_convertor_test_schema(), @@ -858,7 +888,8 @@ mod tests { &ConfigOptions::new(), )?; - assert!(called.load(Ordering::Relaxed)); + assert!(calls.can_be_pushed_down.load(Ordering::Relaxed)); + assert!(!calls.convert.load(Ordering::Relaxed)); assert!(matches!(result.filters.as_slice(), [PushedDown::No])); Ok(()) } @@ -937,31 +968,79 @@ mod tests { #[test] fn format_plumbs_expression_convertor() -> anyhow::Result<()> { let session = VortexSession::default(); - let called = Arc::new(AtomicBool::new(false)); - let convertor = Arc::new(RejectingExpressionConvertor::new( + let calls = Arc::new(ExpressionConvertorCalls::default()); + let convertor = Arc::new(TestExpressionConvertor::new( session.clone(), - Arc::clone(&called), + PushdownMode::Reject, + Arc::clone(&calls), )); let format = VortexFormat::new(session).with_expression_convertor(convertor); - assert_uses_expression_convertor(&format, &called) + assert_rejects_pushdown_with_expression_convertor(&format, &calls) } #[test] fn factory_plumbs_expression_convertor() -> anyhow::Result<()> { - let called = Arc::new(AtomicBool::new(false)); - let convertor = Arc::new(RejectingExpressionConvertor::new( + let calls = Arc::new(ExpressionConvertorCalls::default()); + let convertor = Arc::new(TestExpressionConvertor::new( VortexSession::default(), - Arc::clone(&called), + PushdownMode::Reject, + Arc::clone(&calls), )); let factory = VortexFormatFactory::new().with_expression_convertor(convertor); let ctx = TestSessionContext::default(); let format = factory.create(&ctx.session.state(), &Default::default())?; - assert_uses_expression_convertor(format.as_ref(), &called)?; + assert_rejects_pushdown_with_expression_convertor(format.as_ref(), &calls)?; - called.store(false, Ordering::Relaxed); + calls.reset(); let format = FileFormatFactory::default(&factory); - assert_uses_expression_convertor(format.as_ref(), &called) + assert_rejects_pushdown_with_expression_convertor(format.as_ref(), &calls) + } + + #[tokio::test] + async fn external_table_query_uses_factory_expression_convertor() -> anyhow::Result<()> { + let calls = Arc::new(ExpressionConvertorCalls::default()); + let convertor = Arc::new(TestExpressionConvertor::new( + VortexSession::default(), + PushdownMode::Delegate, + Arc::clone(&calls), + )); + let factory = Arc::new(VortexFormatFactory::new().with_expression_convertor(convertor)); + let ctx = TestSessionContext::new_with_factory(factory); + + ctx.session + .sql( + "CREATE EXTERNAL TABLE numbers (a INT NOT NULL) \ + STORED AS vortex LOCATION '/expression-convertor/'", + ) + .await?; + ctx.session + .sql("INSERT INTO numbers VALUES (1), (2), (3)") + .await? + .collect() + .await?; + + calls.reset(); + let batches = ctx + .session + .sql("SELECT a FROM numbers WHERE a > 1 ORDER BY a") + .await? + .collect() + .await?; + + assert!(calls.can_be_pushed_down.load(Ordering::Relaxed)); + assert!(calls.convert.load(Ordering::Relaxed)); + let mut values = Vec::new(); + for batch in batches { + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("expected Int32 result column"))?; + values.extend(array.values().iter().copied()); + } + assert_eq!(values, vec![2, 3]); + Ok(()) } }