diff --git a/Cargo.lock b/Cargo.lock index f7d3103c9d5..fdfd4062858 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10899,6 +10899,7 @@ dependencies = [ "tracing", "url", "vortex", + "vortex-array", "vortex-arrow", "vortex-utils", ] diff --git a/docs/developer-guide/integrations/datafusion.md b/docs/developer-guide/integrations/datafusion.md index 93f45be0a8d..8ccfb8cdbb2 100644 --- a/docs/developer-guide/integrations/datafusion.md +++ b/docs/developer-guide/integrations/datafusion.md @@ -40,7 +40,7 @@ of parallelism. ## Filter and Projection Pushdown The integration converts DataFusion physical expressions into Vortex expressions using an -`ExpressionConvertor` trait. Supported predicates (comparisons, LIKE, IS NULL, IN lists, casts) +`ExpressionConverter` trait. Supported predicates (comparisons, LIKE, IS NULL, IN lists, casts) are pushed into the Vortex scan where they participate in pruning and filter evaluation at the layout level. Unsupported predicates remain in the DataFusion plan and are evaluated after the scan. @@ -52,7 +52,7 @@ efficiently is pushed into the per-file scan for row-level filtering. Projection pushdown maps DataFusion's requested column indices to Vortex field names and passes them as a projection expression to the scan. Only the requested columns are read from storage. -The integration supports pluggable expression conversion via a custom `ExpressionConvertor`, +The integration supports pluggable expression conversion via a custom `ExpressionConverter`, allowing engine-specific rewrites or schema adaptation when file schemas diverge from the table schema. diff --git a/vortex-arrow/src/dtype.rs b/vortex-arrow/src/dtype.rs index f58cdf1cafb..8d539d0fde6 100644 --- a/vortex-arrow/src/dtype.rs +++ b/vortex-arrow/src/dtype.rs @@ -171,7 +171,7 @@ pub(crate) fn from_arrow_data_type( | DataType::Decimal64(precision, scale) | DataType::Decimal128(precision, scale) | DataType::Decimal256(precision, scale) => { - DType::Decimal(DecimalDType::new(*precision, *scale), nullability) + DType::Decimal(DecimalDType::try_new(*precision, *scale)?, nullability) } DataType::Boolean => DType::Bool(nullability), DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DType::Utf8(nullability), @@ -612,6 +612,17 @@ mod test { assert_eq!(dtype.to_arrow_dtype().unwrap(), expected); } + #[rstest] + #[case::decimal32(DataType::Decimal32(1, 2))] + #[case::decimal64(DataType::Decimal64(1, 2))] + #[case::decimal128(DataType::Decimal128(1, 2))] + #[case::decimal256(DataType::Decimal256(1, 2))] + #[case::zero_precision(DataType::Decimal128(0, 0))] + #[case::excessive_precision(DataType::Decimal256(77, 0))] + fn test_malformed_decimal_dtype_from_arrow(#[case] data_type: DataType) { + assert!(from_arrow_data_type(&data_type, Nullability::Nullable).is_err()); + } + #[test] fn test_variant_dtype_to_arrow_dtype_errors() { let err = DType::Variant(Nullability::NonNullable) diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml index 49fb22d4f59..cb60afa7c14 100644 --- a/vortex-datafusion/Cargo.toml +++ b/vortex-datafusion/Cargo.toml @@ -48,6 +48,7 @@ rstest = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "fs"] } url = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } [lints] workspace = true diff --git a/vortex-datafusion/src/convert/exprs.rs b/vortex-datafusion/src/convert/exprs.rs deleted file mode 100644 index 2ab975ecfd7..00000000000 --- a/vortex-datafusion/src/convert/exprs.rs +++ /dev/null @@ -1,1314 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use arrow_schema::DataType; -use arrow_schema::Field; -use arrow_schema::Schema; -use datafusion_common::Result as DFResult; -use datafusion_common::ScalarValue; -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_functions::string::octet_length::OctetLengthFunc; -use datafusion_functions_nested::length::ArrayLength; -use datafusion_physical_expr::DynamicFilterTracking; -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_plan::expressions as df_expr; -use itertools::Itertools; -use vortex::VortexSessionDefault; -use vortex::dtype::Nullability; -use vortex::expr::Expression; -use vortex::expr::and_collect; -use vortex::expr::byte_length; -use vortex::expr::cast; -use vortex::expr::get_item; -use vortex::expr::is_not_null; -use vortex::expr::is_null; -use vortex::expr::list_contains; -use vortex::expr::list_length; -use vortex::expr::lit; -use vortex::expr::nested_case_when; -use vortex::expr::not; -use vortex::expr::pack; -use vortex::expr::root; -use vortex::scalar::Scalar; -use vortex::scalar_fn::ScalarFnVTableExt; -use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::like::Like; -use vortex::scalar_fn::fns::like::LikeOptions; -use vortex::scalar_fn::fns::operators::Operator; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; - -use crate::convert::scalar_from_df; - -/// 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, -} - -/// 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], -) -> DFResult> { - let exprs = predicate - .iter() - .map(|e| expr_convertor.convert(e.as_ref())) - .collect::>>()?; - - Ok(and_collect(exprs)) -} - -/// 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; - - /// Try and convert a DataFusion [`PhysicalExpr`] into a Vortex [`Expression`]. - fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult; - - /// 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; - - /// Create a projection that reads only the required columns without pushing down - /// any expressions. All projection logic is applied after the scan. - fn no_pushdown_projection( - &self, - source_projection: ProjectionExprs, - input_schema: &Schema, - ) -> DFResult { - // Get all unique column indices referenced by the projection - let column_indices = source_projection.column_indices(); - - // Create scan projection that reads the required columns - let scan_columns: Vec<(String, Expression)> = column_indices - .into_iter() - .map(|idx| { - let field = input_schema.field(idx); - let name = field.name().clone(); - (name.clone(), get_item(name, root())) - }) - .collect(); - - Ok(ProcessedProjection { - scan_projection: pack(scan_columns, Nullability::NonNullable), - leftover_projection: source_projection, - }) - } -} - -/// The default [`ExpressionConvertor`] implementation. -pub struct DefaultExpressionConvertor { - /// Session used to resolve Arrow → Vortex dtypes through the extension - /// plugin registry, so registered extension types (e.g. UUID ⇄ - /// `FixedSizeBinary[16]`) convert correctly instead of hitting the static, - /// non-plugin-aware `DType::from_arrow`. - session: VortexSession, -} - -impl Default for DefaultExpressionConvertor { - fn default() -> Self { - Self { - session: VortexSession::default(), - } - } -} - -impl DefaultExpressionConvertor { - /// Create a convertor that resolves Arrow extension types using `session`'s - /// dtype registry. - pub fn new(session: VortexSession) -> Self { - Self { session } - } - - /// Attempts to convert DataFusion's `octet_length` function to Vortex `byte_length`. - fn try_convert_octet_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult { - let [input] = scalar_fn.args() else { - return Err(exec_datafusion_err!( - "octet_length requires exactly one argument" - )); - }; - - let input = self.convert(input.as_ref())?; - let return_dtype = self - .session - .arrow() - .from_arrow_field(&Field::new( - "", - scalar_fn.return_type().clone(), - scalar_fn.nullable(), - )) - .map_err(|e| exec_datafusion_err!("Failed to convert return type to dtype: {e}"))?; - Ok(cast(byte_length(input), return_dtype)) - } - - /// Attempts to convert DataFusion's `array_length` function (aliased as `list_length`) to - /// Vortex `list_length`. - /// - /// Supports the single-argument form `array_length(arr)` and the equivalent two-argument - /// form with an explicit first dimension `array_length(arr, 1)`. - fn try_convert_array_length(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult { - let Some(input) = array_length_input(scalar_fn) else { - return Err(exec_datafusion_err!( - "array_length pushdown supports only the one-argument form or an explicit first \ - dimension" - )); - }; - - let input = self.convert(input.as_ref())?; - let return_dtype = self - .session - .arrow() - .from_arrow_field(&Field::new( - "", - scalar_fn.return_type().clone(), - scalar_fn.nullable(), - )) - .map_err(|e| exec_datafusion_err!("Failed to convert return type to dtype: {e}"))?; - Ok(cast(list_length(input), return_dtype)) - } - - /// Attempts to convert a DataFusion ScalarFunctionExpr to a Vortex expression. - fn try_convert_scalar_function(&self, scalar_fn: &ScalarFunctionExpr) -> DFResult { - if let Some(octet_length_fn) = - ScalarFunctionExpr::try_downcast_func::(scalar_fn) - { - return self.try_convert_octet_length(octet_length_fn); - } - - if let Some(array_length_fn) = - ScalarFunctionExpr::try_downcast_func::(scalar_fn) - { - return self.try_convert_array_length(array_length_fn); - } - - if let Some(get_field_fn) = ScalarFunctionExpr::try_downcast_func::(scalar_fn) - { - // DataFusion's GetFieldFunc flattens nested field access into a single call - // with multiple field name arguments. For example, `outer.inner.leaf` becomes - // get_field(Column("outer"), "inner", "leaf"). We build a chain of get_item - // calls for each field name in the path. - let (source_expr, field_names) = get_field_fn - .args() - .split_first() - .ok_or_else(|| exec_datafusion_err!("get_field missing source expression"))?; - - let mut result = self.convert(source_expr.as_ref())?; - for expr in field_names { - let field_name = expr - .downcast_ref::() - .ok_or_else(|| exec_datafusion_err!("get_field field name must be a literal"))? - .value() - .try_as_str() - .flatten() - .ok_or_else(|| { - exec_datafusion_err!("get_field field name must be a UTF-8 string") - })?; - result = get_item(field_name.to_string(), result); - } - return Ok(result); - } - - Err(exec_datafusion_err!( - "Unsupported ScalarFunctionExpr: {}", - scalar_fn.name() - )) - } - - /// Attempts to convert a DataFusion CaseExpr to a Vortex expression. - fn try_convert_case_expr(&self, case_expr: &df_expr::CaseExpr) -> DFResult { - // DataFusion CaseExpr has: - // - expr(): Optional base expression (for "CASE expr WHEN ..." form) - // - when_then_expr(): Vec of (when, then) pairs - // - else_expr(): Optional else expression - - // We don't support the "CASE expr WHEN value1 THEN result1" form yet - if case_expr.expr().is_some() { - return Err(exec_datafusion_err!( - "CASE expr WHEN form is not yet supported, only searched CASE is supported" - )); - } - - let when_then_pairs = case_expr.when_then_expr(); - if when_then_pairs.is_empty() { - return Err(exec_datafusion_err!( - "CASE expression must have at least one WHEN clause" - )); - } - - // Convert all when/then pairs to (condition, value) tuples - let mut pairs = Vec::with_capacity(when_then_pairs.len()); - for (when_expr, then_expr) in when_then_pairs { - let condition = self.convert(when_expr.as_ref())?; - let value = self.convert(then_expr.as_ref())?; - pairs.push((condition, value)); - } - - // Convert optional else expression - let else_value = case_expr - .else_expr() - .map(|e| self.convert(e.as_ref())) - .transpose()?; - - // Build a single n-ary CASE WHEN expression from DataFusion WHEN/THEN pairs - Ok(nested_case_when(pairs, else_value)) - } -} - -impl ExpressionConvertor for DefaultExpressionConvertor { - fn can_be_pushed_down(&self, expr: &Arc, schema: &Schema) -> bool { - can_be_pushed_down_impl(expr, schema) - } - - fn convert(&self, df: &dyn PhysicalExpr) -> DFResult { - // 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.downcast_ref::() { - let left = self.convert(binary_expr.left().as_ref())?; - let right = self.convert(binary_expr.right().as_ref())?; - let operator = try_operator_from_df(binary_expr.op())?; - - return Ok(Binary.new_expr(operator, [left, right])); - } - - if let Some(col_expr) = df.downcast_ref::() { - return Ok(get_item(col_expr.name().to_owned(), root())); - } - - if let Some(like) = df.downcast_ref::() { - let child = self.convert(like.expr().as_ref())?; - let pattern = self.convert(like.pattern().as_ref())?; - return Ok(Like.new_expr( - LikeOptions { - negated: like.negated(), - case_insensitive: like.case_insensitive(), - }, - [child, pattern], - )); - } - - if let Some(literal) = df.downcast_ref::() { - let value = scalar_from_df(literal.value(), &self.session); - return Ok(lit(value)); - } - - if let Some(cast_expr) = df.downcast_ref::() { - let cast_dtype = self - .session - .arrow() - .from_arrow_field(cast_expr.target_field().as_ref()) - .map_err(|e| exec_datafusion_err!("Failed to convert cast target to dtype: {e}"))?; - let child = self.convert(cast_expr.expr().as_ref())?; - return Ok(cast(child, cast_dtype)); - } - - if let Some(is_null_expr) = df.downcast_ref::() { - let arg = self.convert(is_null_expr.arg().as_ref())?; - return Ok(is_null(arg)); - } - - if let Some(is_not_null_expr) = df.downcast_ref::() { - let arg = self.convert(is_not_null_expr.arg().as_ref())?; - return Ok(is_not_null(arg)); - } - - if let Some(in_list) = df.downcast_ref::() { - let value = self.convert(in_list.expr().as_ref())?; - let list_elements: Vec<_> = in_list - .list() - .iter() - .map(|e| { - if let Some(lit) = e.downcast_ref::() { - Ok(scalar_from_df(lit.value(), &self.session)) - } else { - Err(exec_datafusion_err!("Failed to cast sub-expression")) - } - }) - .try_collect()?; - - let list = Scalar::list( - list_elements[0].dtype().clone(), - list_elements, - Nullability::Nullable, - ); - let expr = list_contains(lit(list), value); - - return Ok(if in_list.negated() { not(expr) } else { expr }); - } - - if let Some(scalar_fn) = df.downcast_ref::() { - return self.try_convert_scalar_function(scalar_fn); - } - - if let Some(case_expr) = df.downcast_ref::() { - return self.try_convert_case_expr(case_expr); - } - - 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 { - let mut scan_projection = vec![]; - let mut leftover_projection: Vec = 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.downcast_ref::() - && !can_scalar_fn_be_pushed_down(scalar_fn_expr, 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); - } - - // 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.downcast_ref::() - && binary_expr.op().is_numerical_operators() - && binary_expr.left().data_type(input_schema)?.is_decimal() - && binary_expr.right().data_type(input_schema)?.is_decimal() - { - 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) -> DFResult { - match value { - DFOperator::Eq => Ok(Operator::Eq), - DFOperator::NotEq => Ok(Operator::NotEq), - DFOperator::Lt => Ok(Operator::Lt), - DFOperator::LtEq => Ok(Operator::Lte), - DFOperator::Gt => Ok(Operator::Gt), - DFOperator::GtEq => Ok(Operator::Gte), - DFOperator::And => Ok(Operator::And), - DFOperator::Or => Ok(Operator::Or), - DFOperator::Plus => Ok(Operator::Add), - DFOperator::Minus => Ok(Operator::Sub), - DFOperator::Multiply => Ok(Operator::Mul), - DFOperator::Divide => Ok(Operator::Div), - DFOperator::IsDistinctFrom - | DFOperator::IsNotDistinctFrom - | DFOperator::RegexMatch - | DFOperator::RegexIMatch - | DFOperator::RegexNotMatch - | DFOperator::RegexNotIMatch - | DFOperator::LikeMatch - | DFOperator::ILikeMatch - | DFOperator::NotLikeMatch - | DFOperator::NotILikeMatch - | DFOperator::BitwiseAnd - | DFOperator::BitwiseOr - | DFOperator::BitwiseXor - | DFOperator::BitwiseShiftRight - | DFOperator::BitwiseShiftLeft - | DFOperator::StringConcat - | DFOperator::AtArrow - | DFOperator::ArrowAt - | DFOperator::Modulo - | DFOperator::Arrow - | DFOperator::LongArrow - | DFOperator::HashArrow - | DFOperator::HashLongArrow - | DFOperator::AtAt - | DFOperator::IntegerDivide - | DFOperator::HashMinus - | DFOperator::AtQuestion - | DFOperator::Question - | DFOperator::QuestionAnd - | DFOperator::QuestionPipe - | DFOperator::Colon => { - tracing::debug!(operator = %value, "Can't pushdown binary_operator operator"); - Err(exec_datafusion_err!( - "Unsupported datafusion operator {value}" - )) - } - } -} - -fn can_be_pushed_down_impl(expr: &Arc, schema: &Schema) -> bool { - // We currently do not support pushdown of dynamic expressions in DF. - // See issue: https://github.com/vortex-data/vortex/issues/4034 - if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { - return false; - } - - if let Some(binary) = expr.downcast_ref::() { - can_binary_be_pushed_down(binary, schema) - } else if let Some(col) = expr.downcast_ref::() { - schema - .field_with_name(col.name()) - .is_ok_and(|field| supported_data_types(field.data_type())) - } else if let Some(like) = expr.downcast_ref::() { - can_be_pushed_down_impl(like.expr(), schema) - && can_be_pushed_down_impl(like.pattern(), schema) - } else if let Some(lit) = expr.downcast_ref::() { - supported_data_types(&lit.value().data_type()) - } else if let Some(cast_expr) = expr.downcast_ref::() { - // CastExpr child must be an expression type that convert() can handle - is_convertible_expr(cast_expr.expr()) - } else if let Some(is_null) = expr.downcast_ref::() { - can_be_pushed_down_impl(is_null.arg(), schema) - } else if let Some(is_not_null) = expr.downcast_ref::() { - can_be_pushed_down_impl(is_not_null.arg(), schema) - } else if let Some(in_list) = expr.downcast_ref::() { - can_be_pushed_down_impl(in_list.expr(), schema) - && in_list - .list() - .iter() - .all(|e| can_be_pushed_down_impl(e, schema)) - } else if let Some(scalar_fn) = expr.downcast_ref::() { - can_scalar_fn_be_pushed_down(scalar_fn, schema) - } else if let Some(case_expr) = expr.downcast_ref::() { - can_case_be_pushed_down(case_expr, schema) - } else { - tracing::debug!(%expr, "DataFusion expression can't be pushed down"); - false - } -} - -/// Checks if an expression type is one that convert() can handle. -/// This is less restrictive than can_be_pushed_down since it only checks -/// expression types, not data type support. -fn is_convertible_expr(expr: &Arc) -> bool { - // Expression types that convert() handles - expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - || expr - .downcast_ref::() - .is_some_and(|e| is_convertible_expr(e.expr())) - || expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some() - || expr.downcast_ref::().is_some_and(|sf| { - ScalarFunctionExpr::try_downcast_func::(sf).is_some() - || ScalarFunctionExpr::try_downcast_func::(sf).is_some() - || ScalarFunctionExpr::try_downcast_func::(sf).is_some() - }) -} - -fn can_binary_be_pushed_down(binary: &df_expr::BinaryExpr, schema: &Schema) -> bool { - let is_op_supported = try_operator_from_df(binary.op()).is_ok(); - is_op_supported - && can_be_pushed_down_impl(binary.left(), schema) - && can_be_pushed_down_impl(binary.right(), schema) -} - -fn can_case_be_pushed_down(case_expr: &df_expr::CaseExpr, schema: &Schema) -> bool { - // We only support the "searched CASE" form (CASE WHEN cond THEN result ...) - // not the "simple CASE" form (CASE expr WHEN value THEN result ...) - if case_expr.expr().is_some() { - return false; - } - - // Check all when/then pairs - for (when_expr, then_expr) in case_expr.when_then_expr() { - if !can_be_pushed_down_impl(when_expr, schema) - || !can_be_pushed_down_impl(then_expr, schema) - { - return false; - } - } - - // Check the optional else clause - if let Some(else_expr) = case_expr.else_expr() - && !can_be_pushed_down_impl(else_expr, schema) - { - return false; - } - - true -} - -fn supported_data_types(dt: &DataType) -> bool { - use DataType::*; - - // For dictionary types, check if the value type is supported. - if let Dictionary(_, value_type) = dt { - return supported_data_types(value_type.as_ref()); - } - - let is_supported = dt.is_null() - || dt.is_numeric() - || dt.is_binary() - || dt.is_string() - || matches!( - dt, - Boolean | Date32 | Date64 | Timestamp(_, _) | Time32(_) | Time64(_) - ); - - if !is_supported { - tracing::debug!("DataFusion data type {dt:?} is not supported"); - } - - is_supported -} - -/// Checks if a scalar function can be pushed down. -/// Currently GetFieldFunc, OctetLengthFunc, and ArrayLength are supported. -fn can_scalar_fn_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { - if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { - return true; - } - - if ScalarFunctionExpr::try_downcast_func::(scalar_fn) - .is_some_and(|octet_length| can_octet_length_be_pushed_down(octet_length, schema)) - { - return true; - } - - ScalarFunctionExpr::try_downcast_func::(scalar_fn) - .is_some_and(|array_length| can_array_length_be_pushed_down(array_length, schema)) -} - -fn can_octet_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { - let [input] = scalar_fn.args() else { - return false; - }; - - input.data_type(schema).as_ref().is_ok_and(|data_type| { - let dt = if let DataType::Dictionary(_, value_type) = data_type { - value_type.as_ref() - } else { - data_type - }; - - dt.is_binary() || dt.is_string() - }) && can_be_pushed_down_impl(input, schema) -} - -fn can_array_length_be_pushed_down(scalar_fn: &ScalarFunctionExpr, schema: &Schema) -> bool { - let Some(input) = array_length_input(scalar_fn) else { - return false; - }; - - // The argument must resolve to a list type. We gate on the resolved data type rather than - // `can_be_pushed_down_impl`, since list columns are intentionally rejected there. We still - // require the argument to be a convertible expression (e.g. a column or struct field access). - input.data_type(schema).as_ref().is_ok_and(|data_type| { - matches!( - data_type, - DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) - ) - }) && is_convertible_expr(input) -} - -/// Returns the list argument of an `array_length` call if the call is a form we can rewrite to -/// `list_length`: either the single-argument form `array_length(arr)`, or the two-argument form -/// with an explicit first dimension `array_length(arr, 1)`, which is equivalent. Higher -/// dimensions recurse into nested lists and are not supported. -fn array_length_input(scalar_fn: &ScalarFunctionExpr) -> Option<&Arc> { - match scalar_fn.args() { - [input] => Some(input), - [input, dimension] if is_dimension_one(dimension) => Some(input), - _ => None, - } -} - -/// Returns true if `expr` is an `Int64` literal equal to 1. DataFusion coerces the `array_length` -/// dimension argument to `Int64`, so that is the only form we need to recognize; any other literal -/// simply isn't pushed down. -fn is_dimension_one(expr: &Arc) -> bool { - expr.downcast_ref::() - .is_some_and(|literal| matches!(literal.value(), ScalarValue::Int64(Some(1)))) -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use arrow_schema::DataType; - use arrow_schema::Field; - use arrow_schema::Schema; - use arrow_schema::TimeUnit as ArrowTimeUnit; - use datafusion::arrow::array::AsArray; - use datafusion::arrow::datatypes::Int32Type; - use datafusion_common::ScalarValue; - use datafusion_common::config::ConfigOptions; - use datafusion_expr::Operator as DFOperator; - use datafusion_expr::ScalarUDF; - use datafusion_physical_expr::PhysicalExpr; - use datafusion_physical_plan::expressions as df_expr; - use insta::assert_snapshot; - use rstest::rstest; - - use super::*; - use crate::common_tests::TestSessionContext; - - #[rstest::fixture] - fn test_schema() -> Schema { - Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("name", DataType::Utf8, true), - Field::new("score", DataType::Float64, true), - Field::new("active", DataType::Boolean, false), - Field::new( - "created_at", - DataType::Timestamp(ArrowTimeUnit::Millisecond, None), - true, - ), - Field::new( - "tags", - DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), - true, - ), - ]) - } - - fn octet_length_expr(input: Arc, schema: &Schema) -> Arc { - Arc::new( - ScalarFunctionExpr::try_new( - Arc::new(ScalarUDF::from(OctetLengthFunc::new())), - vec![input], - schema, - Arc::new(ConfigOptions::new()), - ) - .unwrap(), - ) - } - - fn array_length_expr( - args: Vec>, - schema: &Schema, - ) -> Arc { - Arc::new( - ScalarFunctionExpr::try_new( - Arc::new(ScalarUDF::from(ArrayLength::new())), - args, - schema, - Arc::new(ConfigOptions::new()), - ) - .unwrap(), - ) - } - - #[test] - fn test_make_vortex_predicate_empty() { - let expr_convertor = DefaultExpressionConvertor::default(); - let result = make_vortex_predicate(&expr_convertor, &[]).unwrap(); - assert!(result.is_none()); - } - - #[test] - fn test_make_vortex_predicate_single() { - let expr_convertor = DefaultExpressionConvertor::default(); - let col_expr = Arc::new(df_expr::Column::new("test", 0)) as Arc; - let result = make_vortex_predicate(&expr_convertor, &[col_expr]).unwrap(); - assert!(result.is_some()); - } - - #[test] - fn test_make_vortex_predicate_multiple() { - let expr_convertor = DefaultExpressionConvertor::default(); - let col1 = Arc::new(df_expr::Column::new("col1", 0)) as Arc; - let col2 = Arc::new(df_expr::Column::new("col2", 1)) as Arc; - let result = make_vortex_predicate(&expr_convertor, &[col1, col2]).unwrap(); - assert!(result.is_some()); - // Result should be an AND expression combining the two columns - } - - #[rstest] - #[case::eq(DFOperator::Eq, Operator::Eq)] - #[case::not_eq(DFOperator::NotEq, Operator::NotEq)] - #[case::lt(DFOperator::Lt, Operator::Lt)] - #[case::lte(DFOperator::LtEq, Operator::Lte)] - #[case::gt(DFOperator::Gt, Operator::Gt)] - #[case::gte(DFOperator::GtEq, Operator::Gte)] - #[case::and(DFOperator::And, Operator::And)] - #[case::or(DFOperator::Or, Operator::Or)] - #[case::plus(DFOperator::Plus, Operator::Add)] - #[case::plus(DFOperator::Minus, Operator::Sub)] - #[case::plus(DFOperator::Multiply, Operator::Mul)] - #[case::plus(DFOperator::Divide, Operator::Div)] - fn test_operator_conversion_supported( - #[case] df_op: DFOperator, - #[case] expected_vortex_op: Operator, - ) { - let result = try_operator_from_df(&df_op).unwrap(); - assert_eq!(result, expected_vortex_op); - } - - #[rstest] - #[case::modulo(DFOperator::Modulo)] - #[case::bitwise_and(DFOperator::BitwiseAnd)] - #[case::regex_match(DFOperator::RegexMatch)] - #[case::like_match(DFOperator::LikeMatch)] - fn test_operator_conversion_unsupported(#[case] df_op: DFOperator) { - let result = try_operator_from_df(&df_op); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("Unsupported datafusion operator") - ); - } - - #[test] - fn test_expr_from_df_column() { - let col_expr = df_expr::Column::new("test_column", 0); - let result = DefaultExpressionConvertor::default() - .convert(&col_expr) - .unwrap(); - - assert_snapshot!(result.display_tree().to_string(), @r" - vortex.get_item(test_column) - └── input: vortex.root() - "); - } - - #[test] - fn test_expr_from_df_literal() { - let literal_expr = df_expr::Literal::new(ScalarValue::Int32(Some(42))); - let result = DefaultExpressionConvertor::default() - .convert(&literal_expr) - .unwrap(); - - assert_snapshot!(result.display_tree().to_string(), @"vortex.literal(42i32)"); - } - - #[test] - fn test_expr_from_df_binary() { - let left = Arc::new(df_expr::Column::new("left", 0)) as Arc; - let right = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; - let binary_expr = df_expr::BinaryExpr::new(left, DFOperator::Eq, right); - - let result = DefaultExpressionConvertor::default() - .convert(&binary_expr) - .unwrap(); - - assert_snapshot!(result.display_tree().to_string(), @r" - vortex.binary(=) - ├── lhs: vortex.get_item(left) - │ └── input: vortex.root() - └── rhs: vortex.literal(42i32) - "); - } - - #[rstest] - #[case::like_normal(false, false)] - #[case::like_negated(true, false)] - #[case::like_case_insensitive(false, true)] - #[case::like_negated_case_insensitive(true, true)] - fn test_expr_from_df_like(#[case] negated: bool, #[case] case_insensitive: bool) { - let expr = Arc::new(df_expr::Column::new("text_col", 0)) as Arc; - let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( - "test%".to_string(), - )))) as Arc; - let like_expr = df_expr::LikeExpr::new(negated, case_insensitive, expr, pattern); - - let result = DefaultExpressionConvertor::default() - .convert(&like_expr) - .unwrap(); - let like_opts = result.as_::(); - assert_eq!( - like_opts, - &LikeOptions { - negated, - case_insensitive - } - ); - } - - #[rstest] - fn test_expr_from_df_octet_length(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; - let octet_length = octet_length_expr(expr, &test_schema); - - let result = DefaultExpressionConvertor::default() - .convert(octet_length.as_ref()) - .unwrap(); - - assert_snapshot!(result.display_tree().to_string(), @r" - vortex.cast(i32?) - └── input: vortex.byte_length() - └── input: vortex.get_item(name) - └── input: vortex.root() - "); - } - - #[rstest] - fn test_expr_from_df_array_length(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let array_length = array_length_expr(vec![expr], &test_schema); - - let result = DefaultExpressionConvertor::default() - .convert(array_length.as_ref()) - .unwrap(); - - assert_snapshot!(result.display_tree().to_string(), @r" - vortex.cast(u64?) - └── input: vortex.list.length() - └── input: vortex.get_item(tags) - └── input: vortex.root() - "); - } - - #[rstest] - // Supported types - #[case::null(DataType::Null, true)] - #[case::boolean(DataType::Boolean, true)] - #[case::int8(DataType::Int8, true)] - #[case::int16(DataType::Int16, true)] - #[case::int32(DataType::Int32, true)] - #[case::int64(DataType::Int64, true)] - #[case::uint8(DataType::UInt8, true)] - #[case::uint16(DataType::UInt16, true)] - #[case::uint32(DataType::UInt32, true)] - #[case::uint64(DataType::UInt64, true)] - #[case::float32(DataType::Float32, true)] - #[case::float64(DataType::Float64, true)] - #[case::utf8(DataType::Utf8, true)] - #[case::utf8_view(DataType::Utf8View, true)] - #[case::binary(DataType::Binary, true)] - #[case::binary_view(DataType::BinaryView, true)] - #[case::date32(DataType::Date32, true)] - #[case::date64(DataType::Date64, true)] - #[case::timestamp_ms(DataType::Timestamp(ArrowTimeUnit::Millisecond, None), true)] - #[case::timestamp_us( - DataType::Timestamp(ArrowTimeUnit::Microsecond, Some(Arc::from("UTC"))), - true - )] - #[case::time32_s(DataType::Time32(ArrowTimeUnit::Second), true)] - #[case::time64_ns(DataType::Time64(ArrowTimeUnit::Nanosecond), true)] - // Unsupported types - #[case::list( - DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), - false - )] - #[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into() - ), false)] - // Dictionary types - should be supported if value type is supported - #[case::dict_utf8( - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true - )] - #[case::dict_int32( - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Int32)), - true - )] - #[case::dict_unsupported( - DataType::Dictionary( - Box::new(DataType::UInt32), - Box::new(DataType::List(Arc::new(Field::new("item", DataType::Int32, true)))) - ), - false - )] - fn test_supported_data_types(#[case] data_type: DataType, #[case] expected: bool) { - assert_eq!(supported_data_types(&data_type), expected); - } - - #[rstest] - fn test_can_be_pushed_down_column_supported(test_schema: Schema) { - let col_expr = Arc::new(df_expr::Column::new("id", 0)) as Arc; - - assert!(can_be_pushed_down_impl(&col_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_column_unsupported_type(test_schema: Schema) { - let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - - assert!(!can_be_pushed_down_impl(&col_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_column_not_found(test_schema: Schema) { - let col_expr = Arc::new(df_expr::Column::new("nonexistent", 99)) as Arc; - - assert!(!can_be_pushed_down_impl(&col_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_literal_supported(test_schema: Schema) { - let lit_expr = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; - - assert!(can_be_pushed_down_impl(&lit_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_literal_unsupported(test_schema: Schema) { - // Use a simpler unsupported type - Duration is not supported - let unsupported_literal = ScalarValue::DurationSecond(Some(42)); - let lit_expr = - Arc::new(df_expr::Literal::new(unsupported_literal)) as Arc; - - assert!(!can_be_pushed_down_impl(&lit_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_binary_supported(test_schema: Schema) { - let left = Arc::new(df_expr::Column::new("id", 0)) as Arc; - let right = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; - let binary_expr = Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right)) - as Arc; - - assert!(can_be_pushed_down_impl(&binary_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_binary_unsupported_operator(test_schema: Schema) { - let left = Arc::new(df_expr::Column::new("id", 0)) as Arc; - let right = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; - let binary_expr = Arc::new(df_expr::BinaryExpr::new( - left, - DFOperator::AtQuestion, - right, - )) as Arc; - - assert!(!can_be_pushed_down_impl(&binary_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_binary_unsupported_operand(test_schema: Schema) { - let left = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let right = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; - let binary_expr = Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right)) - as Arc; - - assert!(!can_be_pushed_down_impl(&binary_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_like_supported(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; - let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( - "test%".to_string(), - )))) as Arc; - let like_expr = - Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc; - - assert!(can_be_pushed_down_impl(&like_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_like_unsupported_operand(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( - "test%".to_string(), - )))) as Arc; - let like_expr = - Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc; - - assert!(!can_be_pushed_down_impl(&like_expr, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_octet_length_supported(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; - let octet_length = octet_length_expr(expr, &test_schema); - - assert!(can_be_pushed_down_impl(&octet_length, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let octet_length = Arc::new(ScalarFunctionExpr::new( - "octet_length", - Arc::new(ScalarUDF::from(OctetLengthFunc::new())), - vec![expr], - Arc::new(Field::new("octet_length", DataType::Int32, true)), - Arc::new(ConfigOptions::new()), - )) as Arc; - - assert!(!can_be_pushed_down_impl(&octet_length, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_array_length_supported(test_schema: Schema) { - let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let array_length = array_length_expr(vec![expr], &test_schema); - - assert!(can_be_pushed_down_impl(&array_length, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_array_length_unsupported_operand(test_schema: Schema) { - // `array_length` over a non-list column cannot be pushed down. - let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; - let array_length = Arc::new(ScalarFunctionExpr::new( - "array_length", - Arc::new(ScalarUDF::from(ArrayLength::new())), - vec![expr], - Arc::new(Field::new("array_length", DataType::UInt64, true)), - Arc::new(ConfigOptions::new()), - )) as Arc; - - assert!(!can_be_pushed_down_impl(&array_length, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_array_length_dimension_one_supported(test_schema: Schema) { - // `array_length(arr, 1)` is the first-dimension length, equivalent to `list_length`. - let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let dimension = - Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(1)))) as Arc; - let array_length = array_length_expr(vec![list, dimension], &test_schema); - - assert!(can_be_pushed_down_impl(&array_length, &test_schema)); - } - - #[rstest] - fn test_can_be_pushed_down_array_length_higher_dimension_not_supported(test_schema: Schema) { - // Dimensions other than 1 recurse into nested lists, which `list_length` does not model, - // so they must not be pushed down. - let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - let dimension = - Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(2)))) as Arc; - let array_length = array_length_expr(vec![list, dimension], &test_schema); - - assert!(!can_be_pushed_down_impl(&array_length, &test_schema)); - } - - // https://github.com/vortex-data/vortex/issues/6211 - #[tokio::test] - async fn test_cast_int_to_string() -> anyhow::Result<()> { - let ctx = TestSessionContext::default(); - - ctx.session - .sql(r#"copy (select 1 as id) to 'example.vortex'"#) - .await? - .show() - .await?; - - ctx.session - .sql(r#"select cast(id as string) as sid from 'example.vortex' where id > 0"#) - .await? - .show() - .await?; - - ctx.session - .sql(r#"select id from 'example.vortex' where cast (id as string) == '1'"#) - .await? - .show() - .await?; - - // This fails as it pushes string cast to the scan - ctx.session - .sql(r#"select cast(id as string) from 'example.vortex'"#) - .await? - .collect() - .await?; - - Ok(()) - } - - /// A cast whose target is a UUID-tagged `FixedSizeBinary(16)` must resolve - /// through the dtype extension registry (UUID is registered on the default - /// session) instead of the static, non-plugin-aware `DType::from_arrow`, - /// which does not support `FixedSizeBinary` and previously panicked here. - #[test] - fn test_cast_to_uuid_resolves_via_registry() -> anyhow::Result<()> { - use arrow_schema::extension::Uuid; - - let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), true); - uuid_field.try_with_extension_type(Uuid)?; - - let child = Arc::new(df_expr::Column::new("id", 0)) as Arc; - let cast = df_expr::CastExpr::new_with_target_field(child, Arc::new(uuid_field), None); - - // Must convert without panicking — the static path would `unimplemented!()`. - DefaultExpressionConvertor::default().convert(&cast)?; - Ok(()) - } - - /// Test that applying a CASE expression to an Arrow RecordBatch using DataFusion - /// matches the result of applying the converted Vortex expression. - #[test] - fn test_case_when_datafusion_vortex_equivalence() { - use datafusion::arrow::array::Int32Array; - use datafusion::arrow::array::RecordBatch; - use datafusion_physical_expr::expressions::CaseExpr; - use vortex::VortexSessionDefault; - use vortex::array::ArrayRef; - use vortex::array::Canonical; - use vortex::array::VortexSessionExecute as _; - use vortex::session::VortexSession; - - // Create test data - let values = Arc::new(Int32Array::from(vec![1, 5, 10, 15, 20])); - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new(schema, vec![values]).unwrap(); - - // Build a DataFusion CASE expression: - // CASE WHEN value > 10 THEN 100 WHEN value > 5 THEN 50 ELSE 0 END - let col_value = Arc::new(df_expr::Column::new("value", 0)) as Arc; - let lit_10 = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(10)))) as Arc; - let lit_5 = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(5)))) as Arc; - let lit_100 = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(100)))) as Arc; - let lit_50 = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(50)))) as Arc; - let lit_0 = - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))) as Arc; - - // WHEN value > 10 THEN 100 - let when1 = Arc::new(df_expr::BinaryExpr::new( - Arc::clone(&col_value), - DFOperator::Gt, - lit_10, - )) as Arc; - // WHEN value > 5 THEN 50 - let when2 = Arc::new(df_expr::BinaryExpr::new(col_value, DFOperator::Gt, lit_5)) - as Arc; - - let case_expr = - CaseExpr::try_new(None, vec![(when1, lit_100), (when2, lit_50)], Some(lit_0)).unwrap(); - - // Apply DataFusion expression - let df_result = case_expr.evaluate(&batch).unwrap(); - let df_array = df_result.into_array(batch.num_rows()).unwrap(); - - // Convert to Vortex expression - let expr_convertor = DefaultExpressionConvertor::default(); - let vortex_expr = expr_convertor.try_convert_case_expr(&case_expr).unwrap(); - - // Convert batch to Vortex array - let session = VortexSession::default(); - let vortex_array: ArrayRef = session - .arrow() - .from_arrow_record_batch(batch.clone(), &batch.schema()) - .unwrap(); - - // Apply Vortex expression - let mut ctx = session.create_execution_ctx(); - let vortex_result = vortex_array - .apply(&vortex_expr) - .unwrap() - .execute::(&mut ctx) - .unwrap(); - - // Convert back to Arrow for comparison - let vortex_as_arrow = vortex_result.into_primitive().as_slice::().to_vec(); - - // Convert DataFusion result to Vec for comparison - let df_as_arrow: Vec = df_array.as_primitive::().values().to_vec(); - - // Compare results - // Expected: [0, 0, 50, 100, 100] for values [1, 5, 10, 15, 20] - // value=1: not > 10, not > 5 -> ELSE 0 - // value=5: not > 10, not > 5 -> ELSE 0 - // value=10: not > 10, > 5 -> 50 - // value=15: > 10 -> 100 - // value=20: > 10 -> 100 - assert_eq!(df_as_arrow, vec![0, 0, 50, 100, 100]); - assert_eq!(vortex_as_arrow, df_as_arrow); - } -} diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs new file mode 100644 index 00000000000..b358c5f3ed9 --- /dev/null +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -0,0 +1,711 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::Schema; +use datafusion_common::DataFusionError; +use datafusion_common::Result as DFResult; +use datafusion_common::ScalarValue; +use datafusion_common::exec_datafusion_err; +use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; +use datafusion_expr::Operator as DFOperator; +use datafusion_functions::core::getfield::GetFieldFunc; +use datafusion_functions::string::octet_length::OctetLengthFunc; +use datafusion_functions_nested::length::ArrayLength; +use datafusion_physical_expr::DynamicFilterTracking; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::ScalarFunctionExpr; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_plan::expressions as df_expr; +use itertools::Itertools; +use vortex::VortexSessionDefault; +use vortex::array::VortexSessionExecute; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +use vortex::expr::Expression; +use vortex::expr::analysis::label_infallible; +use vortex::expr::and_collect; +use vortex::expr::byte_length; +use vortex::expr::cast; +use vortex::expr::get_item; +use vortex::expr::is_not_null; +use vortex::expr::is_null; +use vortex::expr::list_length; +use vortex::expr::lit; +use vortex::expr::nested_case_when; +use vortex::expr::or_collect; +use vortex::expr::pack; +use vortex::expr::root; +use vortex::scalar_fn::ScalarFnVTableExt; +use vortex::scalar_fn::fns::binary::Binary; +use vortex::scalar_fn::fns::like::Like; +use vortex::scalar_fn::fns::like::LikeOptions; +use vortex::scalar_fn::fns::operators::Operator; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +#[cfg(test)] +mod tests; + +/// 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, + /// Arrow reference types and metadata for the scan output. + pub scan_reference_schema: Schema, + /// Projection evaluated by DataFusion after the Vortex scan. + pub leftover_projection: ProjectionExprs, +} + +/// Trait for converting DataFusion expressions to Vortex ones. +/// +/// Custom converters implement a single schema-aware decision. Conversion should preserve +/// DataFusion values, nulls, and evaluation errors; see [`DefaultExpressionConverter`] for +/// the temporary arithmetic exception. Unsupported expressions remain in DataFusion, +/// including when a file's schema adapter introduces them. +/// +/// # Implementing a custom converter +/// +/// ``` +/// use std::sync::Arc; +/// +/// use arrow_schema::Schema; +/// use datafusion_common::Result as DFResult; +/// use datafusion_physical_expr::PhysicalExpr; +/// use vortex::expr::Expression; +/// use vortex_datafusion::convert::DefaultExpressionConverter; +/// use vortex_datafusion::convert::ExpressionConverter; +/// +/// struct CustomExpressionConverter(DefaultExpressionConverter); +/// +/// impl ExpressionConverter for CustomExpressionConverter { +/// fn try_convert( +/// &self, +/// expr: &Arc, +/// schema: &Schema, +/// ) -> DFResult> { +/// self.0.try_convert(expr, schema) +/// } +/// } +/// +/// let _converter: Arc = Arc::new(CustomExpressionConverter( +/// DefaultExpressionConverter::default(), +/// )); +/// ``` +pub trait ExpressionConverter: Send + Sync { + /// Convert an expression for native evaluation against this schema. + /// + /// Returns None for valid but unsupported expressions. Malformed expressions and + /// conversion failures return errors. Callers must retain unsupported exact predicates + /// for DataFusion evaluation. + fn try_convert( + &self, + expr: &Arc, + schema: &Schema, + ) -> DFResult>; + + /// Split a projection into Vortex expressions that can be pushed down and leftover + /// DataFusion projections that need to be evaluated after the scan. + /// + /// If any expression is unsupported, evaluate the complete projection in DataFusion + /// over deduplicated raw inputs. This avoids mixing input names with output aliases. + fn split_projection( + &self, + source_projection: ProjectionExprs, + input_schema: &Schema, + output_schema: &Schema, + ) -> DFResult { + // Duplicate output names cannot identify native pack fields unambiguously. + if !source_projection + .iter() + .map(|projection| &projection.alias) + .all_unique() + { + return self.no_pushdown_projection(source_projection, input_schema); + } + let mut scan_projection = Vec::with_capacity(source_projection.as_ref().len()); + for projection in source_projection.as_ref() { + let Some(expr) = self.try_convert(&projection.expr, input_schema)? else { + return self.no_pushdown_projection(source_projection, input_schema); + }; + scan_projection.push((projection.alias.clone(), expr)); + } + // The output schema names its fields after the projection aliases. + let output_indices = (0..scan_projection.len()).collect_vec(); + Ok(ProcessedProjection { + scan_projection: pack(scan_projection, Nullability::NonNullable), + scan_reference_schema: output_schema.clone(), + leftover_projection: ProjectionExprs::from_indices(&output_indices, output_schema), + }) + } + + /// Create a projection that reads only the required columns without pushing down + /// any expressions. All projection logic is applied after the scan. + fn no_pushdown_projection( + &self, + source_projection: ProjectionExprs, + input_schema: &Schema, + ) -> DFResult { + let (scan_projection, scan_reference_schema) = + raw_projection(&source_projection.column_indices(), input_schema)?; + Ok(ProcessedProjection { + scan_projection, + scan_reference_schema, + leftover_projection: source_projection, + }) + } +} + +/// Read the raw columns at `indices` in file order, without involving custom expression +/// conversion. Returns the scan projection and the Arrow schema of its output. +pub(crate) fn raw_projection( + indices: &[usize], + input_schema: &Schema, +) -> DFResult<(Expression, Schema)> { + let schema = input_schema.project(indices)?; + let scan_columns = schema.fields().iter().map(|field| { + ( + field.name().clone(), + get_item(field.name().as_str(), root()), + ) + }); + Ok((pack(scan_columns, Nullability::NonNullable), schema)) +} + +/// Why an expression was not converted. +enum Unconverted { + /// A valid expression that Vortex cannot evaluate; it stays in DataFusion. + Unsupported, + /// A malformed expression or a failed conversion. + Failed(DataFusionError), +} + +impl From for Unconverted { + fn from(error: DataFusionError) -> Self { + Self::Failed(error) + } +} + +/// Conversion result where `?` propagates both unsupported expressions and errors. +type Conversion = Result; + +/// The default [`ExpressionConverter`] implementation. +/// +/// Supported arithmetic is pushed down using Vortex semantics, including its checked +/// integer arithmetic. Matching DataFusion's overflow behavior is deferred to a future patch. +/// Other expressions require compatible SQL semantics or remain in DataFusion. +pub struct DefaultExpressionConverter { + /// Session used to resolve Arrow → Vortex dtypes through the extension + /// plugin registry, so registered extension types (e.g. UUID ⇄ + /// `FixedSizeBinary[16]`) convert correctly instead of hitting the static, + /// non-plugin-aware `DType::from_arrow`. + session: VortexSession, +} + +impl Default for DefaultExpressionConverter { + fn default() -> Self { + Self { + session: VortexSession::default(), + } + } +} + +impl DefaultExpressionConverter { + /// Create a converter that resolves Arrow extension types using `session`'s + /// dtype registry. + pub fn new(session: VortexSession) -> Self { + Self { session } + } + + /// Resolve an Arrow field through the session registry; unknown types are unsupported. + fn arrow_dtype(&self, field: &Field) -> Conversion { + self.session + .arrow() + .from_arrow_field(field) + .map_err(|_| Unconverted::Unsupported) + } + + /// Convert `expr` and check that it returns the dtype DataFusion expects. + fn convert_checked( + &self, + expr: &Arc, + schema: &Schema, + ) -> Conversion { + let columns = collect_columns(expr); + let mut column_indices = Vec::with_capacity(columns.len()); + for column in columns { + let field = schema.fields().get(column.index()).ok_or_else(|| { + exec_datafusion_err!( + "Column {}@{} is out of bounds", + column.name(), + column.index() + ) + })?; + if field.name() != column.name() { + return Err(exec_datafusion_err!( + "Column {}@{} refers to field {}", + column.name(), + column.index(), + field.name() + ) + .into()); + } + column_indices.push(column.index()); + } + column_indices.sort_unstable(); + column_indices.dedup(); + let input_dtype = self + .session + .arrow() + .from_arrow_schema( + &schema + .project(&column_indices) + .map_err(DataFusionError::from)?, + ) + .map_err(|_| Unconverted::Unsupported)?; + let converted = self.convert_expr(expr, schema, &input_dtype)?; + let expected_dtype = self.arrow_dtype(expr.return_field(schema)?.as_ref())?; + if !converted_dtype(&converted, &input_dtype)?.eq_ignore_nullability(&expected_dtype) { + return Err(Unconverted::Unsupported); + } + Ok(converted) + } + + fn convert_expr( + &self, + expr: &Arc, + schema: &Schema, + input_dtype: &DType, + ) -> Conversion { + if let Some(binary_expr) = expr.downcast_ref::() { + let operator = + try_operator_from_df(binary_expr.op()).ok_or(Unconverted::Unsupported)?; + let boolean_operator = matches!(operator, Operator::And | Operator::Or); + let left_type = binary_expr.left().data_type(schema)?; + let right_type = binary_expr.right().data_type(schema)?; + if boolean_operator { + if left_type != DataType::Boolean || right_type != DataType::Boolean { + return Err(exec_datafusion_err!( + "Boolean operator requires Boolean operands: {expr}" + ) + .into()); + } + } else if !supported_data_types(&left_type) || !supported_data_types(&right_type) { + return Err(Unconverted::Unsupported); + } + let left = self.convert_expr(binary_expr.left(), schema, input_dtype)?; + let right = self.convert_expr(binary_expr.right(), schema, input_dtype)?; + // DataFusion may evaluate the RHS only on rows selected by the LHS. + if boolean_operator && !(is_infallible(&left) && is_infallible(&right)) { + return Err(Unconverted::Unsupported); + } + if !converted_dtype(&left, input_dtype)? + .eq_ignore_nullability(&converted_dtype(&right, input_dtype)?) + { + return Err(Unconverted::Unsupported); + } + return Ok(Binary.new_expr(operator, [left, right])); + } + + if let Some(col_expr) = expr.downcast_ref::() { + return Ok(get_item(col_expr.name(), root())); + } + + if let Some(literal) = expr.downcast_ref::() { + let field = literal.return_field(schema)?; + let array = literal.value().to_array()?; + if array.len() != 1 { + return Err(exec_datafusion_err!( + "Literal must contain exactly one value, found {}", + array.len() + ) + .into()); + } + // Literals of unknown Arrow types stay in DataFusion; conversion failures are errors. + self.arrow_dtype(&field)?; + let scalar = self + .session + .arrow() + .from_arrow_array(array, &field) + .and_then(|array| array.execute_scalar(0, &mut self.session.create_execution_ctx())) + .map_err(|e| exec_datafusion_err!("Failed to convert literal {expr}: {e}"))?; + return Ok(lit(scalar)); + } + + if let Some(cast_expr) = expr.downcast_ref::() { + if !supported_cast(cast_expr, schema)? { + return Err(Unconverted::Unsupported); + } + let child = self.convert_expr(cast_expr.expr(), schema, input_dtype)?; + // DataFusion casts preserve input nulls even when the target field is declared + // non-nullable. Match the expression's runtime nullability rather than narrowing + // to the target field's metadata. + let cast_dtype = self + .arrow_dtype(cast_expr.target_field())? + .with_nullability(Nullability::from(cast_expr.nullable(schema)?)); + // Matching Arrow storage types do not imply matching extension semantics. + let child_dtype = converted_dtype(&child, input_dtype)?; + if (child_dtype.is_extension() || cast_dtype.is_extension()) + && !child_dtype.eq_ignore_nullability(&cast_dtype) + { + return Err(Unconverted::Unsupported); + } + return Ok(cast(child, cast_dtype)); + } + + if let Some(is_null_expr) = expr.downcast_ref::() { + let arg = self.convert_expr(is_null_expr.arg(), schema, input_dtype)?; + return Ok(is_null(arg)); + } + + if let Some(is_not_null_expr) = expr.downcast_ref::() { + let arg = self.convert_expr(is_not_null_expr.arg(), schema, input_dtype)?; + return Ok(is_not_null(arg)); + } + + if let Some(like) = expr.downcast_ref::() { + if !like.expr().data_type(schema)?.is_string() + || !like.pattern().data_type(schema)?.is_string() + { + return Err(Unconverted::Unsupported); + } + let child = self.convert_expr(like.expr(), schema, input_dtype)?; + let pattern = self.convert_expr(like.pattern(), schema, input_dtype)?; + return Ok(Like.new_expr( + LikeOptions { + negated: like.negated(), + case_insensitive: like.case_insensitive(), + }, + [child, pattern], + )); + } + + if let Some(in_list) = expr.downcast_ref::() { + return self.convert_in_list(in_list, schema, input_dtype); + } + + if let Some(scalar_fn) = expr.downcast_ref::() { + return self.convert_scalar_function(scalar_fn, schema, input_dtype); + } + + if let Some(case_expr) = expr.downcast_ref::() { + return self.convert_case_expr(case_expr, schema, input_dtype); + } + + Err(Unconverted::Unsupported) + } + + fn convert_in_list( + &self, + in_list: &df_expr::InListExpr, + schema: &Schema, + input_dtype: &DType, + ) -> Conversion { + if in_list.is_empty() + || !in_list + .list() + .iter() + .all(|expr| expr.is::()) + || !supported_data_types(&in_list.expr().data_type(schema)?) + { + return Err(Unconverted::Unsupported); + } + let value = self.convert_expr(in_list.expr(), schema, input_dtype)?; + // Boolean rewrites may skip evaluating the input, particularly for all-null lists. + if !is_infallible(&value) { + return Err(Unconverted::Unsupported); + } + let value_dtype = converted_dtype(&value, input_dtype)?; + let operator = if in_list.negated() { + Operator::NotEq + } else { + Operator::Eq + }; + let mut comparisons = Vec::with_capacity(in_list.len()); + for element in in_list.list() { + let element = self.convert_expr(element, schema, input_dtype)?; + let element_dtype = converted_dtype(&element, input_dtype)?; + if element_dtype == DType::Null { + comparisons.push(lit(None::)); + } else if value_dtype.eq_ignore_nullability(&element_dtype) { + comparisons.push(Binary.new_expr(operator, [value.clone(), element])); + } else { + return Err(Unconverted::Unsupported); + } + } + // Kleene AND/OR preserve SQL IN/NOT IN nulls; list_contains does not. + let membership = if in_list.negated() { + and_collect(comparisons) + } else { + or_collect(comparisons) + }; + membership.ok_or(Unconverted::Unsupported) + } + + fn convert_scalar_function( + &self, + scalar_fn: &ScalarFunctionExpr, + schema: &Schema, + input_dtype: &DType, + ) -> Conversion { + if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { + // DataFusion's GetFieldFunc flattens nested field access into a single call + // with multiple field name arguments, e.g. `outer.inner.leaf` becomes + // get_field(Column("outer"), "inner", "leaf"). + let [source, paths @ ..] = scalar_fn.args() else { + return Err( + exec_datafusion_err!("get_field requires a source and field path").into(), + ); + }; + if paths.is_empty() { + return Err(exec_datafusion_err!("get_field requires a field path").into()); + } + let mut source_type = source.data_type(schema)?; + let mut nullable = source.nullable(schema)?; + let mut nullable_struct = false; + let mut names = Vec::with_capacity(paths.len()); + for path in paths { + let name = path + .downcast_ref::() + .and_then(|literal| literal.value().try_as_str().flatten()) + .ok_or_else(|| { + exec_datafusion_err!("get_field path must be a non-null string literal") + })?; + let DataType::Struct(fields) = &source_type else { + return Err(Unconverted::Unsupported); + }; + nullable_struct |= nullable; + let (_, field) = fields.find(name).ok_or_else(|| { + exec_datafusion_err!("get_field references missing field {name}") + })?; + nullable = field.is_nullable(); + source_type = field.data_type().clone(); + names.push(name); + } + // DataFusion extracts the child without applying parent struct validity; + // Vortex get_item masks the child when its parent is null. + if nullable_struct { + return Err(Unconverted::Unsupported); + } + let mut result = self.convert_expr(source, schema, input_dtype)?; + for name in names { + result = get_item(name, result); + } + return Ok(result); + } + + let (input, length): (_, fn(Expression) -> Expression) = + if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { + let [input] = scalar_fn.args() else { + return Err( + exec_datafusion_err!("octet_length requires exactly one argument").into(), + ); + }; + let data_type = input.data_type(schema)?; + let data_type = match &data_type { + DataType::Dictionary(_, value) => value.as_ref(), + data_type => data_type, + }; + if !data_type.is_binary() && !data_type.is_string() { + return Err(Unconverted::Unsupported); + } + (input, byte_length) + } else if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { + let input = array_length_input(scalar_fn)?.ok_or(Unconverted::Unsupported)?; + if !matches!( + input.data_type(schema)?, + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ) { + return Err(Unconverted::Unsupported); + } + (input, list_length) + } else { + return Err(Unconverted::Unsupported); + }; + let input = self.convert_expr(input, schema, input_dtype)?; + let return_dtype = self.arrow_dtype(&Field::new( + "", + scalar_fn.return_type().clone(), + scalar_fn.nullable(), + ))?; + Ok(cast(length(input), return_dtype)) + } + + fn convert_case_expr( + &self, + case_expr: &df_expr::CaseExpr, + schema: &Schema, + input_dtype: &DType, + ) -> Conversion { + // Only the searched form (CASE WHEN cond THEN value ...) is supported. + if case_expr.expr().is_some() { + return Err(Unconverted::Unsupported); + } + let mut pairs = Vec::with_capacity(case_expr.when_then_expr().len()); + for (when_expr, then_expr) in case_expr.when_then_expr() { + if when_expr.data_type(schema)? != DataType::Boolean { + return Err(exec_datafusion_err!("CASE WHEN must be Boolean").into()); + } + let condition = self.convert_expr(when_expr, schema, input_dtype)?; + let value = self.convert_expr(then_expr, schema, input_dtype)?; + pairs.push((condition, value)); + } + let else_value = case_expr + .else_expr() + .map(|e| self.convert_expr(e, schema, input_dtype)) + .transpose()?; + let case = nested_case_when(pairs, else_value); + // Vortex may evaluate branch values on rows excluded by the condition. + if !is_infallible(&case) { + return Err(Unconverted::Unsupported); + } + Ok(case) + } +} + +impl ExpressionConverter for DefaultExpressionConverter { + fn try_convert( + &self, + expr: &Arc, + schema: &Schema, + ) -> DFResult> { + // We currently do not support pushdown of dynamic expressions in DF. + // See issue: https://github.com/vortex-data/vortex/issues/4034 + if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { + return Ok(None); + } + match self.convert_checked(expr, schema) { + Ok(converted) => Ok(Some(converted)), + Err(Unconverted::Unsupported) => Ok(None), + Err(Unconverted::Failed(error)) => Err(error), + } + } +} + +/// Whether evaluating `expr` cannot fail, so Vortex may evaluate it on rows DataFusion skips. +fn is_infallible(expr: &Expression) -> bool { + label_infallible(expr).get(expr) == Some(&true) +} + +/// The Vortex return dtype of a converted expression; unresolvable dtypes are unsupported. +fn converted_dtype(expr: &Expression, input_dtype: &DType) -> Conversion { + expr.return_dtype(input_dtype) + .map_err(|_| Unconverted::Unsupported) +} + +fn supported_cast(cast: &df_expr::CastExpr, schema: &Schema) -> DFResult { + use DataType::*; + let options = cast.cast_options(); + if options.safe || options.format_options != DEFAULT_FORMAT_OPTIONS { + return Ok(false); + } + let source = cast.expr().data_type(schema)?; + let target = cast.cast_type(); + Ok(source == *target + || matches!( + (&source, target), + (Int8, Int16 | Int32 | Int64) + | (Int16, Int32 | Int64) + | (Int32, Int64) + | (UInt8, UInt16 | UInt32 | UInt64) + | (UInt16, UInt32 | UInt64) + | (UInt32, UInt64) + | (Float32, Float64) + )) +} + +fn try_operator_from_df(value: &DFOperator) -> Option { + match value { + DFOperator::Eq => Some(Operator::Eq), + DFOperator::NotEq => Some(Operator::NotEq), + DFOperator::Lt => Some(Operator::Lt), + DFOperator::LtEq => Some(Operator::Lte), + DFOperator::Gt => Some(Operator::Gt), + DFOperator::GtEq => Some(Operator::Gte), + DFOperator::And => Some(Operator::And), + DFOperator::Or => Some(Operator::Or), + DFOperator::Plus => Some(Operator::Add), + DFOperator::Minus => Some(Operator::Sub), + DFOperator::Multiply => Some(Operator::Mul), + DFOperator::Divide => Some(Operator::Div), + DFOperator::IsDistinctFrom + | DFOperator::IsNotDistinctFrom + | DFOperator::RegexMatch + | DFOperator::RegexIMatch + | DFOperator::RegexNotMatch + | DFOperator::RegexNotIMatch + | DFOperator::LikeMatch + | DFOperator::ILikeMatch + | DFOperator::NotLikeMatch + | DFOperator::NotILikeMatch + | DFOperator::BitwiseAnd + | DFOperator::BitwiseOr + | DFOperator::BitwiseXor + | DFOperator::BitwiseShiftRight + | DFOperator::BitwiseShiftLeft + | DFOperator::StringConcat + | DFOperator::AtArrow + | DFOperator::ArrowAt + | DFOperator::Modulo + | DFOperator::Arrow + | DFOperator::LongArrow + | DFOperator::HashArrow + | DFOperator::HashLongArrow + | DFOperator::AtAt + | DFOperator::IntegerDivide + | DFOperator::HashMinus + | DFOperator::AtQuestion + | DFOperator::Question + | DFOperator::QuestionAnd + | DFOperator::QuestionPipe + | DFOperator::Colon => None, + } +} + +fn supported_data_types(dt: &DataType) -> bool { + use DataType::*; + + // For dictionary types, check if the value type is supported. + if let Dictionary(_, value_type) = dt { + return supported_data_types(value_type.as_ref()); + } + + let is_supported = dt.is_null() + || dt.is_numeric() + || dt.is_binary() + || dt.is_string() + || matches!( + dt, + Boolean | Date32 | Date64 | Timestamp(_, _) | Time32(_) | Time64(_) + ); + + if !is_supported { + tracing::debug!("DataFusion data type {dt:?} is not supported"); + } + + is_supported +} + +/// Returns the list argument of an `array_length` call that can be rewritten to `list_length`: +/// the single-argument form `array_length(arr)` or the equivalent explicit first dimension +/// `array_length(arr, 1)`. Other dimensions are unsupported (`None`); other arities are errors. +fn array_length_input(scalar_fn: &ScalarFunctionExpr) -> DFResult>> { + match scalar_fn.args() { + [input] => Ok(Some(input)), + [input, dimension] + if dimension + .downcast_ref::() + .is_some_and(|literal| matches!(literal.value(), ScalarValue::Int64(Some(1)))) => + { + Ok(Some(input)) + } + [_, _] => Ok(None), + _ => Err(exec_datafusion_err!( + "array_length requires one or two arguments" + )), + } +} diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs new file mode 100644 index 00000000000..120cafde885 --- /dev/null +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -0,0 +1,1670 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::Array; +use arrow_array::FixedSizeBinaryArray; +use arrow_array::RecordBatch; +use arrow_array::StructArray; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::IntervalUnit; +use arrow_schema::Schema; +use arrow_schema::TimeUnit as ArrowTimeUnit; +use arrow_schema::extension::Uuid as ArrowUuid; +use datafusion_common::ScalarValue; +use datafusion_common::arrow::buffer::NullBuffer; +use datafusion_common::arrow::datatypes::i256 as arrow_i256; +use datafusion_common::config::ConfigOptions; +use datafusion_common::metadata::FieldMetadata; +use datafusion_expr::Operator as DFOperator; +use datafusion_expr::ScalarUDF; +use datafusion_functions::core::coalesce::CoalesceFunc; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion_physical_expr::projection::ProjectionExpr; +use datafusion_physical_plan::expressions as df_expr; +use insta::assert_snapshot; +use rstest::rstest; +use vortex::array::IntoArray; +use vortex::array::arrays::ConstantArray; +use vortex::extension::uuid::Uuid; +use vortex::scalar::Scalar; +use vortex::scalar_fn::fns::literal::Literal; +use vortex_array::assert_arrays_eq; + +use super::*; +use crate::common_tests::TestSessionContext; +use crate::convert::TryToDataFusion; + +#[rstest::fixture] +fn test_schema() -> Schema { + Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("score", DataType::Float64, true), + Field::new("active", DataType::Boolean, false), + Field::new( + "created_at", + DataType::Timestamp(ArrowTimeUnit::Millisecond, None), + true, + ), + Field::new( + "tags", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ), + ]) +} + +fn octet_length_expr(input: Arc, schema: &Schema) -> Arc { + Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(OctetLengthFunc::new())), + vec![input], + schema, + Arc::new(ConfigOptions::new()), + ) + .unwrap(), + ) +} + +fn array_length_expr(args: Vec>, schema: &Schema) -> Arc { + Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(ArrayLength::new())), + args, + schema, + Arc::new(ConfigOptions::new()), + ) + .unwrap(), + ) +} + +/// Whether the default converter accepts `expr` for native evaluation against `schema`. +fn converts(expr: &Arc, schema: &Schema) -> DFResult { + Ok(DefaultExpressionConverter::default() + .try_convert(expr, schema)? + .is_some()) +} + +/// Convert `expr` natively, failing the test if the converter declines it. +fn convert(expr: Arc, schema: &Schema) -> DFResult { + DefaultExpressionConverter::default() + .try_convert(&expr, schema)? + .ok_or_else(|| exec_datafusion_err!("Expected native conversion for {expr}")) +} + +/// Convert an `IN` list expression and compare native evaluation against DataFusion. +fn assert_in_list_matches( + value: Arc, + list: Vec>, + negated: bool, + batch: RecordBatch, +) -> anyhow::Result<()> { + let expr = df_expr::InListExpr::try_new(value, list, negated, &batch.schema())?; + expr.evaluate(&batch)?; + assert_native_matches(Arc::new(expr), batch) +} + +fn int_literals(values: impl IntoIterator>) -> Vec> { + values + .into_iter() + .map(|value| Arc::new(df_expr::Literal::new(ScalarValue::Int32(value))) as _) + .collect() +} + +struct FailingConverter; + +impl ExpressionConverter for FailingConverter { + fn try_convert( + &self, + _expr: &Arc, + _schema: &Schema, + ) -> DFResult> { + Err(exec_datafusion_err!("Expression conversion must not run")) + } +} + +#[test] +fn test_duplicate_aliases_fall_back_before_conversion() -> DFResult<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let projection = ProjectionExprs::from(vec![ + ProjectionExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("b", 1)), + DFOperator::Plus, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), + )), + "duplicate", + ), + ProjectionExpr::new(Arc::new(df_expr::Column::new("a", 0)), "duplicate"), + ProjectionExpr::new(Arc::new(df_expr::Column::new("b", 1)), "duplicate"), + ]); + let output_schema = projection.project_schema(&schema)?; + let processed = + FailingConverter.split_projection(projection.clone(), &schema, &output_schema)?; + assert_eq!( + processed.scan_projection, + pack( + [("a", get_item("a", root())), ("b", get_item("b", root()))], + Nullability::NonNullable, + ) + ); + assert_eq!(processed.scan_reference_schema, schema); + assert_eq!(processed.leftover_projection, projection); + Ok(()) +} + +#[rstest] +fn test_predicate_rejects_cast_over_modulo(test_schema: Schema) -> DFResult<()> { + let modulo = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("id", 0)), + DFOperator::Modulo, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), + )); + let expr: Arc = + Arc::new(df_expr::CastExpr::new(modulo, DataType::Int64, None)); + assert!(!converts(&expr, &test_schema)?); + Ok(()) +} + +#[rstest] +#[case::empty(false)] +#[case::column(true)] +fn test_predicate_rejects_unsupported_in_list( + test_schema: Schema, + #[case] nonempty: bool, + #[values(false, true)] negated: bool, +) -> DFResult<()> { + let column: Arc = Arc::new(df_expr::Column::new("id", 0)); + let list = if nonempty { + vec![Arc::clone(&column)] + } else { + vec![] + }; + let expr: Arc = Arc::new(df_expr::InListExpr::try_new( + column, + list, + negated, + &test_schema, + )?); + assert!(!converts(&expr, &test_schema)?); + Ok(()) +} + +#[rstest] +#[case::values(vec![Some(1), Some(3)])] +#[case::with_null(vec![Some(1), None])] +#[case::null_only(vec![None])] +#[case::singleton(vec![Some(1)])] +#[case::duplicates(vec![Some(1), Some(1), None, None])] +fn test_native_in_list( + #[case] list: Vec>, + #[values(false, true)] negated: bool, +) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3), None]))?; + assert_in_list_matches( + Arc::new(df_expr::Column::new("a", 0)), + int_literals(list), + negated, + batch, + ) +} + +#[rstest] +#[case::boolean(ScalarValue::Boolean(Some(true)), ScalarValue::Boolean(Some(false)))] +#[case::unsigned(ScalarValue::UInt64(Some(u64::MAX)), ScalarValue::UInt64(Some(0)))] +#[case::utf8(ScalarValue::Utf8(Some("a".into())), ScalarValue::Utf8(Some("b".into())))] +#[case::utf8_view(ScalarValue::Utf8View(Some("a".into())), ScalarValue::Utf8View(Some("b".into())))] +#[case::large_utf8(ScalarValue::LargeUtf8(Some("a".into())), ScalarValue::LargeUtf8(Some("b".into())))] +#[case::binary(ScalarValue::Binary(Some(vec![0])), ScalarValue::Binary(Some(vec![1])))] +#[case::binary_view(ScalarValue::BinaryView(Some(vec![0])), ScalarValue::BinaryView(Some(vec![1])))] +#[case::large_binary(ScalarValue::LargeBinary(Some(vec![0])), ScalarValue::LargeBinary(Some(vec![1])))] +#[case::decimal32( + ScalarValue::Decimal32(Some(1234), 5, 2), + ScalarValue::Decimal32(Some(5678), 5, 2) +)] +#[case::decimal64( + ScalarValue::Decimal64(Some(1234), 10, 2), + ScalarValue::Decimal64(Some(5678), 10, 2) +)] +#[case::decimal128( + ScalarValue::Decimal128(Some(1234), 20, 2), + ScalarValue::Decimal128(Some(5678), 20, 2) +)] +#[case::decimal256( + ScalarValue::Decimal256(Some(arrow_i256::from_i128(1234)), 50, 2), + ScalarValue::Decimal256(Some(arrow_i256::from_i128(5678)), 50, 2) +)] +#[case::date(ScalarValue::Date32(Some(1)), ScalarValue::Date32(Some(2)))] +#[case::time( + ScalarValue::Time64Microsecond(Some(1)), + ScalarValue::Time64Microsecond(Some(2)) +)] +#[case::timestamp(ScalarValue::TimestampNanosecond(Some(1), Some("UTC".into())), ScalarValue::TimestampNanosecond(Some(2), Some("UTC".into())))] +#[case::dictionary(ScalarValue::Dictionary(Box::new(DataType::Int8), Box::new(ScalarValue::Utf8(Some("a".into())))), ScalarValue::Dictionary(Box::new(DataType::Int8), Box::new(ScalarValue::Utf8(Some("b".into())))))] +fn test_native_in_list_data_types( + #[case] member: ScalarValue, + #[case] absent: ScalarValue, + #[values(false, true)] negated: bool, +) -> anyhow::Result<()> { + let null = ScalarValue::try_new_null(&member.data_type())?; + let batch = RecordBatch::try_from_iter([( + "a", + ScalarValue::iter_to_array([member.clone(), absent, null.clone()])?, + )])?; + assert_in_list_matches( + Arc::new(df_expr::Column::new("a", 0)), + vec![ + Arc::new(df_expr::Literal::new(member)), + Arc::new(df_expr::Literal::new(null)), + ], + negated, + batch, + ) +} + +#[rstest] +fn test_native_in_list_untyped_null( + #[values(false, true)] negated: bool, + #[values(false, true)] mixed: bool, + #[values(false, true)] literal_input: bool, +) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(("a", Int32, vec![Some(1), Some(2), None]))?; + let mut list: Vec> = + vec![Arc::new(df_expr::Literal::new(ScalarValue::Null))]; + if mixed { + list.push(Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1))))); + } + let value: Arc = if literal_input { + Arc::new(df_expr::Literal::new(ScalarValue::Int32(None))) + } else { + Arc::new(df_expr::Column::new("a", 0)) + }; + assert_in_list_matches(value, list, negated, batch) +} + +#[rstest] +fn test_native_in_list_large(#[values(false, true)] negated: bool) -> anyhow::Result<()> { + let batch = + arrow_array::record_batch!(("a", Int32, vec![Some(0), Some(1023), Some(1024), None]))?; + assert_in_list_matches( + Arc::new(df_expr::Column::new("a", 0)), + int_literals((0..1024).map(Some)), + negated, + batch, + ) +} + +#[rstest] +fn test_native_in_list_float( + #[values(false, true)] negated: bool, + #[values(DataType::Float32, DataType::Float64)] data_type: DataType, +) -> anyhow::Result<()> { + let values = [ + Some(-0.0), + Some(0.0), + Some(f64::NAN), + Some(f64::from_bits(f64::NAN.to_bits() + (1 << 29))), + Some(f64::INFINITY), + None, + ]; + let scalar = |value| ScalarValue::Float64(value).cast_to(&data_type); + let batch = RecordBatch::try_from_iter([( + "a", + ScalarValue::iter_to_array( + values + .into_iter() + .map(scalar) + .collect::>>()?, + )?, + )])?; + assert_in_list_matches( + Arc::new(df_expr::Column::new("a", 0)), + vec![ + Arc::new(df_expr::Literal::new(scalar(Some(-0.0))?)), + Arc::new(df_expr::Literal::new(scalar(Some(f64::NAN))?)), + ], + negated, + batch, + ) +} + +#[rstest] +#[case::unsupported(DFOperator::Modulo)] +#[case::fallible(DFOperator::Divide)] +fn test_in_list_unsupported_input( + #[case] operator: DFOperator, + #[values(false, true)] negated: bool, +) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(("a", Int32, vec![0, 1]))?; + let expr: Arc = Arc::new(df_expr::InListExpr::try_new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(12)))), + operator, + Arc::new(df_expr::Column::new("a", 0)), + )), + vec![Arc::new(df_expr::Literal::new(ScalarValue::Null))], + negated, + &batch.schema(), + )?); + assert!(expr.evaluate(&batch).is_err()); + assert!(!converts(&expr, &batch.schema())?); + Ok(()) +} + +#[test] +fn test_in_list_unsupported_type() -> DFResult<()> { + let value: Arc = + Arc::new(df_expr::Literal::new(ScalarValue::DurationSecond(Some(1)))); + let expr: Arc = Arc::new(df_expr::InListExpr::try_new( + Arc::clone(&value), + vec![value], + false, + &Schema::empty(), + )?); + assert!(!converts(&expr, &Schema::empty())?); + Ok(()) +} + +#[test] +fn test_in_list_malformed_literal() -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Decimal128(10, 2), false)]); + let column: Arc = Arc::new(df_expr::Column::new("a", 0)); + let expr = Arc::new(df_expr::InListExpr::try_new( + Arc::clone(&column), + vec![Arc::new(df_expr::Literal::new(ScalarValue::Decimal128( + Some(1), + 10, + 2, + )))], + false, + &schema, + )?); + let expr = expr.with_new_children(vec![ + column, + Arc::new(df_expr::Literal::new(ScalarValue::Decimal128( + Some(1), + 0, + 0, + ))), + ])?; + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &schema) + .is_err() + ); + Ok(()) +} + +#[rstest] +#[case::eq(DFOperator::Eq, Operator::Eq)] +#[case::not_eq(DFOperator::NotEq, Operator::NotEq)] +#[case::lt(DFOperator::Lt, Operator::Lt)] +#[case::lte(DFOperator::LtEq, Operator::Lte)] +#[case::gt(DFOperator::Gt, Operator::Gt)] +#[case::gte(DFOperator::GtEq, Operator::Gte)] +#[case::and(DFOperator::And, Operator::And)] +#[case::or(DFOperator::Or, Operator::Or)] +#[case::plus(DFOperator::Plus, Operator::Add)] +#[case::plus(DFOperator::Minus, Operator::Sub)] +#[case::plus(DFOperator::Multiply, Operator::Mul)] +#[case::plus(DFOperator::Divide, Operator::Div)] +fn test_operator_conversion_supported( + #[case] df_op: DFOperator, + #[case] expected_vortex_op: Operator, +) { + assert_eq!(try_operator_from_df(&df_op), Some(expected_vortex_op)); +} + +#[rstest] +#[case::modulo(DFOperator::Modulo)] +#[case::bitwise_and(DFOperator::BitwiseAnd)] +#[case::regex_match(DFOperator::RegexMatch)] +#[case::like_match(DFOperator::LikeMatch)] +fn test_operator_conversion_unsupported(#[case] df_op: DFOperator) { + assert_eq!(try_operator_from_df(&df_op), None); +} + +#[test] +fn test_expr_from_df_column() -> DFResult<()> { + let col_expr = df_expr::Column::new("test_column", 0); + let result = convert( + Arc::new(col_expr), + &Schema::new(vec![Field::new("test_column", DataType::Int32, false)]), + )?; + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.get_item(test_column) + └── input: vortex.root() + "); + Ok(()) +} + +#[test] +fn test_expr_from_df_literal() -> DFResult<()> { + let literal_expr = df_expr::Literal::new(ScalarValue::Int32(Some(42))); + let result = convert(Arc::new(literal_expr), &Schema::empty())?; + + assert_snapshot!(result.display_tree().to_string(), @"vortex.literal(42i32)"); + Ok(()) +} + +fn convert_literal(expr: Arc) -> anyhow::Result { + convert(expr, &Schema::empty())? + .as_opt::() + .cloned() + .ok_or_else(|| anyhow::anyhow!("Expected a literal expression")) +} + +#[rstest] +#[case::null(ScalarValue::Null)] +#[case::boolean(ScalarValue::Boolean(Some(true)))] +#[case::false_value(ScalarValue::Boolean(Some(false)))] +#[case::u32(ScalarValue::UInt32(Some(42)))] +#[case::i32(ScalarValue::Int32(Some(-42)))] +#[case::i64(ScalarValue::Int64(Some(-123)))] +#[case::f64(ScalarValue::Float64(Some(2.5)))] +#[case::utf8(ScalarValue::Utf8(Some("test string".into())))] +#[case::binary(ScalarValue::Binary(Some(vec![1, 2, 3])))] +#[case::decimal32(ScalarValue::Decimal32(Some(1234), 5, 2))] +#[case::decimal64(ScalarValue::Decimal64(Some(12345), 10, 2))] +#[case::decimal128(ScalarValue::Decimal128(Some(12345), 20, 2))] +#[case::decimal256(ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10))] +#[case::date32(ScalarValue::Date32(Some(18628)))] +#[case::date64(ScalarValue::Date64(Some(1609459200000)))] +#[case::time32_second(ScalarValue::Time32Second(Some(3661)))] +#[case::time32_millisecond(ScalarValue::Time32Millisecond(Some(3661000)))] +#[case::time64_microsecond(ScalarValue::Time64Microsecond(Some(3661000000)))] +#[case::time64_nanosecond(ScalarValue::Time64Nanosecond(Some(3661000000000)))] +#[case::timestamp_second(ScalarValue::TimestampSecond(Some(1609459200), None))] +#[case::timestamp_millisecond(ScalarValue::TimestampMillisecond(Some(1609459200000), None))] +#[case::timestamp_microsecond(ScalarValue::TimestampMicrosecond(Some(1609459200000000), None))] +#[case::timestamp_nanosecond(ScalarValue::TimestampNanosecond(Some(1609459200000000000), None))] +#[case::timestamp_timezone(ScalarValue::TimestampNanosecond(Some(1609459200000000000), Some("UTC".into())))] +fn test_literal_round_trip( + #[case] value: ScalarValue, + #[values(false, true)] null: bool, +) -> anyhow::Result<()> { + let value = if null { + ScalarValue::try_from(&value.data_type())? + } else { + value + }; + let converted = convert_literal(Arc::new(df_expr::Literal::new(value.clone())))?; + assert_eq!(converted.is_null(), value.is_null()); + assert_eq!(converted.try_to_df()?, value); + Ok(()) +} + +#[rstest] +#[case::utf8_view(ScalarValue::Utf8View(Some("test string".into())), Scalar::from("test string"))] +#[case::large_utf8(ScalarValue::LargeUtf8(Some("test string".into())), Scalar::from("test string"))] +#[case::binary_view(ScalarValue::BinaryView(Some(vec![1, 2, 3])), Scalar::binary(vec![1, 2, 3], Nullability::NonNullable))] +#[case::large_binary(ScalarValue::LargeBinary(Some(vec![1, 2, 3])), Scalar::binary(vec![1, 2, 3], Nullability::NonNullable))] +#[case::dictionary(ScalarValue::Dictionary(Box::new(DataType::Int8), Box::new(ScalarValue::Utf8(Some("test string".into())))), Scalar::from("test string"))] +fn test_literal_storage_variants( + #[case] value: ScalarValue, + #[case] expected: Scalar, + #[values(false, true)] null: bool, +) -> anyhow::Result<()> { + let (value, expected) = if null { + ( + ScalarValue::try_from(&value.data_type())?, + Scalar::null(expected.dtype().as_nullable()), + ) + } else { + (value, expected) + }; + let converted = convert_literal(Arc::new(df_expr::Literal::new(value)))?; + assert!( + converted.eq_ignore_nullability(&expected), + "{converted} != {expected}" + ); + Ok(()) +} + +#[rstest] +fn test_struct_literal_preserves_extension_child( + #[values(false, true)] null: bool, +) -> anyhow::Result<()> { + let mut id_field = Field::new("id", DataType::FixedSizeBinary(16), false); + id_field.try_with_extension_type(ArrowUuid)?; + let fields = vec![Arc::new(id_field)].into(); + let array = if null { + StructArray::new_null(fields, 1) + } else { + let ids = FixedSizeBinaryArray::try_from_iter([*b"0123456789abcdef"].into_iter())?; + StructArray::try_new(fields, vec![Arc::new(ids)], None)? + }; + let converted = convert_literal(Arc::new(df_expr::Literal::new(ScalarValue::Struct( + Arc::new(array), + ))))?; + assert_eq!(converted.is_null(), null); + let id_dtype = converted + .dtype() + .as_struct_fields() + .field_by_index(0) + .ok_or_else(|| anyhow::anyhow!("Expected the id field"))?; + assert!(id_dtype.as_extension().is::()); + assert!(!id_dtype.is_nullable()); + Ok(()) +} + +#[rstest] +fn test_literal_preserves_extension_metadata( + #[values(false, true)] null: bool, +) -> anyhow::Result<()> { + let mut field = Field::new("id", DataType::FixedSizeBinary(16), null); + field.try_with_extension_type(ArrowUuid)?; + let value = ScalarValue::FixedSizeBinary(16, (!null).then(|| b"0123456789abcdef".to_vec())); + let converted = convert_literal(Arc::new(df_expr::Literal::new_with_metadata( + value, + Some(FieldMetadata::new_from_field(&field)), + )))?; + assert_eq!(converted.is_null(), null); + assert!(converted.dtype().as_extension().is::()); + Ok(()) +} + +#[test] +fn test_expr_from_df_binary() -> DFResult<()> { + let left = Arc::new(df_expr::Column::new("left", 0)) as Arc; + let right = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; + let binary_expr = df_expr::BinaryExpr::new(left, DFOperator::Eq, right); + + let result = convert( + Arc::new(binary_expr), + &Schema::new(vec![Field::new("left", DataType::Int32, false)]), + )?; + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.binary(=) + ├── lhs: vortex.get_item(left) + │ └── input: vortex.root() + └── rhs: vortex.literal(42i32) + "); + Ok(()) +} + +#[rstest] +#[case::like_normal(false, false)] +#[case::like_negated(true, false)] +#[case::like_case_insensitive(false, true)] +#[case::like_negated_case_insensitive(true, true)] +fn test_expr_from_df_like(#[case] negated: bool, #[case] case_insensitive: bool) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("text_col", 0)) as Arc; + let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "test%".to_string(), + )))) as Arc; + let like_expr = df_expr::LikeExpr::new(negated, case_insensitive, expr, pattern); + + let result = convert( + Arc::new(like_expr), + &Schema::new(vec![Field::new("text_col", DataType::Utf8, true)]), + )?; + let like_opts = result.as_::(); + assert_eq!( + like_opts, + &LikeOptions { + negated, + case_insensitive + } + ); + Ok(()) +} + +#[test] +fn test_like_preserves_nullable_input_through_non_nullable_cast() -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(( + "text_col", + Utf8View, + vec![Some("google"), None, Some("example")] + ))?; + let column = Arc::new(df_expr::Column::new("text_col", 0)) as Arc; + let cast = Arc::new(df_expr::CastExpr::new_with_target_field( + column, + Arc::new(Field::new("text_col", DataType::Utf8View, false)), + None, + )) as Arc; + let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8View(Some( + "%google%".to_string(), + )))) as Arc; + let like = Arc::new(df_expr::LikeExpr::new(false, false, cast, pattern)); + + assert_native_matches(like, batch) +} + +#[rstest] +fn test_expr_from_df_octet_length(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let octet_length = octet_length_expr(expr, &test_schema); + + let result = convert(octet_length, &test_schema)?; + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.cast(i32?) + └── input: vortex.byte_length() + └── input: vortex.get_item(name) + └── input: vortex.root() + "); + Ok(()) +} + +#[rstest] +fn test_expr_from_df_array_length(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let array_length = array_length_expr(vec![expr], &test_schema); + + let result = convert(array_length, &test_schema)?; + + assert_snapshot!(result.display_tree().to_string(), @r" + vortex.cast(u64?) + └── input: vortex.list.length() + └── input: vortex.get_item(tags) + └── input: vortex.root() + "); + Ok(()) +} + +#[rstest] +// Supported types +#[case::null(DataType::Null, true)] +#[case::boolean(DataType::Boolean, true)] +#[case::int8(DataType::Int8, true)] +#[case::int16(DataType::Int16, true)] +#[case::int32(DataType::Int32, true)] +#[case::int64(DataType::Int64, true)] +#[case::uint8(DataType::UInt8, true)] +#[case::uint16(DataType::UInt16, true)] +#[case::uint32(DataType::UInt32, true)] +#[case::uint64(DataType::UInt64, true)] +#[case::float32(DataType::Float32, true)] +#[case::float64(DataType::Float64, true)] +#[case::utf8(DataType::Utf8, true)] +#[case::utf8_view(DataType::Utf8View, true)] +#[case::binary(DataType::Binary, true)] +#[case::binary_view(DataType::BinaryView, true)] +#[case::date32(DataType::Date32, true)] +#[case::date64(DataType::Date64, true)] +#[case::timestamp_ms(DataType::Timestamp(ArrowTimeUnit::Millisecond, None), true)] +#[case::timestamp_us( + DataType::Timestamp(ArrowTimeUnit::Microsecond, Some(Arc::from("UTC"))), + true +)] +#[case::time32_s(DataType::Time32(ArrowTimeUnit::Second), true)] +#[case::time64_ns(DataType::Time64(ArrowTimeUnit::Nanosecond), true)] +// Unsupported types +#[case::list( + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + false +)] +#[case::struct_type(DataType::Struct(vec![Field::new("field", DataType::Int32, true)].into() +), false)] +// Dictionary types - should be supported if value type is supported +#[case::dict_utf8( + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true +)] +#[case::dict_int32( + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Int32)), + true +)] +#[case::dict_unsupported( + DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::List(Arc::new(Field::new("item", DataType::Int32, true)))) + ), + false +)] +fn test_supported_data_types(#[case] data_type: DataType, #[case] expected: bool) { + assert_eq!(supported_data_types(&data_type), expected); +} + +#[rstest] +fn test_can_be_pushed_down_column_supported(test_schema: Schema) -> DFResult<()> { + let col_expr = Arc::new(df_expr::Column::new("id", 0)) as Arc; + + assert!(converts(&col_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +#[case::duration(DataType::Duration(ArrowTimeUnit::Second))] +#[case::interval(DataType::Interval(IntervalUnit::YearMonth))] +#[case::fixed_size_binary(DataType::FixedSizeBinary(16))] +#[case::decimal32(DataType::Decimal32(1, 2))] +#[case::decimal64(DataType::Decimal64(1, 2))] +#[case::decimal128(DataType::Decimal128(1, 2))] +#[case::decimal256(DataType::Decimal256(1, 2))] +fn test_unsupported_field_is_declined_only_when_referenced( + #[case] unsupported: DataType, + #[values(false, true)] referenced: bool, +) -> DFResult<()> { + let schema = Schema::new(vec![ + Field::new("unsupported", unsupported, true), + Field::new("id", DataType::Int32, false), + ]); + if referenced { + let expr: Arc = Arc::new(df_expr::IsNotNullExpr::new(Arc::new( + df_expr::Column::new("unsupported", 0), + ))); + assert!(!converts(&expr, &schema)?); + } else { + let expr: Arc = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("id", 1)), + DFOperator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))), + )); + assert_eq!( + convert(expr, &schema)?, + Binary.new_expr(Operator::Eq, [get_item("id", root()), lit(42i32)]), + ); + } + Ok(()) +} + +#[test] +fn test_literal_ignores_unreferenced_malformed_decimal() -> DFResult<()> { + let schema = Schema::new(vec![Field::new( + "invalid", + DataType::Decimal128(1, 2), + true, + )]); + let expr: Arc = Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))); + assert_eq!(convert(expr, &schema)?, lit(42i32)); + Ok(()) +} + +#[test] +fn test_projection_ignores_unreferenced_unsupported_field() -> DFResult<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new( + "unsupported", + DataType::Duration(ArrowTimeUnit::Second), + true, + ), + Field::new("b", DataType::Int32, false), + ]); + let projection = ProjectionExprs::from(vec![ProjectionExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("b", 2)), + DFOperator::Plus, + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + DFOperator::Plus, + Arc::new(df_expr::Column::new("b", 2)), + )), + )), + "sum", + )]); + let output_schema = projection.project_schema(&schema)?; + let processed = DefaultExpressionConverter::default().split_projection( + projection, + &schema, + &output_schema, + )?; + let a = get_item("a", root()); + let b = get_item("b", root()); + assert_eq!( + processed.scan_projection, + pack( + [( + "sum", + Binary.new_expr( + Operator::Add, + [b.clone(), Binary.new_expr(Operator::Add, [a, b])] + ) + )], + Nullability::NonNullable, + ), + ); + assert_eq!(processed.scan_reference_schema, output_schema); + assert_eq!( + processed.leftover_projection, + ProjectionExprs::from(vec![ProjectionExpr::new( + Arc::new(df_expr::Column::new("sum", 0)), + "sum", + )]), + ); + Ok(()) +} + +#[rstest] +fn test_referenced_field_preserves_extension_metadata( + #[values(false, true)] nested: bool, +) -> DFResult<()> { + let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), true); + uuid_field.try_with_extension_type(ArrowUuid)?; + let field = if nested { + Field::new("nested", DataType::Struct(vec![uuid_field].into()), false) + } else { + uuid_field + }; + let expr: Arc = Arc::new(df_expr::Column::new(field.name(), 1)); + let expected = get_item(field.name().as_str(), root()); + let schema = Schema::new(vec![ + Field::new("unsupported", DataType::FixedSizeBinary(16), true), + field, + ]); + assert_eq!(convert(expr, &schema)?, expected); + Ok(()) +} + +#[rstest] +fn test_nested_column_conversion(test_schema: Schema) -> DFResult<()> { + let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + + assert!(converts(&col_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_column_not_found(test_schema: Schema) { + let col_expr = Arc::new(df_expr::Column::new("nonexistent", 99)) as Arc; + + assert!( + DefaultExpressionConverter::default() + .try_convert(&col_expr, &test_schema) + .is_err() + ); +} + +#[rstest] +fn test_can_be_pushed_down_literal_supported(test_schema: Schema) -> DFResult<()> { + let lit_expr = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; + + assert!(converts(&lit_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +#[case::duration(ScalarValue::DurationSecond(Some(42)))] +#[case::null_duration(ScalarValue::DurationSecond(None))] +#[case::interval(ScalarValue::IntervalYearMonth(Some(1)))] +#[case::fixed_size_binary(ScalarValue::FixedSizeBinary(5, Some(vec![1, 2, 3, 4, 5])))] +#[case::null_fixed_size_binary(ScalarValue::FixedSizeBinary(5, None))] +fn test_can_be_pushed_down_literal_unsupported( + test_schema: Schema, + #[case] value: ScalarValue, +) -> DFResult<()> { + let lit_expr = Arc::new(df_expr::Literal::new(value)) as Arc; + assert!(!converts(&lit_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_binary_supported(test_schema: Schema) -> DFResult<()> { + let left = Arc::new(df_expr::Column::new("id", 0)) as Arc; + let right = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; + let binary_expr = + Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right)) as Arc; + + assert!(converts(&binary_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_binary_unsupported_operator(test_schema: Schema) -> DFResult<()> { + let left = Arc::new(df_expr::Column::new("id", 0)) as Arc; + let right = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; + let binary_expr = Arc::new(df_expr::BinaryExpr::new( + left, + DFOperator::AtQuestion, + right, + )) as Arc; + + assert!(!converts(&binary_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_binary_unsupported_operand(test_schema: Schema) -> DFResult<()> { + let left = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let right = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))) as Arc; + let binary_expr = + Arc::new(df_expr::BinaryExpr::new(left, DFOperator::Eq, right)) as Arc; + + assert!(!converts(&binary_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_like_supported(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "test%".to_string(), + )))) as Arc; + let like_expr = + Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc; + + assert!(converts(&like_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_like_unsupported_operand(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let pattern = Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "test%".to_string(), + )))) as Arc; + let like_expr = + Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc; + + assert!(!converts(&like_expr, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_octet_length_supported(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let octet_length = octet_length_expr(expr, &test_schema); + + assert!(converts(&octet_length, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let octet_length = Arc::new(ScalarFunctionExpr::new( + "octet_length", + Arc::new(ScalarUDF::from(OctetLengthFunc::new())), + vec![expr], + Arc::new(Field::new("octet_length", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + )) as Arc; + + assert!(!converts(&octet_length, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_array_length_supported(test_schema: Schema) -> DFResult<()> { + let expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let array_length = array_length_expr(vec![expr], &test_schema); + + assert!(converts(&array_length, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_array_length_unsupported_operand(test_schema: Schema) -> DFResult<()> { + // `array_length` over a non-list column cannot be pushed down. + let expr = Arc::new(df_expr::Column::new("name", 1)) as Arc; + let array_length = Arc::new(ScalarFunctionExpr::new( + "array_length", + Arc::new(ScalarUDF::from(ArrayLength::new())), + vec![expr], + Arc::new(Field::new("array_length", DataType::UInt64, true)), + Arc::new(ConfigOptions::new()), + )) as Arc; + + assert!(!converts(&array_length, &test_schema)?); + Ok(()) +} + +#[rstest] +fn test_can_be_pushed_down_array_length_dimension_one_supported( + test_schema: Schema, +) -> DFResult<()> { + // `array_length(arr, 1)` is the first-dimension length, equivalent to `list_length`. + let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let dimension = + Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(1)))) as Arc; + let array_length = array_length_expr(vec![list, dimension], &test_schema); + + assert!(converts(&array_length, &test_schema)?); + Ok(()) +} + +#[rstest] +#[case::higher(Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(2)))))] +#[case::zero(Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(0)))))] +#[case::negative(Arc::new(df_expr::Literal::new(ScalarValue::Int64(Some(-1)))))] +#[case::null(Arc::new(df_expr::Literal::new(ScalarValue::Int64(None))))] +#[case::dynamic(Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("id", 0)), + DataType::Int64, + None, +)))] +fn test_can_be_pushed_down_array_length_higher_dimension_not_supported( + test_schema: Schema, + #[case] dimension: Arc, +) -> DFResult<()> { + let list = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + let array_length = array_length_expr(vec![list, dimension], &test_schema); + assert!(!converts(&array_length, &test_schema)?); + Ok(()) +} + +#[rstest] +#[case::octet_zero(OctetLengthFunc::new().into(), 0)] +#[case::octet_two(OctetLengthFunc::new().into(), 2)] +#[case::array_zero(ArrayLength::new().into(), 0)] +#[case::array_three(ArrayLength::new().into(), 3)] +fn test_length_function_invalid_arity(#[case] function: ScalarUDF, #[case] arity: usize) { + let input = Arc::new(df_expr::Literal::new(ScalarValue::Null)) as Arc; + let name = function.name().to_owned(); + let expr: Arc = Arc::new(ScalarFunctionExpr::new( + &name, + Arc::new(function), + vec![input; arity], + Arc::new(Field::new("", DataType::Int64, true)), + Arc::new(ConfigOptions::new()), + )); + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &Schema::empty()) + .is_err() + ); +} + +// https://github.com/vortex-data/vortex/issues/6211 +#[tokio::test] +async fn test_cast_int_to_string() -> anyhow::Result<()> { + let ctx = TestSessionContext::default(); + + ctx.session + .sql(r#"copy (select 1 as id) to 'example.vortex'"#) + .await? + .show() + .await?; + + ctx.session + .sql(r#"select cast(id as string) as sid from 'example.vortex' where id > 0"#) + .await? + .show() + .await?; + + ctx.session + .sql(r#"select id from 'example.vortex' where cast (id as string) == '1'"#) + .await? + .show() + .await?; + + // This fails as it pushes string cast to the scan + ctx.session + .sql(r#"select cast(id as string) from 'example.vortex'"#) + .await? + .collect() + .await?; + + Ok(()) +} + +/// A cast whose target is a UUID-tagged `FixedSizeBinary(16)` must resolve +/// through the dtype extension registry (UUID is registered on the default +/// session) instead of the static, non-plugin-aware `DType::from_arrow`, +/// which does not support `FixedSizeBinary` and previously panicked here. +#[test] +fn test_cast_to_uuid_resolves_via_registry() -> anyhow::Result<()> { + let mut uuid_field = Field::new("id", DataType::FixedSizeBinary(16), true); + uuid_field.try_with_extension_type(ArrowUuid)?; + + let child = Arc::new(df_expr::Column::new("id", 0)) as Arc; + let schema = Schema::new(vec![uuid_field.clone()]); + let cast: Arc = Arc::new(df_expr::CastExpr::new_with_target_field( + child, + Arc::new(uuid_field), + None, + )); + + // Must convert without panicking — the static path would `unimplemented!()`. + assert!(converts(&cast, &schema)?); + Ok(()) +} + +/// Test that applying a CASE expression to an Arrow RecordBatch using DataFusion +/// matches the result of applying the converted Vortex expression. +#[test] +fn test_case_when_datafusion_vortex_equivalence() -> anyhow::Result<()> { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])), + vec![Arc::new(arrow_array::Int32Array::from(vec![ + 1, 5, 10, 15, 20, + ]))], + )?; + let value: Arc = Arc::new(df_expr::Column::new("value", 0)); + let int32 = + |v| Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(v)))) as Arc; + let greater_than = |threshold| { + Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&value), + DFOperator::Gt, + int32(threshold), + )) as Arc + }; + + // CASE WHEN value > 10 THEN 100 WHEN value > 5 THEN 50 ELSE 0 END + let case_expr = df_expr::CaseExpr::try_new( + None, + vec![(greater_than(10), int32(100)), (greater_than(5), int32(50))], + Some(int32(0)), + )?; + assert_native_matches(Arc::new(case_expr), batch) +} + +fn assert_native_matches(expr: Arc, batch: RecordBatch) -> anyhow::Result<()> { + let session = VortexSession::default(); + let converted = DefaultExpressionConverter::new(session.clone()) + .try_convert(&expr, &batch.schema())? + .ok_or_else(|| anyhow::anyhow!("Expected native conversion for {expr}"))?; + for constant in [false, true] { + let batches = if constant { + (0..batch.num_rows()).map(|i| batch.slice(i, 1)).collect() + } else { + vec![batch.clone()] + }; + for batch in batches { + let mut ctx = session.create_execution_ctx(); + let mut input = session + .arrow() + .from_arrow_record_batch(batch.clone(), &batch.schema())?; + if constant { + input = ConstantArray::new(input.execute_scalar(0, &mut ctx)?, 1).into_array(); + } + let actual = input + .apply(&converted)? + .execute::(&mut ctx); + let expected = expr + .evaluate(&batch) + .and_then(|value| value.into_array(batch.num_rows())); + match (actual, expected) { + (Ok(actual), Ok(expected)) => { + let expected = session + .arrow() + .from_arrow_array(expected, expr.return_field(&batch.schema())?.as_ref())?; + assert_arrays_eq!(actual, expected, &mut ctx); + } + (Err(_), Err(_)) => {} + (actual, expected) => anyhow::bail!( + "Evaluation mismatch for {expr}: Vortex {actual:?}, DataFusion {expected:?}" + ), + } + } + } + Ok(()) +} + +#[rstest] +#[case::ordinary(vec![Some(i32::MIN), Some(i32::MAX), Some(-1), Some(0), None], vec![2, 2, 2, 2, 0])] +#[case::zero(vec![Some(1)], vec![0])] +#[case::overflow(vec![Some(i32::MIN)], vec![-1])] +fn test_native_integer_division( + #[case] values: Vec>, + #[case] denominators: Vec, +) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(("a", Int32, values), ("b", Int32, denominators))?; + assert_native_matches( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + DFOperator::Divide, + Arc::new(df_expr::Column::new("b", 1)), + )), + batch, + ) +} + +#[rstest] +#[case::add(DFOperator::Plus)] +#[case::sub(DFOperator::Minus)] +#[case::mul(DFOperator::Multiply)] +#[case::div(DFOperator::Divide)] +fn test_native_float_arithmetic(#[case] operator: DFOperator) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!( + ( + "a", + Float64, + vec![ + Some(f64::NAN), + Some(f64::INFINITY), + Some(-0.0), + Some(f64::MAX), + None + ] + ), + ("b", Float64, vec![0.0, f64::NEG_INFINITY, 0.0, 2.0, 0.0]) + )?; + assert_native_matches( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + operator, + Arc::new(df_expr::Column::new("b", 1)), + )), + batch, + ) +} + +#[rstest] +#[case::identity(DataType::Int32)] +#[case::widen(DataType::Int64)] +fn test_native_integer_cast(#[case] target: DataType) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(( + "a", + Int32, + vec![Some(i32::MIN), Some(i32::MAX), Some(0), None] + ))?; + assert_native_matches( + Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + target, + None, + )), + batch, + ) +} + +#[test] +fn test_native_float_cast() -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(( + "a", + Float32, + vec![ + Some(f32::NAN), + Some(f32::INFINITY), + Some(-0.0), + Some(f32::MAX), + None + ] + ))?; + assert_native_matches( + Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + DataType::Float64, + None, + )), + batch, + ) +} + +#[rstest] +#[case::narrow(DataType::Int8)] +#[case::signedness(DataType::UInt32)] +#[case::string(DataType::Utf8)] +#[case::decimal(DataType::Decimal128(10, 2))] +fn test_cast_falls_back(#[case] target: DataType) -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let expr: Arc = Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + target, + None, + )); + assert!(!converts(&expr, &schema)?); + Ok(()) +} + +#[test] +fn test_cast_options_fall_back() -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let expr: Arc = Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + DataType::Int64, + Some(datafusion_common::arrow::compute::CastOptions { + safe: true, + format_options: DEFAULT_FORMAT_OPTIONS, + }), + )); + assert!(!converts(&expr, &schema)?); + Ok(()) +} + +#[rstest] +#[case::add(DFOperator::Plus)] +#[case::sub(DFOperator::Minus)] +#[case::mul(DFOperator::Multiply)] +#[case::div(DFOperator::Divide)] +fn test_integer_arithmetic_pushdown_ignores_overflow_mode( + #[case] op: DFOperator, + #[values(false, true)] checked: bool, +) -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let expr: Arc = Arc::new( + df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + op, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), + ) + .with_fail_on_overflow(checked), + ); + assert!(converts(&expr, &schema)?); + Ok(()) +} + +#[rstest] +#[case::add(DFOperator::Plus)] +#[case::sub(DFOperator::Minus)] +#[case::mul(DFOperator::Multiply)] +#[case::div(DFOperator::Divide)] +fn test_native_integer_arithmetic( + #[case] op: DFOperator, + #[values(false, true)] checked: bool, +) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!( + ("a", Int32, vec![Some(-12), Some(0), Some(12), None]), + ("b", Int32, vec![-2, 2, 2, 0]) + )?; + assert_native_matches( + Arc::new( + df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + op, + Arc::new(df_expr::Column::new("b", 1)), + ) + .with_fail_on_overflow(checked), + ), + batch, + ) +} + +#[rstest] +#[case::add(DFOperator::Plus)] +#[case::sub(DFOperator::Minus)] +#[case::mul(DFOperator::Multiply)] +#[case::div(DFOperator::Divide)] +fn test_native_decimal_arithmetic(#[case] op: DFOperator) -> anyhow::Result<()> { + let batch = RecordBatch::try_from_iter([ + ( + "a", + Arc::new( + arrow_array::Decimal128Array::from(vec![Some(-1200), Some(0), Some(1200), None]) + .with_precision_and_scale(12, 2)?, + ) as arrow_array::ArrayRef, + ), + ( + "b", + Arc::new( + arrow_array::Decimal128Array::from(vec![-200, 200, 200, 0]) + .with_precision_and_scale(12, 2)?, + ) as arrow_array::ArrayRef, + ), + ])?; + assert_native_matches( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + op, + Arc::new(df_expr::Column::new("b", 1)), + )), + batch, + ) +} + +#[test] +fn test_fallible_case_branch_is_residual() -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let a: Arc = Arc::new(df_expr::Column::new("a", 0)); + let zero: Arc = Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))); + let expr: Arc = Arc::new(df_expr::CaseExpr::try_new( + None, + vec![( + Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&a), + DFOperator::NotEq, + Arc::clone(&zero), + )), + Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&zero), + DFOperator::Divide, + a, + )), + )], + Some(zero), + )?); + assert!(!converts(&expr, &schema)?); + Ok(()) +} + +#[rstest] +fn test_nested_functions_reject_unknown_children( + #[values(false, true)] list: bool, +) -> DFResult<()> { + let input_type = if list { + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))) + } else { + DataType::Struct(vec![Field::new("leaf", DataType::Int32, true)].into()) + }; + let field = Arc::new(Field::new("a", input_type, true)); + let schema = Schema::new(vec![Arc::clone(&field)]); + let unknown: Arc = Arc::new(ScalarFunctionExpr::new( + "coalesce", + Arc::new(ScalarUDF::from(CoalesceFunc::new())), + vec![Arc::new(df_expr::Column::new("a", 0))], + field, + Arc::new(ConfigOptions::new()), + )); + let expr = if list { + array_length_expr(vec![unknown], &schema) + } else { + Arc::new(ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + vec![ + unknown, + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "leaf".into(), + )))), + ], + &schema, + Arc::new(ConfigOptions::new()), + )?) as Arc + }; + assert!(!converts(&expr, &schema)?); + Ok(()) +} + +#[rstest] +#[case::no_args(vec![])] +#[case::no_path(vec![Arc::new(df_expr::Column::new("a", 0)) as Arc])] +#[case::null_path(vec![ + Arc::new(df_expr::Column::new("a", 0)) as Arc, + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(None))), +])] +#[case::column_path(vec![ + Arc::new(df_expr::Column::new("a", 0)) as Arc, + Arc::new(df_expr::Column::new("a", 0)), +])] +fn test_malformed_get_field_returns_error( + #[case] args: Vec>, +) -> DFResult<()> { + let schema = Schema::new(vec![Field::new( + "a", + DataType::Struct(vec![Field::new("leaf", DataType::Int32, true)].into()), + true, + )]); + let expr: Arc = Arc::new(ScalarFunctionExpr::new( + "get_field", + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + args, + Arc::new(Field::new("", DataType::Int32, true)), + Arc::new(ConfigOptions::new()), + )); + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &schema) + .is_err() + ); + Ok(()) +} + +#[test] +fn test_column_identity_validation_skips_dynamic_filters() -> DFResult<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]); + let expr: Arc = Arc::new(df_expr::Column::new("a", 1)); + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &schema) + .is_err() + ); + let expr = Arc::new(DynamicFilterPhysicalExpr::new( + vec![expr], + Arc::new(df_expr::Literal::new(ScalarValue::Boolean(Some(true)))), + )) as Arc; + + assert!(!converts(&expr, &schema)?); + Ok(()) +} + +#[rstest] +fn test_native_nested_list_length( + #[values(false, true)] nullable_parent: bool, +) -> anyhow::Result<()> { + let mut lists = + arrow_array::builder::ListBuilder::new(arrow_array::builder::Int32Builder::new()); + lists.values().append_value(1); + lists.append(true); + lists.append(true); + lists.append(false); + let lists = Arc::new(lists.finish()); + let field = Arc::new(Field::new("items", lists.data_type().clone(), true)); + let payload = Arc::new(StructArray::new( + vec![field].into(), + vec![lists], + nullable_parent.then(|| NullBuffer::from(vec![true, false, true])), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "payload", + payload.data_type().clone(), + nullable_parent, + )])); + let batch = RecordBatch::try_new(schema, vec![payload])?; + let get_field: Arc = Arc::new(ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(df_expr::Column::new("payload", 0)), + Arc::new(df_expr::Literal::new(ScalarValue::Utf8(Some( + "items".into(), + )))), + ], + &batch.schema(), + Arc::new(ConfigOptions::new()), + )?); + let length = array_length_expr(vec![get_field], &batch.schema()); + if nullable_parent { + assert!(!converts(&length, &batch.schema())?); + Ok(()) + } else { + assert_native_matches(length, batch) + } +} + +#[test] +fn test_native_case_null_conditions() -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(("a", Int32, vec![Some(1), None, Some(3)]))?; + let a: Arc = Arc::new(df_expr::Column::new("a", 0)); + let expr = Arc::new(df_expr::CaseExpr::try_new( + None, + vec![( + Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&a), + DFOperator::Gt, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), + )) as Arc, + a, + )], + None, + )?); + assert_native_matches(expr, batch) +} + +#[test] +fn test_simple_case_falls_back() -> DFResult<()> { + let value = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))) as Arc; + let expr: Arc = Arc::new(df_expr::CaseExpr::try_new( + Some(Arc::clone(&value)), + vec![(Arc::clone(&value), value)], + None, + )?); + assert!(!converts(&expr, &Schema::empty())?); + Ok(()) +} + +#[test] +fn test_case_non_boolean_condition_returns_error() -> DFResult<()> { + let value = + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))) as Arc; + let expr: Arc = Arc::new(df_expr::CaseExpr::try_new( + None, + vec![(Arc::clone(&value), value)], + None, + )?); + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &Schema::empty()) + .is_err() + ); + Ok(()) +} + +#[test] +fn test_malformed_literal_returns_error() { + let value = ScalarValue::Decimal128(Some(1), 0, 0); + let expr: Arc = Arc::new(df_expr::Literal::new(value)); + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &Schema::empty()) + .is_err() + ); +} + +#[rstest] +fn test_literal_invalid_row_count(#[values(0, 2)] len: usize) { + let array = StructArray::new_empty_fields(len, None); + let expr: Arc = + Arc::new(df_expr::Literal::new(ScalarValue::Struct(Arc::new(array)))); + assert!( + DefaultExpressionConverter::default() + .try_convert(&expr, &Schema::empty()) + .is_err() + ); +} + +#[rstest] +#[case::and(DFOperator::And, DFOperator::NotEq)] +#[case::or(DFOperator::Or, DFOperator::Eq)] +fn test_fallible_boolean_rhs_is_residual( + #[case] operator: DFOperator, + #[case] comparison: DFOperator, +) -> anyhow::Result<()> { + let batch = arrow_array::record_batch!(("a", Int32, vec![0, 0, 0, 0, 2]))?; + let a: Arc = Arc::new(df_expr::Column::new("a", 0)); + let zero: Arc = Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))); + let expr: Arc = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&a), + comparison, + Arc::clone(&zero), + )), + operator, + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(12)))), + DFOperator::Divide, + a, + )), + DFOperator::Gt, + zero, + )), + )); + // DataFusion selects the one row that needs RHS evaluation at this selectivity. + expr.evaluate(&batch)?; + assert!(!converts(&expr, &batch.schema())?); + Ok(()) +} diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 476f32afa49..a52486a4c9e 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -4,7 +4,7 @@ //! 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 +//! [`ExpressionConverter`] - 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. @@ -16,8 +16,8 @@ mod scalars; pub(crate) mod schema; pub(crate) mod stats; -pub use exprs::DefaultExpressionConvertor; -pub use exprs::ExpressionConvertor; +pub use exprs::DefaultExpressionConverter; +pub use exprs::ExpressionConverter; pub use exprs::ProcessedProjection; pub use scalars::scalar_from_df; diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 6d05144e315..948b5a32c44 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -68,7 +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::ExpressionConverter; use crate::convert::TryToDataFusion; use crate::convert::stats::is_constant_to_distinct_count; @@ -124,7 +124,7 @@ const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usi pub struct VortexFormat { session: VortexSession, opts: VortexTableOptions, - expression_convertor: Option>, + expression_converter: Option>, } impl Debug for VortexFormat { @@ -132,8 +132,8 @@ impl Debug for VortexFormat { f.debug_struct("VortexFormat") .field("opts", &self.opts) .field( - "has_expression_convertor", - &self.expression_convertor.is_some(), + "has_expression_converter", + &self.expression_converter.is_some(), ) .finish() } @@ -279,7 +279,7 @@ impl ConfigExtension for VortexTableOptions { pub struct VortexFormatFactory { session: VortexSession, options: Option, - expression_convertor: Option>, + expression_converter: Option>, } impl Debug for VortexFormatFactory { @@ -288,8 +288,8 @@ impl Debug for VortexFormatFactory { .field("session", &self.session) .field("options", &self.options) .field( - "has_expression_convertor", - &self.expression_convertor.is_some(), + "has_expression_converter", + &self.expression_converter.is_some(), ) .finish() } @@ -316,7 +316,7 @@ impl VortexFormatFactory { Self { session: VortexSession::default(), options: None, - expression_convertor: None, + expression_converter: None, } } @@ -329,7 +329,7 @@ impl VortexFormatFactory { Self { session, options: None, - expression_convertor: None, + expression_converter: None, } } @@ -344,7 +344,7 @@ impl VortexFormatFactory { Self { session, options: Some(options), - expression_convertor: None, + expression_converter: None, } } @@ -372,12 +372,12 @@ impl VortexFormatFactory { self } - /// Sets the [`ExpressionConvertor`] used by formats and sources created by this factory. - pub fn with_expression_convertor( + /// Sets the [`ExpressionConverter`] used by formats and sources created by this factory. + pub fn with_expression_converter( mut self, - expression_convertor: Arc, + expression_converter: Arc, ) -> Self { - self.expression_convertor = Some(expression_convertor); + self.expression_converter = Some(expression_converter); self } } @@ -418,16 +418,16 @@ impl FileFormatFactory for VortexFormatFactory { } 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)); + if let Some(expression_converter) = &self.expression_converter { + format = format.with_expression_converter(Arc::clone(expression_converter)); } Ok(Arc::new(format)) } fn default(&self) -> Arc { 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)); + if let Some(expression_converter) = &self.expression_converter { + format = format.with_expression_converter(Arc::clone(expression_converter)); } Arc::new(format) } @@ -449,7 +449,7 @@ impl VortexFormat { Self { session, opts, - expression_convertor: None, + expression_converter: None, } } @@ -459,12 +459,12 @@ impl VortexFormat { &self.opts } - /// Sets the [`ExpressionConvertor`] used by every [`VortexSource`] created by this format. - pub fn with_expression_convertor( + /// Sets the [`ExpressionConverter`] used by every [`VortexSource`] created by this format. + pub fn with_expression_converter( mut self, - expression_convertor: Arc, + expression_converter: Arc, ) -> Self { - self.expression_convertor = Some(expression_convertor); + self.expression_converter = Some(expression_converter); self } } @@ -771,8 +771,8 @@ impl FileFormat for VortexFormat { fn file_source(&self, table_schema: TableSchema) -> Arc { 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)); + if let Some(expression_converter) = &self.expression_converter { + source = source.with_expression_converter(Arc::clone(expression_converter)); } Arc::new(source) as _ } @@ -817,7 +817,7 @@ mod tests { use super::*; use crate::common_tests::TestSessionContext; - use crate::convert::DefaultExpressionConvertor; + use crate::convert::DefaultExpressionConverter; use crate::convert::ProcessedProjection; #[derive(Clone, Copy)] @@ -827,52 +827,49 @@ mod tests { } #[derive(Default)] - struct ExpressionConvertorCalls { - can_be_pushed_down: AtomicBool, - convert: AtomicBool, + struct ExpressionConverterCalls { + try_convert: AtomicBool, } - impl ExpressionConvertorCalls { + impl ExpressionConverterCalls { fn reset(&self) { - self.can_be_pushed_down.store(false, Ordering::Relaxed); - self.convert.store(false, Ordering::Relaxed); + self.try_convert.store(false, Ordering::Relaxed); } } - struct TestExpressionConvertor { - inner: DefaultExpressionConvertor, + struct TestExpressionConverter { + inner: DefaultExpressionConverter, pushdown_mode: PushdownMode, - calls: Arc, + calls: Arc, } - impl TestExpressionConvertor { + impl TestExpressionConverter { fn new( session: VortexSession, pushdown_mode: PushdownMode, - calls: Arc, + calls: Arc, ) -> Self { Self { - inner: DefaultExpressionConvertor::new(session), + inner: DefaultExpressionConverter::new(session), pushdown_mode, calls, } } } - 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); + impl ExpressionConverter for TestExpressionConverter { + fn try_convert( + &self, + expr: &Arc, + schema: &Schema, + ) -> DFResult> { + self.calls.try_convert.store(true, Ordering::Relaxed); match self.pushdown_mode { - PushdownMode::Reject => false, - PushdownMode::Delegate => self.inner.can_be_pushed_down(expr, schema), + PushdownMode::Reject => Ok(None), + PushdownMode::Delegate => self.inner.try_convert(expr, schema), } } - fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { - self.calls.convert.store(true, Ordering::Relaxed); - self.inner.convert(expr) - } - fn split_projection( &self, source_projection: ProjectionExprs, @@ -884,29 +881,28 @@ mod tests { } } - fn expression_convertor_test_schema() -> Arc { + fn expression_converter_test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) } - fn expression_convertor_test_filter() -> Arc { + fn expression_converter_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_rejects_pushdown_with_expression_convertor( + fn assert_rejects_pushdown_with_expression_converter( format: &dyn FileFormat, - calls: &ExpressionConvertorCalls, + calls: &ExpressionConverterCalls, ) -> anyhow::Result<()> { - let source = format.file_source(TableSchema::from(expression_convertor_test_schema())); + let source = format.file_source(TableSchema::from(expression_converter_test_schema())); let result = source.try_pushdown_filters( - vec![expression_convertor_test_filter()], + vec![expression_converter_test_filter()], &ConfigOptions::new(), )?; - assert!(calls.can_be_pushed_down.load(Ordering::Relaxed)); - assert!(!calls.convert.load(Ordering::Relaxed)); + assert!(calls.try_convert.load(Ordering::Relaxed)); assert!(matches!(result.filters.as_slice(), [PushedDown::No])); Ok(()) } @@ -983,53 +979,53 @@ mod tests { } #[test] - fn format_plumbs_expression_convertor() -> anyhow::Result<()> { + fn format_plumbs_expression_converter() -> anyhow::Result<()> { let session = VortexSession::default(); - let calls = Arc::new(ExpressionConvertorCalls::default()); - let convertor = Arc::new(TestExpressionConvertor::new( + let calls = Arc::new(ExpressionConverterCalls::default()); + let converter = Arc::new(TestExpressionConverter::new( session.clone(), PushdownMode::Reject, Arc::clone(&calls), )); - let format = VortexFormat::new(session).with_expression_convertor(convertor); + let format = VortexFormat::new(session).with_expression_converter(converter); - assert_rejects_pushdown_with_expression_convertor(&format, &calls) + assert_rejects_pushdown_with_expression_converter(&format, &calls) } #[test] - fn factory_plumbs_expression_convertor() -> anyhow::Result<()> { - let calls = Arc::new(ExpressionConvertorCalls::default()); - let convertor = Arc::new(TestExpressionConvertor::new( + fn factory_plumbs_expression_converter() -> anyhow::Result<()> { + let calls = Arc::new(ExpressionConverterCalls::default()); + let converter = Arc::new(TestExpressionConverter::new( VortexSession::default(), PushdownMode::Reject, Arc::clone(&calls), )); - let factory = VortexFormatFactory::new().with_expression_convertor(convertor); + let factory = VortexFormatFactory::new().with_expression_converter(converter); let ctx = TestSessionContext::default(); let format = factory.create(&ctx.session.state(), &Default::default())?; - assert_rejects_pushdown_with_expression_convertor(format.as_ref(), &calls)?; + assert_rejects_pushdown_with_expression_converter(format.as_ref(), &calls)?; calls.reset(); let format = FileFormatFactory::default(&factory); - assert_rejects_pushdown_with_expression_convertor(format.as_ref(), &calls) + assert_rejects_pushdown_with_expression_converter(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( + async fn external_table_query_uses_factory_expression_converter() -> anyhow::Result<()> { + let calls = Arc::new(ExpressionConverterCalls::default()); + let converter = Arc::new(TestExpressionConverter::new( VortexSession::default(), PushdownMode::Delegate, Arc::clone(&calls), )); - let factory = Arc::new(VortexFormatFactory::new().with_expression_convertor(convertor)); + let factory = Arc::new(VortexFormatFactory::new().with_expression_converter(converter)); let ctx = TestSessionContext::new_with_factory(factory); ctx.session .sql( "CREATE EXTERNAL TABLE numbers (a INT NOT NULL) \ - STORED AS vortex LOCATION '/expression-convertor/'", + STORED AS vortex LOCATION '/expression-converter/'", ) .await?; ctx.session @@ -1046,8 +1042,7 @@ mod tests { .collect() .await?; - assert!(calls.can_be_pushed_down.load(Ordering::Relaxed)); - assert!(calls.convert.load(Ordering::Relaxed)); + assert!(calls.try_convert.load(Ordering::Relaxed)); let mut values = Vec::new(); for batch in batches { let array = batch diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 0f000452f17..8d7696d6c2b 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -1,13 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::future::ready; use std::ops::Range; use std::sync::Arc; use std::sync::Weak; use arrow_array::RecordBatchOptions; +use arrow_schema::DataType; use arrow_schema::Field; -use arrow_schema::Schema; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; use datafusion_common::ScalarValue; @@ -26,9 +27,11 @@ use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; use datafusion_physical_expr::split_conjunction; use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::utils::conjunction_opt; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_adapter::replace_columns_with_literals; +use datafusion_physical_plan::filter::batch_filter; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::metrics::MetricBuilder; use datafusion_physical_plan::metrics::MetricCategory; @@ -42,6 +45,7 @@ use tracing::Instrument; use vortex::array::VortexSessionExecute; use vortex::error::VortexError; use vortex::error::VortexExpect; +use vortex::expr::and_collect; use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; @@ -54,9 +58,9 @@ use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; use crate::VortexAccessPlan; -use crate::convert::exprs::ExpressionConvertor; +use crate::convert::exprs::ExpressionConverter; use crate::convert::exprs::ProcessedProjection; -use crate::convert::exprs::make_vortex_predicate; +use crate::convert::exprs::raw_projection; use crate::convert::schema::calculate_physical_schema; use crate::metrics::PARTITION_LABEL; use crate::metrics::PATH_LABEL; @@ -73,9 +77,8 @@ pub(crate) struct VortexOpener { /// Optional table schema projection. The indices are w.r.t. the `table_schema`, which is /// all fields in the final scan result not including the partition columns. pub projection: ProjectionExprs, - /// Filter expression optimized for pushdown into Vortex scan operations. - /// This may be a subset of file_pruning_predicate containing only expressions - /// that Vortex can efficiently evaluate. + /// Exact filter accepted during planning. Per-file adaptation may move parts + /// of it to DataFusion residual evaluation before projection and limits. pub filter: Option, /// Filter expression used by DataFusion's FilePruner to eliminate files based on /// statistics and partition values without opening them. @@ -100,7 +103,7 @@ pub(crate) struct VortexOpener { /// Whether the query has output ordering specified pub has_output_ordering: bool, - pub expression_convertor: Arc, + pub expression_converter: Arc, pub file_metadata_cache: Option>, /// Whether to enable expression pushdown into the underlying Vortex scan. pub projection_pushdown: bool, @@ -141,7 +144,7 @@ impl FileOpener for VortexOpener { let has_output_ordering = self.has_output_ordering; let scan_concurrency = self.scan_concurrency; - let expr_convertor = Arc::clone(&self.expression_convertor); + let expr_converter = Arc::clone(&self.expression_converter); let projection_pushdown = self.projection_pushdown; let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) @@ -278,23 +281,72 @@ impl FileOpener for VortexOpener { // another simplification pass. simplifier.simplify(expr_adapter.rewrite(filter)?) }) - .transpose()?; + .transpose() + .map_err(|e| { + exec_datafusion_err!("Failed to adapt filter in {}: {e}", file.path()) + })?; let projection = projection.try_map_exprs(|p| simplifier.simplify(expr_adapter.rewrite(p)?))?; + // Split the adapted filter into conjuncts Vortex evaluates natively and residual + // conjuncts DataFusion evaluates on the raw scan output. + let mut native_filters = Vec::new(); + let mut residual_filters = Vec::new(); + if let Some(filter) = &filter { + if filter.data_type(&this_file_schema)? != DataType::Boolean { + return Err(exec_datafusion_err!( + "Filter must be Boolean in {}: {filter}", + file.path() + )); + } + for expr in split_conjunction(filter) { + let converted = expr_converter + .try_convert(expr, &this_file_schema) + .map_err(|e| { + exec_datafusion_err!( + "Failed to convert filter {expr} in {}: {e}", + file.path() + ) + })?; + match converted { + Some(expr) => native_filters.push(expr), + None => residual_filters.push(Arc::clone(expr)), + } + } + } + let residual_filter = conjunction_opt(residual_filters); + let native_filter = and_collect(native_filters) + .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) + .transpose() + .map_err(|e| { + exec_datafusion_err!("Couldn't bind Vortex scan filter in {}: {e}", file.path()) + })?; + + // Residual filters must see raw inputs before computed projections or aliases. let ProcessedProjection { scan_projection, + scan_reference_schema, leftover_projection, - } = if projection_pushdown { - expr_convertor.split_projection( - projection.clone(), + } = if let Some(residual) = &residual_filter { + let mut indices = projection.column_indices(); + indices.extend(collect_columns(residual).iter().map(|column| column.index())); + indices.sort_unstable(); + indices.dedup(); + let (scan_projection, scan_reference_schema) = + raw_projection(&indices, &this_file_schema)?; + ProcessedProjection { + scan_projection, + scan_reference_schema, + leftover_projection: projection, + } + } else if projection_pushdown { + expr_converter.split_projection( + projection, &this_file_schema, output_schema.as_ref(), )? } else { - // When projection pushdown is disabled, read only the required columns - // and apply the full projection after the scan. - expr_convertor.no_pushdown_projection(projection.clone(), &this_file_schema)? + expr_converter.no_pushdown_projection(projection, &this_file_schema)? }; // The schema of the stream returned from the vortex scan. @@ -307,24 +359,14 @@ impl FileOpener for VortexOpener { })?; let scan_dtype = scan_projection.dtype().clone(); - // When projection pushdown is enabled, the scan outputs the projected columns. - // When disabled, the scan outputs raw columns and the projection is applied after. - let scan_reference_schema = if projection_pushdown { - (*output_schema).clone() - } else { - // Build schema from the raw columns being read - let column_indices = projection.column_indices(); - let fields: Vec<_> = column_indices - .into_iter() - .map(|idx| this_file_schema.field(idx).clone()) - .collect(); - Schema::new_with_metadata(fields, this_file_schema.metadata().clone()) - }; let stream_schema = calculate_physical_schema(&scan_dtype, &scan_reference_schema, &session.arrow())?; let leftover_projection = leftover_projection .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; + let residual_filter = residual_filter + .map(|expr| reassign_expr_columns(expr, &stream_schema)) + .transpose()?; let projector = leftover_projection.make_projector(&stream_schema)?; // We share our layout readers with others partitions in the scan, so we can only need to read each layout in each file once. @@ -361,44 +403,9 @@ impl FileOpener for VortexOpener { scan_builder = vortex_plan.apply_to_builder(scan_builder); } - let filter = filter - .and_then(|f| { - // Verify that all filters we've accepted from DataFusion get pushed down. - // This will only fail if the user has not configured a suitable - // PhysicalExprAdapterFactory on the file source to handle rewriting the - // expression to handle missing/reordered columns in the Vortex file. - let (pushed, unpushed): (Vec, Vec) = - split_conjunction(&f) - .into_iter() - .cloned() - .partition(|expr| { - expr_convertor.can_be_pushed_down(expr, &this_file_schema) - }); - - if !unpushed.is_empty() { - return Some(Err(exec_datafusion_err!( - r#"VortexSource accepted but failed to push {} filters. - This should never happen if you have a properly configured - PhysicalExprAdapterFactory configured on the source. - - Failed filters: - - {unpushed:#?} - "#, - unpushed.len() - ))); - } - - make_vortex_predicate(expr_convertor.as_ref(), &pushed).transpose() - }) - .transpose()?; - let filter = filter - .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) - .transpose() - .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; - if let Some(limit) = limit - && filter.is_none() + && native_filter.is_none() + && residual_filter.is_none() { scan_builder = scan_builder.with_limit(limit); } @@ -411,9 +418,9 @@ impl FileOpener for VortexOpener { // the fields the scan's projection and filter reference. scan_builder = scan_builder .with_projection(scan_projection) - .with_some_filter(filter); + .with_some_filter(native_filter); - if let Some(file_range) = file.range { + if let Some(file_range) = &file.range { let byte_range = Range { start: u64::try_from(file_range.start) .map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?, @@ -445,6 +452,7 @@ impl FileOpener for VortexOpener { } let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); + let residual_path = file.path().clone(); let stream = scan_builder .with_metrics_registry(metrics_registry) .with_ordered(has_output_ordering) @@ -466,21 +474,29 @@ impl FileOpener for VortexOpener { file.object_meta.location )))) }) - .map(move |batch| { - let batch = if projector.projection().as_ref().is_empty() { - batch - } else { - batch.and_then(|b| projector.project_batch(&b)) - }?; + .map(move |batch| -> DFResult> { + let mut batch = batch?; + if let Some(residual) = &residual_filter { + batch = batch_filter(&batch, residual).map_err(|e| { + exec_datafusion_err!( + "Failed to evaluate residual filter {residual} in {residual_path}: {e}" + ) + })?; + // Do not evaluate the projection on batches the residual filter emptied. + if batch.num_rows() == 0 { + return Ok(None); + } + } + let batch = projector.project_batch(&batch)?; let (_, columns, row_count) = batch.into_parts(); - RecordBatch::try_new_with_options( + Ok(Some(RecordBatch::try_new_with_options( Arc::clone(&output_schema), columns, &RecordBatchOptions::new().with_row_count(Some(row_count)), - ) - .map_err(Into::into) + )?)) }) + .filter_map(|batch| ready(batch.transpose())) .boxed(); if let Some(file_pruner) = file_pruner @@ -663,6 +679,8 @@ mod tests { use datafusion::physical_expr::planner::logical2physical; use datafusion::physical_expr_adapter::DefaultPhysicalExprAdapterFactory; use datafusion::scalar::ScalarValue; + use datafusion_common::arrow::compute::concat_batches; + use datafusion_common::assert_batches_eq; use datafusion_common::stats::Precision; use datafusion_execution::cache::default_cache::DefaultCache; use datafusion_expr::Operator; @@ -687,7 +705,7 @@ mod tests { use super::*; use crate::VortexAccessPlan; - use crate::convert::exprs::DefaultExpressionConvertor; + use crate::convert::exprs::DefaultExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; static SESSION: LazyLock = LazyLock::new(VortexSession::default); @@ -868,7 +886,7 @@ mod tests { layout_readers: Default::default(), natural_splits: Default::default(), has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + expression_converter: Arc::new(DefaultExpressionConverter::default()), file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, @@ -927,6 +945,51 @@ mod tests { Ok(()) } + #[rstest] + #[tokio::test] + async fn test_residual_filter_unprojected_column( + #[values(false, true)] projection_pushdown: bool, + #[values(false, true)] reordered_schema: bool, + ) -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let batch = record_batch!( + ("a", Int32, vec![10, 20, 30, 40]), + ("unused", Int32, vec![0, 0, 0, 0]), + ("b", Int32, vec![Some(1), Some(2), None, Some(4)]) + )?; + let size = + write_arrow_to_vortex(Arc::clone(&store), "residual.vortex", batch.clone()).await?; + let schema = Arc::new(batch.schema().project(if reordered_schema { + &[2, 0, 1] + } else { + &[0, 1, 2] + })?); + let modulo: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("b", schema.index_of("b")?)), + Operator::Modulo, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), + )); + let filter = Arc::new(df_expr::BinaryExpr::new( + modulo, + Operator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))), + )); + let mut opener = make_opener(store, TableSchema::from(Arc::clone(&schema)), Some(filter)); + opener.projection = ProjectionExprs::from_indices(&[schema.index_of("a")?], &schema); + opener.projection_pushdown = projection_pushdown; + opener.limit = Some(1); + let batches = opener + .open(PartitionedFile::new("residual.vortex", size))? + .await? + .try_collect::>() + .await?; + assert_batches_eq!( + ["+----+", "| a |", "+----+", "| 20 |", "| 40 |", "+----+"], + &batches + ); + Ok(()) + } + #[tokio::test] async fn test_open_preserves_declared_schema_metadata() -> anyhow::Result<()> { let object_store = Arc::new(InMemory::new()) as Arc; @@ -1193,7 +1256,7 @@ mod tests { layout_readers: Default::default(), natural_splits: Default::default(), has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + expression_converter: Arc::new(DefaultExpressionConverter::default()), file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, @@ -1280,7 +1343,7 @@ mod tests { layout_readers: Default::default(), natural_splits: Default::default(), has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + expression_converter: Arc::new(DefaultExpressionConverter::default()), file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, @@ -1434,7 +1497,7 @@ mod tests { layout_readers: Default::default(), natural_splits: Default::default(), has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + expression_converter: Arc::new(DefaultExpressionConverter::default()), file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, @@ -1494,7 +1557,7 @@ mod tests { layout_readers: Default::default(), natural_splits: Default::default(), has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + expression_converter: Arc::new(DefaultExpressionConverter::default()), file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, @@ -1701,7 +1764,7 @@ mod tests { layout_readers: Default::default(), natural_splits: Default::default(), has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + expression_converter: Arc::new(DefaultExpressionConverter::default()), file_metadata_cache: None, projection_pushdown: false, scan_concurrency: None, @@ -1776,4 +1839,203 @@ mod tests { ); Ok(()) } + + #[rstest] + #[tokio::test] + async fn test_temporal_adapter_residual( + #[values(false, true)] projection_pushdown: bool, + ) -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new( + "ts", + DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])), + Arc::new(arrow_array::TimestampMillisecondArray::from(vec![ + Some(1_000), + None, + Some(2_000), + ])), + ], + )?; + let size = write_arrow_to_vortex(Arc::clone(&store), "temporal.vortex", batch).await?; + let logical = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new( + "ts", + DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None), + true, + ), + ])); + let filter: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("ts", 1)), + Operator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::TimestampMicrosecond( + Some(1_000_000), + None, + ))), + )); + assert!( + DefaultExpressionConverter::default() + .try_convert(&filter, &logical)? + .is_some() + ); + let mut opener = make_opener(store, TableSchema::from(logical), Some(filter)); + opener.projection_pushdown = projection_pushdown; + let batches = opener + .open(PartitionedFile::new("temporal.vortex", size))? + .await? + .try_collect::>() + .await?; + assert_batches_eq!(["+----+", "| a |", "+----+", "| 10 |", "+----+"], &batches); + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn test_residual_literal_retains_zero_column_rows( + #[values(Some(true), Some(false), None)] value: Option, + ) -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let batch = record_batch!(("a", Int32, vec![1, 2, 3]))?; + let size = + write_arrow_to_vortex(Arc::clone(&store), "literal.vortex", batch.clone()).await?; + let mut opener = make_opener( + store, + TableSchema::from(batch.schema()), + Some(Arc::new(df_expr::Literal::new(ScalarValue::Boolean(value)))), + ); + opener.projection = Vec::::new().into(); + // Force a supported physical expression to use the residual path. + opener.expression_converter = Arc::new(ResidualConverter); + let batches = opener + .open(PartitionedFile::new("literal.vortex", size))? + .await? + .try_collect::>() + .await?; + assert!(batches.iter().all(|batch| batch.num_columns() == 0)); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + if value == Some(true) { 3 } else { 0 } + ); + Ok(()) + } + + struct ResidualConverter; + impl ExpressionConverter for ResidualConverter { + fn try_convert( + &self, + _expr: &PhysicalExprRef, + _schema: &Schema, + ) -> DFResult> { + Ok(None) + } + } + + #[rstest] + #[tokio::test] + async fn test_native_and_residual_composition( + #[values(Operator::And, Operator::Or)] operator: Operator, + ) -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let batch = record_batch!( + ("a", Int32, vec![10, 20, 30, 40]), + ("b", Int32, vec![Some(1), None, Some(2), Some(4)]) + )?; + let size = + write_arrow_to_vortex(Arc::clone(&store), "composition.vortex", batch.clone()).await?; + let native: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + Operator::Gt, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(35)))), + )); + let residual: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("b", 1)), + Operator::Modulo, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), + )), + Operator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))), + )); + let filter: PhysicalExprRef = + Arc::new(df_expr::BinaryExpr::new(native, operator, residual)); + let expected = batch_filter(&batch, &filter)?.project(&[0])?; + let opener = make_opener(store, TableSchema::from(batch.schema()), Some(filter)); + let actual = opener + .open(PartitionedFile::new("composition.vortex", size))? + .await? + .try_collect::>() + .await?; + let actual = concat_batches(&expected.schema(), &actual)?; + assert_eq!(actual, expected); + Ok(()) + } + + #[tokio::test] + async fn test_residual_error_has_file_and_predicate() -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let batch = record_batch!(("a", Int32, vec![1, 2]))?; + let size = + write_arrow_to_vortex(Arc::clone(&store), "filter-error.vortex", batch.clone()).await?; + let filter: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + Operator::Divide, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))), + )), + Operator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), + )); + let mut opener = make_opener(store, TableSchema::from(batch.schema()), Some(filter)); + opener.expression_converter = Arc::new(ResidualConverter); + let result = opener + .open(PartitionedFile::new("filter-error.vortex", size))? + .await? + .try_collect::>() + .await; + let error = result + .err() + .ok_or_else(|| anyhow::anyhow!("Expected division by zero"))? + .to_string(); + assert!(error.contains("filter-error.vortex"), "{error}"); + assert!(error.contains("residual filter"), "{error}"); + Ok(()) + } + + #[tokio::test] + async fn test_residual_rejects_rows_before_projection_evaluation() -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let batch = record_batch!(("a", Int32, vec![1, 2]))?; + let size = + write_arrow_to_vortex(Arc::clone(&store), "empty-projection.vortex", batch.clone()) + .await?; + let mut opener = make_opener( + store, + TableSchema::from(batch.schema()), + Some(Arc::new(df_expr::Literal::new(ScalarValue::Boolean(Some( + false, + ))))), + ); + opener.expression_converter = Arc::new(ResidualConverter); + opener.projection = vec![ProjectionExpr { + expr: Arc::new(SnapshotErrorExpr), + alias: "failure".into(), + }] + .into(); + let actual = opener + .open(PartitionedFile::new("empty-projection.vortex", size))? + .await? + .try_collect::>() + .await?; + assert_eq!(actual.iter().map(RecordBatch::num_rows).sum::(), 0); + Ok(()) + } } diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index fbd37f04ded..16940997eed 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -39,8 +39,8 @@ use vortex_utils::aliases::dash_map::DashMap; use super::opener::NaturalSplits; use super::opener::VortexOpener; use crate::VortexTableOptions; -use crate::convert::exprs::DefaultExpressionConvertor; -use crate::convert::exprs::ExpressionConvertor; +use crate::convert::exprs::DefaultExpressionConverter; +use crate::convert::exprs::ExpressionConverter; use crate::persistent::reader::DefaultVortexReaderFactory; use crate::persistent::reader::VortexReaderFactory; @@ -129,23 +129,25 @@ use crate::persistent::reader::VortexReaderFactory; /// /// - `full_predicate`, which is used by DataFusion's `FilePruner` to skip whole /// files before they are opened, -/// - `vortex_predicate`, which contains only the expressions Vortex can evaluate -/// during the scan. +/// - `vortex_predicate`, which contains accepted exact filters. After per-file +/// adaptation, these run either natively or as DataFusion residual filters. /// /// Projection handling depends on /// [`VortexTableOptions::projection_pushdown`]: /// /// - when disabled, `VortexSource` still prunes unreferenced top-level columns, /// but DataFusion applies the full projection after the scan, -/// - when enabled, the scan can evaluate a Vortex-native projection and leave -/// only unsupported expressions for DataFusion. +/// - when enabled, the default converter evaluates fully supported projections +/// natively. Otherwise, DataFusion evaluates the full projection over raw columns. /// /// Predicate handling depends on [`VortexTableOptions::predicate_pushdown`]: /// /// - when disabled, `VortexSource` still keeps the full predicate for /// DataFusion file pruning, but reports filters as not pushed down so /// DataFusion evaluates them after the scan, -/// - when enabled, supported filters are pushed into the Vortex scan. +/// - when enabled, supported filters are pushed into the Vortex scan. If file +/// adaptation requires residual filtering, DataFusion filters raw scan batches +/// before the final projection and limit. /// /// # Observability /// @@ -187,8 +189,7 @@ pub struct VortexSource { /// Combined predicate expression containing all filters from DataFusion query planning. /// Used with FilePruner to skip files based on statistics and partition values. pub(crate) full_predicate: Option, - /// Subset of predicates that can be pushed down into Vortex scan operations. - /// These are expressions that Vortex can efficiently evaluate during scanning. + /// Accepted exact predicates, evaluated natively or as per-file residuals. pub(crate) vortex_predicate: Option, /// DataFusion-native metrics exposed through `DataSourceExec`. df_metrics: ExecutionPlanMetricsSet, @@ -198,7 +199,7 @@ pub struct VortexSource { layout_readers: Arc>>, /// Shared full-file natural splits keyed by path. natural_splits: Arc>>, - expression_convertor: Arc, + expression_converter: Arc, pub(crate) vortex_reader_factory: Option>, pub(crate) ordered: bool, vx_metrics_registry: Arc, @@ -220,7 +221,7 @@ impl VortexSource { let full_schema = table_schema.table_schema(); let indices = (0..full_schema.fields().len()).collect::>(); let projection = ProjectionExprs::from_indices(&indices, full_schema); - let expression_convertor = Arc::new(DefaultExpressionConvertor::new(session.clone())); + let expression_converter = Arc::new(DefaultExpressionConverter::new(session.clone())); Self { session, @@ -231,7 +232,7 @@ impl VortexSource { df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), natural_splits: Arc::new(DashMap::default()), - expression_convertor, + expression_converter, vortex_reader_factory: None, vx_metrics_registry: Arc::new(DefaultMetricsRegistry::default()), file_metadata_cache: None, @@ -259,16 +260,16 @@ impl VortexSource { self } - /// Sets the [`ExpressionConvertor`] used to translate DataFusion expressions + /// Sets the [`ExpressionConverter`] used to translate DataFusion expressions /// into Vortex expressions. /// /// Override this when the default converter is insufficient for an engine /// integration or for a custom schema-adaptation strategy. - pub fn with_expression_convertor( + pub fn with_expression_converter( mut self, - expr_convertor: Arc, + expr_converter: Arc, ) -> Self { - self.expression_convertor = expr_convertor; + self.expression_converter = expr_converter; self } @@ -356,7 +357,7 @@ impl VortexSource { layout_readers: Arc::clone(&self.layout_readers), natural_splits: Arc::clone(&self.natural_splits), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, - expression_convertor: Arc::clone(&self.expression_convertor), + expression_converter: Arc::clone(&self.expression_converter), file_metadata_cache: self.file_metadata_cache.clone(), projection_pushdown: self.options.projection_pushdown, scan_concurrency: self.options.scan_concurrency, @@ -473,15 +474,16 @@ impl FileSource for VortexSource { .into_iter() .map(|expr| { if self - .expression_convertor - .can_be_pushed_down(&expr, self.table_schema.file_schema()) + .expression_converter + .try_convert(&expr, self.table_schema.table_schema())? + .is_some() { - PushedDownPredicate::supported(expr) + Ok(PushedDownPredicate::supported(expr)) } else { - PushedDownPredicate::unsupported(expr) + Ok(PushedDownPredicate::unsupported(expr)) } }) - .collect::>(); + .collect::>>()?; if supported_filters .iter() @@ -568,17 +570,17 @@ mod tests { use super::*; use crate::convert::exprs::ProcessedProjection; - struct TrackingExpressionConvertor { - inner: DefaultExpressionConvertor, + struct TrackingExpressionConverter { + inner: DefaultExpressionConverter, } - impl ExpressionConvertor for TrackingExpressionConvertor { - fn can_be_pushed_down(&self, expr: &PhysicalExprRef, schema: &Schema) -> bool { - self.inner.can_be_pushed_down(expr, schema) - } - - fn convert(&self, expr: &dyn PhysicalExpr) -> DFResult { - self.inner.convert(expr) + impl ExpressionConverter for TrackingExpressionConverter { + fn try_convert( + &self, + expr: &PhysicalExprRef, + schema: &Schema, + ) -> DFResult> { + self.inner.try_convert(expr, schema) } fn split_projection( @@ -662,14 +664,14 @@ mod tests { } #[test] - fn create_vortex_opener_preserves_expression_convertor() -> anyhow::Result<()> { + fn create_vortex_opener_preserves_expression_converter() -> anyhow::Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let expression_convertor = Arc::new(TrackingExpressionConvertor { - inner: DefaultExpressionConvertor::default(), - }) as Arc; + let expression_converter = Arc::new(TrackingExpressionConverter { + inner: DefaultExpressionConverter::default(), + }) as Arc; let source = VortexSource::new(TableSchema::from(file_schema), VortexSession::default()) - .with_expression_convertor(Arc::clone(&expression_convertor)); + .with_expression_converter(Arc::clone(&expression_converter)); let config = FileScanConfigBuilder::new( ObjectStoreUrl::local_filesystem(), @@ -684,8 +686,8 @@ mod tests { )?; assert!(Arc::ptr_eq( - &opener.expression_convertor, - &expression_convertor + &opener.expression_converter, + &expression_converter )); Ok(()) } diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 35dd745461d..66d9bf88eb6 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -13,6 +13,7 @@ use datafusion::execution::SessionStateBuilder; use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; use datafusion_common::GetExt; +use datafusion_common::arrow::compute::concat_batches; use datafusion_physical_plan::display::DisplayableExecutionPlan; use insta::assert_snapshot; use object_store::ObjectStore; @@ -620,3 +621,93 @@ async fn arrow_uuid_extension_roundtrip_nested_struct() -> anyhow::Result<()> { Ok(()) } + +#[rstest] +#[case::modulo("id, a", "CAST(a % 2 AS BIGINT) = 0")] +#[case::in_list("id, a IN (1, 4, 5, 6, 7, NULL) AS member", "TRUE")] +#[case::not_in("id, a", "a NOT IN (1, 4, 5, 6, 7, 8)")] +#[case::null_list("id, a", "a NOT IN (1, 4, 5, 6, 7, NULL)")] +#[case::column_list("id, a", "a IN (b, 4, 5, 6, 7, 8)")] +// Overflow equivalence is deferred until Vortex matches DataFusion's arithmetic semantics. +#[case::add("id, a + CAST(1 AS INT) AS n", "id > 1")] +#[case::subtract("id, a - CAST(1 AS INT) AS n", "id > 1")] +#[case::multiply("id, a * CAST(2 AS INT) AS n", "id > 1")] +#[case::arithmetic_filter("id", "b * CAST(2 AS INT) = 4")] +#[case::narrow_cast("id, CAST(a AS TINYINT) AS n", "TRUE")] +#[case::case_branch("id, CASE WHEN b <> 0 THEN 12 / b ELSE 0 END AS n", "TRUE")] +#[case::case_filter("id", "CASE WHEN b <> 0 THEN 12 / b ELSE 0 END = 6")] +#[case::projection_alias("id, a + CAST(1 AS INT) AS a, abs(a) AS b", "id > 1")] +#[case::unprojected("id", "CAST(a AS VARCHAR) = '1'")] +#[case::decimal( + "id, CAST(a AS DECIMAL(12, 2)) + CAST(b AS DECIMAL(12, 2)) AS n", + "TRUE" +)] +#[tokio::test] +async fn test_predicate_memtable_oracle( + #[case] projection: &str, + #[case] predicate: &str, + #[values(false, true)] pushdown: bool, +) -> anyhow::Result<()> { + let options = crate::VortexTableOptions { + projection_pushdown: pushdown, + predicate_pushdown: pushdown, + ..Default::default() + }; + let ctx = TestSessionContext::new_with_factory(Arc::new( + VortexFormatFactory::new().with_options(options), + )); + let batch = arrow_array::record_batch!( + ("id", Int32, vec![0, 1, 2, 3, 4, 5]), + ( + "a", + Int32, + vec![ + Some(i32::MAX), + Some(i32::MIN), + Some(1), + Some(2), + None, + Some(4) + ] + ), + ("b", Int32, vec![0, 0, 1, 2, 0, 2]) + )?; + ctx.write_arrow_batch("oracle.vortex", &batch).await?; + let actual_table = ctx + .table_provider("actual", "/oracle.vortex", batch.schema().as_ref().clone()) + .await?; + ctx.session.register_table("actual", actual_table)?; + ctx.session.register_table( + "oracle", + Arc::new(datafusion::datasource::MemTable::try_new( + batch.schema(), + vec![vec![batch]], + )?), + )?; + let actual = ctx + .session + .sql(&format!( + "SELECT {projection} FROM actual WHERE {predicate} ORDER BY id" + )) + .await?; + let expected = ctx + .session + .sql(&format!( + "SELECT {projection} FROM oracle WHERE {predicate} ORDER BY id" + )) + .await?; + assert_eq!(actual.schema().as_arrow(), expected.schema().as_arrow()); + let schema = Arc::new(expected.schema().as_arrow().clone()); + match (actual.collect().await, expected.collect().await) { + (Ok(actual), Ok(expected)) => { + let actual = concat_batches(&schema, &actual)?; + let expected = concat_batches(&schema, &expected)?; + assert_eq!(actual, expected); + } + (Err(_), Err(_)) => {} + (actual, expected) => anyhow::bail!( + "Vortex/MemTable mismatch for {projection} WHERE {predicate}: {actual:?} / {expected:?}" + ), + } + Ok(()) +} diff --git a/vortex-datafusion/src/v2/mod.rs b/vortex-datafusion/src/v2/mod.rs index 5c2e6aa542f..3a506abc622 100644 --- a/vortex-datafusion/src/v2/mod.rs +++ b/vortex-datafusion/src/v2/mod.rs @@ -36,5 +36,8 @@ mod source; mod table; +#[cfg(test)] +mod tests; + pub use source::VortexDataSource; pub use table::VortexTable; diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 6437cb24dc5..706e3bbc7ef 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -118,10 +118,9 @@ use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; use vortex_utils::parallelism::get_available_parallelism; -use crate::convert::exprs::DefaultExpressionConvertor; -use crate::convert::exprs::ExpressionConvertor; +use crate::convert::exprs::DefaultExpressionConverter; +use crate::convert::exprs::ExpressionConverter; use crate::convert::exprs::ProcessedProjection; -use crate::convert::exprs::make_vortex_predicate; use crate::convert::stats::stats_set_to_df; /// Builder for [`VortexDataSource`]. @@ -545,17 +544,18 @@ impl DataSource for VortexDataSource { projection ); - let convertor = DefaultExpressionConvertor::default(); + let converter = DefaultExpressionConverter::default(); let input_schema = self.initial_schema.as_ref(); let projected_schema = projection.project_schema(input_schema)?; - // Use the shared ExpressionConvertor to split the projection into a Vortex + // Use the shared ExpressionConverter to split the projection into a Vortex // scan_projection and a leftover DataFusion projection for expressions that // can't be pushed down (e.g., unsupported scalar functions, decimal binary). let ProcessedProjection { scan_projection, leftover_projection, - } = convertor.split_projection(projection.clone(), input_schema, &projected_schema)?; + .. + } = converter.split_projection(projection.clone(), input_schema, &projected_schema)?; // Compose with the initial projection so the scan operates on the original // source columns, not the initial projection's output columns. @@ -602,15 +602,28 @@ impl DataSource for VortexDataSource { )); } - let convertor = DefaultExpressionConvertor::default(); - let input_schema = self.initial_schema.as_ref(); + let converter = DefaultExpressionConverter::default(); + let filters = filters + .into_iter() + .map(|filter| { + let filter = match &self.leftover_projection { + Some(projection) => projection.unproject_expr(&filter)?, + None => filter, + }; + reassign_expr_columns(filter, &self.projected_schema) + }) + .collect::>>()?; // Classify each filter: pushable filters are passed into the ScanRequest in open(), // so we can safely claim PushedDown::Yes for them. - let pushdown_results: Vec = filters + let converted = filters + .iter() + .map(|expr| converter.try_convert(expr, &self.projected_schema)) + .collect::>>()?; + let pushdown_results: Vec = converted .iter() .map(|expr| { - if convertor.can_be_pushed_down(expr, input_schema) { + if expr.is_some() { PushedDown::Yes } else { PushedDown::No @@ -625,18 +638,9 @@ impl DataSource for VortexDataSource { )); } - // Collect the pushable filter expressions. - let pushable: Vec> = filters - .iter() - .zip(pushdown_results.iter()) - .filter_map(|(expr, pushed)| match pushed { - PushedDown::Yes => Some(Arc::clone(expr)), - PushedDown::No => None, - }) - .collect(); - // Convert to Vortex conjunction. - let vortex_pred = make_vortex_predicate(&convertor, &pushable)?; + let vortex_pred = vortex::expr::and_collect(converted.into_iter().flatten()) + .map(|expr| replace(expr, &root(), self.projected_projection.clone())); // Combine with existing filter. let new_filter = match (&self.filter, vortex_pred) { diff --git a/vortex-datafusion/src/v2/tests.rs b/vortex-datafusion/src/v2/tests.rs new file mode 100644 index 00000000000..a8815191ac8 --- /dev/null +++ b/vortex-datafusion/src/v2/tests.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use anyhow::anyhow; +use datafusion::assert_batches_eq; +use datafusion::prelude::SessionContext; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_datasource::source::DataSource; +use datafusion_datasource::source::DataSourceExec; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::projection::ProjectionExpr; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expressions as df_expr; +use datafusion_physical_plan::filter_pushdown::PushedDown; +use vortex::VortexSessionDefault; +use vortex::buffer::ByteBufferMut; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::WriteOptionsSessionExt; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; + +use super::VortexDataSource; + +async fn vortex_source() -> anyhow::Result { + let batch = + arrow_array::record_batch!(("a", Int32, vec![1, 2, 3]), ("b", Int32, vec![30, 10, 20]))?; + let session = VortexSession::default(); + let schema = batch.schema(); + let array = session.arrow().from_arrow_record_batch(batch, &schema)?; + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buffer, array.to_array_stream()) + .await?; + let file = session.open_options().open_buffer(buffer)?; + Ok(VortexDataSource::builder(file.data_source()?, session) + .with_arrow_schema(schema) + .build() + .await?) +} + +#[tokio::test] +async fn projection_filters_use_scan_schema() -> anyhow::Result<()> { + let source = vortex_source().await?; + let projection = ProjectionExprs::from(vec![ + ProjectionExpr::new(Arc::new(df_expr::Column::new("b", 1)), "b"), + ProjectionExpr::new( + Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + Operator::Modulo, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), + )), + "odd", + ), + ]); + let projected = source + .try_swapping_with_projection(&projection)? + .ok_or_else(|| anyhow!("projection was not swapped"))?; + let supported: Arc = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("b", 0)), + Operator::Gt, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(15)))), + )); + let unsupported: Arc = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("odd", 1)), + Operator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), + )); + + let pushed = + projected.try_pushdown_filters(vec![supported, unsupported], &ConfigOptions::new())?; + assert!(matches!( + pushed.filters.as_slice(), + [PushedDown::Yes, PushedDown::No] + )); + let updated = pushed + .updated_node + .ok_or_else(|| anyhow!("filter was not added to the source"))?; + let plan = Arc::new(DataSourceExec::new(updated)) as Arc; + let ctx = SessionContext::new(); + let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx()).await?; + + assert_batches_eq!( + [ + "+----+-----+", + "| b | odd |", + "+----+-----+", + "| 30 | 1 |", + "| 20 | 1 |", + "+----+-----+", + ], + &batches + ); + Ok(()) +}