Move more things to the convertor

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

Author: AdamGS

.gitignore

@@ -224,3 +224,6 @@ compile_commands.json
 # cargo sweep output
 sweep.timestamp
 
+# Pending cargo-insta snapshots
+*.pending-snap
+

vortex-datafusion/src/convert/exprs.rs

@@ -5,20 +5,24 @@ use std::sync::Arc;
 
 use arrow_schema::DataType;
 use arrow_schema::Schema;
+use datafusion_common::Result as DFResult;
+use datafusion_common::exec_datafusion_err;
+use datafusion_common::tree_node::TreeNode;
+use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_expr::Operator as DFOperator;
 use datafusion_functions::core::getfield::GetFieldFunc;
 use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_expr::ScalarFunctionExpr;
+use datafusion_physical_expr::projection::ProjectionExpr;
+use datafusion_physical_expr::projection::ProjectionExprs;
+use datafusion_physical_expr::utils::collect_columns;
 use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr;
 use datafusion_physical_plan::expressions as df_expr;
 use itertools::Itertools;
 use vortex::compute::LikeOptions;
 use vortex::dtype::DType;
 use vortex::dtype::Nullability;
 use vortex::dtype::arrow::FromArrowType;
-use vortex::error::VortexResult;
-use vortex::error::vortex_bail;
-use vortex::error::vortex_err;
 use vortex::expr::Binary;
 use vortex::expr::Expression;
 use vortex::expr::Like;
@@ -31,20 +35,27 @@ use vortex::expr::is_null;
 use vortex::expr::list_contains;
 use vortex::expr::lit;
 use vortex::expr::not;
+use vortex::expr::pack;
 use vortex::expr::root;
 use vortex::scalar::Scalar;
 
 use crate::convert::FromDataFusion;
 
+/// Result of splitting a projection into Vortex expressions and leftover DataFusion projections.
+pub struct ProcessedProjection {
+    pub scan_projection: Expression,
+    pub leftover_projection: ProjectionExprs,
+}
+
 /// 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(
     expr_convertor: &dyn ExpressionConvertor,
     predicate: &[Arc<dyn PhysicalExpr>],
-) -> VortexResult<Option<Expression>> {
+) -> DFResult<Option<Expression>> {
     let exprs = predicate
         .iter()
         .map(|e| expr_convertor.convert(e.as_ref()))
-        .collect::<VortexResult<Vec<_>>>()?;
+        .collect::<DFResult<Vec<_>>>()?;
 
     Ok(exprs.into_iter().reduce(and))
 }
@@ -55,7 +66,16 @@ pub trait ExpressionConvertor: Send + Sync {
     fn can_be_pushed_down(&self, expr: &Arc<dyn PhysicalExpr>, schema: &Schema) -> bool;
 
     /// Try and convert a DataFusion [`PhysicalExpr`] into a Vortex [`Expression`].
-    fn convert(&self, expr: &dyn PhysicalExpr) -> VortexResult<Expression>;
+    fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult<Expression>;
+
+    /// Split a projection into Vortex expressions that can be pushed down and leftover
+    /// DataFusion projections that need to be evaluated after the scan.
+    fn split_projection(
+        &self,
+        source_projection: ProjectionExprs,
+        input_schema: &Schema,
+        output_schema: &Schema,
+    ) -> DFResult<ProcessedProjection>;
 }
 
 /// The default [`ExpressionConvertor`].
@@ -64,37 +84,35 @@ pub struct DefaultExpressionConvertor {}
 
 impl DefaultExpressionConvertor {
     /// Attempts to convert a DataFusion ScalarFunctionExpr to a Vortex expression.
-    fn try_convert_scalar_function(
-        &self,
-        scalar_fn: &ScalarFunctionExpr,
-    ) -> VortexResult<Expression> {
+    fn try_convert_scalar_function(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult<Expression> {
         if let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn)
         {
             let source_expr = get_field_fn
                 .args()
                 .first()
-                .ok_or_else(|| vortex_err!("get_field missing source expression"))?
+                .ok_or_else(|| exec_datafusion_err!("get_field missing source expression"))?
                 .as_ref();
             let field_name_expr = get_field_fn
                 .args()
                 .get(1)
-                .ok_or_else(|| vortex_err!("get_field missing field name argument"))?;
+                .ok_or_else(|| exec_datafusion_err!("get_field missing field name argument"))?;
             let field_name = field_name_expr
                 .as_any()
                 .downcast_ref::<df_expr::Literal>()
-                .ok_or_else(|| vortex_err!("get_field field name must be a literal"))?
+                .ok_or_else(|| exec_datafusion_err!("get_field field name must be a literal"))?
                 .value()
                 .try_as_str()
                 .flatten()
-                .ok_or_else(|| vortex_err!("get_field field name must be a UTF-8 string"))?;
+                .ok_or_else(|| {
+                    exec_datafusion_err!("get_field field name must be a UTF-8 string")
+                })?;
             return Ok(get_item(field_name.to_string(), self.convert(source_expr)?));
         }
 
-        tracing::debug!(
-            function_name = scalar_fn.name(),
-            "Unsupported ScalarFunctionExpr"
-        );
-        vortex_bail!("Unsupported ScalarFunctionExpr: {}", scalar_fn.name())
+        Err(exec_datafusion_err!(
+            "Unsupported ScalarFunctionExpr: {}",
+            scalar_fn.name()
+        ))
     }
 }
 
@@ -103,7 +121,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor {
         can_be_pushed_down_impl(expr, schema)
     }
 
-    fn convert(&self, df: &dyn PhysicalExpr) -> VortexResult<Expression> {
+    fn convert(&self, df: &dyn PhysicalExpr) -> DFResult<Expression> {
         // 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.
         if let Some(binary_expr) = df.as_any().downcast_ref::<df_expr::BinaryExpr>() {
@@ -168,7 +186,7 @@ impl ExpressionConvertor for DefaultExpressionConvertor {
                     if let Some(lit) = e.as_any().downcast_ref::<df_expr::Literal>() {
                         Ok(Scalar::from_df(lit.value()))
                     } else {
-                        Err(vortex_err!("Failed to cast sub-expression"))
+                        Err(exec_datafusion_err!("Failed to cast sub-expression"))
                     }
                 })
                 .try_collect()?;
@@ -187,11 +205,93 @@ impl ExpressionConvertor for DefaultExpressionConvertor {
             return self.try_convert_scalar_function(scalar_fn);
         }
 
-        vortex_bail!("Couldn't convert DataFusion physical {df} expression to a vortex expression")
+        Err(exec_datafusion_err!(
+            "Couldn't convert DataFusion physical {df} expression to a vortex expression"
+        ))
+    }
+
+    fn split_projection(
+        &self,
+        source_projection: ProjectionExprs,
+        input_schema: &Schema,
+        output_schema: &Schema,
+    ) -> DFResult<ProcessedProjection> {
+        let mut scan_projection = vec![];
+        let mut leftover_projection: Vec<ProjectionExpr> = vec![];
+
+        for projection_expr in source_projection.iter() {
+            let r = projection_expr.expr.apply(|node| {
+                // We only pull column children of scalar functions that we can't push into the scan.
+                if let Some(scalar_fn_expr) = node.as_any().downcast_ref::<ScalarFunctionExpr>()
+                    && !can_scalar_fn_be_pushed_down(scalar_fn_expr)
+                {
+                    scan_projection.extend(
+                        collect_columns(node)
+                            .into_iter()
+                            .map(|c| (c.name().to_string(), get_item(c.name(), root()))),
+                    );
+
+                    leftover_projection.push(projection_expr.clone());
+                    return Ok(TreeNodeRecursion::Stop);
+                }
+
+                // If the projection contains a `CastColumnExpr` that casts to physical types that don't have a 1:1 mapping
+                // with Vortex's types system, we make sure to pull the input from the file and keep the projection
+                if let Some(cast_expr) = node.as_any().downcast_ref::<df_expr::CastColumnExpr>()
+                    && is_dtype_incompatible(cast_expr.target_field().data_type())
+                {
+                    scan_projection.push((
+                        cast_expr.input_field().name().clone(),
+                        self.convert(cast_expr.expr().as_ref())?,
+                    ));
+                    leftover_projection.push(projection_expr.clone());
+                    return Ok(TreeNodeRecursion::Stop);
+                }
+
+                // DataFusion assumes different decimal types can be coerced.
+                // Vortex expects a perfect match so we don't push it down.
+                if let Some(binary_expr) = node.as_any().downcast_ref::<df_expr::BinaryExpr>()
+                    && binary_expr.op().is_numerical_operators()
+                    && (is_decimal(&binary_expr.left().data_type(input_schema)?)
+                        && is_decimal(&binary_expr.right().data_type(input_schema)?))
+                {
+                    scan_projection.extend(
+                        collect_columns(node)
+                            .into_iter()
+                            .map(|c| (c.name().to_string(), get_item(c.name(), root()))),
+                    );
+
+                    leftover_projection.push(projection_expr.clone());
+                    return Ok(TreeNodeRecursion::Stop);
+                }
+
+                Ok(TreeNodeRecursion::Continue)
+            })?;
+
+            // if we didn't stop early
+            if matches!(r, TreeNodeRecursion::Continue) {
+                scan_projection.push((
+                    projection_expr.alias.clone(),
+                    self.convert(projection_expr.expr.as_ref())?,
+                ));
+                leftover_projection.push(ProjectionExpr {
+                    expr: Arc::new(df_expr::Column::new_with_schema(
+                        projection_expr.alias.as_str(),
+                        output_schema,
+                    )?),
+                    alias: projection_expr.alias.clone(),
+                });
+            }
+        }
+
+        Ok(ProcessedProjection {
+            scan_projection: pack(scan_projection, Nullability::NonNullable),
+            leftover_projection: leftover_projection.into(),
+        })
     }
 }
 
-fn try_operator_from_df(value: &DFOperator) -> VortexResult<Operator> {
+fn try_operator_from_df(value: &DFOperator) -> DFResult<Operator> {
     match value {
         DFOperator::Eq => Ok(Operator::Eq),
         DFOperator::NotEq => Ok(Operator::NotEq),
@@ -236,7 +336,9 @@ fn try_operator_from_df(value: &DFOperator) -> VortexResult<Operator> {
         | DFOperator::QuestionAnd
         | DFOperator::QuestionPipe => {
             tracing::debug!(operator = %value, "Can't pushdown binary_operator operator");
-            Err(vortex_err!("Unsupported datafusion operator {value}"))
+            Err(exec_datafusion_err!(
+                "Unsupported datafusion operator {value}"
+            ))
         }
     }
 }
@@ -328,6 +430,35 @@ fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr) -> bool {
     ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(scalar_fn).is_some()
 }
 
+fn is_dtype_incompatible(dt: &DataType) -> bool {
+    use DataType::*;
+
+    dt.is_run_ends_type()
+        || is_decimal(dt)
+        || matches!(
+            dt,
+            Dictionary(..)
+                | Utf8
+                | LargeUtf8
+                | Binary
+                | LargeBinary
+                | FixedSizeBinary(_)
+                | FixedSizeList(..)
+                | ListView(..)
+                | LargeListView(..)
+        )
+}
+
+fn is_decimal(dt: &DataType) -> bool {
+    matches!(
+        dt,
+        DataType::Decimal32(_, _)
+            | DataType::Decimal64(_, _)
+            | DataType::Decimal128(_, _)
+            | DataType::Decimal256(_, _)
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

vortex-datafusion/src/persistent/mod.rs

@@ -7,7 +7,6 @@ mod cache;
 mod format;
 pub mod metrics;
 mod opener;
-mod projection;
 mod sink;
 mod source;