feat: Support more datafusion features (#4180)

After spending some time today debugging a Datafusion issue with
@joseph-isaacs and @robert3005 (ended up not being related to us), I
figured I'll use the momentum to add some of the missing features:
1. Support for basic arithmetic expression pushdown (which exists in
TPC-DS), including subtraction to `vortex-expr`.
2. Support files with different (but compatible) schemas in the same
table. This is currently only supported if a
`PhysicalExprAdapterFactory` is provided (either through the
`FileScanConfig` or directly).
3. Support hive-style partitioning columns.

Both features aren't fully fleshed out and probably require more tests,
I want to figure out if we can use datafusion's own test suite to close
that gap, they have pretty cool tooling to test SQL.
One of the changes it required is making the source and opener more
datafusion-y, doing most of the conversion work into Vortex in the
actual `FileOpenFuture`. It allows us to utilize more of datafusion's
tools for rewriting expressions over specific files.

---------

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

Cargo.lock

@@ -87,7 +87,13 @@ crossbeam-queue = "0.3.12"
 crossterm = "0.29"
 dashmap = "6.1.0"
 datafusion = { version = "49", default-features = false }
+datafusion-catalog = { version = "49" }
 datafusion-common = { version = "49" }
+datafusion-common-runtime = { version = "49" }
+datafusion-datasource = { version = "49", default-features = false }
+datafusion-execution = { version = "49" }
+datafusion-expr = { version = "49" }
+datafusion-physical-expr = { version = "49" }
 datafusion-physical-plan = { version = "49" }
 dirs = "6.0.0"
 divan = { package = "codspeed-divan-compat", version = "3.0" }

Cargo.toml

@@ -214,8 +214,10 @@ fn execute_queries<B: Benchmark>(
             && query_idx < expected_counts.len()
         {
             assert_eq!(
-                row_count, expected_counts[query_idx],
-                "Row count mismatch for query {query_idx} - duckdb:{format}",
+                row_count,
+                expected_counts[query_idx],
+                "Row count mismatch for query {query_idx} - {}:{format}",
+                engine_ctx.to_engine()
             );
         }
 
@@ -225,10 +227,7 @@ fn execute_queries<B: Benchmark>(
         {
             memory_measurements.push(MemoryMeasurement {
                 query_idx,
-                target: match engine_ctx {
-                    EngineCtx::DataFusion(_) => Target::new(Engine::DataFusion, format),
-                    EngineCtx::DuckDB(_) => Target::new(Engine::DuckDB, format),
-                },
+                target: Target::new(engine_ctx.to_engine(), format),
                 benchmark_dataset: benchmark.dataset(),
                 storage: url_scheme_to_storage(benchmark.data_url())?,
                 physical_memory_delta: memory_result.physical_memory_delta,

bench-vortex/src/benchmark_driver.rs

@@ -14,10 +14,18 @@ rust-version = { workspace = true }
 version = { workspace = true }
 
 [dependencies]
+arrow-schema = { workspace = true }
 async-trait = { workspace = true }
 chrono = { workspace = true }
 dashmap = { workspace = true }
-datafusion = { workspace = true }
+datafusion-catalog = { workspace = true }
+datafusion-common = { workspace = true }
+datafusion-common-runtime = { workspace = true }
+datafusion-datasource = { workspace = true, default-features = false }
+datafusion-execution = { workspace = true }
+datafusion-expr = { workspace = true }
+datafusion-physical-expr = { workspace = true }
+datafusion-physical-plan = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
@@ -30,6 +38,7 @@ vortex = { workspace = true, features = ["object_store", "tokio"] }
 
 [dev-dependencies]
 anyhow = { workspace = true }
+datafusion = { workspace = true }
 rstest = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "fs"] }

vortex-datafusion/Cargo.toml

@@ -1,39 +1,67 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use datafusion::logical_expr::Operator as DFOperator;
-use datafusion::physical_expr::{PhysicalExpr, expressions};
+use std::sync::Arc;
+
+use arrow_schema::{DataType, Schema};
+use datafusion_expr::Operator as DFOperator;
+use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef};
+use datafusion_physical_plan::expressions as df_expr;
 use vortex::error::{VortexResult, vortex_bail, vortex_err};
-use vortex::expr::{BinaryExpr, ExprRef, IntoExpr, LikeExpr, Operator, get_item, lit, root};
+use vortex::expr::{BinaryExpr, ExprRef, LikeExpr, Operator, and, get_item, lit, root};
 use vortex::scalar::Scalar;
 
 use crate::convert::{FromDataFusion, TryFromDataFusion};
 
+const SUPPORTED_BINARY_OPS: &[DFOperator] = &[
+    DFOperator::Eq,
+    DFOperator::NotEq,
+    DFOperator::Gt,
+    DFOperator::GtEq,
+    DFOperator::Lt,
+    DFOperator::LtEq,
+];
+
+/// Tries to convert the expressions into a vortex conjunction. Will return Ok(None) iff the input conjunction is empty.
+pub(crate) fn make_vortex_predicate(
+    predicate: &[&Arc<dyn PhysicalExpr>],
+) -> VortexResult<Option<ExprRef>> {
+    let exprs = predicate
+        .iter()
+        .map(|e| ExprRef::try_from_df(e.as_ref()))
+        .collect::<VortexResult<Vec<_>>>()?;
+
+    Ok(exprs.into_iter().reduce(and))
+}
+
 // TODO(joe): Don't return an error when we have an unsupported node, bubble up "TRUE" as in keep
 //  for that node, up to any `and` or `or` node.
 impl TryFromDataFusion<dyn PhysicalExpr> for ExprRef {
     fn try_from_df(df: &dyn PhysicalExpr) -> VortexResult<Self> {
-        if let Some(binary_expr) = df.as_any().downcast_ref::<expressions::BinaryExpr>() {
+        if let Some(binary_expr) = df.as_any().downcast_ref::<df_expr::BinaryExpr>() {
             let left = ExprRef::try_from_df(binary_expr.left().as_ref())?;
             let right = ExprRef::try_from_df(binary_expr.right().as_ref())?;
             let operator = Operator::try_from_df(binary_expr.op())?;
 
             return Ok(BinaryExpr::new_expr(left, operator, right));
         }
 
-        if let Some(col_expr) = df.as_any().downcast_ref::<expressions::Column>() {
+        if let Some(col_expr) = df.as_any().downcast_ref::<df_expr::Column>() {
             return Ok(get_item(col_expr.name().to_owned(), root()));
         }
 
-        if let Some(like) = df.as_any().downcast_ref::<expressions::LikeExpr>() {
+        if let Some(like) = df.as_any().downcast_ref::<df_expr::LikeExpr>() {
             let child = ExprRef::try_from_df(like.expr().as_ref())?;
             let pattern = ExprRef::try_from_df(like.pattern().as_ref())?;
-            return Ok(
-                LikeExpr::new(child, pattern, like.negated(), like.case_insensitive()).into_expr(),
-            );
+            return Ok(LikeExpr::new_expr(
+                child,
+                pattern,
+                like.negated(),
+                like.case_insensitive(),
+            ));
         }
 
-        if let Some(literal) = df.as_any().downcast_ref::<expressions::Literal>() {
+        if let Some(literal) = df.as_any().downcast_ref::<df_expr::Literal>() {
             let value = Scalar::from_df(literal.value());
             return Ok(lit(value));
         }
@@ -92,3 +120,55 @@ impl TryFromDataFusion<DFOperator> for Operator {
         }
     }
 }
+
+pub(crate) fn can_be_pushed_down(expr: &PhysicalExprRef, schema: &Schema) -> bool {
+    let expr = expr.as_any();
+    if let Some(binary) = expr.downcast_ref::<df_expr::BinaryExpr>() {
+        can_binary_be_pushed_down(binary, schema)
+    } else if let Some(col) = expr.downcast_ref::<df_expr::Column>() {
+        schema
+            .field_with_name(col.name())
+            .ok()
+            .is_some_and(|field| supported_data_types(field.data_type()))
+    } else if let Some(like) = expr.downcast_ref::<df_expr::LikeExpr>() {
+        can_be_pushed_down(like.expr(), schema) && can_be_pushed_down(like.pattern(), schema)
+    } else if let Some(lit) = expr.downcast_ref::<df_expr::Literal>() {
+        supported_data_types(&lit.value().data_type())
+    } else {
+        log::debug!("DataFusion expression can't be pushed down: {expr:?}");
+        false
+    }
+}
+
+fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool {
+    let is_op_supported =
+        binary.op().is_logic_operator() || SUPPORTED_BINARY_OPS.contains(binary.op());
+    is_op_supported
+        && can_be_pushed_down(binary.left(), schema)
+        && can_be_pushed_down(binary.right(), schema)
+}
+
+fn supported_data_types(dt: &DataType) -> bool {
+    use DataType::*;
+    let is_supported = dt.is_null()
+        || dt.is_numeric()
+        || matches!(
+            dt,
+            Boolean
+                | Utf8
+                | Utf8View
+                | Binary
+                | BinaryView
+                | Date32
+                | Date64
+                | Timestamp(_, _)
+                | Time32(_)
+                | Time64(_)
+        );
+
+    if !is_supported {
+        log::debug!("DataFusion data type {dt:?} is not supported");
+    }
+
+    is_supported
+}

vortex-datafusion/src/convert/exprs.rs

@@ -3,7 +3,7 @@
 
 use vortex::error::VortexResult;
 
-mod exprs;
+pub(crate) mod exprs;
 mod scalars;
 
 /// First-party trait for implementing conversion from DataFusion types to Vortex types.

vortex-datafusion/src/convert/mod.rs

@@ -3,7 +3,7 @@
 
 use std::sync::Arc;
 
-use datafusion::common::ScalarValue;
+use datafusion_common::ScalarValue;
 use vortex::buffer::ByteBuffer;
 use vortex::dtype::datetime::arrow::make_temporal_ext_dtype;
 use vortex::dtype::datetime::{TemporalMetadata, TimeUnit, is_temporal_ext_type};

vortex-datafusion/src/convert/scalars.rs

@@ -5,74 +5,14 @@
 #![deny(missing_docs)]
 use std::fmt::Debug;
 
-use datafusion::arrow::datatypes::{DataType, Schema};
-use datafusion::common::stats::Precision as DFPrecision;
-use datafusion::logical_expr::Operator;
-use datafusion::physical_expr::PhysicalExprRef;
-use datafusion::physical_plan::expressions::{BinaryExpr, Column, LikeExpr, Literal};
+use datafusion_common::stats::Precision as DFPrecision;
 use vortex::stats::Precision;
 
 mod convert;
 mod persistent;
 
 pub use persistent::*;
 
-const SUPPORTED_BINARY_OPS: &[Operator] = &[
-    Operator::Eq,
-    Operator::NotEq,
-    Operator::Gt,
-    Operator::GtEq,
-    Operator::Lt,
-    Operator::LtEq,
-];
-
-fn supported_data_types(dt: DataType) -> bool {
-    use DataType::*;
-    let is_supported = dt.is_integer()
-        || dt.is_floating()
-        || dt.is_null()
-        || matches!(
-            dt,
-            Boolean
-                | Utf8
-                | Utf8View
-                | Binary
-                | BinaryView
-                | Date32
-                | Date64
-                | Timestamp(_, _)
-                | Time32(_)
-                | Time64(_)
-        );
-
-    if !is_supported {
-        log::debug!("DataFusion data type {dt:?} is not supported");
-    }
-
-    is_supported
-}
-
-fn can_be_pushed_down(expr: &PhysicalExprRef, schema: &Schema) -> bool {
-    let expr = expr.as_any();
-    if let Some(binary) = expr.downcast_ref::<BinaryExpr>() {
-        (binary.op().is_logic_operator() || SUPPORTED_BINARY_OPS.contains(binary.op()))
-            && can_be_pushed_down(binary.left(), schema)
-            && can_be_pushed_down(binary.right(), schema)
-    } else if let Some(col) = expr.downcast_ref::<Column>() {
-        schema
-            .column_with_name(col.name())
-            .map(|(_, field)| supported_data_types(field.data_type().clone()))
-            .unwrap_or(false)
-    } else if let Some(like) = expr.downcast_ref::<LikeExpr>() {
-        can_be_pushed_down(like.expr(), schema) && can_be_pushed_down(like.pattern(), schema)
-    } else if let Some(lit) = expr.downcast_ref::<Literal>() {
-        supported_data_types(lit.value().data_type())
-    } else {
-        log::debug!("DataFusion expression can't be pushed down: {expr:?}");
-        false
-    }
-}
-
 /// Extension trait to convert our [`Precision`](vortex::stats::Precision) to Datafusion's [`Precision`](datafusion_common::stats::Precision)
 trait PrecisionExt<T>
 where

vortex-datafusion/src/lib.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
-use datafusion::common::ScalarValue;
+use datafusion_common::ScalarValue;
 use moka::future::Cache;
 use object_store::path::Path;
 use object_store::{ObjectMeta, ObjectStore};