From 374e2ab40d29291659738e9dfb184b5f4e99852c Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 7 Sep 2026 12:38:57 +0100 Subject: [PATCH 01/10] Align DF expression conversion and pushdown Signed-off-by: Adam Gutglick --- Cargo.lock | 1 + vortex-datafusion/Cargo.toml | 1 + vortex-datafusion/src/convert/exprs.rs | 1314 ------------ vortex-datafusion/src/convert/exprs/mod.rs | 657 ++++++ vortex-datafusion/src/convert/exprs/tests.rs | 1156 +++++++++++ vortex-datafusion/src/persistent/format.rs | 29 +- vortex-datafusion/src/persistent/opener.rs | 1779 ----------------- .../src/persistent/opener/mod.rs | 669 +++++++ .../src/persistent/opener/tests.rs | 1535 ++++++++++++++ vortex-datafusion/src/persistent/source.rs | 41 +- vortex-datafusion/src/persistent/tests.rs | 88 + vortex-datafusion/src/v2/source.rs | 22 +- 12 files changed, 4151 insertions(+), 3141 deletions(-) delete mode 100644 vortex-datafusion/src/convert/exprs.rs create mode 100644 vortex-datafusion/src/convert/exprs/mod.rs create mode 100644 vortex-datafusion/src/convert/exprs/tests.rs delete mode 100644 vortex-datafusion/src/persistent/opener.rs create mode 100644 vortex-datafusion/src/persistent/opener/mod.rs create mode 100644 vortex-datafusion/src/persistent/opener/tests.rs 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/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..65d6148f7be --- /dev/null +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -0,0 +1,657 @@ +// 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::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::ProjectionExpr; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_plan::expressions as df_expr; +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::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::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; + +/// 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 convertors implement a single schema-aware decision. Successful conversion must +/// preserve DataFusion values, nulls, and evaluation errors. Unsupported expressions remain +/// in DataFusion, including when a file's schema adapter introduces them. +/// +/// # Implementing a custom convertor +/// +/// 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::DefaultExpressionConvertor; +/// use vortex_datafusion::convert::ExpressionConvertor; +/// +/// struct CustomExpressionConvertor(DefaultExpressionConvertor); +/// +/// impl ExpressionConvertor for CustomExpressionConvertor { +/// fn try_convert( +/// &self, +/// expr: &Arc, +/// schema: &Schema, +/// ) -> DFResult> { +/// self.0.try_convert(expr, schema) +/// } +/// } +pub trait ExpressionConvertor: Send + Sync { + /// Convert an expression with equivalent DataFusion behavior for 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 native and DataFusion evaluation. + /// + /// 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 { + let mut scan_projection = Vec::with_capacity(source_projection.as_ref().len()); + for projection in source_projection.iter() { + let Some(expr) = self.try_convert(&projection.expr, input_schema)? else { + return self.no_pushdown_projection(source_projection.clone(), input_schema); + }; + scan_projection.push((projection.alias.clone(), expr)); + } + // Duplicate output names cannot identify native pack fields unambiguously. + let mut names = scan_projection + .iter() + .map(|(name, _)| name) + .collect::>(); + names.sort_unstable(); + if names.windows(2).any(|names| names[0] == names[1]) { + return self.no_pushdown_projection(source_projection, input_schema); + } + Ok(ProcessedProjection { + scan_projection: pack(scan_projection, Nullability::NonNullable), + scan_reference_schema: output_schema.clone(), + leftover_projection: source_projection + .iter() + .enumerate() + .map(|(index, projection)| ProjectionExpr { + expr: Arc::new(df_expr::Column::new(&projection.alias, index)), + alias: projection.alias.clone(), + }) + .collect::>() + .into(), + }) + } + + /// Read the required raw columns and apply the complete projection in DataFusion. + fn no_pushdown_projection( + &self, + source_projection: ProjectionExprs, + input_schema: &Schema, + ) -> DFResult { + raw_projection(source_projection, input_schema) + } +} + +/// Read raw columns in file-index order without involving custom expression conversion. +pub(crate) fn raw_projection( + source_projection: ProjectionExprs, + input_schema: &Schema, +) -> DFResult { + let column_indices = source_projection.column_indices(); + let mut scan_columns = Vec::with_capacity(column_indices.len()); + let mut fields = Vec::with_capacity(column_indices.len()); + for index in column_indices { + let field = input_schema.fields().get(index).ok_or_else(|| { + exec_datafusion_err!("Projection column index {index} is out of bounds") + })?; + scan_columns.push(( + field.name().clone(), + get_item(field.name().as_str(), root()), + )); + fields.push(Arc::clone(field)); + } + Ok(ProcessedProjection { + scan_projection: pack(scan_columns, Nullability::NonNullable), + scan_reference_schema: Schema::new_with_metadata(fields, input_schema.metadata().clone()), + leftover_projection: source_projection, + }) +} + +/// The default schema-aware DataFusion expression convertor. +/// +/// Casts and operators are accepted only for the SQL semantics implemented by Vortex. +/// Other valid expressions are evaluated by DataFusion. +pub struct DefaultExpressionConvertor { + session: VortexSession, +} + +impl Default for DefaultExpressionConvertor { + fn default() -> Self { + Self::new(VortexSession::default()) + } +} + +impl DefaultExpressionConvertor { + /// Create a convertor that resolves Arrow extension types using the session registry. + pub fn new(session: VortexSession) -> Self { + Self { session } + } + + fn convert_expr( + &self, + expr: &Arc, + schema: &Schema, + input_dtype: &DType, + ) -> DFResult> { + let converted = if let Some(binary) = expr.downcast_ref::() { + let Ok(operator) = try_operator_from_df(binary.op()) else { + return Ok(None); + }; + let left_type = binary.left().data_type(schema)?; + let right_type = binary.right().data_type(schema)?; + if operator.is_arithmetic() { + // DataFusion's integer +/-/* can wrap; Vortex arithmetic is checked. + // Decimal coercion and result precision also differ. + let supported = matches!( + (&left_type, &right_type), + (DataType::Float32, DataType::Float32) | (DataType::Float64, DataType::Float64) + ) || (*binary.op() == DFOperator::Divide + && left_type.is_integer() + && left_type == right_type); + if !supported { + return Ok(None); + } + } else if *binary.op() == DFOperator::And || *binary.op() == DFOperator::Or { + if left_type != DataType::Boolean || right_type != DataType::Boolean { + return Err(exec_datafusion_err!( + "Boolean operator requires Boolean operands: {expr}" + )); + } + } else if !supported_data_types(&left_type) || !supported_data_types(&right_type) { + return Ok(None); + } + let (Some(left), Some(right)) = ( + self.convert_expr(binary.left(), schema, input_dtype)?, + self.convert_expr(binary.right(), schema, input_dtype)?, + ) else { + return Ok(None); + }; + if matches!(binary.op(), DFOperator::And | DFOperator::Or) + && (label_infallible(&left).get(&left) != Some(&true) + || label_infallible(&right).get(&right) != Some(&true)) + { + // DataFusion may evaluate the RHS only on rows selected by the LHS. + return Ok(None); + } + let (Ok(left_dtype), Ok(right_dtype)) = ( + left.return_dtype(input_dtype), + right.return_dtype(input_dtype), + ) else { + return Ok(None); + }; + if !left_dtype.eq_ignore_nullability(&right_dtype) { + return Ok(None); + } + Binary.new_expr(operator, [left, right]) + } else if let Some(column) = expr.downcast_ref::() { + get_item(column.name(), root()) + } else if let Some(literal) = expr.downcast_ref::() { + let field = literal.return_field(schema)?; + let array = literal.value().to_array()?; + if self.session.arrow().from_arrow_field(&field).is_err() { + return Ok(None); + } + if array.len() != 1 { + return Err(exec_datafusion_err!( + "Literal must contain exactly one value: {expr}" + )); + } + let array = self + .session + .arrow() + .from_arrow_array(array, &field) + .map_err(|e| exec_datafusion_err!("Failed to convert literal {expr}: {e}"))?; + lit(array + .execute_scalar(0, &mut self.session.create_execution_ctx()) + .map_err(|e| exec_datafusion_err!("Failed to evaluate literal {expr}: {e}"))?) + } else if let Some(cast_expr) = expr.downcast_ref::() { + if !supported_cast(cast_expr, schema)? { + return Ok(None); + } + let Some(child) = self.convert_expr(cast_expr.expr(), schema, input_dtype)? else { + return Ok(None); + }; + let Ok(target) = self + .session + .arrow() + .from_arrow_field(cast_expr.target_field()) + else { + return Ok(None); + }; + let Ok(child_dtype) = child.return_dtype(input_dtype) else { + return Ok(None); + }; + // Matching Arrow storage types do not imply matching extension semantics. + if (child_dtype.is_extension() || target.is_extension()) + && !child_dtype.eq_ignore_nullability(&target) + { + return Ok(None); + } + cast(child, target) + } else if let Some(is_null_expr) = expr.downcast_ref::() { + let Some(child) = self.convert_expr(is_null_expr.arg(), schema, input_dtype)? else { + return Ok(None); + }; + is_null(child) + } else if let Some(is_not_null_expr) = expr.downcast_ref::() { + let Some(child) = self.convert_expr(is_not_null_expr.arg(), schema, input_dtype)? + else { + return Ok(None); + }; + is_not_null(child) + } else if let Some(like) = expr.downcast_ref::() { + if !like.expr().data_type(schema)?.is_string() + || !like.pattern().data_type(schema)?.is_string() + { + return Ok(None); + } + let (Some(child), Some(pattern)) = ( + self.convert_expr(like.expr(), schema, input_dtype)?, + self.convert_expr(like.pattern(), schema, input_dtype)?, + ) else { + return Ok(None); + }; + Like.new_expr( + LikeOptions { + negated: like.negated(), + case_insensitive: like.case_insensitive(), + }, + [child, pattern], + ) + } else if expr.downcast_ref::().is_some() { + // list_contains does not implement SQL IN/NOT IN null semantics. + return Ok(None); + } else if let Some(scalar_fn) = expr.downcast_ref::() { + return self.convert_scalar_function(scalar_fn, schema, input_dtype); + } else if let Some(case_expr) = expr.downcast_ref::() { + if case_expr.expr().is_some() { + return Ok(None); + } + if case_expr.when_then_expr().is_empty() { + return Err(exec_datafusion_err!( + "CASE requires at least one WHEN clause" + )); + } + let mut pairs = Vec::with_capacity(case_expr.when_then_expr().len()); + for (when, then) in case_expr.when_then_expr() { + if when.data_type(schema)? != DataType::Boolean { + return Err(exec_datafusion_err!("CASE WHEN must be Boolean")); + } + let (Some(when), Some(then)) = ( + self.convert_expr(when, schema, input_dtype)?, + self.convert_expr(then, schema, input_dtype)?, + ) else { + return Ok(None); + }; + pairs.push((when, then)); + } + let otherwise = match case_expr.else_expr() { + Some(expr) => { + let Some(expr) = self.convert_expr(expr, schema, input_dtype)? else { + return Ok(None); + }; + Some(expr) + } + None => None, + }; + let case = nested_case_when(pairs, otherwise); + // Vortex may evaluate branch values on rows excluded by the condition. + if label_infallible(&case).get(&case) != Some(&true) { + return Ok(None); + } + case + } else { + return Ok(None); + }; + Ok(Some(converted)) + } + + fn convert_scalar_function( + &self, + scalar_fn: &ScalarFunctionExpr, + schema: &Schema, + input_dtype: &DType, + ) -> DFResult> { + if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { + let [source, paths @ ..] = scalar_fn.args() else { + return Err(exec_datafusion_err!( + "get_field requires a source and field path" + )); + }; + if paths.is_empty() { + return Err(exec_datafusion_err!("get_field requires a field path")); + } + 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 Ok(None); + }; + nullable_struct |= nullable; + let field = fields + .iter() + .find(|field| field.name() == 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 Ok(None); + } + let Some(mut result) = self.convert_expr(source, schema, input_dtype)? else { + return Ok(None); + }; + for name in names { + result = get_item(name, result); + } + return Ok(Some(result)); + } + + let octet_length = + ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some(); + let array_length = + ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some(); + let input = if octet_length { + let [input] = scalar_fn.args() else { + return Err(exec_datafusion_err!( + "octet_length requires exactly one argument" + )); + }; + 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 Ok(None); + } + input + } else if array_length { + if scalar_fn.args().is_empty() || scalar_fn.args().len() > 2 { + return Err(exec_datafusion_err!( + "array_length requires one or two arguments" + )); + } + let Some(input) = array_length_input(scalar_fn) else { + return Ok(None); + }; + if !matches!( + input.data_type(schema)?, + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ) { + return Ok(None); + } + input + } else { + return Ok(None); + }; + let Some(input) = self.convert_expr(input, schema, input_dtype)? else { + return Ok(None); + }; + let Ok(return_dtype) = self.session.arrow().from_arrow_field(&Field::new( + "", + scalar_fn.return_type().clone(), + scalar_fn.nullable(), + )) else { + return Ok(None); + }; + Ok(Some(cast( + if octet_length { + byte_length(input) + } else { + list_length(input) + }, + return_dtype, + ))) + } +} + +impl ExpressionConvertor for DefaultExpressionConvertor { + fn try_convert( + &self, + expr: &Arc, + schema: &Schema, + ) -> DFResult> { + for column in collect_columns(expr) { + 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() + )); + } + } + if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { + return Ok(None); + } + let Ok(input_dtype) = self.session.arrow().from_arrow_schema(schema) else { + return Ok(None); + }; + let Some(converted) = self.convert_expr(expr, schema, &input_dtype)? else { + return Ok(None); + }; + let Ok(expected_dtype) = self + .session + .arrow() + .from_arrow_field(expr.return_field(schema)?.as_ref()) + else { + return Ok(None); + }; + let Ok(actual_dtype) = converted.return_dtype(&input_dtype) else { + return Ok(None); + }; + if !actual_dtype.eq_ignore_nullability(&expected_dtype) { + return Ok(None); + } + Ok(Some(converted)) + } +} + +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 + || (cast.expr().nullable(schema)? && !cast.target_field().is_nullable()) + { + 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) -> 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 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 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; diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs new file mode 100644 index 00000000000..1ffa7fea73e --- /dev/null +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -0,0 +1,1156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::Array; +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::arrow::buffer::NullBuffer; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::Operator as DFOperator; +use datafusion_expr::ScalarUDF; +use datafusion_functions::core::coalesce::CoalesceFunc; +use datafusion_physical_expr::PhysicalExpr; +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_array::assert_arrays_eq; + +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(), + ) +} + +#[rstest] +fn test_predicate_rejects_cast_over_modulo(test_schema: Schema) { + 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!( + !DefaultExpressionConvertor::default() + .try_convert(&expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[rstest] +#[case::empty(false)] +#[case::column(true)] +fn test_predicate_rejects_in_list(test_schema: Schema, #[case] nonempty: 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, + false, + &test_schema, + )?); + assert!( + !DefaultExpressionConvertor::default() + .try_convert(&expr, &test_schema)? + .is_some() + ); + 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, +) { + 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() + .try_convert( + &(Arc::new(col_expr) as Arc), + &Schema::new(vec![Field::new("test_column", DataType::Int32, false)]), + ) + .unwrap() + .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() + .try_convert( + &(Arc::new(literal_expr) as Arc), + &Schema::empty(), + ) + .unwrap() + .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() + .try_convert( + &(Arc::new(binary_expr) as Arc), + &Schema::new(vec![Field::new("left", DataType::Int32, false)]), + ) + .unwrap() + .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() + .try_convert( + &(Arc::new(like_expr) as Arc), + &Schema::new(vec![Field::new("text_col", DataType::Utf8, true)]), + ) + .unwrap() + .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() + .try_convert(&octet_length, &test_schema) + .unwrap() + .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() + .try_convert(&array_length, &test_schema) + .unwrap() + .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!( + DefaultExpressionConvertor::default() + .try_convert(&col_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[rstest] +fn test_nested_column_conversion(test_schema: Schema) { + let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; + + assert!( + DefaultExpressionConvertor::default() + .try_convert(&col_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&col_expr, &test_schema) + .is_err() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&lit_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&lit_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&binary_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&binary_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&binary_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&like_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&like_expr, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&octet_length, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&octet_length, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&array_length, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&array_length, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&array_length, &test_schema) + .unwrap() + .is_some() + ); +} + +#[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!( + !DefaultExpressionConvertor::default() + .try_convert(&array_length, &test_schema) + .unwrap() + .is_some() + ); +} + +// 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 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!( + DefaultExpressionConvertor::default() + .try_convert(&cast, &schema)? + .is_some() + ); + 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( + &(Arc::new(case_expr) as Arc), + &batch.schema(), + ) + .unwrap() + .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); +} + +fn assert_native_matches( + expr: Arc, + batch: arrow_array::RecordBatch, +) -> anyhow::Result<()> { + let session = VortexSession::default(); + let converted = DefaultExpressionConvertor::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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + Ok(()) +} + +#[rstest] +#[case::add(DFOperator::Plus)] +#[case::sub(DFOperator::Minus)] +#[case::mul(DFOperator::Multiply)] +fn test_integer_overflow_modes_fall_back( + #[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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + Ok(()) +} + +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema) + .is_err() + ); + Ok(()) +} + +#[test] +fn test_mismatched_column_identity_returns_error() { + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema) + .is_err() + ); +} + +#[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(arrow_array::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 = arrow_array::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!( + DefaultExpressionConvertor::default() + .try_convert(&length, &batch.schema())? + .is_none() + ); + 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_malformed_literal_returns_error() { + let value = ScalarValue::Decimal128(Some(1), 0, 0); + let expr: Arc = Arc::new(df_expr::Literal::new(value)); + assert!( + DefaultExpressionConvertor::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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &batch.schema())? + .is_none() + ); + Ok(()) +} diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 6d05144e315..2e9afbb173e 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -828,14 +828,12 @@ mod tests { #[derive(Default)] struct ExpressionConvertorCalls { - can_be_pushed_down: AtomicBool, - convert: AtomicBool, + try_convert: AtomicBool, } impl ExpressionConvertorCalls { 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); } } @@ -860,19 +858,18 @@ mod tests { } 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); + 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, @@ -905,8 +902,7 @@ mod tests { &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(()) } @@ -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 deleted file mode 100644 index 0f000452f17..00000000000 --- a/vortex-datafusion/src/persistent/opener.rs +++ /dev/null @@ -1,1779 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::ops::Range; -use std::sync::Arc; -use std::sync::Weak; - -use arrow_array::RecordBatchOptions; -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::Statistics; -use datafusion_common::arrow::array::AsArray; -use datafusion_common::arrow::array::RecordBatch; -use datafusion_common::exec_datafusion_err; -use datafusion_datasource::PartitionedFile; -use datafusion_datasource::TableSchema; -use datafusion_datasource::file_stream::FileOpenFuture; -use datafusion_datasource::file_stream::FileOpener; -use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; -use datafusion_execution::cache::cache_manager::FileMetadataCache; -use datafusion_physical_expr::PhysicalExprRef; -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::reassign_expr_columns; -use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; -use datafusion_physical_expr_adapter::replace_columns_with_literals; -use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; -use datafusion_physical_plan::metrics::MetricBuilder; -use datafusion_physical_plan::metrics::MetricCategory; -use datafusion_pruning::FilePruner; -use futures::FutureExt; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::stream; -use object_store::path::Path; -use tracing::Instrument; -use vortex::array::VortexSessionExecute; -use vortex::error::VortexError; -use vortex::error::VortexExpect; -use vortex::file::OpenOptionsSessionExt; -use vortex::io::InstrumentedReadAt; -use vortex::layout::LayoutReader; -use vortex::layout::scan::scan_builder::ScanBuilder; -use vortex::metrics::Label; -use vortex::metrics::MetricsRegistry; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; -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::ProcessedProjection; -use crate::convert::exprs::make_vortex_predicate; -use crate::convert::schema::calculate_physical_schema; -use crate::metrics::PARTITION_LABEL; -use crate::metrics::PATH_LABEL; -use crate::persistent::cache::CachedVortexMetadata; -use crate::persistent::reader::VortexReaderFactory; -use crate::persistent::stream::PrunableStream; - -#[derive(Clone)] -pub(crate) struct VortexOpener { - /// The partition this opener is assigned to. Only used for labeling metrics. - pub partition: usize, - pub session: VortexSession, - pub vortex_reader_factory: Arc, - /// 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. - pub filter: Option, - /// Filter expression used by DataFusion's FilePruner to eliminate files based on - /// statistics and partition values without opening them. - pub file_pruning_predicate: Option, - pub expr_adapter_factory: Arc, - /// This is the table's schema without partition columns. It may contain fields which do - /// not exist in the file, and are supplied by the `schema_adapter_factory`. - pub table_schema: TableSchema, - /// If provided, the scan will not return more than this many rows. - pub limit: Option, - /// A metrics object for tracking performance of the scan. - pub metrics_registry: Arc, - /// DataFusion-native metrics exposed through `DataSourceExec`. - pub df_metrics: ExecutionPlanMetricsSet, - /// A shared cache of file readers. - /// - /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache - /// a file reader the first time we read a file. - pub layout_readers: Arc>>, - /// Shared full-file natural splits keyed by file path. - pub natural_splits: Arc>>, - /// Whether the query has output ordering specified - pub has_output_ordering: bool, - - pub expression_convertor: Arc, - pub file_metadata_cache: Option>, - /// Whether to enable expression pushdown into the underlying Vortex scan. - pub projection_pushdown: bool, - pub scan_concurrency: Option, -} - -impl FileOpener for VortexOpener { - fn open(&self, file: PartitionedFile) -> DFResult { - // Calculate the output schema before replacing partition columns with literals so it - // retains the table and partition-field metadata declared by the plan. - let output_schema = Arc::new( - self.projection - .project_schema(self.table_schema.table_schema())?, - ); - let session = self.session.clone(); - let metrics_registry = Arc::clone(&self.metrics_registry); - let labels = vec![ - Label::new(PATH_LABEL, file.path().to_string()), - Label::new(PARTITION_LABEL, self.partition.to_string()), - ]; - - let mut projection = self.projection.clone(); - let mut filter = self.filter.clone(); - - let reader = self.vortex_reader_factory.create_reader(&file, &session)?; - - let reader = - InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); - - let mut file_pruning_predicate = self.file_pruning_predicate.clone(); - let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); - let file_metadata_cache = self.file_metadata_cache.clone(); - - let unified_file_schema = Arc::clone(self.table_schema.file_schema()); - let limit = self.limit; - let layout_readers = Arc::clone(&self.layout_readers); - let natural_splits = Arc::clone(&self.natural_splits); - let has_output_ordering = self.has_output_ordering; - let scan_concurrency = self.scan_concurrency; - - let expr_convertor = Arc::clone(&self.expression_convertor); - let projection_pushdown = self.projection_pushdown; - - let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) - .with_category(MetricCategory::Rows) - .global_counter("num_predicate_creation_errors"); - - // Replace column access for partition columns with literals - #[expect(clippy::disallowed_types)] - let literal_value_cols = self - .table_schema - .table_partition_cols() - .iter() - .map(|f| f.name()) - .cloned() - .zip(file.partition_values.clone()) - .collect::>(); - - let predicate_uses_partition_columns = - file_pruning_predicate.as_ref().is_some_and(|predicate| { - collect_columns(predicate) - .iter() - .any(|column| literal_value_cols.contains_key(column.name())) - }); - - if !literal_value_cols.is_empty() { - projection = projection.try_map_exprs(|expr| { - replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) - })?; - filter = filter - .map(|p| replace_columns_with_literals(p, &literal_value_cols)) - .transpose()?; - file_pruning_predicate = file_pruning_predicate - .map(|p| replace_columns_with_literals(p, &literal_value_cols)) - .transpose()?; - } - - Ok(async move { - // FilePruner requires a statistics object even when the rewritten predicate - // only contains partition literals. Supply unknown file-column statistics in - // that case so static and dynamic partition predicates can still prune. - let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) - .then(|| { - file.clone() - .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) - }); - let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); - - let mut file_pruner = file_pruning_predicate - .filter(|_| file.has_statistics() || predicate_uses_partition_columns) - .and_then(|predicate| { - FilePruner::try_new( - Arc::clone(&predicate), - &unified_file_schema, - pruning_file, - predicate_creation_errors, - ) - }); - - // Check if this file should be pruned based on statistics/partition values. - // Returns empty stream if file can be skipped entirely. - if let Some(file_pruner) = file_pruner.as_mut() - && file_pruner.should_prune()? - { - return Ok(stream::empty().boxed()); - } - - let mut open_opts = session - .open_options() - .with_file_size(file.object_meta.size) - .with_metrics_registry(Arc::clone(&metrics_registry)) - .with_labels(labels); - - let cached_footer = file_metadata_cache - .as_ref() - .and_then(|cache| cache.get(file.path())) - .filter(|entry| entry.is_valid_for(&file.object_meta)) - .and_then(|entry| { - entry - .file_metadata - .as_any() - .downcast_ref::() - .map(|vortex_metadata| vortex_metadata.footer().clone()) - }); - let footer_cache_hit = cached_footer.is_some(); - - if let Some(footer) = cached_footer { - open_opts = open_opts.with_footer(footer); - } - - let vxf = open_opts - .open_read(reader) - .await - .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?; - - // On a miss, cache the parsed footer so other partitions and later executions - // skip the footer fetch and parse. `infer_schema`/`infer_stats` also populate - // this cache, but only when planning goes through `VortexFormat`. - if !footer_cache_hit && let Some(cache) = &file_metadata_cache { - cache.put( - file.path(), - CachedFileMetadataEntry::new( - file.object_meta.clone(), - Arc::new(CachedVortexMetadata::new(&vxf)), - ), - ); - } - - // Check if there are rows in this file. If not, we can save - // ourselves some work and return an empty stream. - if vxf.row_count() == 0 { - return Ok(stream::empty().boxed()); - } - - // This is the expected arrow types of the actual columns in the file, which might have different types - // from the unified logical schema or miss - let this_file_schema = Arc::new(calculate_physical_schema( - vxf.dtype(), - &unified_file_schema, - &session.arrow(), - )?); - - let expr_adapter = expr_adapter_factory.create( - Arc::clone(&unified_file_schema), - Arc::clone(&this_file_schema), - )?; - - let simplifier = PhysicalExprSimplifier::new(&this_file_schema); - - // The adapter rewrites the expressions to the local file schema, allowing - // for schema evolution and divergence between the table's schema and individual files. - let filter = filter - .map(|filter| { - // Expression might now reference columns that don't exist in the file, so we can give it - // another simplification pass. - simplifier.simplify(expr_adapter.rewrite(filter)?) - }) - .transpose()?; - let projection = - projection.try_map_exprs(|p| simplifier.simplify(expr_adapter.rewrite(p)?))?; - - let ProcessedProjection { - scan_projection, - leftover_projection, - } = if projection_pushdown { - expr_convertor.split_projection( - projection.clone(), - &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)? - }; - - // The schema of the stream returned from the vortex scan. - // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). - let scan_projection = scan_projection - .optimize_recursive(vxf.dtype()) - .and_then(|projection| projection.bind(vxf.dtype())) - .map_err(|_e| { - exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") - })?; - 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 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. - let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { - Entry::Occupied(mut occupied_entry) => { - if let Some(reader) = occupied_entry.get().upgrade() { - tracing::trace!("reusing layout reader for {}", occupied_entry.key()); - reader - } else { - tracing::trace!("creating layout reader for {}", occupied_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!( - "Failed to create layout reader: {e}" - )) - })?; - occupied_entry.insert(Arc::downgrade(&reader)); - reader - } - } - Entry::Vacant(vacant_entry) => { - tracing::trace!("creating layout reader for {}", vacant_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!("Failed to create layout reader: {e}")) - })?; - vacant_entry.insert(Arc::downgrade(&reader)); - - reader - } - }; - - let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); - - if let Some(vortex_plan) = file.extensions.get::() { - 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() - { - scan_builder = scan_builder.with_limit(limit); - } - - if let Some(concurrency) = scan_concurrency { - scan_builder = scan_builder.with_concurrency(concurrency); - } - - // Set before the byte-range translation below, which computes natural splits for - // the fields the scan's projection and filter reference. - scan_builder = scan_builder - .with_projection(scan_projection) - .with_some_filter(filter); - - 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"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, - }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_splits = natural_splits_for_file( - natural_splits.as_ref(), - &file.object_meta.location, - &scan_builder, - file.object_meta.size, - )?; - - let Some(row_range) = - split_aligned_row_range(byte_range, natural_splits.as_ref()) - else { - return Ok(stream::empty().boxed()); - }; - - scan_builder = scan_builder - .with_row_range(row_range) - // Hand the shared full-file boundaries back to the scan so prepare() - // skips its own layout walk. - .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); - } - } - - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); - let stream = scan_builder - .with_metrics_registry(metrics_registry) - .with_ordered(has_output_ordering) - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) - .into_stream() - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - 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)) - }?; - - let (_, columns, row_count) = batch.into_parts(); - RecordBatch::try_new_with_options( - Arc::clone(&output_schema), - columns, - &RecordBatchOptions::new().with_row_count(Some(row_count)), - ) - .map_err(Into::into) - }) - .boxed(); - - if let Some(file_pruner) = file_pruner - && file_pruner.is_watching() - { - Ok(PrunableStream::new(file_pruner, stream).boxed()) - } else { - Ok(stream) - } - } - .in_current_span() - .boxed()) - } -} - -/// A file's natural split boundaries plus the precomputed byte each split is assigned to, -/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row -/// boundaries with a binary search instead of re-projecting every split per partition. -/// -/// The boundaries are computed for the fields referenced by the scan's projection and filter. -/// All partitions translate through the first opener's cached entry (the cache lives on the -/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges -/// tiling the file's rows exactly once. -#[derive(Debug)] -pub(crate) struct NaturalSplits { - /// Sorted row boundaries of the natural splits; split `i` covers - /// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the - /// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the - /// per-partition layout walk in `prepare`. - row_boundaries: Arc<[u64]>, - /// For each split, the byte a DataFusion byte range must contain to own it (see - /// [`split_assignment_byte`]); one entry per split, sorted because split midpoints - /// increase monotonically under the row-to-byte projection. - assignment_bytes: Box<[u64]>, -} - -impl NaturalSplits { - fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { - let row_count = row_boundaries.last().copied().unwrap_or_default(); - let assignment_bytes = if row_count == 0 { - Box::default() - } else { - row_boundaries - .windows(2) - .enumerate() - .map(|(idx, boundaries)| { - split_assignment_byte( - idx, - &(boundaries[0]..boundaries[1]), - row_count, - total_size, - ) - }) - .collect() - }; - - debug_assert!(assignment_bytes.is_sorted()); - debug_assert_eq!( - assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), - row_boundaries.len() - ); - - Self { - row_boundaries, - assignment_bytes, - } - } -} - -/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. -fn natural_splits_for_file( - natural_splits: &DashMap>, - path: &Path, - scan_builder: &ScanBuilder, - total_size: u64, -) -> DFResult> { - if let Some(splits) = natural_splits.get(path) { - return Ok(Arc::clone(splits.value())); - } - - // Compute while holding the entry so concurrent partitions opening the same file wait - // for the winner instead of all walking the layout tree; the redundant walks contend on - // the lazily-initialized layout children and dominate the cost of the computation itself. - match natural_splits.entry(path.clone()) { - Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), - Entry::Vacant(entry) => { - let splits = compute_natural_splits(scan_builder, total_size)?; - entry.insert(Arc::clone(&splits)); - Ok(splits) - } - } -} - -/// Walk the layout tree to compute the file's full natural split boundaries for the fields -/// referenced by the scan's projection and filter. -fn compute_natural_splits( - scan_builder: &ScanBuilder, - total_size: u64, -) -> DFResult> { - let row_boundaries = scan_builder - .full_file_splits() - .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; - - Ok(Arc::new(NaturalSplits::new( - row_boundaries.into(), - total_size, - ))) -} - -/// Translate a DataFusion byte range to the contiguous natural split ranges it owns. -/// Most splits are assigned by midpoint, but the leading split stays with the range that owns -/// byte 0 so a tiny first byte range still claims the first rows. -fn split_aligned_row_range( - byte_range: Range, - natural_splits: &NaturalSplits, -) -> Option> { - if byte_range.start >= byte_range.end { - return None; - } - - let first_split = natural_splits - .assignment_bytes - .partition_point(|&assignment_byte| assignment_byte < byte_range.start); - let after_last_split = natural_splits - .assignment_bytes - .partition_point(|&assignment_byte| assignment_byte < byte_range.end); - if first_split == after_last_split { - return None; - } - - Some( - natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split], - ) -} - -fn split_assignment_byte( - idx: usize, - split_range: &Range, - row_count: u64, - total_size: u64, -) -> u64 { - if idx == 0 && split_range.start == 0 { - // Byte 0 is the only stable representative for the leading split. A midpoint can fall - // into the next DataFusion byte range and leave the first range with no rows to read. - 0 - } else { - split_midpoint_to_byte(split_range, row_count, total_size) - } -} - -fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64) -> u64 { - let midpoint_row = split_range.start + (split_range.end - split_range.start) / 2; - let midpoint_byte = (u128::from(midpoint_row) * u128::from(total_size)) / u128::from(row_count); - - u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") -} - -#[cfg(test)] -mod tests { - use std::fmt; - use std::sync::Arc; - use std::sync::LazyLock; - - use arrow_array::record_batch; - use arrow_schema::Field; - use arrow_schema::Fields; - use arrow_schema::SchemaRef; - use datafusion::arrow::array::DictionaryArray; - use datafusion::arrow::array::Int32Array; - use datafusion::arrow::array::RecordBatch; - use datafusion::arrow::array::StringArray; - use datafusion::arrow::array::StructArray; - use datafusion::arrow::datatypes::DataType; - use datafusion::arrow::datatypes::Schema; - use datafusion::arrow::datatypes::UInt32Type; - use datafusion::arrow::util::display::FormatOptions; - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - use datafusion::logical_expr::col; - use datafusion::logical_expr::lit; - use datafusion::physical_expr::planner::logical2physical; - use datafusion::physical_expr_adapter::DefaultPhysicalExprAdapterFactory; - use datafusion::scalar::ScalarValue; - use datafusion_common::stats::Precision; - use datafusion_execution::cache::default_cache::DefaultCache; - use datafusion_expr::Operator; - use datafusion_physical_expr::PhysicalExpr; - use datafusion_physical_expr::expressions as df_expr; - use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; - use datafusion_physical_expr::projection::ProjectionExpr; - use insta::assert_snapshot; - use itertools::Itertools; - use object_store::ObjectStore; - use object_store::memory::InMemory; - use rstest::rstest; - use vortex::VortexSessionDefault; - use vortex::buffer::Buffer; - use vortex::file::WriteOptionsSessionExt; - use vortex::io::VortexWrite; - use vortex::io::object_store::ObjectStoreWrite; - use vortex::metrics::DefaultMetricsRegistry; - use vortex::scan::selection::Selection; - use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; - use vortex::session::VortexSession; - - use super::*; - use crate::VortexAccessPlan; - use crate::convert::exprs::DefaultExpressionConvertor; - use crate::persistent::reader::DefaultVortexReaderFactory; - - static SESSION: LazyLock = LazyLock::new(VortexSession::default); - - /// Test-only expr used to test error reporting. - #[derive(Debug, Eq, Hash, PartialEq)] - struct SnapshotErrorExpr; - - impl fmt::Display for SnapshotErrorExpr { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "snapshot_error") - } - } - - impl PhysicalExpr for SnapshotErrorExpr { - fn data_type(&self, _input_schema: &Schema) -> DFResult { - Ok(DataType::Boolean) - } - - fn nullable(&self, _input_schema: &Schema) -> DFResult { - Ok(false) - } - - fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt::Display::fmt(self, f) - } - - fn evaluate(&self, _batch: &RecordBatch) -> DFResult { - Err(DataFusionError::Internal( - "intentional snapshot error".to_owned(), - )) - } - - fn children(&self) -> Vec<&PhysicalExprRef> { - Vec::new() - } - - fn with_new_children( - self: Arc, - children: Vec, - ) -> DFResult { - assert!(children.is_empty()); - Ok(self) - } - - fn snapshot(&self) -> DFResult> { - Err(DataFusionError::Internal( - "intentional snapshot error".to_owned(), - )) - } - } - - fn natural_splits(total_size: u64, split_ranges: &[Range]) -> NaturalSplits { - let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1); - if let Some(first) = split_ranges.first() { - row_boundaries.push(first.start); - row_boundaries.extend(split_ranges.iter().map(|range| range.end)); - } - NaturalSplits::new(row_boundaries.into(), total_size) - } - - #[rstest] - #[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] - #[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] - #[case(1..8, 10, vec![0..1, 1..9, 9..10], Some(1..9))] - #[case(1..4, 16, vec![0..1, 1..2, 2..3, 3..4], None)] - #[case(0..1, 10, vec![0..2, 2..10], Some(0..2))] - #[case(0..2, 2, vec![], None)] - fn test_split_aligned_row_range( - #[case] byte_range: Range, - #[case] total_size: u64, - #[case] split_ranges: Vec>, - #[case] expected: Option>, - ) { - assert_eq!( - split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)), - expected - ); - } - - #[test] - fn test_split_aligned_ranges_cover_splits_exactly_once() { - let split_ranges = vec![0..1, 1..4, 4..10, 10..13]; - let byte_ranges = [0..4, 4..8, 8..12, 12..16]; - let natural_splits = natural_splits(16, &split_ranges); - - let assigned = byte_ranges - .into_iter() - .filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits)) - .collect::>(); - - assert_eq!(assigned, vec![0..4, 4..10, 10..13]); - assert_eq!( - assigned - .iter() - .map(|range| range.end - range.start) - .sum::(), - 13 - ); - - let split_starts = split_ranges - .iter() - .map(|range| range.start) - .collect::>(); - let split_ends = split_ranges - .iter() - .map(|range| range.end) - .collect::>(); - - for range in &assigned { - assert!(split_starts.contains(&range.start)); - assert!(split_ends.contains(&range.end)); - } - - for (left, right) in assigned.iter().tuple_windows() { - assert_eq!(left.end, right.start); - } - } - - #[rstest] - #[case(vec![], 10)] - #[case(vec![0], 10)] - #[case(vec![], 0)] - #[case(vec![0], 0)] - fn test_natural_splits_empty_file(#[case] row_boundaries: Vec, #[case] total_size: u64) { - let splits = NaturalSplits::new(row_boundaries.clone().into(), total_size); - - assert!(splits.assignment_bytes.is_empty()); - assert_eq!(splits.row_boundaries.as_ref(), row_boundaries.as_slice()); - assert_eq!(split_aligned_row_range(0..u64::MAX, &splits), None); - } - - #[test] - fn test_split_aligned_row_range_keeps_colliding_assignments_together() { - let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]); - - assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]); - assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2)); - assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); - } - - async fn write_arrow_to_vortex( - object_store: Arc, - path: &str, - rb: RecordBatch, - ) -> anyhow::Result { - let schema = rb.schema(); - let array = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; - let path = Path::parse(path)?; - - let mut write = ObjectStoreWrite::new(object_store, &path).await?; - let summary = SESSION - .write_options() - .write(&mut write, array.to_array_stream()) - .await?; - write.shutdown().await?; - - Ok(summary.size()) - } - - fn make_opener( - object_store: Arc, - table_schema: TableSchema, - filter: Option, - ) -> VortexOpener { - VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), - projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), - filter, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema, - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - } - } - - #[tokio::test] - async fn test_open() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "part=1/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file_schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.partition_values = vec![ScalarValue::Int32(Some(1))]; - - let table_schema = TableSchema::builder(Arc::clone(&file_schema)) - .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) - .build(); - - // filter matches partition value - let filter = col("part").eq(lit(1)); - let filter = logical2physical(&filter, table_schema.table_schema()); - - let opener = make_opener( - Arc::clone(&object_store), - table_schema.clone(), - Some(filter), - ); - let stream = opener.open(file.clone()).unwrap().await.unwrap(); - - let data = stream.try_collect::>().await?; - let num_batches = data.len(); - let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); - - assert_eq!((num_batches, num_rows), (1, 3)); - - // filter doesn't matches partition value - let filter = col("part").eq(lit(2)); - let filter = logical2physical(&filter, table_schema.table_schema()); - - let opener = make_opener( - Arc::clone(&object_store), - table_schema.clone(), - Some(filter), - ); - let stream = opener.open(file.clone()).unwrap().await.unwrap(); - - let data = stream.try_collect::>().await?; - let num_batches = data.len(); - let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); - assert_eq!((num_batches, num_rows), (0, 0)); - - Ok(()) - } - - #[tokio::test] - async fn test_open_preserves_declared_schema_metadata() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "part=1/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file_schema = Arc::new( - batch.schema().as_ref().clone().with_metadata( - [("table".to_string(), "metadata".to_string())] - .into_iter() - .collect(), - ), - ); - let table_schema = TableSchema::builder(file_schema) - .with_table_partition_cols(vec![Arc::new( - Field::new("part", DataType::Int32, false).with_metadata( - [("partition".to_string(), "metadata".to_string())] - .into_iter() - .collect(), - ), - )]) - .build(); - let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); - let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); - - assert_eq!( - expected_schema.metadata().get("table"), - Some(&"metadata".to_string()) - ); - assert_eq!( - expected_schema.field(1).metadata().get("partition"), - Some(&"metadata".to_string()) - ); - - for projection_pushdown in [false, true] { - let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); - opener.projection = projection.clone(); - opener.projection_pushdown = projection_pushdown; - - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.partition_values = vec![ScalarValue::Int32(Some(1))]; - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert!(!batches.is_empty()); - for batch in batches { - assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); - } - } - - Ok(()) - } - - #[tokio::test] - async fn test_open_all_valid_nullable_columns_with_nonnullable_table_schema() - -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "nullable/file.vortex"; - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), - vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], - )?; - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - let expected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let table_schema = TableSchema::from(Arc::clone(&expected_schema)); - - for projection_pushdown in [false, true] { - let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); - opener.projection_pushdown = projection_pushdown; - - let file = PartitionedFile::new(file_path.to_string(), data_size); - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert_eq!(batches.len(), 1); - assert_eq!(batches[0].schema().as_ref(), expected_schema.as_ref()); - } - - Ok(()) - } - - #[tokio::test] - async fn test_file_pruning_replaces_partition_columns_without_file_statistics() - -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let table_schema = TableSchema::builder(Arc::clone(&file_schema)) - .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) - .build(); - - let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; - let predicate = Arc::new(df_expr::BinaryExpr::new( - Arc::clone(&partition_column), - Operator::Gt, - df_expr::lit(ScalarValue::Int32(Some(1))), - )) as PhysicalExprRef; - let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( - vec![partition_column], - predicate, - )) as PhysicalExprRef; - - let mut opener = make_opener(object_store, table_schema, None); - opener.file_pruning_predicate = Some(dynamic_predicate); - let df_metrics = opener.df_metrics.clone(); - - // The file does not exist and has no statistics. Replacing `part` with 1 - // makes the predicate false, so pruning must happen before any file I/O. - let mut file = PartitionedFile::new("missing.vortex", 1); - file.partition_values = vec![ScalarValue::Int32(Some(1))]; - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert!(batches.is_empty()); - assert_eq!( - df_metrics - .clone_inner() - .sum_by_name("num_predicate_creation_errors") - .map(|metric| metric.as_usize()), - Some(0) - ); - - Ok(()) - } - - #[tokio::test] - async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "metrics/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); - statistics.column_statistics[0].null_count = Precision::Exact(0); - let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); - - let mut opener = make_opener(object_store, TableSchema::from(batch.schema()), None); - opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); - let df_metrics = opener.df_metrics.clone(); - - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); - assert_eq!( - df_metrics - .clone_inner() - .sum_by_name("num_predicate_creation_errors") - .map(|metric| metric.as_usize()), - Some(1) - ); - - Ok(()) - } - - #[tokio::test] - async fn test_open_empty_file() -> anyhow::Result<()> { - use futures::TryStreamExt; - - let object_store = Arc::new(InMemory::new()) as Arc; - let data_batch = record_batch!(("a", Int32, Vec::::new())).unwrap(); - let file_path = "part=1/empty.vortex"; - let file_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, data_batch.clone()).await?; - - let file_schema = data_batch.schema(); - // Parallel scans may attach a byte range even for empty files; the - // opener must return early before attempting split-aligned translation. - let file = - PartitionedFile::new_with_range(file_path.to_string(), file_size, 0, file_size as i64); - - let table_schema = TableSchema::from(Arc::clone(&file_schema)); - - let opener = make_opener(object_store, table_schema, None); - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - assert_eq!(data.len(), 0); - - Ok(()) - } - - #[tokio::test] - async fn test_open_populates_file_metadata_cache() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "cached/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file = PartitionedFile::new(file_path.to_string(), data_size); - let table_schema = TableSchema::from(batch.schema()); - - let cache: Arc = Arc::new( - DefaultCache::::new(64 * 1024 * 1024), - ); - let mut opener = make_opener(Arc::clone(&object_store), table_schema, None); - opener.file_metadata_cache = Some(Arc::clone(&cache)); - - // The first open misses the cache and must write the parsed footer back. - let stream = opener.open(file.clone())?.await?; - stream.try_collect::>().await?; - - let entry = cache - .get(file.path()) - .ok_or_else(|| anyhow::anyhow!("footer was not cached after open"))?; - assert!(entry.is_valid_for(&file.object_meta)); - assert!( - entry - .file_metadata - .as_any() - .downcast_ref::() - .is_some() - ); - - // The second open hits the cache and still returns the same data. - let stream = opener.open(file.clone())?.await?; - let data = stream.try_collect::>().await?; - assert_eq!(data.iter().map(|rb| rb.num_rows()).sum::(), 3); - - Ok(()) - } - - #[rstest] - #[tokio::test] - async fn test_open_files_different_table_schema() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - - let file1 = { - let file1_path = "/path/file1.vortex"; - let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size1 = - write_arrow_to_vortex(Arc::clone(&object_store), file1_path, batch1).await?; - PartitionedFile::new(file1_path.to_string(), data_size1) - }; - - let file2 = { - let file2_path = "/path/file2.vortex"; - let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); - let data_size2 = - write_arrow_to_vortex(Arc::clone(&object_store), file2_path, batch2).await?; - PartitionedFile::new(file2_path.to_string(), data_size2) - }; - - // Table schema has can accommodate both files - let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( - "a", - DataType::Int32, - true, - )]))); - - let make_opener = |filter| VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone( - &object_store, - ))), - projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), - filter: Some(filter), - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: table_schema.clone(), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - let filter = col("a").lt(lit(100_i32)); - let filter = logical2physical(&filter, table_schema.table_schema()); - - let opener1 = make_opener(Arc::clone(&filter)); - let stream = opener1.open(file1)?.await?; - - let format_opts = FormatOptions::new().with_types_info(true); - - let data = stream.try_collect::>().await?; - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+ - | a | - | Int32 | - +-------+ - | 1 | - | 2 | - | 3 | - +-------+ - "); - - let opener2 = make_opener(Arc::clone(&filter)); - let stream = opener2.open(file2)?.await?; - - let data = stream.try_collect::>().await?; - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+ - | a | - | Int32 | - +-------+ - | -1 | - | -2 | - | -3 | - +-------+ - "); - - Ok(()) - } - - #[tokio::test] - // This test verifies that files with different column order than the - // table schema can be opened without errors. The fix ensures that the - // schema mapper is only used for type casting, not for reordering, - // since the vortex projection already handles reordering. - async fn test_schema_different_column_order() -> anyhow::Result<()> { - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - // File has columns in order: c, b, a - let batch = record_batch!( - ("c", Int32, vec![Some(300), Some(301), Some(302)]), - ("b", Int32, vec![Some(200), Some(201), Some(202)]), - ("a", Int32, vec![Some(100), Some(101), Some(102)]) - ) - .unwrap(); - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - let file = PartitionedFile::new(file_path.to_string(), data_size); - - // Table schema has columns in different order: a, b, c - let table_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - ])); - - let opener = VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), - projection: ProjectionExprs::from_indices(&[0, 1, 2], &table_schema), - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: TableSchema::from(Arc::clone(&table_schema)), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - let stream = opener.open(file)?.await?; - - let format_opts = FormatOptions::new().with_types_info(true); - let data = stream.try_collect::>().await?; - - // Verify the output has columns in table schema order (a, b, c) - // not file order (c, b, a) - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+-------+-------+ - | a | b | c | - | Int32 | Int32 | Int32 | - +-------+-------+-------+ - | 100 | 200 | 300 | - | 101 | 201 | 301 | - | 102 | 202 | 302 | - +-------+-------+-------+ - "); - - Ok(()) - } - - #[tokio::test] - // This test verifies that expression rewriting doesn't fail when there is - // a nested schema mismatch between the physical file schema and logical - // table schema. - async fn test_adapter_logical_physical_struct_mismatch() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - let file_struct_fields = Fields::from(vec![ - Field::new("field1", DataType::Utf8, true), - Field::new("field2", DataType::Utf8, true), - ]); - let struct_array = StructArray::new( - file_struct_fields.clone(), - vec![ - Arc::new(StringArray::from(vec!["value1", "value2", "value3"])), - Arc::new(StringArray::from(vec!["a", "b", "c"])), - ], - None, - ); - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new( - "my_struct", - DataType::Struct(file_struct_fields), - true, - )])), - vec![Arc::new(struct_array)], - )?; - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - // Table schema has an extra utf8 field. - let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( - "my_struct", - DataType::Struct(Fields::from(vec![ - Field::new( - "field1", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - Field::new( - "field2", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - Field::new("field3", DataType::Utf8, true), - ])), - true, - )]))); - - let opener = make_opener( - Arc::clone(&object_store), - table_schema.clone(), - // expression references my_struct column which has different fields in each - // field. - Some(logical2physical( - &col("my_struct").is_not_null(), - table_schema.table_schema(), - )), - ); - - // The opener should be able to open the file with a filter on the - // struct column. - let data = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? - .await? - .try_collect::>() - .await?; - - assert_eq!(data.len(), 1); - assert_eq!(data[0].num_rows(), 3); - - Ok(()) - } - - #[tokio::test] - // Minimal reproducing test for the schema projection bug. - // Before the fix, this would fail with a cast error when the file schema - // and table schema have different field orders and we project a subset of columns. - async fn test_projection_bug_minimal_repro() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - // File has columns in order: a, b, c with simple types - let batch = record_batch!( - ("a", Int32, vec![Some(1)]), - ("b", Utf8, vec![Some("test")]), - ("c", Int32, vec![Some(2)]) - ) - .unwrap(); - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - // Table schema has columns in DIFFERENT order: c, a, b - // and different types that require casting (Utf8 -> Dictionary) - let table_schema = TableSchema::from(Arc::new(Schema::new(vec![ - Field::new("c", DataType::Int32, true), - Field::new("a", DataType::Int32, true), - Field::new( - "b", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - ]))); - - // Project columns [0, 2] from table schema, which should give us: c, b - // Before the fix, the schema adapter would get confused about which fields - // to select from the file, causing incorrect type mappings. - let projection = vec![0, 2]; - - let opener = VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone( - &object_store, - ))), - projection: ProjectionExprs::from_indices( - projection.as_ref(), - table_schema.file_schema(), - ), - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: table_schema.clone(), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - // This should succeed and return the correctly projected and cast data - let data = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? - .await? - .try_collect::>() - .await?; - - // Verify the columns are in the right order and have the right values - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - let format_opts = FormatOptions::new().with_types_info(true); - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+--------------------------+ - | c | b | - | Int32 | Dictionary(UInt32, Utf8) | - +-------+--------------------------+ - | 2 | test | - +-------+--------------------------+ - "); - - Ok(()) - } - - fn make_test_batch_with_10_rows() -> RecordBatch { - record_batch!( - ("a", Int32, (0..=9).map(Some).collect::>()), - ( - "b", - Utf8, - (0..=9).map(|i| Some(format!("r{}", i))).collect::>() - ) - ) - .unwrap() - } - - fn make_test_opener( - object_store: Arc, - schema: SchemaRef, - projection: ProjectionExprs, - ) -> VortexOpener { - VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), - projection, - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: TableSchema::from(schema), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - } - } - - #[tokio::test] - // Test that Selection::IncludeByIndex filters to specific row indices. - async fn test_selection_include_by_index() -> anyhow::Result<()> { - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.extensions - .insert( - VortexAccessPlan::default().with_selection(Selection::IncludeByIndex( - StrictSortedBuffer::try_new(Buffer::from_iter(vec![1, 3, 5, 7]))?, - )), - ); - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0, 1], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - let format_opts = FormatOptions::new().with_types_info(true); - - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+------+ - | a | b | - | Int32 | Utf8 | - +-------+------+ - | 1 | r1 | - | 3 | r3 | - | 5 | r5 | - | 7 | r7 | - +-------+------+ - "); - - Ok(()) - } - - #[tokio::test] - // Test that Selection::ExcludeByIndex excludes specific row indices. - async fn test_selection_exclude_by_index() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.extensions - .insert( - VortexAccessPlan::default().with_selection(Selection::ExcludeByIndex( - StrictSortedBuffer::try_new(Buffer::from_iter(vec![0, 2, 4, 6, 8]))?, - )), - ); - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0, 1], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - let format_opts = FormatOptions::new().with_types_info(true); - - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+------+ - | a | b | - | Int32 | Utf8 | - +-------+------+ - | 1 | r1 | - | 3 | r3 | - | 5 | r5 | - | 7 | r7 | - | 9 | r9 | - +-------+------+ - "); - - Ok(()) - } - - #[tokio::test] - // Test that Selection::All returns all rows. - async fn test_selection_all() -> anyhow::Result<()> { - use vortex::scan::selection::Selection; - - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.extensions - .insert(VortexAccessPlan::default().with_selection(Selection::All)); - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); - assert_eq!(total_rows, 10); - - Ok(()) - } - - #[tokio::test] - // Test that when no extensions are provided, all rows are returned (backward compatibility). - async fn test_selection_no_extensions() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let file = PartitionedFile::new(file_path.to_string(), data_size); - // file.extensions is None by default - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); - assert_eq!(total_rows, 10); - - Ok(()) - } - - #[tokio::test] - async fn test_projection_expr_pushdown() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = record_batch!( - ("a", Int32, vec![Some(1), Some(2), Some(3)]), - ("b", Int32, vec![Some(10), Some(20), Some(30)]) - ) - .unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file_schema = batch.schema(); - let table_schema = TableSchema::from(Arc::clone(&file_schema)); - - // Create a projection that includes an arithmetic expression: a + b * 2 - let col_a = df_expr::col("a", &file_schema)?; - let col_b = df_expr::col("b", &file_schema)?; - let two = df_expr::lit(ScalarValue::Int32(Some(2))); - - // b * 2 - let b_times_2 = df_expr::binary(col_b, Operator::Multiply, two, &file_schema)?; - // a + (b * 2) - let a_plus_b_times_2 = df_expr::binary(col_a, Operator::Plus, b_times_2, &file_schema)?; - - let projection = ProjectionExprs::new(vec![ProjectionExpr::new( - a_plus_b_times_2, - "result".to_string(), - )]); - - let opener = VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone( - &object_store, - ))), - projection, - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema, - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - let file = PartitionedFile::new(file_path.to_string(), data_size); - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - // Expected: a + b * 2 - // row 0: 1 + 10 * 2 = 21 - // row 1: 2 + 20 * 2 = 42 - // row 2: 3 + 30 * 2 = 63 - assert_snapshot!(pretty_format_batches_with_options(&data, &FormatOptions::new().with_types_info(true))?.to_string(), @r" - +--------+ - | result | - | Int32 | - +--------+ - | 21 | - | 42 | - | 63 | - +--------+ - "); - - Ok(()) - } - - /// When a Struct contains Dictionary fields, writing to vortex and reading back - /// should preserve the Dictionary type. - #[tokio::test] - async fn test_struct_with_dictionary_roundtrip() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - - let struct_fields = Fields::from(vec![ - Field::new_dictionary("a", DataType::UInt32, DataType::Utf8, true), - Field::new_dictionary("b", DataType::UInt32, DataType::Utf8, true), - ]); - let struct_array = StructArray::new( - struct_fields.clone(), - vec![ - Arc::new(DictionaryArray::::from_iter(["x", "y", "x"])), - Arc::new(DictionaryArray::::from_iter(["p", "p", "q"])), - ], - None, - ); - - let schema = Arc::new(Schema::new(vec![Field::new( - "labels", - DataType::Struct(struct_fields.clone()), - false, - )])); - let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(struct_array)])?; - - let file_path = "/test.vortex"; - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0], &schema), - ); - let data: Vec<_> = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? - .await? - .try_collect() - .await?; - - assert_eq!( - data[0].schema().field(0).data_type(), - &DataType::Struct(struct_fields), - "Struct(Dictionary) type should be preserved" - ); - Ok(()) - } -} diff --git a/vortex-datafusion/src/persistent/opener/mod.rs b/vortex-datafusion/src/persistent/opener/mod.rs new file mode 100644 index 00000000000..c69efe08d19 --- /dev/null +++ b/vortex-datafusion/src/persistent/opener/mod.rs @@ -0,0 +1,669 @@ +// 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::Field; +use datafusion_common::DataFusionError; +use datafusion_common::Result as DFResult; +use datafusion_common::ScalarValue; +use datafusion_common::Statistics; +use datafusion_common::arrow::array::AsArray; +use datafusion_common::arrow::array::RecordBatch; +use datafusion_common::exec_datafusion_err; +use datafusion_common::tree_node::Transformed; +use datafusion_common::tree_node::TreeNode; +use datafusion_datasource::PartitionedFile; +use datafusion_datasource::TableSchema; +use datafusion_datasource::file_stream::FileOpenFuture; +use datafusion_datasource::file_stream::FileOpener; +use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; +use datafusion_execution::cache::cache_manager::FileMetadataCache; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions as df_expr; +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; +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; +use datafusion_pruning::FilePruner; +use futures::FutureExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use object_store::path::Path; +use tracing::Instrument; +use vortex::array::VortexSessionExecute; +use vortex::error::VortexError; +use vortex::error::VortexExpect; +use vortex::file::OpenOptionsSessionExt; +use vortex::io::InstrumentedReadAt; +use vortex::layout::LayoutReader; +use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::metrics::Label; +use vortex::metrics::MetricsRegistry; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; +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::ProcessedProjection; +use crate::convert::exprs::raw_projection; +use crate::convert::schema::calculate_physical_schema; +use crate::metrics::PARTITION_LABEL; +use crate::metrics::PATH_LABEL; +use crate::persistent::cache::CachedVortexMetadata; +use crate::persistent::reader::VortexReaderFactory; +use crate::persistent::stream::PrunableStream; + +#[derive(Clone)] +pub(crate) struct VortexOpener { + /// The partition this opener is assigned to. Only used for labeling metrics. + pub partition: usize, + pub session: VortexSession, + pub vortex_reader_factory: Arc, + /// 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, + /// 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. + pub file_pruning_predicate: Option, + pub expr_adapter_factory: Arc, + /// This is the table's schema without partition columns. It may contain fields which do + /// not exist in the file, and are supplied by the `schema_adapter_factory`. + pub table_schema: TableSchema, + /// If provided, the scan will not return more than this many rows. + pub limit: Option, + /// A metrics object for tracking performance of the scan. + pub metrics_registry: Arc, + /// DataFusion-native metrics exposed through `DataSourceExec`. + pub df_metrics: ExecutionPlanMetricsSet, + /// A shared cache of file readers. + /// + /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache + /// a file reader the first time we read a file. + pub layout_readers: Arc>>, + /// Shared full-file natural splits keyed by file path. + pub natural_splits: Arc>>, + /// Whether the query has output ordering specified + pub has_output_ordering: bool, + + pub expression_convertor: Arc, + pub file_metadata_cache: Option>, + /// Whether to enable expression pushdown into the underlying Vortex scan. + pub projection_pushdown: bool, + pub scan_concurrency: Option, +} + +impl FileOpener for VortexOpener { + fn open(&self, file: PartitionedFile) -> DFResult { + // Calculate the output schema before replacing partition columns with literals so it + // retains the table and partition-field metadata declared by the plan. + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); + let session = self.session.clone(); + let metrics_registry = Arc::clone(&self.metrics_registry); + let labels = vec![ + Label::new(PATH_LABEL, file.path().to_string()), + Label::new(PARTITION_LABEL, self.partition.to_string()), + ]; + + let mut projection = self.projection.clone(); + let mut filter = self.filter.clone(); + + let reader = self.vortex_reader_factory.create_reader(&file, &session)?; + + let reader = + InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); + + let mut file_pruning_predicate = self.file_pruning_predicate.clone(); + let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); + let file_metadata_cache = self.file_metadata_cache.clone(); + + let unified_file_schema = Arc::clone(self.table_schema.file_schema()); + let limit = self.limit; + let layout_readers = Arc::clone(&self.layout_readers); + let natural_splits = Arc::clone(&self.natural_splits); + let has_output_ordering = self.has_output_ordering; + let scan_concurrency = self.scan_concurrency; + + let expr_convertor = Arc::clone(&self.expression_convertor); + let projection_pushdown = self.projection_pushdown; + + let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) + .with_category(MetricCategory::Rows) + .global_counter("num_predicate_creation_errors"); + + // Replace column access for partition columns with literals + #[expect(clippy::disallowed_types)] + let literal_value_cols = self + .table_schema + .table_partition_cols() + .iter() + .map(|f| f.name()) + .cloned() + .zip(file.partition_values.clone()) + .collect::>(); + + let predicate_uses_partition_columns = + file_pruning_predicate.as_ref().is_some_and(|predicate| { + collect_columns(predicate) + .iter() + .any(|column| literal_value_cols.contains_key(column.name())) + }); + + if !literal_value_cols.is_empty() { + projection = projection.try_map_exprs(|expr| { + replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) + })?; + filter = filter + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; + file_pruning_predicate = file_pruning_predicate + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; + } + + Ok(async move { + // FilePruner requires a statistics object even when the rewritten predicate + // only contains partition literals. Supply unknown file-column statistics in + // that case so static and dynamic partition predicates can still prune. + let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) + .then(|| { + file.clone() + .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) + }); + let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); + + let mut file_pruner = file_pruning_predicate + .filter(|_| file.has_statistics() || predicate_uses_partition_columns) + .and_then(|predicate| { + FilePruner::try_new( + Arc::clone(&predicate), + &unified_file_schema, + pruning_file, + predicate_creation_errors, + ) + }); + + // Check if this file should be pruned based on statistics/partition values. + // Returns empty stream if file can be skipped entirely. + if let Some(file_pruner) = file_pruner.as_mut() + && file_pruner.should_prune()? + { + return Ok(stream::empty().boxed()); + } + + let mut open_opts = session + .open_options() + .with_file_size(file.object_meta.size) + .with_metrics_registry(Arc::clone(&metrics_registry)) + .with_labels(labels); + + let cached_footer = file_metadata_cache + .as_ref() + .and_then(|cache| cache.get(file.path())) + .filter(|entry| entry.is_valid_for(&file.object_meta)) + .and_then(|entry| { + entry + .file_metadata + .as_any() + .downcast_ref::() + .map(|vortex_metadata| vortex_metadata.footer().clone()) + }); + let footer_cache_hit = cached_footer.is_some(); + + if let Some(footer) = cached_footer { + open_opts = open_opts.with_footer(footer); + } + + let vxf = open_opts + .open_read(reader) + .await + .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?; + + // On a miss, cache the parsed footer so other partitions and later executions + // skip the footer fetch and parse. `infer_schema`/`infer_stats` also populate + // this cache, but only when planning goes through `VortexFormat`. + if !footer_cache_hit && let Some(cache) = &file_metadata_cache { + cache.put( + file.path(), + CachedFileMetadataEntry::new( + file.object_meta.clone(), + Arc::new(CachedVortexMetadata::new(&vxf)), + ), + ); + } + + // Check if there are rows in this file. If not, we can save + // ourselves some work and return an empty stream. + if vxf.row_count() == 0 { + return Ok(stream::empty().boxed()); + } + + // This is the expected arrow types of the actual columns in the file, which might have different types + // from the unified logical schema or miss + let this_file_schema = Arc::new(calculate_physical_schema( + vxf.dtype(), + &unified_file_schema, + &session.arrow(), + )?); + + let expr_adapter = expr_adapter_factory.create( + Arc::clone(&unified_file_schema), + Arc::clone(&this_file_schema), + )?; + + let simplifier = PhysicalExprSimplifier::new(&this_file_schema); + + // The adapter rewrites the expressions to the local file schema, allowing + // for schema evolution and divergence between the table's schema and individual files. + let filter = filter + .map(|filter| { + // Expression might now reference columns that don't exist in the file, so we can give it + // another simplification pass. + let adapted = expr_adapter.rewrite(Arc::clone(&filter)) + .map_err(|e| exec_datafusion_err!("Failed to adapt filter {filter} in {}: {e}", file.path()))?; + simplifier.simplify(adapted) + .map_err(|e| exec_datafusion_err!("Failed to simplify filter {filter} in {}: {e}", file.path())) + }) + .transpose()?; + let projection = + projection.try_map_exprs(|p| simplifier.simplify(expr_adapter.rewrite(p)?))?; + + let mut native_filters = Vec::new(); + let mut residual_filters = Vec::new(); + if let Some(filter) = &filter { + if filter.data_type(&this_file_schema)? != arrow_schema::DataType::Boolean { + return Err(exec_datafusion_err!("Filter must be Boolean in {}: {filter}", file.path())); + } + for expr in split_conjunction(filter) { + match expr_convertor.try_convert(expr, &this_file_schema) + .map_err(|e| exec_datafusion_err!("Failed to convert filter {expr} in {}: {e}", file.path()))? + { + Some(expr) => native_filters.push(expr), + None => residual_filters.push(Arc::clone(expr)), + } + } + } + let residual_filter = if residual_filters.is_empty() { + None + } else { + Some(conjunction(residual_filters)) + }; + let native_filter = vortex::expr::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 mut residual_columns = None; + let ProcessedProjection { + scan_projection, + scan_reference_schema, + leftover_projection, + } = if let Some(residual) = &residual_filter { + let mut indices = projection.column_indices(); + indices.extend(collect_columns(residual).into_iter().map(|column| column.index())); + indices.sort_unstable(); + indices.dedup(); + let required = ProjectionExprs::from_indices(&indices, &this_file_schema); + let raw = raw_projection(required, &this_file_schema)?; + residual_columns = Some(indices); + ProcessedProjection { + scan_projection: raw.scan_projection, + scan_reference_schema: raw.scan_reference_schema, + leftover_projection: projection, + } + } else if projection_pushdown { + expr_convertor.split_projection( + projection, + &this_file_schema, + output_schema.as_ref(), + )? + } else { + expr_convertor.no_pushdown_projection(projection, &this_file_schema)? + }; + + // The schema of the stream returned from the vortex scan. + // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). + let scan_projection = scan_projection + .optimize_recursive(vxf.dtype()) + .and_then(|projection| projection.bind(vxf.dtype())) + .map_err(|_e| { + exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") + })?; + let scan_dtype = scan_projection.dtype().clone(); + + let stream_schema = + calculate_physical_schema(&scan_dtype, &scan_reference_schema, &session.arrow())?; + + let (leftover_projection, residual_filter) = if let Some(indices) = residual_columns { + ( + leftover_projection.try_map_exprs(|expr| reassign_raw_columns(expr, &indices))?, + residual_filter.map(|expr| reassign_raw_columns(expr, &indices)).transpose()?, + ) + } else { + ( + leftover_projection.try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?, + residual_filter, + ) + }; + 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. + let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { + Entry::Occupied(mut occupied_entry) => { + if let Some(reader) = occupied_entry.get().upgrade() { + tracing::trace!("reusing layout reader for {}", occupied_entry.key()); + reader + } else { + tracing::trace!("creating layout reader for {}", occupied_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + occupied_entry.insert(Arc::downgrade(&reader)); + reader + } + } + Entry::Vacant(vacant_entry) => { + tracing::trace!("creating layout reader for {}", vacant_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!("Failed to create layout reader: {e}")) + })?; + vacant_entry.insert(Arc::downgrade(&reader)); + + reader + } + }; + + let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); + + if let Some(vortex_plan) = file.extensions.get::() { + scan_builder = vortex_plan.apply_to_builder(scan_builder); + } + + if let Some(limit) = limit + && native_filter.is_none() + && residual_filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + + if let Some(concurrency) = scan_concurrency { + scan_builder = scan_builder.with_concurrency(concurrency); + } + + // Set before the byte-range translation below, which computes natural splits for + // the fields the scan's projection and filter reference. + scan_builder = scan_builder + .with_projection(scan_projection) + .with_some_filter(native_filter); + + 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"))?, + end: u64::try_from(file_range.end) + .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + // Full-file scans already cover every natural split. Only translate the + // byte range back into row boundaries when DataFusion has trimmed the file. + let natural_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + &scan_builder, + file.object_meta.size, + )?; + + let Some(row_range) = + split_aligned_row_range(byte_range, natural_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + + scan_builder = scan_builder + .with_row_range(row_range) + // Hand the shared full-file boundaries back to the scan so prepare() + // skips its own layout walk. + .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + } + } + + 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) + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + .into_stream() + .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new(e.with_context(format!( + "Failed to read Vortex file: {}", + file.object_meta.location + )))) + }) + .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}"))?; + } + Ok(batch) + }) + .try_filter(|batch| ready(batch.num_rows() != 0)) + .map(move |batch| { + let batch = projector.project_batch(&batch?)?; + + let (_, columns, row_count) = batch.into_parts(); + RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + ) + .map_err(Into::into) + }) + .boxed(); + + if let Some(file_pruner) = file_pruner { + Ok(PrunableStream::new(file_pruner, stream).boxed()) + } else { + Ok(stream) + } + } + .in_current_span() + .boxed()) + } +} + +/// A file's natural split boundaries plus the precomputed byte each split is assigned to, +/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row +/// boundaries with a binary search instead of re-projecting every split per partition. +/// +/// The boundaries are computed for the fields referenced by the scan's projection and filter. +/// All partitions translate through the first opener's cached entry (the cache lives on the +/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges +/// tiling the file's rows exactly once. +#[derive(Debug)] +pub(crate) struct NaturalSplits { + /// Sorted row boundaries of the natural splits; split `i` covers + /// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the + /// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the + /// per-partition layout walk in `prepare`. + row_boundaries: Arc<[u64]>, + /// For each split, the byte a DataFusion byte range must contain to own it (see + /// [`split_assignment_byte`]); one entry per split, sorted because split midpoints + /// increase monotonically under the row-to-byte projection. + assignment_bytes: Box<[u64]>, +} + +impl NaturalSplits { + fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { + let row_count = row_boundaries.last().copied().unwrap_or_default(); + let assignment_bytes = if row_count == 0 { + Box::default() + } else { + row_boundaries + .windows(2) + .enumerate() + .map(|(idx, boundaries)| { + split_assignment_byte( + idx, + &(boundaries[0]..boundaries[1]), + row_count, + total_size, + ) + }) + .collect() + }; + + debug_assert!(assignment_bytes.is_sorted()); + debug_assert_eq!( + assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), + row_boundaries.len() + ); + + Self { + row_boundaries, + assignment_bytes, + } + } +} + +/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. +fn natural_splits_for_file( + natural_splits: &DashMap>, + path: &Path, + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + if let Some(splits) = natural_splits.get(path) { + return Ok(Arc::clone(splits.value())); + } + + // Compute while holding the entry so concurrent partitions opening the same file wait + // for the winner instead of all walking the layout tree; the redundant walks contend on + // the lazily-initialized layout children and dominate the cost of the computation itself. + match natural_splits.entry(path.clone()) { + Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), + Entry::Vacant(entry) => { + let splits = compute_natural_splits(scan_builder, total_size)?; + entry.insert(Arc::clone(&splits)); + Ok(splits) + } + } +} + +/// Walk the layout tree to compute the file's full natural split boundaries for the fields +/// referenced by the scan's projection and filter. +fn compute_natural_splits( + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + let row_boundaries = scan_builder + .full_file_splits() + .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; + + Ok(Arc::new(NaturalSplits::new( + row_boundaries.into(), + total_size, + ))) +} + +/// Translate a DataFusion byte range to the contiguous natural split ranges it owns. +/// Most splits are assigned by midpoint, but the leading split stays with the range that owns +/// byte 0 so a tiny first byte range still claims the first rows. +fn split_aligned_row_range( + byte_range: Range, + natural_splits: &NaturalSplits, +) -> Option> { + if byte_range.start >= byte_range.end { + return None; + } + + let first_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.start); + let after_last_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.end); + if first_split == after_last_split { + return None; + } + + Some( + natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split], + ) +} + +fn split_assignment_byte( + idx: usize, + split_range: &Range, + row_count: u64, + total_size: u64, +) -> u64 { + if idx == 0 && split_range.start == 0 { + // Byte 0 is the only stable representative for the leading split. A midpoint can fall + // into the next DataFusion byte range and leave the first range with no rows to read. + 0 + } else { + split_midpoint_to_byte(split_range, row_count, total_size) + } +} + +fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64) -> u64 { + let midpoint_row = split_range.start + (split_range.end - split_range.start) / 2; + let midpoint_byte = (u128::from(midpoint_row) * u128::from(total_size)) / u128::from(row_count); + + u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") +} + +/// Remap physical file indices onto the ordered raw columns emitted by the scan. +fn reassign_raw_columns(expr: PhysicalExprRef, indices: &[usize]) -> DFResult { + expr.transform_up(|expr| { + let Some(column) = expr.downcast_ref::() else { + return Ok(Transformed::no(expr)); + }; + let index = indices + .binary_search(&column.index()) + .map_err(|_| exec_datafusion_err!("Missing raw filter/projection column {column}"))?; + Ok(Transformed::yes( + Arc::new(df_expr::Column::new(column.name(), index)) as PhysicalExprRef, + )) + }) + .map(|result| result.data) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-datafusion/src/persistent/opener/tests.rs b/vortex-datafusion/src/persistent/opener/tests.rs new file mode 100644 index 00000000000..2d08ca74478 --- /dev/null +++ b/vortex-datafusion/src/persistent/opener/tests.rs @@ -0,0 +1,1535 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::sync::Arc; +use std::sync::LazyLock; + +use arrow_array::record_batch; +use arrow_schema::Field; +use arrow_schema::Fields; +use arrow_schema::SchemaRef; +use datafusion::arrow::array::DictionaryArray; +use datafusion::arrow::array::Int32Array; +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::array::StringArray; +use datafusion::arrow::array::StructArray; +use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::datatypes::Schema; +use datafusion::arrow::datatypes::UInt32Type; +use datafusion::arrow::util::display::FormatOptions; +use datafusion::arrow::util::pretty::pretty_format_batches_with_options; +use datafusion::logical_expr::col; +use datafusion::logical_expr::lit; +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; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions as df_expr; +use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion_physical_expr::projection::ProjectionExpr; +use insta::assert_snapshot; +use itertools::Itertools; +use object_store::ObjectStore; +use object_store::ObjectStoreExt; +use object_store::memory::InMemory; +use rstest::rstest; +use vortex::VortexSessionDefault; +use vortex::buffer::Buffer; +use vortex::file::WriteOptionsSessionExt; +use vortex::io::VortexWrite; +use vortex::io::object_store::ObjectStoreWrite; +use vortex::metrics::DefaultMetricsRegistry; +use vortex::scan::selection::Selection; +use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; +use vortex::session::VortexSession; + +use super::*; +use crate::VortexAccessPlan; +use crate::convert::exprs::DefaultExpressionConvertor; +use crate::persistent::reader::DefaultVortexReaderFactory; + +static SESSION: LazyLock = LazyLock::new(VortexSession::default); + +/// Test-only expr used to test error reporting. +#[derive(Debug, Eq, Hash, PartialEq)] +struct SnapshotErrorExpr; + +impl fmt::Display for SnapshotErrorExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "snapshot_error") + } +} + +impl PhysicalExpr for SnapshotErrorExpr { + fn data_type(&self, _input_schema: &Schema) -> DFResult { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> DFResult { + Ok(false) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self, f) + } + + fn evaluate(&self, _batch: &RecordBatch) -> DFResult { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + Vec::new() + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> DFResult { + assert!(children.is_empty()); + Ok(self) + } + + fn snapshot(&self) -> DFResult> { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } +} + +fn natural_splits(total_size: u64, split_ranges: &[Range]) -> NaturalSplits { + let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1); + if let Some(first) = split_ranges.first() { + row_boundaries.push(first.start); + row_boundaries.extend(split_ranges.iter().map(|range| range.end)); + } + NaturalSplits::new(row_boundaries.into(), total_size) +} + +#[rstest] +#[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] +#[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] +#[case(1..8, 10, vec![0..1, 1..9, 9..10], Some(1..9))] +#[case(1..4, 16, vec![0..1, 1..2, 2..3, 3..4], None)] +#[case(0..1, 10, vec![0..2, 2..10], Some(0..2))] +#[case(0..2, 2, vec![], None)] +fn test_split_aligned_row_range( + #[case] byte_range: Range, + #[case] total_size: u64, + #[case] split_ranges: Vec>, + #[case] expected: Option>, +) { + assert_eq!( + split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)), + expected + ); +} + +#[test] +fn test_split_aligned_ranges_cover_splits_exactly_once() { + let split_ranges = vec![0..1, 1..4, 4..10, 10..13]; + let byte_ranges = [0..4, 4..8, 8..12, 12..16]; + let natural_splits = natural_splits(16, &split_ranges); + + let assigned = byte_ranges + .into_iter() + .filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits)) + .collect::>(); + + assert_eq!(assigned, vec![0..4, 4..10, 10..13]); + assert_eq!( + assigned + .iter() + .map(|range| range.end - range.start) + .sum::(), + 13 + ); + + let split_starts = split_ranges + .iter() + .map(|range| range.start) + .collect::>(); + let split_ends = split_ranges + .iter() + .map(|range| range.end) + .collect::>(); + + for range in &assigned { + assert!(split_starts.contains(&range.start)); + assert!(split_ends.contains(&range.end)); + } + + for (left, right) in assigned.iter().tuple_windows() { + assert_eq!(left.end, right.start); + } +} + +#[rstest] +#[case(vec![], 10)] +#[case(vec![0], 10)] +#[case(vec![], 0)] +#[case(vec![0], 0)] +fn test_natural_splits_empty_file(#[case] row_boundaries: Vec, #[case] total_size: u64) { + let splits = NaturalSplits::new(row_boundaries.clone().into(), total_size); + + assert!(splits.assignment_bytes.is_empty()); + assert_eq!(splits.row_boundaries.as_ref(), row_boundaries.as_slice()); + assert_eq!(split_aligned_row_range(0..u64::MAX, &splits), None); +} + +#[test] +fn test_split_aligned_row_range_keeps_colliding_assignments_together() { + let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]); + + assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]); + assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2)); + assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); +} + +async fn write_arrow_to_vortex( + object_store: Arc, + path: &str, + rb: RecordBatch, +) -> anyhow::Result { + let schema = rb.schema(); + let array = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; + let path = Path::parse(path)?; + + let mut write = ObjectStoreWrite::new(object_store, &path).await?; + let summary = SESSION + .write_options() + .write(&mut write, array.to_array_stream()) + .await?; + write.shutdown().await?; + + Ok(summary.size()) +} + +fn make_opener( + object_store: Arc, + table_schema: TableSchema, + filter: Option, +) -> VortexOpener { + VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), + projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), + filter, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema, + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + } +} + +#[tokio::test] +async fn test_open() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) + .build(); + + // filter matches partition value + let filter = col("part").eq(lit(1)); + let filter = logical2physical(&filter, table_schema.table_schema()); + + let opener = make_opener( + Arc::clone(&object_store), + table_schema.clone(), + Some(filter), + ); + let stream = opener.open(file.clone()).unwrap().await.unwrap(); + + let data = stream.try_collect::>().await?; + let num_batches = data.len(); + let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); + + assert_eq!((num_batches, num_rows), (1, 3)); + + // filter doesn't matches partition value + let filter = col("part").eq(lit(2)); + let filter = logical2physical(&filter, table_schema.table_schema()); + + let opener = make_opener( + Arc::clone(&object_store), + table_schema.clone(), + Some(filter), + ); + let stream = opener.open(file.clone()).unwrap().await.unwrap(); + + let data = stream.try_collect::>().await?; + let num_batches = data.len(); + let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); + assert_eq!((num_batches, num_rows), (0, 0)); + + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn test_residual_filter_unprojected_column( + #[values(false, true)] projection_pushdown: bool, +) -> 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), Some(2), None, Some(4)]) + )?; + let size = write_arrow_to_vortex(Arc::clone(&store), "residual.vortex", batch.clone()).await?; + let modulo: PhysicalExprRef = 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)))), + )); + 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(batch.schema()), Some(filter)); + 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; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = Arc::new( + batch.schema().as_ref().clone().with_metadata( + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + ); + let table_schema = TableSchema::builder(file_schema) + .with_table_partition_cols(vec![Arc::new( + Field::new("part", DataType::Int32, false).with_metadata( + [("partition".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + )]) + .build(); + let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); + + assert_eq!( + expected_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + assert_eq!( + expected_schema.field(1).metadata().get("partition"), + Some(&"metadata".to_string()) + ); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection = projection.clone(); + opener.projection_pushdown = projection_pushdown; + + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); + } + } + + Ok(()) +} + +#[tokio::test] +async fn test_open_all_valid_nullable_columns_with_nonnullable_table_schema() -> anyhow::Result<()> +{ + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "nullable/file.vortex"; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], + )?; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + let expected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::from(Arc::clone(&expected_schema)); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection_pushdown = projection_pushdown; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema().as_ref(), expected_schema.as_ref()); + } + + Ok(()) +} + +#[tokio::test] +async fn test_file_pruning_replaces_partition_columns_without_file_statistics() -> anyhow::Result<()> +{ + let object_store = Arc::new(InMemory::new()) as Arc; + let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) + .build(); + + let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; + let predicate = Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&partition_column), + Operator::Gt, + df_expr::lit(ScalarValue::Int32(Some(1))), + )) as PhysicalExprRef; + let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![partition_column], + predicate, + )) as PhysicalExprRef; + + let mut opener = make_opener(object_store, table_schema, None); + opener.file_pruning_predicate = Some(dynamic_predicate); + let df_metrics = opener.df_metrics.clone(); + + // The file does not exist and has no statistics. Replacing `part` with 1 + // makes the predicate false, so pruning must happen before any file I/O. + let mut file = PartitionedFile::new("missing.vortex", 1); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(batches.is_empty()); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(0) + ); + + Ok(()) +} + +#[tokio::test] +async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "metrics/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); + statistics.column_statistics[0].null_count = Precision::Exact(0); + let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); + + let mut opener = make_opener(object_store, TableSchema::from(batch.schema()), None); + opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); + let df_metrics = opener.df_metrics.clone(); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(1) + ); + + Ok(()) +} + +#[tokio::test] +async fn test_open_empty_file() -> anyhow::Result<()> { + use futures::TryStreamExt; + + let object_store = Arc::new(InMemory::new()) as Arc; + let data_batch = record_batch!(("a", Int32, Vec::::new())).unwrap(); + let file_path = "part=1/empty.vortex"; + let file_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, data_batch.clone()).await?; + + let file_schema = data_batch.schema(); + // Parallel scans may attach a byte range even for empty files; the + // opener must return early before attempting split-aligned translation. + let file = + PartitionedFile::new_with_range(file_path.to_string(), file_size, 0, file_size as i64); + + let table_schema = TableSchema::from(Arc::clone(&file_schema)); + + let opener = make_opener(object_store, table_schema, None); + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + assert_eq!(data.len(), 0); + + Ok(()) +} + +#[tokio::test] +async fn test_open_populates_file_metadata_cache() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "cached/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let table_schema = TableSchema::from(batch.schema()); + + let cache: Arc = Arc::new( + DefaultCache::::new(64 * 1024 * 1024), + ); + let mut opener = make_opener(Arc::clone(&object_store), table_schema, None); + opener.file_metadata_cache = Some(Arc::clone(&cache)); + + // The first open misses the cache and must write the parsed footer back. + let stream = opener.open(file.clone())?.await?; + stream.try_collect::>().await?; + + let entry = cache + .get(file.path()) + .ok_or_else(|| anyhow::anyhow!("footer was not cached after open"))?; + assert!(entry.is_valid_for(&file.object_meta)); + assert!( + entry + .file_metadata + .as_any() + .downcast_ref::() + .is_some() + ); + + // The second open hits the cache and still returns the same data. + let stream = opener.open(file.clone())?.await?; + let data = stream.try_collect::>().await?; + assert_eq!(data.iter().map(|rb| rb.num_rows()).sum::(), 3); + + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn test_open_files_different_table_schema() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + + let file1 = { + let file1_path = "/path/file1.vortex"; + let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size1 = + write_arrow_to_vortex(Arc::clone(&object_store), file1_path, batch1).await?; + PartitionedFile::new(file1_path.to_string(), data_size1) + }; + + let file2 = { + let file2_path = "/path/file2.vortex"; + let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); + let data_size2 = + write_arrow_to_vortex(Arc::clone(&object_store), file2_path, batch2).await?; + PartitionedFile::new(file2_path.to_string(), data_size2) + }; + + // Table schema has can accommodate both files + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Int32, + true, + )]))); + + let make_opener = |filter| VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone(&object_store))), + projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), + filter: Some(filter), + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: table_schema.clone(), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + let filter = col("a").lt(lit(100_i32)); + let filter = logical2physical(&filter, table_schema.table_schema()); + + let opener1 = make_opener(Arc::clone(&filter)); + let stream = opener1.open(file1)?.await?; + + let format_opts = FormatOptions::new().with_types_info(true); + + let data = stream.try_collect::>().await?; + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+ + | a | + | Int32 | + +-------+ + | 1 | + | 2 | + | 3 | + +-------+ + "); + + let opener2 = make_opener(Arc::clone(&filter)); + let stream = opener2.open(file2)?.await?; + + let data = stream.try_collect::>().await?; + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+ + | a | + | Int32 | + +-------+ + | -1 | + | -2 | + | -3 | + +-------+ + "); + + Ok(()) +} + +#[tokio::test] +// This test verifies that files with different column order than the +// table schema can be opened without errors. The fix ensures that the +// schema mapper is only used for type casting, not for reordering, +// since the vortex projection already handles reordering. +async fn test_schema_different_column_order() -> anyhow::Result<()> { + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + // File has columns in order: c, b, a + let batch = record_batch!( + ("c", Int32, vec![Some(300), Some(301), Some(302)]), + ("b", Int32, vec![Some(200), Some(201), Some(202)]), + ("a", Int32, vec![Some(100), Some(101), Some(102)]) + ) + .unwrap(); + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + let file = PartitionedFile::new(file_path.to_string(), data_size); + + // Table schema has columns in different order: a, b, c + let table_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + + let opener = VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), + projection: ProjectionExprs::from_indices(&[0, 1, 2], &table_schema), + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: TableSchema::from(Arc::clone(&table_schema)), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + let stream = opener.open(file)?.await?; + + let format_opts = FormatOptions::new().with_types_info(true); + let data = stream.try_collect::>().await?; + + // Verify the output has columns in table schema order (a, b, c) + // not file order (c, b, a) + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+-------+-------+ + | a | b | c | + | Int32 | Int32 | Int32 | + +-------+-------+-------+ + | 100 | 200 | 300 | + | 101 | 201 | 301 | + | 102 | 202 | 302 | + +-------+-------+-------+ + "); + + Ok(()) +} + +#[tokio::test] +// This test verifies that expression rewriting doesn't fail when there is +// a nested schema mismatch between the physical file schema and logical +// table schema. +async fn test_adapter_logical_physical_struct_mismatch() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + let file_struct_fields = Fields::from(vec![ + Field::new("field1", DataType::Utf8, true), + Field::new("field2", DataType::Utf8, true), + ]); + let struct_array = StructArray::new( + file_struct_fields.clone(), + vec![ + Arc::new(StringArray::from(vec!["value1", "value2", "value3"])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "my_struct", + DataType::Struct(file_struct_fields), + true, + )])), + vec![Arc::new(struct_array)], + )?; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + // Table schema has an extra utf8 field. + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( + "my_struct", + DataType::Struct(Fields::from(vec![ + Field::new( + "field1", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + Field::new( + "field2", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + Field::new("field3", DataType::Utf8, true), + ])), + true, + )]))); + + let opener = make_opener( + Arc::clone(&object_store), + table_schema.clone(), + // expression references my_struct column which has different fields in each + // field. + Some(logical2physical( + &col("my_struct").is_not_null(), + table_schema.table_schema(), + )), + ); + + // The opener should be able to open the file with a filter on the + // struct column. + let data = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect::>() + .await?; + + assert_eq!(data.len(), 1); + assert_eq!(data[0].num_rows(), 3); + + Ok(()) +} + +#[tokio::test] +// Minimal reproducing test for the schema projection bug. +// Before the fix, this would fail with a cast error when the file schema +// and table schema have different field orders and we project a subset of columns. +async fn test_projection_bug_minimal_repro() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + // File has columns in order: a, b, c with simple types + let batch = record_batch!( + ("a", Int32, vec![Some(1)]), + ("b", Utf8, vec![Some("test")]), + ("c", Int32, vec![Some(2)]) + ) + .unwrap(); + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + // Table schema has columns in DIFFERENT order: c, a, b + // and different types that require casting (Utf8 -> Dictionary) + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![ + Field::new("c", DataType::Int32, true), + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + ]))); + + // Project columns [0, 2] from table schema, which should give us: c, b + // Before the fix, the schema adapter would get confused about which fields + // to select from the file, causing incorrect type mappings. + let projection = vec![0, 2]; + + let opener = VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone(&object_store))), + projection: ProjectionExprs::from_indices(projection.as_ref(), table_schema.file_schema()), + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: table_schema.clone(), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + // This should succeed and return the correctly projected and cast data + let data = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect::>() + .await?; + + // Verify the columns are in the right order and have the right values + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + let format_opts = FormatOptions::new().with_types_info(true); + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+--------------------------+ + | c | b | + | Int32 | Dictionary(UInt32, Utf8) | + +-------+--------------------------+ + | 2 | test | + +-------+--------------------------+ + "); + + Ok(()) +} + +fn make_test_batch_with_10_rows() -> RecordBatch { + record_batch!( + ("a", Int32, (0..=9).map(Some).collect::>()), + ( + "b", + Utf8, + (0..=9).map(|i| Some(format!("r{}", i))).collect::>() + ) + ) + .unwrap() +} + +fn make_test_opener( + object_store: Arc, + schema: SchemaRef, + projection: ProjectionExprs, +) -> VortexOpener { + VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), + projection, + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: TableSchema::from(schema), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + } +} + +#[tokio::test] +// Test that Selection::IncludeByIndex filters to specific row indices. +async fn test_selection_include_by_index() -> anyhow::Result<()> { + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.extensions.insert( + VortexAccessPlan::default().with_selection(Selection::IncludeByIndex( + StrictSortedBuffer::try_new(Buffer::from_iter(vec![1, 3, 5, 7]))?, + )), + ); + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0, 1], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + let format_opts = FormatOptions::new().with_types_info(true); + + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+------+ + | a | b | + | Int32 | Utf8 | + +-------+------+ + | 1 | r1 | + | 3 | r3 | + | 5 | r5 | + | 7 | r7 | + +-------+------+ + "); + + Ok(()) +} + +#[tokio::test] +// Test that Selection::ExcludeByIndex excludes specific row indices. +async fn test_selection_exclude_by_index() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.extensions.insert( + VortexAccessPlan::default().with_selection(Selection::ExcludeByIndex( + StrictSortedBuffer::try_new(Buffer::from_iter(vec![0, 2, 4, 6, 8]))?, + )), + ); + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0, 1], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + let format_opts = FormatOptions::new().with_types_info(true); + + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+------+ + | a | b | + | Int32 | Utf8 | + +-------+------+ + | 1 | r1 | + | 3 | r3 | + | 5 | r5 | + | 7 | r7 | + | 9 | r9 | + +-------+------+ + "); + + Ok(()) +} + +#[tokio::test] +// Test that Selection::All returns all rows. +async fn test_selection_all() -> anyhow::Result<()> { + use vortex::scan::selection::Selection; + + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.extensions + .insert(VortexAccessPlan::default().with_selection(Selection::All)); + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); + assert_eq!(total_rows, 10); + + Ok(()) +} + +#[tokio::test] +// Test that when no extensions are provided, all rows are returned (backward compatibility). +async fn test_selection_no_extensions() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let file = PartitionedFile::new(file_path.to_string(), data_size); + // file.extensions is None by default + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); + assert_eq!(total_rows, 10); + + Ok(()) +} + +#[tokio::test] +async fn test_projection_expr_pushdown() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = record_batch!( + ("a", Int32, vec![Some(1), Some(2), Some(3)]), + ("b", Int32, vec![Some(10), Some(20), Some(30)]) + ) + .unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = batch.schema(); + let table_schema = TableSchema::from(Arc::clone(&file_schema)); + + // Create a projection that includes an arithmetic expression: a + b * 2 + let col_a = df_expr::col("a", &file_schema)?; + let col_b = df_expr::col("b", &file_schema)?; + let two = df_expr::lit(ScalarValue::Int32(Some(2))); + + // b * 2 + let b_times_2 = df_expr::binary(col_b, Operator::Multiply, two, &file_schema)?; + // a + (b * 2) + let a_plus_b_times_2 = df_expr::binary(col_a, Operator::Plus, b_times_2, &file_schema)?; + + let projection = ProjectionExprs::new(vec![ProjectionExpr::new( + a_plus_b_times_2, + "result".to_string(), + )]); + + let opener = VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone(&object_store))), + projection, + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema, + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + // Expected: a + b * 2 + // row 0: 1 + 10 * 2 = 21 + // row 1: 2 + 20 * 2 = 42 + // row 2: 3 + 30 * 2 = 63 + assert_snapshot!(pretty_format_batches_with_options(&data, &FormatOptions::new().with_types_info(true))?.to_string(), @r" + +--------+ + | result | + | Int32 | + +--------+ + | 21 | + | 42 | + | 63 | + +--------+ + "); + + Ok(()) +} + +/// When a Struct contains Dictionary fields, writing to vortex and reading back +/// should preserve the Dictionary type. +#[tokio::test] +async fn test_struct_with_dictionary_roundtrip() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + + let struct_fields = Fields::from(vec![ + Field::new_dictionary("a", DataType::UInt32, DataType::Utf8, true), + Field::new_dictionary("b", DataType::UInt32, DataType::Utf8, true), + ]); + let struct_array = StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(DictionaryArray::::from_iter(["x", "y", "x"])), + Arc::new(DictionaryArray::::from_iter(["p", "p", "q"])), + ], + None, + ); + + let schema = Arc::new(Schema::new(vec![Field::new( + "labels", + DataType::Struct(struct_fields.clone()), + false, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(struct_array)])?; + + let file_path = "/test.vortex"; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0], &schema), + ); + let data: Vec<_> = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect() + .await?; + + assert_eq!( + data[0].schema().field(0).data_type(), + &DataType::Struct(struct_fields), + "Struct(Dictionary) type should be preserved" + ); + Ok(()) +} + +#[derive(Debug)] +struct ModuloAdapterFactory; + +#[derive(Debug)] +struct ModuloAdapter(Arc); + +impl PhysicalExprAdapterFactory for ModuloAdapterFactory { + fn create( + &self, + logical: SchemaRef, + physical: SchemaRef, + ) -> DFResult> { + Ok(Arc::new(ModuloAdapter( + DefaultPhysicalExprAdapterFactory.create(logical, physical)?, + ))) + } +} + +impl datafusion_physical_expr_adapter::PhysicalExprAdapter for ModuloAdapter { + fn rewrite(&self, expr: PhysicalExprRef) -> DFResult { + self.0 + .rewrite(expr)? + .transform_up(|expr| { + if expr + .downcast_ref::() + .is_some_and(|c| c.name() == "b") + { + Ok(Transformed::yes(Arc::new(df_expr::BinaryExpr::new( + expr, + Operator::Modulo, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), + )) as PhysicalExprRef)) + } else { + Ok(Transformed::no(expr)) + } + }) + .map(|result| result.data) + } +} + +struct DelegatingConvertor(DefaultExpressionConvertor); + +impl ExpressionConvertor for DelegatingConvertor { + fn try_convert( + &self, + expr: &PhysicalExprRef, + schema: &Schema, + ) -> DFResult> { + self.0.try_convert(expr, schema) + } +} + +#[rstest] +#[tokio::test] +async fn test_adapted_filter_fallback_and_limit( + #[values(false, true)] projection_pushdown: bool, + #[values(false, true)] custom_convertor: bool, + #[values(false, true)] zero_columns: bool, +) -> anyhow::Result<()> { + let ctx = crate::common_tests::TestSessionContext::new(projection_pushdown); + let batch = record_batch!( + ("a", Int32, vec![10, 20, 30, 40]), + ("b", Int32, vec![Some(1), None, Some(2), Some(4)]) + )?; + ctx.write_arrow_batch("adapted.vortex", &batch).await?; + let metadata = ctx.store.head(&Path::from("adapted.vortex")).await?; + let mut source = crate::VortexSource::new(TableSchema::from(batch.schema()), SESSION.clone()) + .with_projection_pushdown(projection_pushdown); + if custom_convertor { + source = source.with_expression_convertor(Arc::new(DelegatingConvertor( + DefaultExpressionConvertor::default(), + ))); + } + let filter: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("b", 1)), + Operator::Eq, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))), + )); + let accepted = datafusion_datasource::file::FileSource::try_pushdown_filters( + &source, + vec![filter], + &datafusion_common::config::ConfigOptions::new(), + )?; + assert!(matches!( + accepted.filters.as_slice(), + [datafusion_physical_plan::filter_pushdown::PushedDown::Yes] + )); + let mut source = accepted + .updated_node + .ok_or_else(|| anyhow::anyhow!("Expected updated source"))?; + // Alias a computed output to b, which is also an unprojected residual input. + let projection = if zero_columns { + ProjectionExprs::from(Vec::::new()) + } else { + vec![ProjectionExpr { + expr: Arc::new(df_expr::BinaryExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + Operator::Plus, + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), + )), + alias: "b".into(), + }] + .into() + }; + source = source + .try_pushdown_projection(&projection)? + .ok_or_else(|| anyhow::anyhow!("Expected projected source"))?; + let config = datafusion_datasource::file_scan_config::FileScanConfigBuilder::new( + datafusion_execution::object_store::ObjectStoreUrl::local_filesystem(), + source, + ) + .with_expr_adapter(Some(Arc::new(ModuloAdapterFactory))) + .with_limit(Some(1)) + .with_file(PartitionedFile::new("adapted.vortex", metadata.size)) + .build(); + let plan: Arc = Arc::new( + datafusion_datasource::source::DataSourceExec::new(Arc::new(config)), + ); + let batches = datafusion_physical_plan::collect(plan, ctx.session.task_ctx()).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + if zero_columns { + assert!(batches.iter().all(|batch| batch.num_columns() == 0)); + } else { + assert_batches_eq!(["+----+", "| b |", "+----+", "| 31 |", "+----+"], &batches); + } + 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!( + DefaultExpressionConvertor::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_convertor = Arc::new(ResidualConvertor); + 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 ResidualConvertor; +impl ExpressionConvertor for ResidualConvertor { + fn try_convert( + &self, + _expr: &PhysicalExprRef, + _schema: &Schema, + ) -> DFResult> { + Ok(None) + } +} + +#[rstest] +#[tokio::test] +async fn test_physical_in_list_residual( + #[values(false, true)] negated: bool, + #[values(0, 1, 2)] list_kind: usize, +) -> anyhow::Result<()> { + let store: Arc = Arc::new(InMemory::new()); + let batch = record_batch!( + ("id", Int32, vec![0, 1, 2]), + ("a", Int32, vec![Some(1), Some(2), None]), + ("b", Int32, vec![Some(1), None, Some(2)]) + )?; + let size = write_arrow_to_vortex(Arc::clone(&store), "in-list.vortex", batch.clone()).await?; + let list: Vec = match list_kind { + 0 => vec![], + 1 => vec![Arc::new(df_expr::Column::new("b", 2))], + _ => vec![ + Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), + Arc::new(df_expr::Literal::new(ScalarValue::Int32(None))), + ], + }; + let filter: PhysicalExprRef = Arc::new(df_expr::InListExpr::try_new( + Arc::new(df_expr::Column::new("a", 1)), + list, + negated, + &batch.schema(), + )?); + assert!( + DefaultExpressionConvertor::default() + .try_convert(&filter, &batch.schema())? + .is_none() + ); + 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("in-list.vortex", size))? + .await? + .try_collect::>() + .await?; + let actual = concat_batches(&expected.schema(), &actual)?; + assert_eq!(actual, expected); + Ok(()) +} + +#[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_convertor = Arc::new(ResidualConvertor); + 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_convertor = Arc::new(ResidualConvertor); + 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..f8ed18f3d10 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -5,8 +5,10 @@ use std::fmt::Formatter; use std::sync::Arc; use std::sync::Weak; +use arrow_schema::DataType; use datafusion_common::Result as DFResult; use datafusion_common::config::ConfigOptions; +use datafusion_common::exec_datafusion_err; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; @@ -129,23 +131,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 convertor 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 +191,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, @@ -472,16 +475,20 @@ impl FileSource for VortexSource { let supported_filters = filters .into_iter() .map(|expr| { + if expr.data_type(self.table_schema.table_schema())? != DataType::Boolean { + return Err(exec_datafusion_err!("Filter must be Boolean: {expr}")); + } if self .expression_convertor - .can_be_pushed_down(&expr, self.table_schema.file_schema()) + .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() @@ -573,12 +580,12 @@ mod tests { } 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) + fn try_convert( + &self, + expr: &PhysicalExprRef, + schema: &Schema, + ) -> DFResult> { + self.inner.try_convert(expr, schema) } fn split_projection( diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 35dd745461d..9077f7a4a6d 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,90 @@ 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)")] +#[case::overflow("id, a + CAST(1 AS INT) AS n", "TRUE")] +#[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)] projection_pushdown: bool, + #[values(false, true)] predicate_pushdown: bool, +) -> anyhow::Result<()> { + let options = crate::VortexTableOptions { + projection_pushdown, + predicate_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/source.rs b/vortex-datafusion/src/v2/source.rs index 6437cb24dc5..f78da1c2a8f 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -121,7 +121,6 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::convert::exprs::DefaultExpressionConvertor; use crate::convert::exprs::ExpressionConvertor; use crate::convert::exprs::ProcessedProjection; -use crate::convert::exprs::make_vortex_predicate; use crate::convert::stats::stats_set_to_df; /// Builder for [`VortexDataSource`]. @@ -555,6 +554,7 @@ impl DataSource for VortexDataSource { let ProcessedProjection { scan_projection, leftover_projection, + .. } = convertor.split_projection(projection.clone(), input_schema, &projected_schema)?; // Compose with the initial projection so the scan operates on the original @@ -607,10 +607,14 @@ impl DataSource for VortexDataSource { // 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| convertor.try_convert(expr, input_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 +629,8 @@ 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()); // Combine with existing filter. let new_filter = match (&self.filter, vortex_pred) { From 9f509a34c77f2b50c8882dbe0b64dd7349fe3623 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 7 Sep 2026 13:03:52 +0100 Subject: [PATCH 02/10] inline fn Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 65d6148f7be..8ccfff9f8e1 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -45,6 +45,9 @@ 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. @@ -640,18 +643,13 @@ fn supported_data_types(dt: &DataType) -> bool { 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), + [input, dimension] + if dimension + .downcast_ref::() + .is_some_and(|literal| matches!(literal.value(), ScalarValue::Int64(Some(1)))) => + { + 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; From 259c48c2a78908aaee8959d3ad2a0a842bfc6dae Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 7 Sep 2026 16:26:10 +0100 Subject: [PATCH 03/10] restore arithmetic conversion Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 28 +++------ vortex-datafusion/src/convert/exprs/tests.rs | 63 ++++++++++++++++++- .../src/persistent/opener/mod.rs | 2 +- vortex-datafusion/src/persistent/tests.rs | 6 +- 4 files changed, 76 insertions(+), 23 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 8ccfff9f8e1..1d267df6e90 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -60,9 +60,10 @@ pub struct ProcessedProjection { /// Trait for converting DataFusion expressions to Vortex ones. /// -/// Custom convertors implement a single schema-aware decision. Successful conversion must -/// preserve DataFusion values, nulls, and evaluation errors. Unsupported expressions remain -/// in DataFusion, including when a file's schema adapter introduces them. +/// Custom convertors implement a single schema-aware decision. Conversion should preserve +/// DataFusion values, nulls, and evaluation errors; see [`DefaultExpressionConvertor`] for +/// the temporary arithmetic exception. Unsupported expressions remain in DataFusion, +/// including when a file's schema adapter introduces them. /// /// # Implementing a custom convertor /// @@ -87,7 +88,7 @@ pub struct ProcessedProjection { /// } /// } pub trait ExpressionConvertor: Send + Sync { - /// Convert an expression with equivalent DataFusion behavior for this schema. + /// 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 @@ -176,8 +177,9 @@ pub(crate) fn raw_projection( /// The default schema-aware DataFusion expression convertor. /// -/// Casts and operators are accepted only for the SQL semantics implemented by Vortex. -/// Other valid expressions are evaluated by DataFusion. +/// 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 DefaultExpressionConvertor { session: VortexSession, } @@ -206,19 +208,7 @@ impl DefaultExpressionConvertor { }; let left_type = binary.left().data_type(schema)?; let right_type = binary.right().data_type(schema)?; - if operator.is_arithmetic() { - // DataFusion's integer +/-/* can wrap; Vortex arithmetic is checked. - // Decimal coercion and result precision also differ. - let supported = matches!( - (&left_type, &right_type), - (DataType::Float32, DataType::Float32) | (DataType::Float64, DataType::Float64) - ) || (*binary.op() == DFOperator::Divide - && left_type.is_integer() - && left_type == right_type); - if !supported { - return Ok(None); - } - } else if *binary.op() == DFOperator::And || *binary.op() == DFOperator::Or { + if *binary.op() == DFOperator::And || *binary.op() == DFOperator::Or { if left_type != DataType::Boolean || right_type != DataType::Boolean { return Err(exec_datafusion_err!( "Boolean operator requires Boolean operands: {expr}" diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 1ffa7fea73e..43680f86473 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -902,7 +902,8 @@ fn test_cast_options_fall_back() -> DFResult<()> { #[case::add(DFOperator::Plus)] #[case::sub(DFOperator::Minus)] #[case::mul(DFOperator::Multiply)] -fn test_integer_overflow_modes_fall_back( +#[case::div(DFOperator::Divide)] +fn test_integer_arithmetic_pushdown_ignores_overflow_mode( #[case] op: DFOperator, #[values(false, true)] checked: bool, ) -> DFResult<()> { @@ -918,11 +919,69 @@ fn test_integer_overflow_modes_fall_back( assert!( DefaultExpressionConvertor::default() .try_convert(&expr, &schema)? - .is_none() + .is_some() ); 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 = arrow_array::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)]); diff --git a/vortex-datafusion/src/persistent/opener/mod.rs b/vortex-datafusion/src/persistent/opener/mod.rs index c69efe08d19..ea7b5bd6ae2 100644 --- a/vortex-datafusion/src/persistent/opener/mod.rs +++ b/vortex-datafusion/src/persistent/opener/mod.rs @@ -496,7 +496,7 @@ impl FileOpener for VortexOpener { }) .boxed(); - if let Some(file_pruner) = file_pruner { + if let Some(file_pruner) = file_pruner && file_pruner.is_watching() { Ok(PrunableStream::new(file_pruner, stream).boxed()) } else { Ok(stream) diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 9077f7a4a6d..60d5a85b869 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -628,7 +628,11 @@ async fn arrow_uuid_extension_roundtrip_nested_struct() -> anyhow::Result<()> { #[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)")] -#[case::overflow("id, a + CAST(1 AS INT) AS n", "TRUE")] +// 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")] From 95101b9e9e49192dbbfea492e8a7e36577d7cb32 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 8 Sep 2026 13:46:09 +0100 Subject: [PATCH 04/10] Simplify Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 174 ++++--- vortex-datafusion/src/convert/exprs/tests.rs | 332 +++++++++++-- vortex-datafusion/src/convert/mod.rs | 1 - vortex-datafusion/src/convert/scalars.rs | 438 +----------------- .../src/persistent/opener/mod.rs | 57 +-- .../src/persistent/opener/tests.rs | 14 +- 6 files changed, 436 insertions(+), 580 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 1d267df6e90..589e9e0d688 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -21,6 +21,7 @@ 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::array::VortexSessionExecute; use vortex::dtype::DType; @@ -109,22 +110,21 @@ pub trait ExpressionConvertor: Send + Sync { 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.iter() { + 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.clone(), input_schema); + return self.no_pushdown_projection(source_projection, input_schema); }; scan_projection.push((projection.alias.clone(), expr)); } - // Duplicate output names cannot identify native pack fields unambiguously. - let mut names = scan_projection - .iter() - .map(|(name, _)| name) - .collect::>(); - names.sort_unstable(); - if names.windows(2).any(|names| names[0] == names[1]) { - return self.no_pushdown_projection(source_projection, input_schema); - } Ok(ProcessedProjection { scan_projection: pack(scan_projection, Nullability::NonNullable), scan_reference_schema: output_schema.clone(), @@ -203,12 +203,13 @@ impl DefaultExpressionConvertor { input_dtype: &DType, ) -> DFResult> { let converted = if let Some(binary) = expr.downcast_ref::() { - let Ok(operator) = try_operator_from_df(binary.op()) else { + let Some(operator) = try_operator_from_df(binary.op()) else { return Ok(None); }; + let boolean_operator = matches!(operator, Operator::And | Operator::Or); let left_type = binary.left().data_type(schema)?; let right_type = binary.right().data_type(schema)?; - if *binary.op() == DFOperator::And || *binary.op() == DFOperator::Or { + if boolean_operator { if left_type != DataType::Boolean || right_type != DataType::Boolean { return Err(exec_datafusion_err!( "Boolean operator requires Boolean operands: {expr}" @@ -223,7 +224,7 @@ impl DefaultExpressionConvertor { ) else { return Ok(None); }; - if matches!(binary.op(), DFOperator::And | DFOperator::Or) + if boolean_operator && (label_infallible(&left).get(&left) != Some(&true) || label_infallible(&right).get(&right) != Some(&true)) { @@ -245,19 +246,23 @@ impl DefaultExpressionConvertor { } else if let Some(literal) = expr.downcast_ref::() { let field = literal.return_field(schema)?; let array = literal.value().to_array()?; - if self.session.arrow().from_arrow_field(&field).is_err() { - return Ok(None); - } if array.len() != 1 { return Err(exec_datafusion_err!( - "Literal must contain exactly one value: {expr}" + "Literal must contain exactly one value, found {}", + array.len() )); } - let array = self - .session - .arrow() - .from_arrow_array(array, &field) - .map_err(|e| exec_datafusion_err!("Failed to convert literal {expr}: {e}"))?; + let array = match self.session.arrow().from_arrow_array(array, &field) { + Ok(array) => array, + Err(error) => { + if self.session.arrow().from_arrow_field(&field).is_err() { + return Ok(None); + } + return Err(exec_datafusion_err!( + "Failed to convert literal {expr}: {error}" + )); + } + }; lit(array .execute_scalar(0, &mut self.session.create_execution_ctx()) .map_err(|e| exec_datafusion_err!("Failed to evaluate literal {expr}: {e}"))?) @@ -324,11 +329,6 @@ impl DefaultExpressionConvertor { if case_expr.expr().is_some() { return Ok(None); } - if case_expr.when_then_expr().is_empty() { - return Err(exec_datafusion_err!( - "CASE requires at least one WHEN clause" - )); - } let mut pairs = Vec::with_capacity(case_expr.when_then_expr().len()); for (when, then) in case_expr.when_then_expr() { if when.data_type(schema)? != DataType::Boolean { @@ -417,44 +417,36 @@ impl DefaultExpressionConvertor { return Ok(Some(result)); } - let octet_length = - ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some(); - let array_length = - ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some(); - let input = if octet_length { - let [input] = scalar_fn.args() else { - return Err(exec_datafusion_err!( - "octet_length requires exactly one argument" - )); - }; - 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 Ok(None); - } - input - } else if array_length { - if scalar_fn.args().is_empty() || scalar_fn.args().len() > 2 { - return Err(exec_datafusion_err!( - "array_length requires one or two arguments" - )); - } - let Some(input) = array_length_input(scalar_fn) else { + 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" + )); + }; + 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 Ok(None); + } + (input, byte_length) + } else if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { + let Some(input) = array_length_input(scalar_fn)? else { + return Ok(None); + }; + if !matches!( + input.data_type(schema)?, + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ) { + return Ok(None); + } + (input, list_length) + } else { return Ok(None); }; - if !matches!( - input.data_type(schema)?, - DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) - ) { - return Ok(None); - } - input - } else { - return Ok(None); - }; let Some(input) = self.convert_expr(input, schema, input_dtype)? else { return Ok(None); }; @@ -465,14 +457,7 @@ impl DefaultExpressionConvertor { )) else { return Ok(None); }; - Ok(Some(cast( - if octet_length { - byte_length(input) - } else { - list_length(input) - }, - return_dtype, - ))) + Ok(Some(cast(length(input), return_dtype))) } } @@ -549,20 +534,20 @@ fn supported_cast(cast: &df_expr::CastExpr, schema: &Schema) -> DFResult { )) } -fn try_operator_from_df(value: &DFOperator) -> DFResult { +fn try_operator_from_df(value: &DFOperator) -> Option { 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::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 @@ -593,12 +578,7 @@ fn try_operator_from_df(value: &DFOperator) -> DFResult { | 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}" - )) - } + | DFOperator::Colon => None, } } @@ -630,16 +610,20 @@ fn supported_data_types(dt: &DataType) -> bool { /// `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> { +/// Calls with other arities return errors. +fn array_length_input(scalar_fn: &ScalarFunctionExpr) -> DFResult>> { match scalar_fn.args() { - [input] => Some(input), + [input] => Ok(Some(input)), [input, dimension] if dimension .downcast_ref::() .is_some_and(|literal| matches!(literal.value(), ScalarValue::Int64(Some(1)))) => { - Some(input) + Ok(Some(input)) } - _ => None, + [_, _] => 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 index 43680f86473..81f86aadc84 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -4,15 +4,20 @@ use std::sync::Arc; use arrow_array::Array; +use arrow_array::FixedSizeBinaryArray; +use arrow_array::StructArray; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::Schema; use arrow_schema::TimeUnit as ArrowTimeUnit; +use arrow_schema::extension::Uuid as ArrowUuid; use datafusion::arrow::array::AsArray; use datafusion::arrow::datatypes::Int32Type; 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; @@ -22,10 +27,14 @@ 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 { @@ -71,6 +80,79 @@ fn array_length_expr(args: Vec>, schema: &Schema) -> 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 { + expr: 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)))), + )), + alias: "duplicate".into(), + }, + ProjectionExpr { + expr: Arc::new(df_expr::Column::new("a", 0)), + alias: "duplicate".into(), + }, + ProjectionExpr { + expr: Arc::new(df_expr::Column::new("b", 1)), + alias: "duplicate".into(), + }, + ]); + let output_schema = projection.project_schema(&schema)?; + let processed = + FailingConvertor.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] +#[case::past_end(1)] +#[case::max(usize::MAX)] +fn test_raw_projection_out_of_bounds(#[case] index: usize) -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let projection = vec![ProjectionExpr { + expr: Arc::new(df_expr::Column::new("a", index)), + alias: "a".into(), + }] + .into(); + let error = raw_projection(projection, &schema) + .err() + .ok_or_else(|| exec_datafusion_err!("Expected an out-of-bounds error"))?; + assert!( + error + .to_string() + .contains(&format!("Projection column index {index} is out of bounds")), + "{error}" + ); + Ok(()) +} + #[rstest] fn test_predicate_rejects_cast_over_modulo(test_schema: Schema) { let modulo = Arc::new(df_expr::BinaryExpr::new( @@ -129,8 +211,7 @@ 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); + assert_eq!(try_operator_from_df(&df_op), Some(expected_vortex_op)); } #[rstest] @@ -139,14 +220,7 @@ fn test_operator_conversion_supported( #[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") - ); + assert_eq!(try_operator_from_df(&df_op), None); } #[test] @@ -180,6 +254,126 @@ fn test_expr_from_df_literal() { assert_snapshot!(result.display_tree().to_string(), @"vortex.literal(42i32)"); } +fn convert_literal(expr: Arc) -> anyhow::Result { + let converted = DefaultExpressionConvertor::default() + .try_convert(&expr, &Schema::empty())? + .ok_or_else(|| anyhow::anyhow!("Expected native conversion for {expr}"))?; + converted + .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() { let left = Arc::new(df_expr::Column::new("left", 0)) as Arc; @@ -371,17 +565,22 @@ fn test_can_be_pushed_down_literal_supported(test_schema: 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; - +#[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!( - !DefaultExpressionConvertor::default() - .try_convert(&lit_expr, &test_schema) - .unwrap() - .is_some() + DefaultExpressionConvertor::default() + .try_convert(&lit_expr, &test_schema)? + .is_none() ); + Ok(()) } #[rstest] @@ -551,19 +750,48 @@ fn test_can_be_pushed_down_array_length_dimension_one_supported(test_schema: Sch } #[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. +#[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_array_length_unsupported_dimension( + test_schema: Schema, + #[case] dimension: Arc, +) -> DFResult<()> { 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!( + DefaultExpressionConvertor::default() + .try_convert(&array_length, &test_schema)? + .is_none() + ); + 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!( - !DefaultExpressionConvertor::default() - .try_convert(&array_length, &test_schema) - .unwrap() - .is_some() + DefaultExpressionConvertor::default() + .try_convert(&expr, &Schema::empty()) + .is_err() ); } @@ -1112,7 +1340,7 @@ fn test_native_nested_list_length( 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(arrow_array::StructArray::new( + let payload = Arc::new(StructArray::new( vec![field].into(), vec![lists], nullable_parent.then(|| NullBuffer::from(vec![true, false, true])), @@ -1166,6 +1394,40 @@ fn test_native_case_null_conditions() -> anyhow::Result<()> { 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &Schema::empty())? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &Schema::empty()) + .is_err() + ); + Ok(()) +} + #[test] fn test_malformed_literal_returns_error() { let value = ScalarValue::Decimal128(Some(1), 0, 0); @@ -1177,6 +1439,18 @@ fn test_malformed_literal_returns_error() { ); } +#[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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &Schema::empty()) + .is_err() + ); +} + #[rstest] #[case::and(DFOperator::And, DFOperator::NotEq)] #[case::or(DFOperator::Or, DFOperator::Eq)] diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 476f32afa49..72ee2438220 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -19,7 +19,6 @@ pub(crate) mod stats; pub use exprs::DefaultExpressionConvertor; pub use exprs::ExpressionConvertor; pub use exprs::ProcessedProjection; -pub use scalars::scalar_from_df; /// First-party trait for implementing conversion from DataFusion types to Vortex types. pub trait FromDataFusion: Sized { diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index c2b1dfc79b2..a046d1ee005 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -3,33 +3,22 @@ use std::sync::Arc; -use arrow_array::Array; -use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::StructArray; use arrow_schema::Field; use arrow_schema::Fields; use datafusion_common::ScalarValue; -use vortex::array::VortexSessionExecute; -use vortex::buffer::ByteBuffer; use vortex::dtype::DType; -use vortex::dtype::DecimalDType; use vortex::dtype::NativeDecimalType; -use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::half::f16; use vortex::dtype::i256; -use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; -use vortex::error::vortex_panic; use vortex::extension::datetime::AnyTemporal; use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; -use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; use crate::convert::TryToDataFusion; @@ -185,139 +174,6 @@ impl TryToDataFusion for Scalar { } } -/// Converts a DataFusion [`ScalarValue`] to a Vortex [`Scalar`], resolving Arrow types through -/// `session`'s [`ArrowSession`](vortex_arrow::ArrowSession). -pub fn scalar_from_df(value: &ScalarValue, session: &VortexSession) -> Scalar { - let arrow = session.arrow(); - match value { - ScalarValue::Null => Scalar::null(DType::Null), - ScalarValue::Boolean(b) => b - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))), - ScalarValue::Float16(f) => f - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F16, Nullability::Nullable))), - ScalarValue::Float32(f) => f - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable))), - ScalarValue::Float64(f) => f - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))), - ScalarValue::Int8(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable))), - ScalarValue::Int16(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable))), - ScalarValue::Int32(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))), - ScalarValue::Int64(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))), - ScalarValue::UInt8(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable))), - ScalarValue::UInt16(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable))), - ScalarValue::UInt32(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable))), - ScalarValue::UInt64(i) => i - .map(Scalar::from) - .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable))), - ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => s - .as_ref() - .map(|s| Scalar::from(s.as_str())) - .unwrap_or_else(|| Scalar::null(DType::Utf8(Nullability::Nullable))), - ScalarValue::Binary(b) - | ScalarValue::BinaryView(b) - | ScalarValue::LargeBinary(b) - | ScalarValue::FixedSizeBinary(_, b) => b - .as_ref() - .map(|b| Scalar::binary(ByteBuffer::from(b.clone()), Nullability::Nullable)) - .unwrap_or_else(|| Scalar::null(DType::Binary(Nullability::Nullable))), - ScalarValue::Date32(v) - | ScalarValue::Time32Second(v) - | ScalarValue::Time32Millisecond(v) => { - let dtype = arrow - .from_arrow_datatype(&value.data_type(), Nullability::Nullable) - .vortex_expect("arrow data type to dtype"); - Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) - .vortex_expect("unable to create a time `Scalar`") - } - ScalarValue::Date64(v) - | ScalarValue::Time64Microsecond(v) - | ScalarValue::Time64Nanosecond(v) - | ScalarValue::TimestampSecond(v, _) - | ScalarValue::TimestampMillisecond(v, _) - | ScalarValue::TimestampMicrosecond(v, _) - | ScalarValue::TimestampNanosecond(v, _) => { - let dtype = arrow - .from_arrow_datatype(&value.data_type(), Nullability::Nullable) - .vortex_expect("arrow data type to dtype"); - Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) - .vortex_expect("unable to create a time `Scalar`") - } - ScalarValue::Decimal32(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I32(*value), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } - } - ScalarValue::Decimal64(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I64(*value), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } - } - ScalarValue::Decimal128(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I128(*value), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } - } - ScalarValue::Decimal256(decimal, precision, scale) => { - let decimal_dtype = DecimalDType::new(*precision, *scale); - let nullable = Nullability::Nullable; - if let Some(value) = decimal { - Scalar::decimal( - DecimalValue::I256(i256::from_le_bytes(value.to_le_bytes())), - decimal_dtype, - Nullability::Nullable, - ) - } else { - Scalar::null(DType::Decimal(decimal_dtype, nullable)) - } - } - ScalarValue::Dictionary(_, v) => scalar_from_df(v.as_ref(), session), - ScalarValue::Struct(array) => struct_from_df(array, session), - _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"), - } -} - /// Converts a Vortex struct scalar to a DataFusion `ScalarValue::Struct`. fn struct_to_df(scalar: &Scalar) -> VortexResult { let scalar = scalar.as_struct(); @@ -359,44 +215,11 @@ fn struct_to_df(scalar: &Scalar) -> VortexResult { Ok(ScalarValue::Struct(Arc::new(struct_array))) } -/// Converts a DataFusion `ScalarValue::Struct` (a one-row struct array) to a Vortex struct scalar. -/// -/// The struct dtype comes from the Arrow `Fields`, so the children must be converted from those -/// same fields. Going through `ScalarValue` instead would drop each field's `ARROW:extension:name` -/// (and its declared nullability), yielding storage-typed children that -/// [`Scalar::struct_`] rejects against an extension-typed struct dtype. -fn struct_from_df(array: &StructArray, session: &VortexSession) -> Scalar { - let arrow = session.arrow(); - let dtype = arrow - .from_arrow_datatype(array.data_type(), Nullability::Nullable) - .vortex_expect("arrow data type to dtype"); - if array.is_null(0) { - Scalar::null(dtype) - } else { - let mut ctx = session.create_execution_ctx(); - let children = array - .columns() - .iter() - .zip(array.fields().iter()) - .map(|(column, field)| { - arrow - .from_arrow_array(ArrowArrayRef::clone(column), field) - .and_then(|column| column.execute_scalar(0, &mut ctx)) - .unwrap_or_else(|e| { - vortex_panic!("cannot convert struct field to a Vortex scalar: {e}") - }) - }) - .collect::>(); - Scalar::struct_(dtype, children) - } -} - #[cfg(test)] mod tests { use datafusion_common::ScalarValue; use datafusion_common::arrow::datatypes::i256 as arrow_i256; use rstest::rstest; - use vortex::VortexSessionDefault; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::DecimalDType; @@ -411,11 +234,6 @@ mod tests { use super::*; - /// Test shim: convert with a default `VortexSession` passed explicitly. - fn from_df(value: &ScalarValue) -> Scalar { - scalar_from_df(value, &VortexSession::default()) - } - #[rstest] #[case::u8_some(Scalar::from(42u8), ScalarValue::UInt8(Some(42)))] #[case::u8_null( @@ -570,160 +388,6 @@ mod tests { assert_eq!(result, expected_df_scalar); } - #[rstest] - #[case::from_df_null(ScalarValue::Null, Scalar::null(DType::Null))] - #[case::from_df_bool_some(ScalarValue::Boolean(Some(true)), Scalar::from(true))] - #[case::from_df_bool_null( - ScalarValue::Boolean(None), - Scalar::null(DType::Bool(Nullability::Nullable)) - )] - #[case::from_df_i32_some(ScalarValue::Int32(Some(42)), Scalar::from(42i32))] - #[case::from_df_i32_null( - ScalarValue::Int32(None), - Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) - )] - #[case::from_df_f64_some(ScalarValue::Float64(Some(2.5)), Scalar::from(2.5f64))] - #[case::from_df_f64_null( - ScalarValue::Float64(None), - Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)) - )] - #[case::from_df_utf8_some(ScalarValue::Utf8(Some("test".to_string())), Scalar::from("test"))] - #[case::from_df_utf8_null( - ScalarValue::Utf8(None), - Scalar::null(DType::Utf8(Nullability::Nullable)) - )] - #[case::from_df_binary_some(ScalarValue::Binary(Some(vec![1, 2, 3])), Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3]), Nullability::Nullable))] - #[case::from_df_binary_null( - ScalarValue::Binary(None), - Scalar::null(DType::Binary(Nullability::Nullable)) - )] - fn test_from_datafusion_scalars( - #[case] df_scalar: ScalarValue, - #[case] expected_vortex: Scalar, - ) { - let result = from_df(&df_scalar); - assert_eq!(result.dtype(), expected_vortex.dtype()); - assert_eq!(result.is_null(), expected_vortex.is_null()); - - // For non-null values, convert both back to DataFusion for comparison - if !result.is_null() { - let result_df = result.try_to_df().unwrap(); - let expected_df = expected_vortex.try_to_df().unwrap(); - assert_eq!(result_df, expected_df); - } - } - - #[rstest] - #[case::decimal128_some(ScalarValue::Decimal128(Some(12345), 10, 2))] - #[case::decimal128_null(ScalarValue::Decimal128(None, 10, 2))] - #[case::decimal256_some(ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10))] - #[case::decimal256_null(ScalarValue::Decimal256(None, 50, 10))] - fn test_from_datafusion_decimals(#[case] df_scalar: ScalarValue) { - let result = from_df(&df_scalar); - match &df_scalar { - ScalarValue::Decimal128(value, precision, scale) => { - if let DType::Decimal(decimal_type, _) = result.dtype() { - assert_eq!(decimal_type.precision(), *precision); - assert_eq!(decimal_type.scale(), *scale); - if value.is_some() { - assert!(!result.is_null()); - } else { - assert!(result.is_null()); - } - } else { - panic!("Expected decimal type"); - } - } - ScalarValue::Decimal256(value, precision, scale) => { - if let DType::Decimal(decimal_type, _) = result.dtype() { - assert_eq!(decimal_type.precision(), *precision); - assert_eq!(decimal_type.scale(), *scale); - if value.is_some() { - assert!(!result.is_null()); - } else { - assert!(result.is_null()); - } - } else { - panic!("Expected decimal type"); - } - } - _ => panic!("Unexpected scalar type"), - } - } - - #[rstest] - #[case::date32(ScalarValue::Date32(Some(18628)))] // 2021-01-01 - #[case::date64(ScalarValue::Date64(Some(1609459200000)))] // 2021-01-01 in milliseconds - #[case::time32_second(ScalarValue::Time32Second(Some(3661)))] // 01:01:01 - #[case::time32_millisecond(ScalarValue::Time32Millisecond(Some(3661000)))] // 01:01:01 - #[case::time64_microsecond(ScalarValue::Time64Microsecond(Some(3661000000)))] // 01:01:01 - #[case::time64_nanosecond(ScalarValue::Time64Nanosecond(Some(3661000000000)))] // 01:01:01 - #[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 - ))] - fn test_from_datafusion_temporals(#[case] df_scalar: ScalarValue) { - let result = from_df(&df_scalar); - - // All temporal types should convert to extension types - if let DType::Extension(_) = result.dtype() { - assert!(!result.is_null()); - } else { - panic!( - "Expected extension type for temporal scalar, got: {:?}", - result.dtype() - ); - } - } - - #[rstest] - #[case::u32(Scalar::from(42u32))] - #[case::i64(Scalar::from(-123i64))] - #[case::f64(Scalar::from(2.5f64))] - #[case::bool(Scalar::from(true))] - #[case::utf8(Scalar::from("hello world"))] - #[case::null_type(Scalar::null(DType::Null))] - #[case::null_i32(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)))] - #[case::decimal128(Scalar::decimal( - DecimalValue::I128(12345), - DecimalDType::new(10, 2), - Nullability::NonNullable - ))] - #[case::binary(Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4, 5]), Nullability::NonNullable))] - fn test_round_trip_conversions(#[case] original: Scalar) { - let df_scalar = original.try_to_df().unwrap(); - let round_trip = from_df(&df_scalar); - - // Check that core types match (ignoring nullability differences that can occur in round-trip) - assert!( - original.dtype().eq_ignore_nullability(round_trip.dtype()), - "DType mismatch for scalar: {:?} vs {:?}", - original.dtype(), - round_trip.dtype() - ); - - assert_eq!( - original.is_null(), - round_trip.is_null(), - "Null status mismatch for scalar: {:?}", - original - ); - - if !original.is_null() { - // For non-null values, compare by converting both to DataFusion scalars - let original_df = original.try_to_df().unwrap(); - let round_trip_df = round_trip.try_to_df().unwrap(); - assert_eq!( - original_df, round_trip_df, - "Value mismatch for scalar: {:?}", - original - ); - } - } - #[rstest] #[case::null_type(Scalar::null(DType::Null), ScalarValue::Null)] #[case::null_bool( @@ -758,80 +422,16 @@ mod tests { Scalar::null(DType::Decimal(DecimalDType::new(5, 2), Nullability::Nullable)), ScalarValue::Decimal32(None, 5, 2) )] - fn test_null_handling(#[case] vortex_null: Scalar, #[case] expected_df_null: ScalarValue) { - // Test Vortex -> DataFusion - let df_result = vortex_null.try_to_df().unwrap(); - assert_eq!(df_result, expected_df_null); - - // Test DataFusion -> Vortex - let vortex_result = from_df(&expected_df_null); - assert!(vortex_result.is_null()); - assert!( - vortex_result - .dtype() - .eq_ignore_nullability(vortex_null.dtype()) - ); - } - - #[rstest] - #[case::utf8(ScalarValue::Utf8(Some("test string".to_string())))] - #[case::utf8_view(ScalarValue::Utf8View(Some("test string".to_string())))] - #[case::large_utf8(ScalarValue::LargeUtf8(Some("test string".to_string())))] - fn test_utf8_variants(#[case] variant: ScalarValue) { - let result = from_df(&variant); - assert_eq!(result.as_utf8().value().unwrap().as_str(), "test string"); - } - - #[rstest] - #[case::binary(ScalarValue::Binary(Some(vec![1u8, 2, 3, 4, 5])))] - #[case::binary_view(ScalarValue::BinaryView(Some(vec![1u8, 2, 3, 4, 5])))] - #[case::large_binary(ScalarValue::LargeBinary(Some(vec![1u8, 2, 3, 4, 5])))] - #[case::fixed_size_binary(ScalarValue::FixedSizeBinary(5, Some(vec![1u8, 2, 3, 4, 5])))] - fn test_binary_variants(#[case] variant: ScalarValue) { - let result = from_df(&variant); - let result_bytes: Vec = result - .as_binary() - .value() - .cloned() - .unwrap() - .into_bytes() - .into(); - assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]); - } - - /// A DataFusion struct whose child field carries Arrow extension metadata: the struct dtype - /// derived from the Arrow type keeps the extension, so the child scalars must keep it too or - /// `Scalar::struct_` rejects them. - #[test] - fn struct_from_df_preserves_extension_child() -> VortexResult<()> { - use arrow_array::FixedSizeBinaryArray; - use arrow_schema::DataType; - use arrow_schema::extension::Uuid as ArrowUuid; - use vortex::extension::uuid::Uuid; - - let mut id_field = Field::new("id", DataType::FixedSizeBinary(16), false); - id_field.try_with_extension_type(ArrowUuid)?; - let ids = FixedSizeBinaryArray::try_from_iter([*b"0123456789abcdef"].into_iter())?; - let struct_array = StructArray::try_new( - Fields::from(vec![Arc::new(id_field)]), - vec![Arc::new(ids) as ArrowArrayRef], - None, - )?; - - let scalar = from_df(&ScalarValue::Struct(Arc::new(struct_array))); - let DType::Struct(fields, _) = scalar.dtype() else { - panic!("expected a struct dtype, got {}", scalar.dtype()); - }; - let id_dtype = fields.field_by_index(0).vortex_expect("one field"); - assert!( - id_dtype.as_extension().is::(), - "expected a Uuid extension field, got {id_dtype}" - ); + fn test_null_handling( + #[case] vortex_null: Scalar, + #[case] expected_df_null: ScalarValue, + ) -> VortexResult<()> { + assert_eq!(vortex_null.try_to_df()?, expected_df_null); Ok(()) } #[test] - fn struct_scalar_round_trips() -> VortexResult<()> { + fn test_struct_scalar_to_datafusion() -> VortexResult<()> { let dtype = DType::Struct( StructFields::new( FieldNames::from(["x", "y"]), @@ -847,17 +447,27 @@ mod tests { vec![Scalar::from(-111.7610f64), Scalar::from(34.8697f64)], ); - let df = original.try_to_df()?; - assert!(matches!(df, ScalarValue::Struct(_))); - - // Back through `from_df` and out again yields the identical DataFusion struct value. - let back = from_df(&df); - assert_eq!(back.try_to_df()?, df); + let expected = StructArray::try_new( + vec![ + Field::new("x", arrow_schema::DataType::Float64, false), + Field::new("y", arrow_schema::DataType::Float64, false), + ] + .into(), + vec![ + Arc::new(arrow_array::Float64Array::from(vec![-111.7610])), + Arc::new(arrow_array::Float64Array::from(vec![34.8697])), + ], + None, + )?; + assert_eq!( + original.try_to_df()?, + ScalarValue::Struct(Arc::new(expected)) + ); Ok(()) } #[test] - fn null_struct_scalar_round_trips() -> VortexResult<()> { + fn test_null_struct_scalar_to_datafusion() -> VortexResult<()> { let dtype = DType::Struct( StructFields::new( FieldNames::from(["x", "y"]), @@ -871,7 +481,7 @@ mod tests { let df = Scalar::null(dtype).try_to_df()?; assert!(matches!(df, ScalarValue::Struct(_))); - assert!(from_df(&df).is_null()); + assert!(df.is_null()); Ok(()) } diff --git a/vortex-datafusion/src/persistent/opener/mod.rs b/vortex-datafusion/src/persistent/opener/mod.rs index ea7b5bd6ae2..c9bb11f1a37 100644 --- a/vortex-datafusion/src/persistent/opener/mod.rs +++ b/vortex-datafusion/src/persistent/opener/mod.rs @@ -15,8 +15,6 @@ use datafusion_common::Statistics; use datafusion_common::arrow::array::AsArray; use datafusion_common::arrow::array::RecordBatch; use datafusion_common::exec_datafusion_err; -use datafusion_common::tree_node::Transformed; -use datafusion_common::tree_node::TreeNode; use datafusion_datasource::PartitionedFile; use datafusion_datasource::TableSchema; use datafusion_datasource::file_stream::FileOpenFuture; @@ -25,11 +23,12 @@ use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; use datafusion_execution::cache::cache_manager::FileMetadataCache; use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::expressions as df_expr; +use datafusion_physical_expr::projection::ProjectionExpr; 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; +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; @@ -304,18 +303,13 @@ impl FileOpener for VortexOpener { } } } - let residual_filter = if residual_filters.is_empty() { - None - } else { - Some(conjunction(residual_filters)) - }; + let residual_filter = conjunction_opt(residual_filters); let native_filter = vortex::expr::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 mut residual_columns = None; let ProcessedProjection { scan_projection, scan_reference_schema, @@ -325,9 +319,16 @@ impl FileOpener for VortexOpener { indices.extend(collect_columns(residual).into_iter().map(|column| column.index())); indices.sort_unstable(); indices.dedup(); - let required = ProjectionExprs::from_indices(&indices, &this_file_schema); - let raw = raw_projection(required, &this_file_schema)?; - residual_columns = Some(indices); + let required = indices.into_iter().map(|index| { + let field = this_file_schema.fields().get(index).ok_or_else(|| { + exec_datafusion_err!("Projection column index {index} is out of bounds") + })?; + Ok(ProjectionExpr { + expr: Arc::new(df_expr::Column::new(field.name(), index)), + alias: field.name().clone(), + }) + }).collect::>>()?; + let raw = raw_projection(required.into(), &this_file_schema)?; ProcessedProjection { scan_projection: raw.scan_projection, scan_reference_schema: raw.scan_reference_schema, @@ -356,17 +357,11 @@ impl FileOpener for VortexOpener { let stream_schema = calculate_physical_schema(&scan_dtype, &scan_reference_schema, &session.arrow())?; - let (leftover_projection, residual_filter) = if let Some(indices) = residual_columns { - ( - leftover_projection.try_map_exprs(|expr| reassign_raw_columns(expr, &indices))?, - residual_filter.map(|expr| reassign_raw_columns(expr, &indices)).transpose()?, - ) - } else { - ( - leftover_projection.try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?, - residual_filter, - ) - }; + 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. @@ -649,21 +644,5 @@ fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") } -/// Remap physical file indices onto the ordered raw columns emitted by the scan. -fn reassign_raw_columns(expr: PhysicalExprRef, indices: &[usize]) -> DFResult { - expr.transform_up(|expr| { - let Some(column) = expr.downcast_ref::() else { - return Ok(Transformed::no(expr)); - }; - let index = indices - .binary_search(&column.index()) - .map_err(|_| exec_datafusion_err!("Missing raw filter/projection column {column}"))?; - Ok(Transformed::yes( - Arc::new(df_expr::Column::new(column.name(), index)) as PhysicalExprRef, - )) - }) - .map(|result| result.data) -} - #[cfg(test)] mod tests; diff --git a/vortex-datafusion/src/persistent/opener/tests.rs b/vortex-datafusion/src/persistent/opener/tests.rs index 2d08ca74478..1531e973db3 100644 --- a/vortex-datafusion/src/persistent/opener/tests.rs +++ b/vortex-datafusion/src/persistent/opener/tests.rs @@ -27,6 +27,8 @@ use datafusion::scalar::ScalarValue; use datafusion_common::arrow::compute::concat_batches; use datafusion_common::assert_batches_eq; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::Transformed; +use datafusion_common::tree_node::TreeNode; use datafusion_execution::cache::default_cache::DefaultCache; use datafusion_expr::Operator; use datafusion_physical_expr::PhysicalExpr; @@ -295,15 +297,22 @@ async fn test_open() -> anyhow::Result<()> { #[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", 1)), + Arc::new(df_expr::Column::new("b", schema.index_of("b")?)), Operator::Modulo, Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), )); @@ -312,7 +321,8 @@ async fn test_residual_filter_unprojected_column( Operator::Eq, Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))), )); - let mut opener = make_opener(store, TableSchema::from(batch.schema()), Some(filter)); + 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 From 77b26a805621b560993038085ad647b716ac3338 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 8 Sep 2026 16:04:35 +0100 Subject: [PATCH 05/10] convert in list Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 64 ++++- vortex-datafusion/src/convert/exprs/tests.rs | 245 +++++++++++++++++- .../src/persistent/opener/tests.rs | 7 +- 3 files changed, 302 insertions(+), 14 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 589e9e0d688..2ff83308ea6 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -28,6 +28,7 @@ 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; @@ -36,6 +37,7 @@ 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; @@ -320,9 +322,8 @@ impl DefaultExpressionConvertor { }, [child, pattern], ) - } else if expr.downcast_ref::().is_some() { - // list_contains does not implement SQL IN/NOT IN null semantics. - return Ok(None); + } else if let Some(in_list) = expr.downcast_ref::() { + return self.convert_in_list(in_list, schema, input_dtype); } else if let Some(scalar_fn) = expr.downcast_ref::() { return self.convert_scalar_function(scalar_fn, schema, input_dtype); } else if let Some(case_expr) = expr.downcast_ref::() { @@ -363,6 +364,63 @@ impl DefaultExpressionConvertor { Ok(Some(converted)) } + fn convert_in_list( + &self, + in_list: &df_expr::InListExpr, + schema: &Schema, + input_dtype: &DType, + ) -> DFResult> { + if in_list.is_empty() + || !in_list + .list() + .iter() + .all(|expr| expr.is::()) + || !supported_data_types(&in_list.expr().data_type(schema)?) + { + return Ok(None); + } + let Some(value) = self.convert_expr(in_list.expr(), schema, input_dtype)? else { + return Ok(None); + }; + // Boolean rewrites may skip evaluating the input, particularly for all-null lists. + if label_infallible(&value).get(&value) != Some(&true) { + return Ok(None); + } + let Ok(value_dtype) = value.return_dtype(input_dtype) else { + return Ok(None); + }; + 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() { + if !supported_data_types(&element.data_type(schema)?) { + return Ok(None); + } + let Some(element) = self.convert_expr(element, schema, input_dtype)? else { + return Ok(None); + }; + let Ok(element_dtype) = element.return_dtype(input_dtype) else { + return Ok(None); + }; + 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 Ok(None); + } + } + // Kleene AND/OR preserve SQL IN/NOT IN nulls; list_contains does not. + Ok(if in_list.negated() { + and_collect(comparisons) + } else { + or_collect(comparisons) + }) + } + fn convert_scalar_function( &self, scalar_fn: &ScalarFunctionExpr, diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 81f86aadc84..881ce040598 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -5,6 +5,7 @@ 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; @@ -173,7 +174,11 @@ fn test_predicate_rejects_cast_over_modulo(test_schema: Schema) { #[rstest] #[case::empty(false)] #[case::column(true)] -fn test_predicate_rejects_in_list(test_schema: Schema, #[case] nonempty: bool) -> DFResult<()> { +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)] @@ -183,7 +188,7 @@ fn test_predicate_rejects_in_list(test_schema: Schema, #[case] nonempty: bool) - let expr: Arc = Arc::new(df_expr::InListExpr::try_new( column, list, - false, + negated, &test_schema, )?); assert!( @@ -194,6 +199,233 @@ fn test_predicate_rejects_in_list(test_schema: Schema, #[case] nonempty: bool) - 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]))?; + let expr = df_expr::InListExpr::try_new( + Arc::new(df_expr::Column::new("a", 0)), + list.into_iter() + .map(|value| Arc::new(df_expr::Literal::new(ScalarValue::Int32(value))) as _) + .collect(), + negated, + &batch.schema(), + )?; + expr.evaluate(&batch)?; + assert_native_matches(Arc::new(expr), 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()])?, + )])?; + let expr = df_expr::InListExpr::try_new( + 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.schema(), + )?; + expr.evaluate(&batch)?; + assert_native_matches(Arc::new(expr), 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)) + }; + let expr = df_expr::InListExpr::try_new(value, list, negated, &batch.schema())?; + expr.evaluate(&batch)?; + assert_native_matches(Arc::new(expr), 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]))?; + let expr = df_expr::InListExpr::try_new( + Arc::new(df_expr::Column::new("a", 0)), + (0..1024) + .map(|i| Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(i)))) as _) + .collect(), + negated, + &batch.schema(), + )?; + expr.evaluate(&batch)?; + assert_native_matches(Arc::new(expr), 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::>>()?, + )?, + )])?; + let expr = df_expr::InListExpr::try_new( + 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.schema(), + )?; + expr.evaluate(&batch)?; + assert_native_matches(Arc::new(expr), 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &batch.schema())? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &Schema::empty())? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema) + .is_err() + ); + Ok(()) +} + #[rstest] #[case::eq(DFOperator::Eq, Operator::Eq)] #[case::not_eq(DFOperator::NotEq, Operator::NotEq)] @@ -951,10 +1183,7 @@ fn test_case_when_datafusion_vortex_equivalence() { assert_eq!(vortex_as_arrow, df_as_arrow); } -fn assert_native_matches( - expr: Arc, - batch: arrow_array::RecordBatch, -) -> anyhow::Result<()> { +fn assert_native_matches(expr: Arc, batch: RecordBatch) -> anyhow::Result<()> { let session = VortexSession::default(); let converted = DefaultExpressionConvertor::new(session.clone()) .try_convert(&expr, &batch.schema())? @@ -1184,7 +1413,7 @@ fn test_native_integer_arithmetic( #[case::mul(DFOperator::Multiply)] #[case::div(DFOperator::Divide)] fn test_native_decimal_arithmetic(#[case] op: DFOperator) -> anyhow::Result<()> { - let batch = arrow_array::RecordBatch::try_from_iter([ + let batch = RecordBatch::try_from_iter([ ( "a", Arc::new( @@ -1350,7 +1579,7 @@ fn test_native_nested_list_length( payload.data_type().clone(), nullable_parent, )])); - let batch = arrow_array::RecordBatch::try_new(schema, vec![payload])?; + let batch = RecordBatch::try_new(schema, vec![payload])?; let get_field: Arc = Arc::new(ScalarFunctionExpr::try_new( Arc::new(ScalarUDF::from(GetFieldFunc::new())), vec![ diff --git a/vortex-datafusion/src/persistent/opener/tests.rs b/vortex-datafusion/src/persistent/opener/tests.rs index 1531e973db3..c8c8f904151 100644 --- a/vortex-datafusion/src/persistent/opener/tests.rs +++ b/vortex-datafusion/src/persistent/opener/tests.rs @@ -1404,7 +1404,7 @@ impl ExpressionConvertor for ResidualConvertor { #[rstest] #[tokio::test] -async fn test_physical_in_list_residual( +async fn test_physical_in_list( #[values(false, true)] negated: bool, #[values(0, 1, 2)] list_kind: usize, ) -> anyhow::Result<()> { @@ -1429,10 +1429,11 @@ async fn test_physical_in_list_residual( negated, &batch.schema(), )?); - assert!( + assert_eq!( DefaultExpressionConvertor::default() .try_convert(&filter, &batch.schema())? - .is_none() + .is_some(), + list_kind == 2, ); let expected = batch_filter(&batch, &filter)?.project(&[0])?; let opener = make_opener(store, TableSchema::from(batch.schema()), Some(filter)); From 629fd24acaf5d37df65dcd8c5a38aa4e3b9227e7 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 8 Sep 2026 16:17:29 +0100 Subject: [PATCH 06/10] fix v2 issue Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 9 +- vortex-datafusion/src/convert/exprs/tests.rs | 14 ++- vortex-datafusion/src/v2/mod.rs | 3 + vortex-datafusion/src/v2/source.rs | 16 ++- vortex-datafusion/src/v2/tests.rs | 100 +++++++++++++++++++ 5 files changed, 132 insertions(+), 10 deletions(-) create mode 100644 vortex-datafusion/src/v2/tests.rs diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 2ff83308ea6..0edd7be43a3 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -396,9 +396,6 @@ impl DefaultExpressionConvertor { }; let mut comparisons = Vec::with_capacity(in_list.len()); for element in in_list.list() { - if !supported_data_types(&element.data_type(schema)?) { - return Ok(None); - } let Some(element) = self.convert_expr(element, schema, input_dtype)? else { return Ok(None); }; @@ -525,6 +522,9 @@ impl ExpressionConvertor for DefaultExpressionConvertor { expr: &Arc, schema: &Schema, ) -> DFResult> { + if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { + return Ok(None); + } for column in collect_columns(expr) { let field = schema.fields().get(column.index()).ok_or_else(|| { exec_datafusion_err!( @@ -542,9 +542,6 @@ impl ExpressionConvertor for DefaultExpressionConvertor { )); } } - if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { - return Ok(None); - } let Ok(input_dtype) = self.session.arrow().from_arrow_schema(schema) else { return Ok(None); }; diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 881ce040598..b62aebe653f 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -23,6 +23,7 @@ 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_plan::expressions as df_expr; use insta::assert_snapshot; use rstest::rstest; @@ -1544,7 +1545,7 @@ fn test_malformed_get_field_returns_error( } #[test] -fn test_mismatched_column_identity_returns_error() { +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), @@ -1555,6 +1556,17 @@ fn test_mismatched_column_identity_returns_error() { .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!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + Ok(()) } #[rstest] 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 f78da1c2a8f..258a97fc591 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -603,13 +603,22 @@ impl DataSource for VortexDataSource { } let convertor = DefaultExpressionConvertor::default(); - let input_schema = self.initial_schema.as_ref(); + 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 converted = filters .iter() - .map(|expr| convertor.try_convert(expr, input_schema)) + .map(|expr| convertor.try_convert(expr, &self.projected_schema)) .collect::>>()?; let pushdown_results: Vec = converted .iter() @@ -630,7 +639,8 @@ impl DataSource for VortexDataSource { } // Convert to Vortex conjunction. - let vortex_pred = vortex::expr::and_collect(converted.into_iter().flatten()); + 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(()) +} From 085b86f0060094727f75effeb7c2e498c86bcf47 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 8 Sep 2026 16:45:02 +0100 Subject: [PATCH 07/10] more fixes Signed-off-by: Adam Gutglick --- vortex-arrow/src/dtype.rs | 13 +- vortex-datafusion/src/convert/exprs/mod.rs | 10 +- vortex-datafusion/src/convert/exprs/tests.rs | 143 +++++++++++++++++++ 3 files changed, 163 insertions(+), 3 deletions(-) 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/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 0edd7be43a3..baec74365c3 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -525,7 +525,9 @@ impl ExpressionConvertor for DefaultExpressionConvertor { if DynamicFilterTracking::classify(expr).contains_dynamic_filter() { return Ok(None); } - for column in collect_columns(expr) { + 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", @@ -541,8 +543,12 @@ impl ExpressionConvertor for DefaultExpressionConvertor { field.name() )); } + column_indices.push(column.index()); } - let Ok(input_dtype) = self.session.arrow().from_arrow_schema(schema) else { + column_indices.sort_unstable(); + column_indices.dedup(); + let referenced_schema = schema.project(&column_indices)?; + let Ok(input_dtype) = self.session.arrow().from_arrow_schema(&referenced_schema) else { return Ok(None); }; let Some(converted) = self.convert_expr(expr, schema, &input_dtype)? else { diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index b62aebe653f..c290b1a33ee 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -9,6 +9,7 @@ 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; @@ -761,6 +762,148 @@ fn test_can_be_pushed_down_column_supported(test_schema: Schema) { ); } +#[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_filter_ignores_unreferenced_unsupported_field( + #[case] unsupported: DataType, +) -> DFResult<()> { + let schema = Schema::new(vec![ + Field::new("unsupported", unsupported, true), + Field::new("id", DataType::Int32, false), + ]); + 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!( + DefaultExpressionConvertor::default().try_convert(&expr, &schema)?, + Some(Binary.new_expr(Operator::Eq, [get_item("id", root()), lit(42i32)])), + ); + 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_referenced_unsupported_field_is_declined(#[case] unsupported: DataType) -> DFResult<()> { + let schema = Schema::new(vec![Field::new("unsupported", unsupported, true)]); + let expr: Arc = Arc::new(df_expr::IsNotNullExpr::new(Arc::new( + df_expr::Column::new("unsupported", 0), + ))); + assert!( + DefaultExpressionConvertor::default() + .try_convert(&expr, &schema)? + .is_none() + ); + 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!( + DefaultExpressionConvertor::default().try_convert(&expr, &schema)?, + Some(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 { + expr: 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)), + )), + )), + alias: "sum".into(), + }]); + let output_schema = projection.project_schema(&schema)?; + let processed = DefaultExpressionConvertor::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 { + expr: Arc::new(df_expr::Column::new("sum", 0)), + alias: "sum".into(), + }]), + ); + 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!( + DefaultExpressionConvertor::default().try_convert(&expr, &schema)?, + Some(expected), + ); + Ok(()) +} + #[rstest] fn test_nested_column_conversion(test_schema: Schema) { let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; From 2187b65c2ce0502c3e34e009f9c8aa1ffe334bfc Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 10:05:42 +0100 Subject: [PATCH 08/10] things Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 571 ++--- vortex-datafusion/src/convert/exprs/tests.rs | 622 ++--- vortex-datafusion/src/convert/mod.rs | 1 + vortex-datafusion/src/convert/scalars.rs | 438 +++- vortex-datafusion/src/persistent/opener.rs | 2041 +++++++++++++++++ .../src/persistent/opener/mod.rs | 648 ------ .../src/persistent/opener/tests.rs | 1546 ------------- vortex-datafusion/src/persistent/source.rs | 5 - vortex-datafusion/src/persistent/tests.rs | 7 +- 9 files changed, 2953 insertions(+), 2926 deletions(-) create mode 100644 vortex-datafusion/src/persistent/opener.rs delete mode 100644 vortex-datafusion/src/persistent/opener/mod.rs delete mode 100644 vortex-datafusion/src/persistent/opener/tests.rs diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index baec74365c3..b05cba79b1f 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -6,6 +6,7 @@ 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; @@ -17,7 +18,6 @@ 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; @@ -70,26 +70,32 @@ pub struct ProcessedProjection { /// /// # Implementing a custom convertor /// -/// use std::sync::Arc; +/// ``` +/// 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::DefaultExpressionConvertor; -/// use vortex_datafusion::convert::ExpressionConvertor; +/// use arrow_schema::Schema; +/// use datafusion_common::Result as DFResult; +/// use datafusion_physical_expr::PhysicalExpr; +/// use vortex::expr::Expression; +/// use vortex_datafusion::convert::DefaultExpressionConvertor; +/// use vortex_datafusion::convert::ExpressionConvertor; /// -/// struct CustomExpressionConvertor(DefaultExpressionConvertor); +/// struct CustomExpressionConvertor(DefaultExpressionConvertor); /// -/// impl ExpressionConvertor for CustomExpressionConvertor { -/// fn try_convert( -/// &self, -/// expr: &Arc, -/// schema: &Schema, -/// ) -> DFResult> { -/// self.0.try_convert(expr, schema) -/// } +/// impl ExpressionConvertor for CustomExpressionConvertor { +/// fn try_convert( +/// &self, +/// expr: &Arc, +/// schema: &Schema, +/// ) -> DFResult> { +/// self.0.try_convert(expr, schema) /// } +/// } +/// +/// let _convertor: Arc = Arc::new(CustomExpressionConvertor( +/// DefaultExpressionConvertor::default(), +/// )); +/// ``` pub trait ExpressionConvertor: Send + Sync { /// Convert an expression for native evaluation against this schema. /// @@ -102,7 +108,8 @@ pub trait ExpressionConvertor: Send + Sync { schema: &Schema, ) -> DFResult>; - /// Split a projection into native and DataFusion evaluation. + /// 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. @@ -127,241 +134,264 @@ pub trait ExpressionConvertor: Send + Sync { }; 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: source_projection - .iter() - .enumerate() - .map(|(index, projection)| ProjectionExpr { - expr: Arc::new(df_expr::Column::new(&projection.alias, index)), - alias: projection.alias.clone(), - }) - .collect::>() - .into(), + leftover_projection: ProjectionExprs::from_indices(&output_indices, output_schema), }) } - /// Read the required raw columns and apply the complete projection in DataFusion. + /// 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 { - raw_projection(source_projection, input_schema) + 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 raw columns in file-index order without involving custom expression conversion. +/// 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( - source_projection: ProjectionExprs, + indices: &[usize], input_schema: &Schema, -) -> DFResult { - let column_indices = source_projection.column_indices(); - let mut scan_columns = Vec::with_capacity(column_indices.len()); - let mut fields = Vec::with_capacity(column_indices.len()); - for index in column_indices { - let field = input_schema.fields().get(index).ok_or_else(|| { - exec_datafusion_err!("Projection column index {index} is out of bounds") - })?; - scan_columns.push(( +) -> 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()), - )); - fields.push(Arc::clone(field)); + ) + }); + 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) } - Ok(ProcessedProjection { - scan_projection: pack(scan_columns, Nullability::NonNullable), - scan_reference_schema: Schema::new_with_metadata(fields, input_schema.metadata().clone()), - leftover_projection: source_projection, - }) } -/// The default schema-aware DataFusion expression convertor. +/// Conversion result where `?` propagates both unsupported expressions and errors. +type Conversion = Result; + +/// The default [`ExpressionConvertor`] 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 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::new(VortexSession::default()) + Self { + session: VortexSession::default(), + } } } impl DefaultExpressionConvertor { - /// Create a convertor that resolves Arrow extension types using the session registry. + /// Create a convertor 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, - ) -> DFResult> { - let converted = if let Some(binary) = expr.downcast_ref::() { - let Some(operator) = try_operator_from_df(binary.op()) else { - return Ok(None); - }; + ) -> 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.left().data_type(schema)?; - let right_type = binary.right().data_type(schema)?; + 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 Ok(None); + return Err(Unconverted::Unsupported); } - let (Some(left), Some(right)) = ( - self.convert_expr(binary.left(), schema, input_dtype)?, - self.convert_expr(binary.right(), schema, input_dtype)?, - ) else { - return Ok(None); - }; - if boolean_operator - && (label_infallible(&left).get(&left) != Some(&true) - || label_infallible(&right).get(&right) != Some(&true)) - { - // DataFusion may evaluate the RHS only on rows selected by the LHS. - return Ok(None); + 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); } - let (Ok(left_dtype), Ok(right_dtype)) = ( - left.return_dtype(input_dtype), - right.return_dtype(input_dtype), - ) else { - return Ok(None); - }; - if !left_dtype.eq_ignore_nullability(&right_dtype) { - return Ok(None); + if !converted_dtype(&left, input_dtype)? + .eq_ignore_nullability(&converted_dtype(&right, input_dtype)?) + { + return Err(Unconverted::Unsupported); } - Binary.new_expr(operator, [left, right]) - } else if let Some(column) = expr.downcast_ref::() { - get_item(column.name(), root()) - } else if let Some(literal) = expr.downcast_ref::() { + 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() - )); - } - let array = match self.session.arrow().from_arrow_array(array, &field) { - Ok(array) => array, - Err(error) => { - if self.session.arrow().from_arrow_field(&field).is_err() { - return Ok(None); - } - return Err(exec_datafusion_err!( - "Failed to convert literal {expr}: {error}" - )); - } - }; - lit(array - .execute_scalar(0, &mut self.session.create_execution_ctx()) - .map_err(|e| exec_datafusion_err!("Failed to evaluate literal {expr}: {e}"))?) - } else if let Some(cast_expr) = expr.downcast_ref::() { - if !supported_cast(cast_expr, schema)? { - return Ok(None); + ) + .into()); } - let Some(child) = self.convert_expr(cast_expr.expr(), schema, input_dtype)? else { - return Ok(None); - }; - let Ok(target) = self + // Literals of unknown Arrow types stay in DataFusion; conversion failures are errors. + self.arrow_dtype(&field)?; + let scalar = self .session .arrow() - .from_arrow_field(cast_expr.target_field()) - else { - return Ok(None); - }; - let Ok(child_dtype) = child.return_dtype(input_dtype) else { - return Ok(None); - }; + .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)?; + let cast_dtype = self.arrow_dtype(cast_expr.target_field())?; // Matching Arrow storage types do not imply matching extension semantics. - if (child_dtype.is_extension() || target.is_extension()) - && !child_dtype.eq_ignore_nullability(&target) + 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 Ok(None); + return Err(Unconverted::Unsupported); } - cast(child, target) - } else if let Some(is_null_expr) = expr.downcast_ref::() { - let Some(child) = self.convert_expr(is_null_expr.arg(), schema, input_dtype)? else { - return Ok(None); - }; - is_null(child) - } else if let Some(is_not_null_expr) = expr.downcast_ref::() { - let Some(child) = self.convert_expr(is_not_null_expr.arg(), schema, input_dtype)? - else { - return Ok(None); - }; - is_not_null(child) - } else if let Some(like) = expr.downcast_ref::() { + 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 Ok(None); + return Err(Unconverted::Unsupported); } - let (Some(child), Some(pattern)) = ( - self.convert_expr(like.expr(), schema, input_dtype)?, - self.convert_expr(like.pattern(), schema, input_dtype)?, - ) else { - return Ok(None); - }; - Like.new_expr( + 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], - ) - } else if let Some(in_list) = expr.downcast_ref::() { + )); + } + + if let Some(in_list) = expr.downcast_ref::() { return self.convert_in_list(in_list, schema, input_dtype); - } else if let Some(scalar_fn) = expr.downcast_ref::() { + } + + if let Some(scalar_fn) = expr.downcast_ref::() { return self.convert_scalar_function(scalar_fn, schema, input_dtype); - } else if let Some(case_expr) = expr.downcast_ref::() { - if case_expr.expr().is_some() { - return Ok(None); - } - let mut pairs = Vec::with_capacity(case_expr.when_then_expr().len()); - for (when, then) in case_expr.when_then_expr() { - if when.data_type(schema)? != DataType::Boolean { - return Err(exec_datafusion_err!("CASE WHEN must be Boolean")); - } - let (Some(when), Some(then)) = ( - self.convert_expr(when, schema, input_dtype)?, - self.convert_expr(then, schema, input_dtype)?, - ) else { - return Ok(None); - }; - pairs.push((when, then)); - } - let otherwise = match case_expr.else_expr() { - Some(expr) => { - let Some(expr) = self.convert_expr(expr, schema, input_dtype)? else { - return Ok(None); - }; - Some(expr) - } - None => None, - }; - let case = nested_case_when(pairs, otherwise); - // Vortex may evaluate branch values on rows excluded by the condition. - if label_infallible(&case).get(&case) != Some(&true) { - return Ok(None); - } - case - } else { - return Ok(None); - }; - Ok(Some(converted)) + } + + 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( @@ -369,7 +399,7 @@ impl DefaultExpressionConvertor { in_list: &df_expr::InListExpr, schema: &Schema, input_dtype: &DType, - ) -> DFResult> { + ) -> Conversion { if in_list.is_empty() || !in_list .list() @@ -377,18 +407,14 @@ impl DefaultExpressionConvertor { .all(|expr| expr.is::()) || !supported_data_types(&in_list.expr().data_type(schema)?) { - return Ok(None); + return Err(Unconverted::Unsupported); } - let Some(value) = self.convert_expr(in_list.expr(), schema, input_dtype)? else { - return Ok(None); - }; + let value = self.convert_expr(in_list.expr(), schema, input_dtype)?; // Boolean rewrites may skip evaluating the input, particularly for all-null lists. - if label_infallible(&value).get(&value) != Some(&true) { - return Ok(None); + if !is_infallible(&value) { + return Err(Unconverted::Unsupported); } - let Ok(value_dtype) = value.return_dtype(input_dtype) else { - return Ok(None); - }; + let value_dtype = converted_dtype(&value, input_dtype)?; let operator = if in_list.negated() { Operator::NotEq } else { @@ -396,26 +422,23 @@ impl DefaultExpressionConvertor { }; let mut comparisons = Vec::with_capacity(in_list.len()); for element in in_list.list() { - let Some(element) = self.convert_expr(element, schema, input_dtype)? else { - return Ok(None); - }; - let Ok(element_dtype) = element.return_dtype(input_dtype) else { - return Ok(None); - }; + 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 Ok(None); + return Err(Unconverted::Unsupported); } } // Kleene AND/OR preserve SQL IN/NOT IN nulls; list_contains does not. - Ok(if in_list.negated() { + let membership = if in_list.negated() { and_collect(comparisons) } else { or_collect(comparisons) - }) + }; + membership.ok_or(Unconverted::Unsupported) } fn convert_scalar_function( @@ -423,15 +446,18 @@ impl DefaultExpressionConvertor { scalar_fn: &ScalarFunctionExpr, schema: &Schema, input_dtype: &DType, - ) -> DFResult> { + ) -> 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" - )); + 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")); + 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)?; @@ -445,15 +471,12 @@ impl DefaultExpressionConvertor { exec_datafusion_err!("get_field path must be a non-null string literal") })?; let DataType::Struct(fields) = &source_type else { - return Ok(None); + return Err(Unconverted::Unsupported); }; nullable_struct |= nullable; - let field = fields - .iter() - .find(|field| field.name() == name) - .ok_or_else(|| { - exec_datafusion_err!("get_field references missing field {name}") - })?; + 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); @@ -461,23 +484,21 @@ impl DefaultExpressionConvertor { // DataFusion extracts the child without applying parent struct validity; // Vortex get_item masks the child when its parent is null. if nullable_struct { - return Ok(None); + return Err(Unconverted::Unsupported); } - let Some(mut result) = self.convert_expr(source, schema, input_dtype)? else { - return Ok(None); - }; + let mut result = self.convert_expr(source, schema, input_dtype)?; for name in names { result = get_item(name, result); } - return Ok(Some(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" - )); + 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 { @@ -485,34 +506,59 @@ impl DefaultExpressionConvertor { data_type => data_type, }; if !data_type.is_binary() && !data_type.is_string() { - return Ok(None); + return Err(Unconverted::Unsupported); } (input, byte_length) } else if ScalarFunctionExpr::try_downcast_func::(scalar_fn).is_some() { - let Some(input) = array_length_input(scalar_fn)? else { - return Ok(None); - }; + let input = array_length_input(scalar_fn)?.ok_or(Unconverted::Unsupported)?; if !matches!( input.data_type(schema)?, DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) ) { - return Ok(None); + return Err(Unconverted::Unsupported); } (input, list_length) } else { - return Ok(None); + return Err(Unconverted::Unsupported); }; - let Some(input) = self.convert_expr(input, schema, input_dtype)? else { - return Ok(None); - }; - let Ok(return_dtype) = self.session.arrow().from_arrow_field(&Field::new( + 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(), - )) else { - return Ok(None); - }; - Ok(Some(cast(length(input), return_dtype))) + ))?; + 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) } } @@ -522,55 +568,30 @@ impl ExpressionConvertor for DefaultExpressionConvertor { 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); } - 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() - )); - } - column_indices.push(column.index()); - } - column_indices.sort_unstable(); - column_indices.dedup(); - let referenced_schema = schema.project(&column_indices)?; - let Ok(input_dtype) = self.session.arrow().from_arrow_schema(&referenced_schema) else { - return Ok(None); - }; - let Some(converted) = self.convert_expr(expr, schema, &input_dtype)? else { - return Ok(None); - }; - let Ok(expected_dtype) = self - .session - .arrow() - .from_arrow_field(expr.return_field(schema)?.as_ref()) - else { - return Ok(None); - }; - let Ok(actual_dtype) = converted.return_dtype(&input_dtype) else { - return Ok(None); - }; - if !actual_dtype.eq_ignore_nullability(&expected_dtype) { - 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), } - Ok(Some(converted)) } } +/// 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(); @@ -667,11 +688,9 @@ fn supported_data_types(dt: &DataType) -> bool { is_supported } -/// 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. -/// Calls with other arities return errors. +/// 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)), diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index c290b1a33ee..e7e0e3bfc93 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -13,8 +13,6 @@ use arrow_schema::IntervalUnit; use arrow_schema::Schema; use arrow_schema::TimeUnit as ArrowTimeUnit; use arrow_schema::extension::Uuid as ArrowUuid; -use datafusion::arrow::array::AsArray; -use datafusion::arrow::datatypes::Int32Type; use datafusion_common::ScalarValue; use datafusion_common::arrow::buffer::NullBuffer; use datafusion_common::arrow::datatypes::i256 as arrow_i256; @@ -25,6 +23,7 @@ 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; @@ -83,6 +82,39 @@ fn array_length_expr(args: Vec>, schema: &Schema) -> Arc, schema: &Schema) -> DFResult { + Ok(DefaultExpressionConvertor::default() + .try_convert(expr, schema)? + .is_some()) +} + +/// Convert `expr` natively, failing the test if the convertor declines it. +fn convert(expr: Arc, schema: &Schema) -> DFResult { + DefaultExpressionConvertor::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 FailingConvertor; impl ExpressionConvertor for FailingConvertor { @@ -102,22 +134,16 @@ fn test_duplicate_aliases_fall_back_before_conversion() -> DFResult<()> { Field::new("b", DataType::Int32, false), ]); let projection = ProjectionExprs::from(vec![ - ProjectionExpr { - expr: Arc::new(df_expr::BinaryExpr::new( + 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)))), )), - alias: "duplicate".into(), - }, - ProjectionExpr { - expr: Arc::new(df_expr::Column::new("a", 0)), - alias: "duplicate".into(), - }, - ProjectionExpr { - expr: Arc::new(df_expr::Column::new("b", 1)), - alias: "duplicate".into(), - }, + "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 = @@ -135,29 +161,7 @@ fn test_duplicate_aliases_fall_back_before_conversion() -> DFResult<()> { } #[rstest] -#[case::past_end(1)] -#[case::max(usize::MAX)] -fn test_raw_projection_out_of_bounds(#[case] index: usize) -> DFResult<()> { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - let projection = vec![ProjectionExpr { - expr: Arc::new(df_expr::Column::new("a", index)), - alias: "a".into(), - }] - .into(); - let error = raw_projection(projection, &schema) - .err() - .ok_or_else(|| exec_datafusion_err!("Expected an out-of-bounds error"))?; - assert!( - error - .to_string() - .contains(&format!("Projection column index {index} is out of bounds")), - "{error}" - ); - Ok(()) -} - -#[rstest] -fn test_predicate_rejects_cast_over_modulo(test_schema: Schema) { +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, @@ -165,12 +169,8 @@ fn test_predicate_rejects_cast_over_modulo(test_schema: Schema) { )); let expr: Arc = Arc::new(df_expr::CastExpr::new(modulo, DataType::Int64, None)); - assert!( - !DefaultExpressionConvertor::default() - .try_convert(&expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(!converts(&expr, &test_schema)?); + Ok(()) } #[rstest] @@ -193,11 +193,7 @@ fn test_predicate_rejects_unsupported_in_list( negated, &test_schema, )?); - assert!( - !DefaultExpressionConvertor::default() - .try_convert(&expr, &test_schema)? - .is_some() - ); + assert!(!converts(&expr, &test_schema)?); Ok(()) } @@ -212,16 +208,12 @@ fn test_native_in_list( #[values(false, true)] negated: bool, ) -> anyhow::Result<()> { let batch = arrow_array::record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3), None]))?; - let expr = df_expr::InListExpr::try_new( + assert_in_list_matches( Arc::new(df_expr::Column::new("a", 0)), - list.into_iter() - .map(|value| Arc::new(df_expr::Literal::new(ScalarValue::Int32(value))) as _) - .collect(), + int_literals(list), negated, - &batch.schema(), - )?; - expr.evaluate(&batch)?; - assert_native_matches(Arc::new(expr), batch) + batch, + ) } #[rstest] @@ -266,17 +258,15 @@ fn test_native_in_list_data_types( "a", ScalarValue::iter_to_array([member.clone(), absent, null.clone()])?, )])?; - let expr = df_expr::InListExpr::try_new( + 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.schema(), - )?; - expr.evaluate(&batch)?; - assert_native_matches(Arc::new(expr), batch) + batch, + ) } #[rstest] @@ -296,25 +286,19 @@ fn test_native_in_list_untyped_null( } else { Arc::new(df_expr::Column::new("a", 0)) }; - let expr = df_expr::InListExpr::try_new(value, list, negated, &batch.schema())?; - expr.evaluate(&batch)?; - assert_native_matches(Arc::new(expr), batch) + 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]))?; - let expr = df_expr::InListExpr::try_new( + assert_in_list_matches( Arc::new(df_expr::Column::new("a", 0)), - (0..1024) - .map(|i| Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(i)))) as _) - .collect(), + int_literals((0..1024).map(Some)), negated, - &batch.schema(), - )?; - expr.evaluate(&batch)?; - assert_native_matches(Arc::new(expr), batch) + batch, + ) } #[rstest] @@ -340,17 +324,15 @@ fn test_native_in_list_float( .collect::>>()?, )?, )])?; - let expr = df_expr::InListExpr::try_new( + 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.schema(), - )?; - expr.evaluate(&batch)?; - assert_native_matches(Arc::new(expr), batch) + batch, + ) } #[rstest] @@ -372,11 +354,7 @@ fn test_in_list_unsupported_input( &batch.schema(), )?); assert!(expr.evaluate(&batch).is_err()); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &batch.schema())? - .is_none() - ); + assert!(!converts(&expr, &batch.schema())?); Ok(()) } @@ -390,11 +368,7 @@ fn test_in_list_unsupported_type() -> DFResult<()> { false, &Schema::empty(), )?); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &Schema::empty())? - .is_none() - ); + assert!(!converts(&expr, &Schema::empty())?); Ok(()) } @@ -458,41 +432,31 @@ fn test_operator_conversion_unsupported(#[case] df_op: DFOperator) { } #[test] -fn test_expr_from_df_column() { +fn test_expr_from_df_column() -> DFResult<()> { let col_expr = df_expr::Column::new("test_column", 0); - let result = DefaultExpressionConvertor::default() - .try_convert( - &(Arc::new(col_expr) as Arc), - &Schema::new(vec![Field::new("test_column", DataType::Int32, false)]), - ) - .unwrap() - .unwrap(); + 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() { +fn test_expr_from_df_literal() -> DFResult<()> { let literal_expr = df_expr::Literal::new(ScalarValue::Int32(Some(42))); - let result = DefaultExpressionConvertor::default() - .try_convert( - &(Arc::new(literal_expr) as Arc), - &Schema::empty(), - ) - .unwrap() - .unwrap(); + 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 { - let converted = DefaultExpressionConvertor::default() - .try_convert(&expr, &Schema::empty())? - .ok_or_else(|| anyhow::anyhow!("Expected native conversion for {expr}"))?; - converted + convert(expr, &Schema::empty())? .as_opt::() .cloned() .ok_or_else(|| anyhow::anyhow!("Expected a literal expression")) @@ -609,19 +573,16 @@ fn test_literal_preserves_extension_metadata( } #[test] -fn test_expr_from_df_binary() { +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 = DefaultExpressionConvertor::default() - .try_convert( - &(Arc::new(binary_expr) as Arc), - &Schema::new(vec![Field::new("left", DataType::Int32, false)]), - ) - .unwrap() - .unwrap(); + 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(=) @@ -629,6 +590,7 @@ fn test_expr_from_df_binary() { │ └── input: vortex.root() └── rhs: vortex.literal(42i32) "); + Ok(()) } #[rstest] @@ -636,20 +598,17 @@ fn test_expr_from_df_binary() { #[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) { +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 = DefaultExpressionConvertor::default() - .try_convert( - &(Arc::new(like_expr) as Arc), - &Schema::new(vec![Field::new("text_col", DataType::Utf8, true)]), - ) - .unwrap() - .unwrap(); + 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, @@ -658,17 +617,15 @@ fn test_expr_from_df_like(#[case] negated: bool, #[case] case_insensitive: bool) case_insensitive } ); + Ok(()) } #[rstest] -fn test_expr_from_df_octet_length(test_schema: Schema) { +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 = DefaultExpressionConvertor::default() - .try_convert(&octet_length, &test_schema) - .unwrap() - .unwrap(); + let result = convert(octet_length, &test_schema)?; assert_snapshot!(result.display_tree().to_string(), @r" vortex.cast(i32?) @@ -676,17 +633,15 @@ fn test_expr_from_df_octet_length(test_schema: Schema) { └── input: vortex.get_item(name) └── input: vortex.root() "); + Ok(()) } #[rstest] -fn test_expr_from_df_array_length(test_schema: Schema) { +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 = DefaultExpressionConvertor::default() - .try_convert(&array_length, &test_schema) - .unwrap() - .unwrap(); + let result = convert(array_length, &test_schema)?; assert_snapshot!(result.display_tree().to_string(), @r" vortex.cast(u64?) @@ -694,6 +649,7 @@ fn test_expr_from_df_array_length(test_schema: Schema) { └── input: vortex.get_item(tags) └── input: vortex.root() "); + Ok(()) } #[rstest] @@ -751,15 +707,11 @@ fn test_supported_data_types(#[case] data_type: DataType, #[case] expected: bool } #[rstest] -fn test_can_be_pushed_down_column_supported(test_schema: Schema) { +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!( - DefaultExpressionConvertor::default() - .try_convert(&col_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&col_expr, &test_schema)?); + Ok(()) } #[rstest] @@ -770,43 +722,30 @@ fn test_can_be_pushed_down_column_supported(test_schema: Schema) { #[case::decimal64(DataType::Decimal64(1, 2))] #[case::decimal128(DataType::Decimal128(1, 2))] #[case::decimal256(DataType::Decimal256(1, 2))] -fn test_filter_ignores_unreferenced_unsupported_field( +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), ]); - 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!( - DefaultExpressionConvertor::default().try_convert(&expr, &schema)?, - Some(Binary.new_expr(Operator::Eq, [get_item("id", root()), lit(42i32)])), - ); - 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_referenced_unsupported_field_is_declined(#[case] unsupported: DataType) -> DFResult<()> { - let schema = Schema::new(vec![Field::new("unsupported", unsupported, true)]); - let expr: Arc = Arc::new(df_expr::IsNotNullExpr::new(Arc::new( - df_expr::Column::new("unsupported", 0), - ))); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_none() - ); + 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(()) } @@ -818,10 +757,7 @@ fn test_literal_ignores_unreferenced_malformed_decimal() -> DFResult<()> { true, )]); let expr: Arc = Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(42)))); - assert_eq!( - DefaultExpressionConvertor::default().try_convert(&expr, &schema)?, - Some(lit(42i32)), - ); + assert_eq!(convert(expr, &schema)?, lit(42i32)); Ok(()) } @@ -836,8 +772,8 @@ fn test_projection_ignores_unreferenced_unsupported_field() -> DFResult<()> { ), Field::new("b", DataType::Int32, false), ]); - let projection = ProjectionExprs::from(vec![ProjectionExpr { - expr: Arc::new(df_expr::BinaryExpr::new( + 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( @@ -846,8 +782,8 @@ fn test_projection_ignores_unreferenced_unsupported_field() -> DFResult<()> { Arc::new(df_expr::Column::new("b", 2)), )), )), - alias: "sum".into(), - }]); + "sum", + )]); let output_schema = projection.project_schema(&schema)?; let processed = DefaultExpressionConvertor::default().split_projection( projection, @@ -872,10 +808,10 @@ fn test_projection_ignores_unreferenced_unsupported_field() -> DFResult<()> { assert_eq!(processed.scan_reference_schema, output_schema); assert_eq!( processed.leftover_projection, - ProjectionExprs::from(vec![ProjectionExpr { - expr: Arc::new(df_expr::Column::new("sum", 0)), - alias: "sum".into(), - }]), + ProjectionExprs::from(vec![ProjectionExpr::new( + Arc::new(df_expr::Column::new("sum", 0)), + "sum", + )]), ); Ok(()) } @@ -897,23 +833,16 @@ fn test_referenced_field_preserves_extension_metadata( Field::new("unsupported", DataType::FixedSizeBinary(16), true), field, ]); - assert_eq!( - DefaultExpressionConvertor::default().try_convert(&expr, &schema)?, - Some(expected), - ); + assert_eq!(convert(expr, &schema)?, expected); Ok(()) } #[rstest] -fn test_nested_column_conversion(test_schema: Schema) { +fn test_nested_column_conversion(test_schema: Schema) -> DFResult<()> { let col_expr = Arc::new(df_expr::Column::new("tags", 5)) as Arc; - assert!( - DefaultExpressionConvertor::default() - .try_convert(&col_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&col_expr, &test_schema)?); + Ok(()) } #[rstest] @@ -928,16 +857,12 @@ fn test_can_be_pushed_down_column_not_found(test_schema: Schema) { } #[rstest] -fn test_can_be_pushed_down_literal_supported(test_schema: Schema) { +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!( - DefaultExpressionConvertor::default() - .try_convert(&lit_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&lit_expr, &test_schema)?); + Ok(()) } #[rstest] @@ -951,32 +876,24 @@ fn test_can_be_pushed_down_literal_unsupported( #[case] value: ScalarValue, ) -> DFResult<()> { let lit_expr = Arc::new(df_expr::Literal::new(value)) as Arc; - assert!( - DefaultExpressionConvertor::default() - .try_convert(&lit_expr, &test_schema)? - .is_none() - ); + assert!(!converts(&lit_expr, &test_schema)?); Ok(()) } #[rstest] -fn test_can_be_pushed_down_binary_supported(test_schema: Schema) { +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!( - DefaultExpressionConvertor::default() - .try_convert(&binary_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&binary_expr, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_binary_unsupported_operator(test_schema: Schema) { +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; @@ -986,32 +903,24 @@ fn test_can_be_pushed_down_binary_unsupported_operator(test_schema: Schema) { right, )) as Arc; - assert!( - !DefaultExpressionConvertor::default() - .try_convert(&binary_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(!converts(&binary_expr, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_binary_unsupported_operand(test_schema: Schema) { +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!( - !DefaultExpressionConvertor::default() - .try_convert(&binary_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(!converts(&binary_expr, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_like_supported(test_schema: Schema) { +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(), @@ -1019,16 +928,12 @@ fn test_can_be_pushed_down_like_supported(test_schema: Schema) { let like_expr = Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc; - assert!( - DefaultExpressionConvertor::default() - .try_convert(&like_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&like_expr, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_like_unsupported_operand(test_schema: Schema) { +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(), @@ -1036,29 +941,21 @@ fn test_can_be_pushed_down_like_unsupported_operand(test_schema: Schema) { let like_expr = Arc::new(df_expr::LikeExpr::new(false, false, expr, pattern)) as Arc; - assert!( - !DefaultExpressionConvertor::default() - .try_convert(&like_expr, &test_schema) - .unwrap() - .is_some() - ); + assert!(!converts(&like_expr, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_octet_length_supported(test_schema: Schema) { +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!( - DefaultExpressionConvertor::default() - .try_convert(&octet_length, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&octet_length, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) { +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", @@ -1068,29 +965,21 @@ fn test_can_be_pushed_down_octet_length_unsupported_operand(test_schema: Schema) Arc::new(ConfigOptions::new()), )) as Arc; - assert!( - !DefaultExpressionConvertor::default() - .try_convert(&octet_length, &test_schema) - .unwrap() - .is_some() - ); + assert!(!converts(&octet_length, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_array_length_supported(test_schema: Schema) { +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!( - DefaultExpressionConvertor::default() - .try_convert(&array_length, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&array_length, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_array_length_unsupported_operand(test_schema: Schema) { +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( @@ -1101,28 +990,22 @@ fn test_can_be_pushed_down_array_length_unsupported_operand(test_schema: Schema) Arc::new(ConfigOptions::new()), )) as Arc; - assert!( - !DefaultExpressionConvertor::default() - .try_convert(&array_length, &test_schema) - .unwrap() - .is_some() - ); + assert!(!converts(&array_length, &test_schema)?); + Ok(()) } #[rstest] -fn test_can_be_pushed_down_array_length_dimension_one_supported(test_schema: Schema) { +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!( - DefaultExpressionConvertor::default() - .try_convert(&array_length, &test_schema) - .unwrap() - .is_some() - ); + assert!(converts(&array_length, &test_schema)?); + Ok(()) } #[rstest] @@ -1135,17 +1018,13 @@ fn test_can_be_pushed_down_array_length_dimension_one_supported(test_schema: Sch DataType::Int64, None, )))] -fn test_array_length_unsupported_dimension( +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!( - DefaultExpressionConvertor::default() - .try_convert(&array_length, &test_schema)? - .is_none() - ); + assert!(!converts(&array_length, &test_schema)?); Ok(()) } @@ -1210,10 +1089,8 @@ async fn test_cast_int_to_string() -> anyhow::Result<()> { /// 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)?; + 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()]); @@ -1224,107 +1101,42 @@ fn test_cast_to_uuid_resolves_via_registry() -> anyhow::Result<()> { )); // Must convert without panicking — the static path would `unimplemented!()`. - assert!( - DefaultExpressionConvertor::default() - .try_convert(&cast, &schema)? - .is_some() - ); + 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() { - 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(); +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 + }; - // 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( - &(Arc::new(case_expr) as Arc), - &batch.schema(), - ) - .unwrap() - .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); + 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<()> { @@ -1472,11 +1284,7 @@ fn test_cast_falls_back(#[case] target: DataType) -> DFResult<()> { target, None, )); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_none() - ); + assert!(!converts(&expr, &schema)?); Ok(()) } @@ -1491,11 +1299,7 @@ fn test_cast_options_fall_back() -> DFResult<()> { format_options: DEFAULT_FORMAT_OPTIONS, }), )); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_none() - ); + assert!(!converts(&expr, &schema)?); Ok(()) } @@ -1517,11 +1321,7 @@ fn test_integer_arithmetic_pushdown_ignores_overflow_mode( ) .with_fail_on_overflow(checked), ); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_some() - ); + assert!(converts(&expr, &schema)?); Ok(()) } @@ -1604,11 +1404,7 @@ fn test_fallible_case_branch_is_residual() -> DFResult<()> { )], Some(zero), )?); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_none() - ); + assert!(!converts(&expr, &schema)?); Ok(()) } @@ -1645,11 +1441,7 @@ fn test_nested_functions_reject_unknown_children( Arc::new(ConfigOptions::new()), )?) as Arc }; - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_none() - ); + assert!(!converts(&expr, &schema)?); Ok(()) } @@ -1704,11 +1496,7 @@ fn test_column_identity_validation_skips_dynamic_filters() -> DFResult<()> { Arc::new(df_expr::Literal::new(ScalarValue::Boolean(Some(true)))), )) as Arc; - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &schema)? - .is_none() - ); + assert!(!converts(&expr, &schema)?); Ok(()) } @@ -1748,11 +1536,7 @@ fn test_native_nested_list_length( )?); let length = array_length_expr(vec![get_field], &batch.schema()); if nullable_parent { - assert!( - DefaultExpressionConvertor::default() - .try_convert(&length, &batch.schema())? - .is_none() - ); + assert!(!converts(&length, &batch.schema())?); Ok(()) } else { assert_native_matches(length, batch) @@ -1787,11 +1571,7 @@ fn test_simple_case_falls_back() -> DFResult<()> { vec![(Arc::clone(&value), value)], None, )?); - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &Schema::empty())? - .is_none() - ); + assert!(!converts(&expr, &Schema::empty())?); Ok(()) } @@ -1864,10 +1644,6 @@ fn test_fallible_boolean_rhs_is_residual( )); // DataFusion selects the one row that needs RHS evaluation at this selectivity. expr.evaluate(&batch)?; - assert!( - DefaultExpressionConvertor::default() - .try_convert(&expr, &batch.schema())? - .is_none() - ); + assert!(!converts(&expr, &batch.schema())?); Ok(()) } diff --git a/vortex-datafusion/src/convert/mod.rs b/vortex-datafusion/src/convert/mod.rs index 72ee2438220..476f32afa49 100644 --- a/vortex-datafusion/src/convert/mod.rs +++ b/vortex-datafusion/src/convert/mod.rs @@ -19,6 +19,7 @@ pub(crate) mod stats; pub use exprs::DefaultExpressionConvertor; pub use exprs::ExpressionConvertor; pub use exprs::ProcessedProjection; +pub use scalars::scalar_from_df; /// First-party trait for implementing conversion from DataFusion types to Vortex types. pub trait FromDataFusion: Sized { diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index a046d1ee005..c2b1dfc79b2 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -3,22 +3,33 @@ use std::sync::Arc; +use arrow_array::Array; +use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::StructArray; use arrow_schema::Field; use arrow_schema::Fields; use datafusion_common::ScalarValue; +use vortex::array::VortexSessionExecute; +use vortex::buffer::ByteBuffer; use vortex::dtype::DType; +use vortex::dtype::DecimalDType; use vortex::dtype::NativeDecimalType; +use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::half::f16; use vortex::dtype::i256; +use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; +use vortex::error::vortex_panic; use vortex::extension::datetime::AnyTemporal; use vortex::extension::datetime::TemporalMetadata; use vortex::extension::datetime::TimeUnit; +use vortex::scalar::DecimalValue; use vortex::scalar::Scalar; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; use crate::convert::TryToDataFusion; @@ -174,6 +185,139 @@ impl TryToDataFusion for Scalar { } } +/// Converts a DataFusion [`ScalarValue`] to a Vortex [`Scalar`], resolving Arrow types through +/// `session`'s [`ArrowSession`](vortex_arrow::ArrowSession). +pub fn scalar_from_df(value: &ScalarValue, session: &VortexSession) -> Scalar { + let arrow = session.arrow(); + match value { + ScalarValue::Null => Scalar::null(DType::Null), + ScalarValue::Boolean(b) => b + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))), + ScalarValue::Float16(f) => f + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F16, Nullability::Nullable))), + ScalarValue::Float32(f) => f + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable))), + ScalarValue::Float64(f) => f + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))), + ScalarValue::Int8(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I8, Nullability::Nullable))), + ScalarValue::Int16(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I16, Nullability::Nullable))), + ScalarValue::Int32(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))), + ScalarValue::Int64(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))), + ScalarValue::UInt8(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable))), + ScalarValue::UInt16(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U16, Nullability::Nullable))), + ScalarValue::UInt32(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable))), + ScalarValue::UInt64(i) => i + .map(Scalar::from) + .unwrap_or_else(|| Scalar::null(DType::Primitive(PType::U64, Nullability::Nullable))), + ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => s + .as_ref() + .map(|s| Scalar::from(s.as_str())) + .unwrap_or_else(|| Scalar::null(DType::Utf8(Nullability::Nullable))), + ScalarValue::Binary(b) + | ScalarValue::BinaryView(b) + | ScalarValue::LargeBinary(b) + | ScalarValue::FixedSizeBinary(_, b) => b + .as_ref() + .map(|b| Scalar::binary(ByteBuffer::from(b.clone()), Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(DType::Binary(Nullability::Nullable))), + ScalarValue::Date32(v) + | ScalarValue::Time32Second(v) + | ScalarValue::Time32Millisecond(v) => { + let dtype = arrow + .from_arrow_datatype(&value.data_type(), Nullability::Nullable) + .vortex_expect("arrow data type to dtype"); + Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) + .vortex_expect("unable to create a time `Scalar`") + } + ScalarValue::Date64(v) + | ScalarValue::Time64Microsecond(v) + | ScalarValue::Time64Nanosecond(v) + | ScalarValue::TimestampSecond(v, _) + | ScalarValue::TimestampMillisecond(v, _) + | ScalarValue::TimestampMicrosecond(v, _) + | ScalarValue::TimestampNanosecond(v, _) => { + let dtype = arrow + .from_arrow_datatype(&value.data_type(), Nullability::Nullable) + .vortex_expect("arrow data type to dtype"); + Scalar::try_new(dtype, v.map(vortex::scalar::ScalarValue::from)) + .vortex_expect("unable to create a time `Scalar`") + } + ScalarValue::Decimal32(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I32(*value), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) + } + } + ScalarValue::Decimal64(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I64(*value), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) + } + } + ScalarValue::Decimal128(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I128(*value), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) + } + } + ScalarValue::Decimal256(decimal, precision, scale) => { + let decimal_dtype = DecimalDType::new(*precision, *scale); + let nullable = Nullability::Nullable; + if let Some(value) = decimal { + Scalar::decimal( + DecimalValue::I256(i256::from_le_bytes(value.to_le_bytes())), + decimal_dtype, + Nullability::Nullable, + ) + } else { + Scalar::null(DType::Decimal(decimal_dtype, nullable)) + } + } + ScalarValue::Dictionary(_, v) => scalar_from_df(v.as_ref(), session), + ScalarValue::Struct(array) => struct_from_df(array, session), + _ => unimplemented!("Can't convert {value:?} value to a Vortex scalar"), + } +} + /// Converts a Vortex struct scalar to a DataFusion `ScalarValue::Struct`. fn struct_to_df(scalar: &Scalar) -> VortexResult { let scalar = scalar.as_struct(); @@ -215,11 +359,44 @@ fn struct_to_df(scalar: &Scalar) -> VortexResult { Ok(ScalarValue::Struct(Arc::new(struct_array))) } +/// Converts a DataFusion `ScalarValue::Struct` (a one-row struct array) to a Vortex struct scalar. +/// +/// The struct dtype comes from the Arrow `Fields`, so the children must be converted from those +/// same fields. Going through `ScalarValue` instead would drop each field's `ARROW:extension:name` +/// (and its declared nullability), yielding storage-typed children that +/// [`Scalar::struct_`] rejects against an extension-typed struct dtype. +fn struct_from_df(array: &StructArray, session: &VortexSession) -> Scalar { + let arrow = session.arrow(); + let dtype = arrow + .from_arrow_datatype(array.data_type(), Nullability::Nullable) + .vortex_expect("arrow data type to dtype"); + if array.is_null(0) { + Scalar::null(dtype) + } else { + let mut ctx = session.create_execution_ctx(); + let children = array + .columns() + .iter() + .zip(array.fields().iter()) + .map(|(column, field)| { + arrow + .from_arrow_array(ArrowArrayRef::clone(column), field) + .and_then(|column| column.execute_scalar(0, &mut ctx)) + .unwrap_or_else(|e| { + vortex_panic!("cannot convert struct field to a Vortex scalar: {e}") + }) + }) + .collect::>(); + Scalar::struct_(dtype, children) + } +} + #[cfg(test)] mod tests { use datafusion_common::ScalarValue; use datafusion_common::arrow::datatypes::i256 as arrow_i256; use rstest::rstest; + use vortex::VortexSessionDefault; use vortex::buffer::ByteBuffer; use vortex::dtype::DType; use vortex::dtype::DecimalDType; @@ -234,6 +411,11 @@ mod tests { use super::*; + /// Test shim: convert with a default `VortexSession` passed explicitly. + fn from_df(value: &ScalarValue) -> Scalar { + scalar_from_df(value, &VortexSession::default()) + } + #[rstest] #[case::u8_some(Scalar::from(42u8), ScalarValue::UInt8(Some(42)))] #[case::u8_null( @@ -388,6 +570,160 @@ mod tests { assert_eq!(result, expected_df_scalar); } + #[rstest] + #[case::from_df_null(ScalarValue::Null, Scalar::null(DType::Null))] + #[case::from_df_bool_some(ScalarValue::Boolean(Some(true)), Scalar::from(true))] + #[case::from_df_bool_null( + ScalarValue::Boolean(None), + Scalar::null(DType::Bool(Nullability::Nullable)) + )] + #[case::from_df_i32_some(ScalarValue::Int32(Some(42)), Scalar::from(42i32))] + #[case::from_df_i32_null( + ScalarValue::Int32(None), + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) + )] + #[case::from_df_f64_some(ScalarValue::Float64(Some(2.5)), Scalar::from(2.5f64))] + #[case::from_df_f64_null( + ScalarValue::Float64(None), + Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)) + )] + #[case::from_df_utf8_some(ScalarValue::Utf8(Some("test".to_string())), Scalar::from("test"))] + #[case::from_df_utf8_null( + ScalarValue::Utf8(None), + Scalar::null(DType::Utf8(Nullability::Nullable)) + )] + #[case::from_df_binary_some(ScalarValue::Binary(Some(vec![1, 2, 3])), Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3]), Nullability::Nullable))] + #[case::from_df_binary_null( + ScalarValue::Binary(None), + Scalar::null(DType::Binary(Nullability::Nullable)) + )] + fn test_from_datafusion_scalars( + #[case] df_scalar: ScalarValue, + #[case] expected_vortex: Scalar, + ) { + let result = from_df(&df_scalar); + assert_eq!(result.dtype(), expected_vortex.dtype()); + assert_eq!(result.is_null(), expected_vortex.is_null()); + + // For non-null values, convert both back to DataFusion for comparison + if !result.is_null() { + let result_df = result.try_to_df().unwrap(); + let expected_df = expected_vortex.try_to_df().unwrap(); + assert_eq!(result_df, expected_df); + } + } + + #[rstest] + #[case::decimal128_some(ScalarValue::Decimal128(Some(12345), 10, 2))] + #[case::decimal128_null(ScalarValue::Decimal128(None, 10, 2))] + #[case::decimal256_some(ScalarValue::Decimal256(Some(arrow_i256::from_i128(12345)), 50, 10))] + #[case::decimal256_null(ScalarValue::Decimal256(None, 50, 10))] + fn test_from_datafusion_decimals(#[case] df_scalar: ScalarValue) { + let result = from_df(&df_scalar); + match &df_scalar { + ScalarValue::Decimal128(value, precision, scale) => { + if let DType::Decimal(decimal_type, _) = result.dtype() { + assert_eq!(decimal_type.precision(), *precision); + assert_eq!(decimal_type.scale(), *scale); + if value.is_some() { + assert!(!result.is_null()); + } else { + assert!(result.is_null()); + } + } else { + panic!("Expected decimal type"); + } + } + ScalarValue::Decimal256(value, precision, scale) => { + if let DType::Decimal(decimal_type, _) = result.dtype() { + assert_eq!(decimal_type.precision(), *precision); + assert_eq!(decimal_type.scale(), *scale); + if value.is_some() { + assert!(!result.is_null()); + } else { + assert!(result.is_null()); + } + } else { + panic!("Expected decimal type"); + } + } + _ => panic!("Unexpected scalar type"), + } + } + + #[rstest] + #[case::date32(ScalarValue::Date32(Some(18628)))] // 2021-01-01 + #[case::date64(ScalarValue::Date64(Some(1609459200000)))] // 2021-01-01 in milliseconds + #[case::time32_second(ScalarValue::Time32Second(Some(3661)))] // 01:01:01 + #[case::time32_millisecond(ScalarValue::Time32Millisecond(Some(3661000)))] // 01:01:01 + #[case::time64_microsecond(ScalarValue::Time64Microsecond(Some(3661000000)))] // 01:01:01 + #[case::time64_nanosecond(ScalarValue::Time64Nanosecond(Some(3661000000000)))] // 01:01:01 + #[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 + ))] + fn test_from_datafusion_temporals(#[case] df_scalar: ScalarValue) { + let result = from_df(&df_scalar); + + // All temporal types should convert to extension types + if let DType::Extension(_) = result.dtype() { + assert!(!result.is_null()); + } else { + panic!( + "Expected extension type for temporal scalar, got: {:?}", + result.dtype() + ); + } + } + + #[rstest] + #[case::u32(Scalar::from(42u32))] + #[case::i64(Scalar::from(-123i64))] + #[case::f64(Scalar::from(2.5f64))] + #[case::bool(Scalar::from(true))] + #[case::utf8(Scalar::from("hello world"))] + #[case::null_type(Scalar::null(DType::Null))] + #[case::null_i32(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)))] + #[case::decimal128(Scalar::decimal( + DecimalValue::I128(12345), + DecimalDType::new(10, 2), + Nullability::NonNullable + ))] + #[case::binary(Scalar::binary(ByteBuffer::from(vec![1u8, 2, 3, 4, 5]), Nullability::NonNullable))] + fn test_round_trip_conversions(#[case] original: Scalar) { + let df_scalar = original.try_to_df().unwrap(); + let round_trip = from_df(&df_scalar); + + // Check that core types match (ignoring nullability differences that can occur in round-trip) + assert!( + original.dtype().eq_ignore_nullability(round_trip.dtype()), + "DType mismatch for scalar: {:?} vs {:?}", + original.dtype(), + round_trip.dtype() + ); + + assert_eq!( + original.is_null(), + round_trip.is_null(), + "Null status mismatch for scalar: {:?}", + original + ); + + if !original.is_null() { + // For non-null values, compare by converting both to DataFusion scalars + let original_df = original.try_to_df().unwrap(); + let round_trip_df = round_trip.try_to_df().unwrap(); + assert_eq!( + original_df, round_trip_df, + "Value mismatch for scalar: {:?}", + original + ); + } + } + #[rstest] #[case::null_type(Scalar::null(DType::Null), ScalarValue::Null)] #[case::null_bool( @@ -422,16 +758,80 @@ mod tests { Scalar::null(DType::Decimal(DecimalDType::new(5, 2), Nullability::Nullable)), ScalarValue::Decimal32(None, 5, 2) )] - fn test_null_handling( - #[case] vortex_null: Scalar, - #[case] expected_df_null: ScalarValue, - ) -> VortexResult<()> { - assert_eq!(vortex_null.try_to_df()?, expected_df_null); + fn test_null_handling(#[case] vortex_null: Scalar, #[case] expected_df_null: ScalarValue) { + // Test Vortex -> DataFusion + let df_result = vortex_null.try_to_df().unwrap(); + assert_eq!(df_result, expected_df_null); + + // Test DataFusion -> Vortex + let vortex_result = from_df(&expected_df_null); + assert!(vortex_result.is_null()); + assert!( + vortex_result + .dtype() + .eq_ignore_nullability(vortex_null.dtype()) + ); + } + + #[rstest] + #[case::utf8(ScalarValue::Utf8(Some("test string".to_string())))] + #[case::utf8_view(ScalarValue::Utf8View(Some("test string".to_string())))] + #[case::large_utf8(ScalarValue::LargeUtf8(Some("test string".to_string())))] + fn test_utf8_variants(#[case] variant: ScalarValue) { + let result = from_df(&variant); + assert_eq!(result.as_utf8().value().unwrap().as_str(), "test string"); + } + + #[rstest] + #[case::binary(ScalarValue::Binary(Some(vec![1u8, 2, 3, 4, 5])))] + #[case::binary_view(ScalarValue::BinaryView(Some(vec![1u8, 2, 3, 4, 5])))] + #[case::large_binary(ScalarValue::LargeBinary(Some(vec![1u8, 2, 3, 4, 5])))] + #[case::fixed_size_binary(ScalarValue::FixedSizeBinary(5, Some(vec![1u8, 2, 3, 4, 5])))] + fn test_binary_variants(#[case] variant: ScalarValue) { + let result = from_df(&variant); + let result_bytes: Vec = result + .as_binary() + .value() + .cloned() + .unwrap() + .into_bytes() + .into(); + assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]); + } + + /// A DataFusion struct whose child field carries Arrow extension metadata: the struct dtype + /// derived from the Arrow type keeps the extension, so the child scalars must keep it too or + /// `Scalar::struct_` rejects them. + #[test] + fn struct_from_df_preserves_extension_child() -> VortexResult<()> { + use arrow_array::FixedSizeBinaryArray; + use arrow_schema::DataType; + use arrow_schema::extension::Uuid as ArrowUuid; + use vortex::extension::uuid::Uuid; + + let mut id_field = Field::new("id", DataType::FixedSizeBinary(16), false); + id_field.try_with_extension_type(ArrowUuid)?; + let ids = FixedSizeBinaryArray::try_from_iter([*b"0123456789abcdef"].into_iter())?; + let struct_array = StructArray::try_new( + Fields::from(vec![Arc::new(id_field)]), + vec![Arc::new(ids) as ArrowArrayRef], + None, + )?; + + let scalar = from_df(&ScalarValue::Struct(Arc::new(struct_array))); + let DType::Struct(fields, _) = scalar.dtype() else { + panic!("expected a struct dtype, got {}", scalar.dtype()); + }; + let id_dtype = fields.field_by_index(0).vortex_expect("one field"); + assert!( + id_dtype.as_extension().is::(), + "expected a Uuid extension field, got {id_dtype}" + ); Ok(()) } #[test] - fn test_struct_scalar_to_datafusion() -> VortexResult<()> { + fn struct_scalar_round_trips() -> VortexResult<()> { let dtype = DType::Struct( StructFields::new( FieldNames::from(["x", "y"]), @@ -447,27 +847,17 @@ mod tests { vec![Scalar::from(-111.7610f64), Scalar::from(34.8697f64)], ); - let expected = StructArray::try_new( - vec![ - Field::new("x", arrow_schema::DataType::Float64, false), - Field::new("y", arrow_schema::DataType::Float64, false), - ] - .into(), - vec![ - Arc::new(arrow_array::Float64Array::from(vec![-111.7610])), - Arc::new(arrow_array::Float64Array::from(vec![34.8697])), - ], - None, - )?; - assert_eq!( - original.try_to_df()?, - ScalarValue::Struct(Arc::new(expected)) - ); + let df = original.try_to_df()?; + assert!(matches!(df, ScalarValue::Struct(_))); + + // Back through `from_df` and out again yields the identical DataFusion struct value. + let back = from_df(&df); + assert_eq!(back.try_to_df()?, df); Ok(()) } #[test] - fn test_null_struct_scalar_to_datafusion() -> VortexResult<()> { + fn null_struct_scalar_round_trips() -> VortexResult<()> { let dtype = DType::Struct( StructFields::new( FieldNames::from(["x", "y"]), @@ -481,7 +871,7 @@ mod tests { let df = Scalar::null(dtype).try_to_df()?; assert!(matches!(df, ScalarValue::Struct(_))); - assert!(df.is_null()); + assert!(from_df(&df).is_null()); Ok(()) } diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs new file mode 100644 index 00000000000..267b58ae9fd --- /dev/null +++ b/vortex-datafusion/src/persistent/opener.rs @@ -0,0 +1,2041 @@ +// 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 datafusion_common::DataFusionError; +use datafusion_common::Result as DFResult; +use datafusion_common::ScalarValue; +use datafusion_common::Statistics; +use datafusion_common::arrow::array::AsArray; +use datafusion_common::arrow::array::RecordBatch; +use datafusion_common::exec_datafusion_err; +use datafusion_datasource::PartitionedFile; +use datafusion_datasource::TableSchema; +use datafusion_datasource::file_stream::FileOpenFuture; +use datafusion_datasource::file_stream::FileOpener; +use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; +use datafusion_execution::cache::cache_manager::FileMetadataCache; +use datafusion_physical_expr::PhysicalExprRef; +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; +use datafusion_pruning::FilePruner; +use futures::FutureExt; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use object_store::path::Path; +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; +use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::metrics::Label; +use vortex::metrics::MetricsRegistry; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; +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::ProcessedProjection; +use crate::convert::exprs::raw_projection; +use crate::convert::schema::calculate_physical_schema; +use crate::metrics::PARTITION_LABEL; +use crate::metrics::PATH_LABEL; +use crate::persistent::cache::CachedVortexMetadata; +use crate::persistent::reader::VortexReaderFactory; +use crate::persistent::stream::PrunableStream; + +#[derive(Clone)] +pub(crate) struct VortexOpener { + /// The partition this opener is assigned to. Only used for labeling metrics. + pub partition: usize, + pub session: VortexSession, + pub vortex_reader_factory: Arc, + /// 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, + /// 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. + pub file_pruning_predicate: Option, + pub expr_adapter_factory: Arc, + /// This is the table's schema without partition columns. It may contain fields which do + /// not exist in the file, and are supplied by the `schema_adapter_factory`. + pub table_schema: TableSchema, + /// If provided, the scan will not return more than this many rows. + pub limit: Option, + /// A metrics object for tracking performance of the scan. + pub metrics_registry: Arc, + /// DataFusion-native metrics exposed through `DataSourceExec`. + pub df_metrics: ExecutionPlanMetricsSet, + /// A shared cache of file readers. + /// + /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache + /// a file reader the first time we read a file. + pub layout_readers: Arc>>, + /// Shared full-file natural splits keyed by file path. + pub natural_splits: Arc>>, + /// Whether the query has output ordering specified + pub has_output_ordering: bool, + + pub expression_convertor: Arc, + pub file_metadata_cache: Option>, + /// Whether to enable expression pushdown into the underlying Vortex scan. + pub projection_pushdown: bool, + pub scan_concurrency: Option, +} + +impl FileOpener for VortexOpener { + fn open(&self, file: PartitionedFile) -> DFResult { + // Calculate the output schema before replacing partition columns with literals so it + // retains the table and partition-field metadata declared by the plan. + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); + let session = self.session.clone(); + let metrics_registry = Arc::clone(&self.metrics_registry); + let labels = vec![ + Label::new(PATH_LABEL, file.path().to_string()), + Label::new(PARTITION_LABEL, self.partition.to_string()), + ]; + + let mut projection = self.projection.clone(); + let mut filter = self.filter.clone(); + + let reader = self.vortex_reader_factory.create_reader(&file, &session)?; + + let reader = + InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); + + let mut file_pruning_predicate = self.file_pruning_predicate.clone(); + let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); + let file_metadata_cache = self.file_metadata_cache.clone(); + + let unified_file_schema = Arc::clone(self.table_schema.file_schema()); + let limit = self.limit; + let layout_readers = Arc::clone(&self.layout_readers); + let natural_splits = Arc::clone(&self.natural_splits); + let has_output_ordering = self.has_output_ordering; + let scan_concurrency = self.scan_concurrency; + + let expr_convertor = Arc::clone(&self.expression_convertor); + let projection_pushdown = self.projection_pushdown; + + let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) + .with_category(MetricCategory::Rows) + .global_counter("num_predicate_creation_errors"); + + // Replace column access for partition columns with literals + #[expect(clippy::disallowed_types)] + let literal_value_cols = self + .table_schema + .table_partition_cols() + .iter() + .map(|f| f.name()) + .cloned() + .zip(file.partition_values.clone()) + .collect::>(); + + let predicate_uses_partition_columns = + file_pruning_predicate.as_ref().is_some_and(|predicate| { + collect_columns(predicate) + .iter() + .any(|column| literal_value_cols.contains_key(column.name())) + }); + + if !literal_value_cols.is_empty() { + projection = projection.try_map_exprs(|expr| { + replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) + })?; + filter = filter + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; + file_pruning_predicate = file_pruning_predicate + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; + } + + Ok(async move { + // FilePruner requires a statistics object even when the rewritten predicate + // only contains partition literals. Supply unknown file-column statistics in + // that case so static and dynamic partition predicates can still prune. + let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) + .then(|| { + file.clone() + .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) + }); + let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); + + let mut file_pruner = file_pruning_predicate + .filter(|_| file.has_statistics() || predicate_uses_partition_columns) + .and_then(|predicate| { + FilePruner::try_new( + Arc::clone(&predicate), + &unified_file_schema, + pruning_file, + predicate_creation_errors, + ) + }); + + // Check if this file should be pruned based on statistics/partition values. + // Returns empty stream if file can be skipped entirely. + if let Some(file_pruner) = file_pruner.as_mut() + && file_pruner.should_prune()? + { + return Ok(stream::empty().boxed()); + } + + let mut open_opts = session + .open_options() + .with_file_size(file.object_meta.size) + .with_metrics_registry(Arc::clone(&metrics_registry)) + .with_labels(labels); + + let cached_footer = file_metadata_cache + .as_ref() + .and_then(|cache| cache.get(file.path())) + .filter(|entry| entry.is_valid_for(&file.object_meta)) + .and_then(|entry| { + entry + .file_metadata + .as_any() + .downcast_ref::() + .map(|vortex_metadata| vortex_metadata.footer().clone()) + }); + let footer_cache_hit = cached_footer.is_some(); + + if let Some(footer) = cached_footer { + open_opts = open_opts.with_footer(footer); + } + + let vxf = open_opts + .open_read(reader) + .await + .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?; + + // On a miss, cache the parsed footer so other partitions and later executions + // skip the footer fetch and parse. `infer_schema`/`infer_stats` also populate + // this cache, but only when planning goes through `VortexFormat`. + if !footer_cache_hit && let Some(cache) = &file_metadata_cache { + cache.put( + file.path(), + CachedFileMetadataEntry::new( + file.object_meta.clone(), + Arc::new(CachedVortexMetadata::new(&vxf)), + ), + ); + } + + // Check if there are rows in this file. If not, we can save + // ourselves some work and return an empty stream. + if vxf.row_count() == 0 { + return Ok(stream::empty().boxed()); + } + + // This is the expected arrow types of the actual columns in the file, which might have different types + // from the unified logical schema or miss + let this_file_schema = Arc::new(calculate_physical_schema( + vxf.dtype(), + &unified_file_schema, + &session.arrow(), + )?); + + let expr_adapter = expr_adapter_factory.create( + Arc::clone(&unified_file_schema), + Arc::clone(&this_file_schema), + )?; + + let simplifier = PhysicalExprSimplifier::new(&this_file_schema); + + // The adapter rewrites the expressions to the local file schema, allowing + // for schema evolution and divergence between the table's schema and individual files. + let filter = filter + .map(|filter| { + // Expression might now reference columns that don't exist in the file, so we can give it + // another simplification pass. + simplifier.simplify(expr_adapter.rewrite(filter)?) + }) + .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_convertor + .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 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_convertor.split_projection( + projection, + &this_file_schema, + output_schema.as_ref(), + )? + } else { + expr_convertor.no_pushdown_projection(projection, &this_file_schema)? + }; + + // The schema of the stream returned from the vortex scan. + // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). + let scan_projection = scan_projection + .optimize_recursive(vxf.dtype()) + .and_then(|projection| projection.bind(vxf.dtype())) + .map_err(|_e| { + exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") + })?; + let scan_dtype = scan_projection.dtype().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. + let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { + Entry::Occupied(mut occupied_entry) => { + if let Some(reader) = occupied_entry.get().upgrade() { + tracing::trace!("reusing layout reader for {}", occupied_entry.key()); + reader + } else { + tracing::trace!("creating layout reader for {}", occupied_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create layout reader: {e}" + )) + })?; + occupied_entry.insert(Arc::downgrade(&reader)); + reader + } + } + Entry::Vacant(vacant_entry) => { + tracing::trace!("creating layout reader for {}", vacant_entry.key()); + let reader = vxf.layout_reader().map_err(|e| { + DataFusionError::Execution(format!("Failed to create layout reader: {e}")) + })?; + vacant_entry.insert(Arc::downgrade(&reader)); + + reader + } + }; + + let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); + + if let Some(vortex_plan) = file.extensions.get::() { + scan_builder = vortex_plan.apply_to_builder(scan_builder); + } + + if let Some(limit) = limit + && native_filter.is_none() + && residual_filter.is_none() + { + scan_builder = scan_builder.with_limit(limit); + } + + if let Some(concurrency) = scan_concurrency { + scan_builder = scan_builder.with_concurrency(concurrency); + } + + // Set before the byte-range translation below, which computes natural splits for + // the fields the scan's projection and filter reference. + scan_builder = scan_builder + .with_projection(scan_projection) + .with_some_filter(native_filter); + + 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"))?, + end: u64::try_from(file_range.end) + .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, + }; + if byte_range.start != 0 || byte_range.end != file.object_meta.size { + // Full-file scans already cover every natural split. Only translate the + // byte range back into row boundaries when DataFusion has trimmed the file. + let natural_splits = natural_splits_for_file( + natural_splits.as_ref(), + &file.object_meta.location, + &scan_builder, + file.object_meta.size, + )?; + + let Some(row_range) = + split_aligned_row_range(byte_range, natural_splits.as_ref()) + else { + return Ok(stream::empty().boxed()); + }; + + scan_builder = scan_builder + .with_row_range(row_range) + // Hand the shared full-file boundaries back to the scan so prepare() + // skips its own layout walk. + .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); + } + } + + 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) + .map(move |chunk| { + let mut ctx = session.create_execution_ctx(); + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) + .into_stream() + .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? + .map_err(move |e: VortexError| { + DataFusionError::External(Box::new(e.with_context(format!( + "Failed to read Vortex file: {}", + file.object_meta.location + )))) + }) + .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(); + Ok(Some(RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(row_count)), + )?)) + }) + .filter_map(|batch| ready(batch.transpose())) + .boxed(); + + if let Some(file_pruner) = file_pruner + && file_pruner.is_watching() + { + Ok(PrunableStream::new(file_pruner, stream).boxed()) + } else { + Ok(stream) + } + } + .in_current_span() + .boxed()) + } +} + +/// A file's natural split boundaries plus the precomputed byte each split is assigned to, +/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row +/// boundaries with a binary search instead of re-projecting every split per partition. +/// +/// The boundaries are computed for the fields referenced by the scan's projection and filter. +/// All partitions translate through the first opener's cached entry (the cache lives on the +/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges +/// tiling the file's rows exactly once. +#[derive(Debug)] +pub(crate) struct NaturalSplits { + /// Sorted row boundaries of the natural splits; split `i` covers + /// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the + /// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the + /// per-partition layout walk in `prepare`. + row_boundaries: Arc<[u64]>, + /// For each split, the byte a DataFusion byte range must contain to own it (see + /// [`split_assignment_byte`]); one entry per split, sorted because split midpoints + /// increase monotonically under the row-to-byte projection. + assignment_bytes: Box<[u64]>, +} + +impl NaturalSplits { + fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { + let row_count = row_boundaries.last().copied().unwrap_or_default(); + let assignment_bytes = if row_count == 0 { + Box::default() + } else { + row_boundaries + .windows(2) + .enumerate() + .map(|(idx, boundaries)| { + split_assignment_byte( + idx, + &(boundaries[0]..boundaries[1]), + row_count, + total_size, + ) + }) + .collect() + }; + + debug_assert!(assignment_bytes.is_sorted()); + debug_assert_eq!( + assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), + row_boundaries.len() + ); + + Self { + row_boundaries, + assignment_bytes, + } + } +} + +/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. +fn natural_splits_for_file( + natural_splits: &DashMap>, + path: &Path, + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + if let Some(splits) = natural_splits.get(path) { + return Ok(Arc::clone(splits.value())); + } + + // Compute while holding the entry so concurrent partitions opening the same file wait + // for the winner instead of all walking the layout tree; the redundant walks contend on + // the lazily-initialized layout children and dominate the cost of the computation itself. + match natural_splits.entry(path.clone()) { + Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), + Entry::Vacant(entry) => { + let splits = compute_natural_splits(scan_builder, total_size)?; + entry.insert(Arc::clone(&splits)); + Ok(splits) + } + } +} + +/// Walk the layout tree to compute the file's full natural split boundaries for the fields +/// referenced by the scan's projection and filter. +fn compute_natural_splits( + scan_builder: &ScanBuilder, + total_size: u64, +) -> DFResult> { + let row_boundaries = scan_builder + .full_file_splits() + .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; + + Ok(Arc::new(NaturalSplits::new( + row_boundaries.into(), + total_size, + ))) +} + +/// Translate a DataFusion byte range to the contiguous natural split ranges it owns. +/// Most splits are assigned by midpoint, but the leading split stays with the range that owns +/// byte 0 so a tiny first byte range still claims the first rows. +fn split_aligned_row_range( + byte_range: Range, + natural_splits: &NaturalSplits, +) -> Option> { + if byte_range.start >= byte_range.end { + return None; + } + + let first_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.start); + let after_last_split = natural_splits + .assignment_bytes + .partition_point(|&assignment_byte| assignment_byte < byte_range.end); + if first_split == after_last_split { + return None; + } + + Some( + natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split], + ) +} + +fn split_assignment_byte( + idx: usize, + split_range: &Range, + row_count: u64, + total_size: u64, +) -> u64 { + if idx == 0 && split_range.start == 0 { + // Byte 0 is the only stable representative for the leading split. A midpoint can fall + // into the next DataFusion byte range and leave the first range with no rows to read. + 0 + } else { + split_midpoint_to_byte(split_range, row_count, total_size) + } +} + +fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64) -> u64 { + let midpoint_row = split_range.start + (split_range.end - split_range.start) / 2; + let midpoint_byte = (u128::from(midpoint_row) * u128::from(total_size)) / u128::from(row_count); + + u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") +} + +#[cfg(test)] +mod tests { + use std::fmt; + use std::sync::Arc; + use std::sync::LazyLock; + + use arrow_array::record_batch; + use arrow_schema::Field; + use arrow_schema::Fields; + use arrow_schema::SchemaRef; + use datafusion::arrow::array::DictionaryArray; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::array::RecordBatch; + use datafusion::arrow::array::StringArray; + use datafusion::arrow::array::StructArray; + use datafusion::arrow::datatypes::DataType; + use datafusion::arrow::datatypes::Schema; + use datafusion::arrow::datatypes::UInt32Type; + use datafusion::arrow::util::display::FormatOptions; + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + use datafusion::logical_expr::col; + use datafusion::logical_expr::lit; + 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; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions as df_expr; + use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; + use datafusion_physical_expr::projection::ProjectionExpr; + use insta::assert_snapshot; + use itertools::Itertools; + use object_store::ObjectStore; + use object_store::memory::InMemory; + use rstest::rstest; + use vortex::VortexSessionDefault; + use vortex::buffer::Buffer; + use vortex::file::WriteOptionsSessionExt; + use vortex::io::VortexWrite; + use vortex::io::object_store::ObjectStoreWrite; + use vortex::metrics::DefaultMetricsRegistry; + use vortex::scan::selection::Selection; + use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; + use vortex::session::VortexSession; + + use super::*; + use crate::VortexAccessPlan; + use crate::convert::exprs::DefaultExpressionConvertor; + use crate::persistent::reader::DefaultVortexReaderFactory; + + static SESSION: LazyLock = LazyLock::new(VortexSession::default); + + /// Test-only expr used to test error reporting. + #[derive(Debug, Eq, Hash, PartialEq)] + struct SnapshotErrorExpr; + + impl fmt::Display for SnapshotErrorExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "snapshot_error") + } + } + + impl PhysicalExpr for SnapshotErrorExpr { + fn data_type(&self, _input_schema: &Schema) -> DFResult { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> DFResult { + Ok(false) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self, f) + } + + fn evaluate(&self, _batch: &RecordBatch) -> DFResult { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + Vec::new() + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> DFResult { + assert!(children.is_empty()); + Ok(self) + } + + fn snapshot(&self) -> DFResult> { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + } + + fn natural_splits(total_size: u64, split_ranges: &[Range]) -> NaturalSplits { + let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1); + if let Some(first) = split_ranges.first() { + row_boundaries.push(first.start); + row_boundaries.extend(split_ranges.iter().map(|range| range.end)); + } + NaturalSplits::new(row_boundaries.into(), total_size) + } + + #[rstest] + #[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] + #[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] + #[case(1..8, 10, vec![0..1, 1..9, 9..10], Some(1..9))] + #[case(1..4, 16, vec![0..1, 1..2, 2..3, 3..4], None)] + #[case(0..1, 10, vec![0..2, 2..10], Some(0..2))] + #[case(0..2, 2, vec![], None)] + fn test_split_aligned_row_range( + #[case] byte_range: Range, + #[case] total_size: u64, + #[case] split_ranges: Vec>, + #[case] expected: Option>, + ) { + assert_eq!( + split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)), + expected + ); + } + + #[test] + fn test_split_aligned_ranges_cover_splits_exactly_once() { + let split_ranges = vec![0..1, 1..4, 4..10, 10..13]; + let byte_ranges = [0..4, 4..8, 8..12, 12..16]; + let natural_splits = natural_splits(16, &split_ranges); + + let assigned = byte_ranges + .into_iter() + .filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits)) + .collect::>(); + + assert_eq!(assigned, vec![0..4, 4..10, 10..13]); + assert_eq!( + assigned + .iter() + .map(|range| range.end - range.start) + .sum::(), + 13 + ); + + let split_starts = split_ranges + .iter() + .map(|range| range.start) + .collect::>(); + let split_ends = split_ranges + .iter() + .map(|range| range.end) + .collect::>(); + + for range in &assigned { + assert!(split_starts.contains(&range.start)); + assert!(split_ends.contains(&range.end)); + } + + for (left, right) in assigned.iter().tuple_windows() { + assert_eq!(left.end, right.start); + } + } + + #[rstest] + #[case(vec![], 10)] + #[case(vec![0], 10)] + #[case(vec![], 0)] + #[case(vec![0], 0)] + fn test_natural_splits_empty_file(#[case] row_boundaries: Vec, #[case] total_size: u64) { + let splits = NaturalSplits::new(row_boundaries.clone().into(), total_size); + + assert!(splits.assignment_bytes.is_empty()); + assert_eq!(splits.row_boundaries.as_ref(), row_boundaries.as_slice()); + assert_eq!(split_aligned_row_range(0..u64::MAX, &splits), None); + } + + #[test] + fn test_split_aligned_row_range_keeps_colliding_assignments_together() { + let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]); + + assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]); + assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2)); + assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); + } + + async fn write_arrow_to_vortex( + object_store: Arc, + path: &str, + rb: RecordBatch, + ) -> anyhow::Result { + let schema = rb.schema(); + let array = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; + let path = Path::parse(path)?; + + let mut write = ObjectStoreWrite::new(object_store, &path).await?; + let summary = SESSION + .write_options() + .write(&mut write, array.to_array_stream()) + .await?; + write.shutdown().await?; + + Ok(summary.size()) + } + + fn make_opener( + object_store: Arc, + table_schema: TableSchema, + filter: Option, + ) -> VortexOpener { + VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), + projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), + filter, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema, + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + } + } + + #[tokio::test] + async fn test_open() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) + .build(); + + // filter matches partition value + let filter = col("part").eq(lit(1)); + let filter = logical2physical(&filter, table_schema.table_schema()); + + let opener = make_opener( + Arc::clone(&object_store), + table_schema.clone(), + Some(filter), + ); + let stream = opener.open(file.clone()).unwrap().await.unwrap(); + + let data = stream.try_collect::>().await?; + let num_batches = data.len(); + let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); + + assert_eq!((num_batches, num_rows), (1, 3)); + + // filter doesn't matches partition value + let filter = col("part").eq(lit(2)); + let filter = logical2physical(&filter, table_schema.table_schema()); + + let opener = make_opener( + Arc::clone(&object_store), + table_schema.clone(), + Some(filter), + ); + let stream = opener.open(file.clone()).unwrap().await.unwrap(); + + let data = stream.try_collect::>().await?; + let num_batches = data.len(); + let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); + assert_eq!((num_batches, num_rows), (0, 0)); + + 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; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = Arc::new( + batch.schema().as_ref().clone().with_metadata( + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + ); + let table_schema = TableSchema::builder(file_schema) + .with_table_partition_cols(vec![Arc::new( + Field::new("part", DataType::Int32, false).with_metadata( + [("partition".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + )]) + .build(); + let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); + + assert_eq!( + expected_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + assert_eq!( + expected_schema.field(1).metadata().get("partition"), + Some(&"metadata".to_string()) + ); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection = projection.clone(); + opener.projection_pushdown = projection_pushdown; + + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); + } + } + + Ok(()) + } + + #[tokio::test] + async fn test_open_all_valid_nullable_columns_with_nonnullable_table_schema() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "nullable/file.vortex"; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], + )?; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + let expected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::from(Arc::clone(&expected_schema)); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection_pushdown = projection_pushdown; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema().as_ref(), expected_schema.as_ref()); + } + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_replaces_partition_columns_without_file_statistics() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) + .build(); + + let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; + let predicate = Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&partition_column), + Operator::Gt, + df_expr::lit(ScalarValue::Int32(Some(1))), + )) as PhysicalExprRef; + let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![partition_column], + predicate, + )) as PhysicalExprRef; + + let mut opener = make_opener(object_store, table_schema, None); + opener.file_pruning_predicate = Some(dynamic_predicate); + let df_metrics = opener.df_metrics.clone(); + + // The file does not exist and has no statistics. Replacing `part` with 1 + // makes the predicate false, so pruning must happen before any file I/O. + let mut file = PartitionedFile::new("missing.vortex", 1); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(batches.is_empty()); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(0) + ); + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "metrics/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); + statistics.column_statistics[0].null_count = Precision::Exact(0); + let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); + + let mut opener = make_opener(object_store, TableSchema::from(batch.schema()), None); + opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); + let df_metrics = opener.df_metrics.clone(); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(1) + ); + + Ok(()) + } + + #[tokio::test] + async fn test_open_empty_file() -> anyhow::Result<()> { + use futures::TryStreamExt; + + let object_store = Arc::new(InMemory::new()) as Arc; + let data_batch = record_batch!(("a", Int32, Vec::::new())).unwrap(); + let file_path = "part=1/empty.vortex"; + let file_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, data_batch.clone()).await?; + + let file_schema = data_batch.schema(); + // Parallel scans may attach a byte range even for empty files; the + // opener must return early before attempting split-aligned translation. + let file = + PartitionedFile::new_with_range(file_path.to_string(), file_size, 0, file_size as i64); + + let table_schema = TableSchema::from(Arc::clone(&file_schema)); + + let opener = make_opener(object_store, table_schema, None); + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + assert_eq!(data.len(), 0); + + Ok(()) + } + + #[tokio::test] + async fn test_open_populates_file_metadata_cache() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "cached/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let table_schema = TableSchema::from(batch.schema()); + + let cache: Arc = Arc::new( + DefaultCache::::new(64 * 1024 * 1024), + ); + let mut opener = make_opener(Arc::clone(&object_store), table_schema, None); + opener.file_metadata_cache = Some(Arc::clone(&cache)); + + // The first open misses the cache and must write the parsed footer back. + let stream = opener.open(file.clone())?.await?; + stream.try_collect::>().await?; + + let entry = cache + .get(file.path()) + .ok_or_else(|| anyhow::anyhow!("footer was not cached after open"))?; + assert!(entry.is_valid_for(&file.object_meta)); + assert!( + entry + .file_metadata + .as_any() + .downcast_ref::() + .is_some() + ); + + // The second open hits the cache and still returns the same data. + let stream = opener.open(file.clone())?.await?; + let data = stream.try_collect::>().await?; + assert_eq!(data.iter().map(|rb| rb.num_rows()).sum::(), 3); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn test_open_files_different_table_schema() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + + let file1 = { + let file1_path = "/path/file1.vortex"; + let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size1 = + write_arrow_to_vortex(Arc::clone(&object_store), file1_path, batch1).await?; + PartitionedFile::new(file1_path.to_string(), data_size1) + }; + + let file2 = { + let file2_path = "/path/file2.vortex"; + let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); + let data_size2 = + write_arrow_to_vortex(Arc::clone(&object_store), file2_path, batch2).await?; + PartitionedFile::new(file2_path.to_string(), data_size2) + }; + + // Table schema has can accommodate both files + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Int32, + true, + )]))); + + let make_opener = |filter| VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone( + &object_store, + ))), + projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), + filter: Some(filter), + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: table_schema.clone(), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + let filter = col("a").lt(lit(100_i32)); + let filter = logical2physical(&filter, table_schema.table_schema()); + + let opener1 = make_opener(Arc::clone(&filter)); + let stream = opener1.open(file1)?.await?; + + let format_opts = FormatOptions::new().with_types_info(true); + + let data = stream.try_collect::>().await?; + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+ + | a | + | Int32 | + +-------+ + | 1 | + | 2 | + | 3 | + +-------+ + "); + + let opener2 = make_opener(Arc::clone(&filter)); + let stream = opener2.open(file2)?.await?; + + let data = stream.try_collect::>().await?; + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+ + | a | + | Int32 | + +-------+ + | -1 | + | -2 | + | -3 | + +-------+ + "); + + Ok(()) + } + + #[tokio::test] + // This test verifies that files with different column order than the + // table schema can be opened without errors. The fix ensures that the + // schema mapper is only used for type casting, not for reordering, + // since the vortex projection already handles reordering. + async fn test_schema_different_column_order() -> anyhow::Result<()> { + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + // File has columns in order: c, b, a + let batch = record_batch!( + ("c", Int32, vec![Some(300), Some(301), Some(302)]), + ("b", Int32, vec![Some(200), Some(201), Some(202)]), + ("a", Int32, vec![Some(100), Some(101), Some(102)]) + ) + .unwrap(); + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + let file = PartitionedFile::new(file_path.to_string(), data_size); + + // Table schema has columns in different order: a, b, c + let table_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + + let opener = VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), + projection: ProjectionExprs::from_indices(&[0, 1, 2], &table_schema), + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: TableSchema::from(Arc::clone(&table_schema)), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + let stream = opener.open(file)?.await?; + + let format_opts = FormatOptions::new().with_types_info(true); + let data = stream.try_collect::>().await?; + + // Verify the output has columns in table schema order (a, b, c) + // not file order (c, b, a) + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+-------+-------+ + | a | b | c | + | Int32 | Int32 | Int32 | + +-------+-------+-------+ + | 100 | 200 | 300 | + | 101 | 201 | 301 | + | 102 | 202 | 302 | + +-------+-------+-------+ + "); + + Ok(()) + } + + #[tokio::test] + // This test verifies that expression rewriting doesn't fail when there is + // a nested schema mismatch between the physical file schema and logical + // table schema. + async fn test_adapter_logical_physical_struct_mismatch() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + let file_struct_fields = Fields::from(vec![ + Field::new("field1", DataType::Utf8, true), + Field::new("field2", DataType::Utf8, true), + ]); + let struct_array = StructArray::new( + file_struct_fields.clone(), + vec![ + Arc::new(StringArray::from(vec!["value1", "value2", "value3"])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "my_struct", + DataType::Struct(file_struct_fields), + true, + )])), + vec![Arc::new(struct_array)], + )?; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + // Table schema has an extra utf8 field. + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( + "my_struct", + DataType::Struct(Fields::from(vec![ + Field::new( + "field1", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + Field::new( + "field2", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + Field::new("field3", DataType::Utf8, true), + ])), + true, + )]))); + + let opener = make_opener( + Arc::clone(&object_store), + table_schema.clone(), + // expression references my_struct column which has different fields in each + // field. + Some(logical2physical( + &col("my_struct").is_not_null(), + table_schema.table_schema(), + )), + ); + + // The opener should be able to open the file with a filter on the + // struct column. + let data = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect::>() + .await?; + + assert_eq!(data.len(), 1); + assert_eq!(data[0].num_rows(), 3); + + Ok(()) + } + + #[tokio::test] + // Minimal reproducing test for the schema projection bug. + // Before the fix, this would fail with a cast error when the file schema + // and table schema have different field orders and we project a subset of columns. + async fn test_projection_bug_minimal_repro() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + // File has columns in order: a, b, c with simple types + let batch = record_batch!( + ("a", Int32, vec![Some(1)]), + ("b", Utf8, vec![Some("test")]), + ("c", Int32, vec![Some(2)]) + ) + .unwrap(); + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + // Table schema has columns in DIFFERENT order: c, a, b + // and different types that require casting (Utf8 -> Dictionary) + let table_schema = TableSchema::from(Arc::new(Schema::new(vec![ + Field::new("c", DataType::Int32, true), + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + ), + ]))); + + // Project columns [0, 2] from table schema, which should give us: c, b + // Before the fix, the schema adapter would get confused about which fields + // to select from the file, causing incorrect type mappings. + let projection = vec![0, 2]; + + let opener = VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone( + &object_store, + ))), + projection: ProjectionExprs::from_indices( + projection.as_ref(), + table_schema.file_schema(), + ), + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: table_schema.clone(), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + // This should succeed and return the correctly projected and cast data + let data = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect::>() + .await?; + + // Verify the columns are in the right order and have the right values + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + let format_opts = FormatOptions::new().with_types_info(true); + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+--------------------------+ + | c | b | + | Int32 | Dictionary(UInt32, Utf8) | + +-------+--------------------------+ + | 2 | test | + +-------+--------------------------+ + "); + + Ok(()) + } + + fn make_test_batch_with_10_rows() -> RecordBatch { + record_batch!( + ("a", Int32, (0..=9).map(Some).collect::>()), + ( + "b", + Utf8, + (0..=9).map(|i| Some(format!("r{}", i))).collect::>() + ) + ) + .unwrap() + } + + fn make_test_opener( + object_store: Arc, + schema: SchemaRef, + projection: ProjectionExprs, + ) -> VortexOpener { + VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), + projection, + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema: TableSchema::from(schema), + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + } + } + + #[tokio::test] + // Test that Selection::IncludeByIndex filters to specific row indices. + async fn test_selection_include_by_index() -> anyhow::Result<()> { + use datafusion::arrow::util::pretty::pretty_format_batches_with_options; + + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.extensions + .insert( + VortexAccessPlan::default().with_selection(Selection::IncludeByIndex( + StrictSortedBuffer::try_new(Buffer::from_iter(vec![1, 3, 5, 7]))?, + )), + ); + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0, 1], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + let format_opts = FormatOptions::new().with_types_info(true); + + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+------+ + | a | b | + | Int32 | Utf8 | + +-------+------+ + | 1 | r1 | + | 3 | r3 | + | 5 | r5 | + | 7 | r7 | + +-------+------+ + "); + + Ok(()) + } + + #[tokio::test] + // Test that Selection::ExcludeByIndex excludes specific row indices. + async fn test_selection_exclude_by_index() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.extensions + .insert( + VortexAccessPlan::default().with_selection(Selection::ExcludeByIndex( + StrictSortedBuffer::try_new(Buffer::from_iter(vec![0, 2, 4, 6, 8]))?, + )), + ); + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0, 1], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + let format_opts = FormatOptions::new().with_types_info(true); + + assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" + +-------+------+ + | a | b | + | Int32 | Utf8 | + +-------+------+ + | 1 | r1 | + | 3 | r3 | + | 5 | r5 | + | 7 | r7 | + | 9 | r9 | + +-------+------+ + "); + + Ok(()) + } + + #[tokio::test] + // Test that Selection::All returns all rows. + async fn test_selection_all() -> anyhow::Result<()> { + use vortex::scan::selection::Selection; + + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.extensions + .insert(VortexAccessPlan::default().with_selection(Selection::All)); + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); + assert_eq!(total_rows, 10); + + Ok(()) + } + + #[tokio::test] + // Test that when no extensions are provided, all rows are returned (backward compatibility). + async fn test_selection_no_extensions() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = make_test_batch_with_10_rows(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let schema = batch.schema(); + let file = PartitionedFile::new(file_path.to_string(), data_size); + // file.extensions is None by default + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0], &schema), + ); + + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); + assert_eq!(total_rows, 10); + + Ok(()) + } + + #[tokio::test] + async fn test_projection_expr_pushdown() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "/path/file.vortex"; + + let batch = record_batch!( + ("a", Int32, vec![Some(1), Some(2), Some(3)]), + ("b", Int32, vec![Some(10), Some(20), Some(30)]) + ) + .unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = batch.schema(); + let table_schema = TableSchema::from(Arc::clone(&file_schema)); + + // Create a projection that includes an arithmetic expression: a + b * 2 + let col_a = df_expr::col("a", &file_schema)?; + let col_b = df_expr::col("b", &file_schema)?; + let two = df_expr::lit(ScalarValue::Int32(Some(2))); + + // b * 2 + let b_times_2 = df_expr::binary(col_b, Operator::Multiply, two, &file_schema)?; + // a + (b * 2) + let a_plus_b_times_2 = df_expr::binary(col_a, Operator::Plus, b_times_2, &file_schema)?; + + let projection = ProjectionExprs::new(vec![ProjectionExpr::new( + a_plus_b_times_2, + "result".to_string(), + )]); + + let opener = VortexOpener { + partition: 1, + session: SESSION.clone(), + vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone( + &object_store, + ))), + projection, + filter: None, + file_pruning_predicate: None, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), + table_schema, + limit: None, + metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), + layout_readers: Default::default(), + natural_splits: Default::default(), + has_output_ordering: false, + expression_convertor: Arc::new(DefaultExpressionConvertor::default()), + file_metadata_cache: None, + projection_pushdown: false, + scan_concurrency: None, + }; + + let file = PartitionedFile::new(file_path.to_string(), data_size); + let stream = opener.open(file)?.await?; + let data = stream.try_collect::>().await?; + + // Expected: a + b * 2 + // row 0: 1 + 10 * 2 = 21 + // row 1: 2 + 20 * 2 = 42 + // row 2: 3 + 30 * 2 = 63 + assert_snapshot!(pretty_format_batches_with_options(&data, &FormatOptions::new().with_types_info(true))?.to_string(), @r" + +--------+ + | result | + | Int32 | + +--------+ + | 21 | + | 42 | + | 63 | + +--------+ + "); + + Ok(()) + } + + /// When a Struct contains Dictionary fields, writing to vortex and reading back + /// should preserve the Dictionary type. + #[tokio::test] + async fn test_struct_with_dictionary_roundtrip() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + + let struct_fields = Fields::from(vec![ + Field::new_dictionary("a", DataType::UInt32, DataType::Utf8, true), + Field::new_dictionary("b", DataType::UInt32, DataType::Utf8, true), + ]); + let struct_array = StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(DictionaryArray::::from_iter(["x", "y", "x"])), + Arc::new(DictionaryArray::::from_iter(["p", "p", "q"])), + ], + None, + ); + + let schema = Arc::new(Schema::new(vec![Field::new( + "labels", + DataType::Struct(struct_fields.clone()), + false, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(struct_array)])?; + + let file_path = "/test.vortex"; + let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; + + let opener = make_test_opener( + Arc::clone(&object_store), + Arc::clone(&schema), + ProjectionExprs::from_indices(&[0], &schema), + ); + let data: Vec<_> = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect() + .await?; + + assert_eq!( + data[0].schema().field(0).data_type(), + &DataType::Struct(struct_fields), + "Struct(Dictionary) type should be preserved" + ); + 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!( + DefaultExpressionConvertor::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_convertor = Arc::new(ResidualConvertor); + 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 ResidualConvertor; + impl ExpressionConvertor for ResidualConvertor { + 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_convertor = Arc::new(ResidualConvertor); + 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_convertor = Arc::new(ResidualConvertor); + 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/opener/mod.rs b/vortex-datafusion/src/persistent/opener/mod.rs deleted file mode 100644 index c9bb11f1a37..00000000000 --- a/vortex-datafusion/src/persistent/opener/mod.rs +++ /dev/null @@ -1,648 +0,0 @@ -// 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::Field; -use datafusion_common::DataFusionError; -use datafusion_common::Result as DFResult; -use datafusion_common::ScalarValue; -use datafusion_common::Statistics; -use datafusion_common::arrow::array::AsArray; -use datafusion_common::arrow::array::RecordBatch; -use datafusion_common::exec_datafusion_err; -use datafusion_datasource::PartitionedFile; -use datafusion_datasource::TableSchema; -use datafusion_datasource::file_stream::FileOpenFuture; -use datafusion_datasource::file_stream::FileOpener; -use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; -use datafusion_execution::cache::cache_manager::FileMetadataCache; -use datafusion_physical_expr::PhysicalExprRef; -use datafusion_physical_expr::expressions as df_expr; -use datafusion_physical_expr::projection::ProjectionExpr; -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; -use datafusion_pruning::FilePruner; -use futures::FutureExt; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::stream; -use object_store::path::Path; -use tracing::Instrument; -use vortex::array::VortexSessionExecute; -use vortex::error::VortexError; -use vortex::error::VortexExpect; -use vortex::file::OpenOptionsSessionExt; -use vortex::io::InstrumentedReadAt; -use vortex::layout::LayoutReader; -use vortex::layout::scan::scan_builder::ScanBuilder; -use vortex::metrics::Label; -use vortex::metrics::MetricsRegistry; -use vortex::session::VortexSession; -use vortex_arrow::ArrowSessionExt; -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::ProcessedProjection; -use crate::convert::exprs::raw_projection; -use crate::convert::schema::calculate_physical_schema; -use crate::metrics::PARTITION_LABEL; -use crate::metrics::PATH_LABEL; -use crate::persistent::cache::CachedVortexMetadata; -use crate::persistent::reader::VortexReaderFactory; -use crate::persistent::stream::PrunableStream; - -#[derive(Clone)] -pub(crate) struct VortexOpener { - /// The partition this opener is assigned to. Only used for labeling metrics. - pub partition: usize, - pub session: VortexSession, - pub vortex_reader_factory: Arc, - /// 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, - /// 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. - pub file_pruning_predicate: Option, - pub expr_adapter_factory: Arc, - /// This is the table's schema without partition columns. It may contain fields which do - /// not exist in the file, and are supplied by the `schema_adapter_factory`. - pub table_schema: TableSchema, - /// If provided, the scan will not return more than this many rows. - pub limit: Option, - /// A metrics object for tracking performance of the scan. - pub metrics_registry: Arc, - /// DataFusion-native metrics exposed through `DataSourceExec`. - pub df_metrics: ExecutionPlanMetricsSet, - /// A shared cache of file readers. - /// - /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache - /// a file reader the first time we read a file. - pub layout_readers: Arc>>, - /// Shared full-file natural splits keyed by file path. - pub natural_splits: Arc>>, - /// Whether the query has output ordering specified - pub has_output_ordering: bool, - - pub expression_convertor: Arc, - pub file_metadata_cache: Option>, - /// Whether to enable expression pushdown into the underlying Vortex scan. - pub projection_pushdown: bool, - pub scan_concurrency: Option, -} - -impl FileOpener for VortexOpener { - fn open(&self, file: PartitionedFile) -> DFResult { - // Calculate the output schema before replacing partition columns with literals so it - // retains the table and partition-field metadata declared by the plan. - let output_schema = Arc::new( - self.projection - .project_schema(self.table_schema.table_schema())?, - ); - let session = self.session.clone(); - let metrics_registry = Arc::clone(&self.metrics_registry); - let labels = vec![ - Label::new(PATH_LABEL, file.path().to_string()), - Label::new(PARTITION_LABEL, self.partition.to_string()), - ]; - - let mut projection = self.projection.clone(); - let mut filter = self.filter.clone(); - - let reader = self.vortex_reader_factory.create_reader(&file, &session)?; - - let reader = - InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); - - let mut file_pruning_predicate = self.file_pruning_predicate.clone(); - let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); - let file_metadata_cache = self.file_metadata_cache.clone(); - - let unified_file_schema = Arc::clone(self.table_schema.file_schema()); - let limit = self.limit; - let layout_readers = Arc::clone(&self.layout_readers); - let natural_splits = Arc::clone(&self.natural_splits); - let has_output_ordering = self.has_output_ordering; - let scan_concurrency = self.scan_concurrency; - - let expr_convertor = Arc::clone(&self.expression_convertor); - let projection_pushdown = self.projection_pushdown; - - let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) - .with_category(MetricCategory::Rows) - .global_counter("num_predicate_creation_errors"); - - // Replace column access for partition columns with literals - #[expect(clippy::disallowed_types)] - let literal_value_cols = self - .table_schema - .table_partition_cols() - .iter() - .map(|f| f.name()) - .cloned() - .zip(file.partition_values.clone()) - .collect::>(); - - let predicate_uses_partition_columns = - file_pruning_predicate.as_ref().is_some_and(|predicate| { - collect_columns(predicate) - .iter() - .any(|column| literal_value_cols.contains_key(column.name())) - }); - - if !literal_value_cols.is_empty() { - projection = projection.try_map_exprs(|expr| { - replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) - })?; - filter = filter - .map(|p| replace_columns_with_literals(p, &literal_value_cols)) - .transpose()?; - file_pruning_predicate = file_pruning_predicate - .map(|p| replace_columns_with_literals(p, &literal_value_cols)) - .transpose()?; - } - - Ok(async move { - // FilePruner requires a statistics object even when the rewritten predicate - // only contains partition literals. Supply unknown file-column statistics in - // that case so static and dynamic partition predicates can still prune. - let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) - .then(|| { - file.clone() - .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) - }); - let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); - - let mut file_pruner = file_pruning_predicate - .filter(|_| file.has_statistics() || predicate_uses_partition_columns) - .and_then(|predicate| { - FilePruner::try_new( - Arc::clone(&predicate), - &unified_file_schema, - pruning_file, - predicate_creation_errors, - ) - }); - - // Check if this file should be pruned based on statistics/partition values. - // Returns empty stream if file can be skipped entirely. - if let Some(file_pruner) = file_pruner.as_mut() - && file_pruner.should_prune()? - { - return Ok(stream::empty().boxed()); - } - - let mut open_opts = session - .open_options() - .with_file_size(file.object_meta.size) - .with_metrics_registry(Arc::clone(&metrics_registry)) - .with_labels(labels); - - let cached_footer = file_metadata_cache - .as_ref() - .and_then(|cache| cache.get(file.path())) - .filter(|entry| entry.is_valid_for(&file.object_meta)) - .and_then(|entry| { - entry - .file_metadata - .as_any() - .downcast_ref::() - .map(|vortex_metadata| vortex_metadata.footer().clone()) - }); - let footer_cache_hit = cached_footer.is_some(); - - if let Some(footer) = cached_footer { - open_opts = open_opts.with_footer(footer); - } - - let vxf = open_opts - .open_read(reader) - .await - .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?; - - // On a miss, cache the parsed footer so other partitions and later executions - // skip the footer fetch and parse. `infer_schema`/`infer_stats` also populate - // this cache, but only when planning goes through `VortexFormat`. - if !footer_cache_hit && let Some(cache) = &file_metadata_cache { - cache.put( - file.path(), - CachedFileMetadataEntry::new( - file.object_meta.clone(), - Arc::new(CachedVortexMetadata::new(&vxf)), - ), - ); - } - - // Check if there are rows in this file. If not, we can save - // ourselves some work and return an empty stream. - if vxf.row_count() == 0 { - return Ok(stream::empty().boxed()); - } - - // This is the expected arrow types of the actual columns in the file, which might have different types - // from the unified logical schema or miss - let this_file_schema = Arc::new(calculate_physical_schema( - vxf.dtype(), - &unified_file_schema, - &session.arrow(), - )?); - - let expr_adapter = expr_adapter_factory.create( - Arc::clone(&unified_file_schema), - Arc::clone(&this_file_schema), - )?; - - let simplifier = PhysicalExprSimplifier::new(&this_file_schema); - - // The adapter rewrites the expressions to the local file schema, allowing - // for schema evolution and divergence between the table's schema and individual files. - let filter = filter - .map(|filter| { - // Expression might now reference columns that don't exist in the file, so we can give it - // another simplification pass. - let adapted = expr_adapter.rewrite(Arc::clone(&filter)) - .map_err(|e| exec_datafusion_err!("Failed to adapt filter {filter} in {}: {e}", file.path()))?; - simplifier.simplify(adapted) - .map_err(|e| exec_datafusion_err!("Failed to simplify filter {filter} in {}: {e}", file.path())) - }) - .transpose()?; - let projection = - projection.try_map_exprs(|p| simplifier.simplify(expr_adapter.rewrite(p)?))?; - - let mut native_filters = Vec::new(); - let mut residual_filters = Vec::new(); - if let Some(filter) = &filter { - if filter.data_type(&this_file_schema)? != arrow_schema::DataType::Boolean { - return Err(exec_datafusion_err!("Filter must be Boolean in {}: {filter}", file.path())); - } - for expr in split_conjunction(filter) { - match expr_convertor.try_convert(expr, &this_file_schema) - .map_err(|e| exec_datafusion_err!("Failed to convert filter {expr} in {}: {e}", file.path()))? - { - Some(expr) => native_filters.push(expr), - None => residual_filters.push(Arc::clone(expr)), - } - } - } - let residual_filter = conjunction_opt(residual_filters); - let native_filter = vortex::expr::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 let Some(residual) = &residual_filter { - let mut indices = projection.column_indices(); - indices.extend(collect_columns(residual).into_iter().map(|column| column.index())); - indices.sort_unstable(); - indices.dedup(); - let required = indices.into_iter().map(|index| { - let field = this_file_schema.fields().get(index).ok_or_else(|| { - exec_datafusion_err!("Projection column index {index} is out of bounds") - })?; - Ok(ProjectionExpr { - expr: Arc::new(df_expr::Column::new(field.name(), index)), - alias: field.name().clone(), - }) - }).collect::>>()?; - let raw = raw_projection(required.into(), &this_file_schema)?; - ProcessedProjection { - scan_projection: raw.scan_projection, - scan_reference_schema: raw.scan_reference_schema, - leftover_projection: projection, - } - } else if projection_pushdown { - expr_convertor.split_projection( - projection, - &this_file_schema, - output_schema.as_ref(), - )? - } else { - expr_convertor.no_pushdown_projection(projection, &this_file_schema)? - }; - - // The schema of the stream returned from the vortex scan. - // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). - let scan_projection = scan_projection - .optimize_recursive(vxf.dtype()) - .and_then(|projection| projection.bind(vxf.dtype())) - .map_err(|_e| { - exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") - })?; - let scan_dtype = scan_projection.dtype().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. - let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) { - Entry::Occupied(mut occupied_entry) => { - if let Some(reader) = occupied_entry.get().upgrade() { - tracing::trace!("reusing layout reader for {}", occupied_entry.key()); - reader - } else { - tracing::trace!("creating layout reader for {}", occupied_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!( - "Failed to create layout reader: {e}" - )) - })?; - occupied_entry.insert(Arc::downgrade(&reader)); - reader - } - } - Entry::Vacant(vacant_entry) => { - tracing::trace!("creating layout reader for {}", vacant_entry.key()); - let reader = vxf.layout_reader().map_err(|e| { - DataFusionError::Execution(format!("Failed to create layout reader: {e}")) - })?; - vacant_entry.insert(Arc::downgrade(&reader)); - - reader - } - }; - - let mut scan_builder = ScanBuilder::new(session.clone(), Arc::clone(&layout_reader)); - - if let Some(vortex_plan) = file.extensions.get::() { - scan_builder = vortex_plan.apply_to_builder(scan_builder); - } - - if let Some(limit) = limit - && native_filter.is_none() - && residual_filter.is_none() - { - scan_builder = scan_builder.with_limit(limit); - } - - if let Some(concurrency) = scan_concurrency { - scan_builder = scan_builder.with_concurrency(concurrency); - } - - // Set before the byte-range translation below, which computes natural splits for - // the fields the scan's projection and filter reference. - scan_builder = scan_builder - .with_projection(scan_projection) - .with_some_filter(native_filter); - - 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"))?, - end: u64::try_from(file_range.end) - .map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?, - }; - if byte_range.start != 0 || byte_range.end != file.object_meta.size { - // Full-file scans already cover every natural split. Only translate the - // byte range back into row boundaries when DataFusion has trimmed the file. - let natural_splits = natural_splits_for_file( - natural_splits.as_ref(), - &file.object_meta.location, - &scan_builder, - file.object_meta.size, - )?; - - let Some(row_range) = - split_aligned_row_range(byte_range, natural_splits.as_ref()) - else { - return Ok(stream::empty().boxed()); - }; - - scan_builder = scan_builder - .with_row_range(row_range) - // Hand the shared full-file boundaries back to the scan so prepare() - // skips its own layout walk. - .with_natural_splits(Arc::clone(&natural_splits.row_boundaries)); - } - } - - 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) - .map(move |chunk| { - let mut ctx = session.create_execution_ctx(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) - .into_stream() - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - file.object_meta.location - )))) - }) - .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}"))?; - } - Ok(batch) - }) - .try_filter(|batch| ready(batch.num_rows() != 0)) - .map(move |batch| { - let batch = projector.project_batch(&batch?)?; - - let (_, columns, row_count) = batch.into_parts(); - RecordBatch::try_new_with_options( - Arc::clone(&output_schema), - columns, - &RecordBatchOptions::new().with_row_count(Some(row_count)), - ) - .map_err(Into::into) - }) - .boxed(); - - if let Some(file_pruner) = file_pruner && file_pruner.is_watching() { - Ok(PrunableStream::new(file_pruner, stream).boxed()) - } else { - Ok(stream) - } - } - .in_current_span() - .boxed()) - } -} - -/// A file's natural split boundaries plus the precomputed byte each split is assigned to, -/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row -/// boundaries with a binary search instead of re-projecting every split per partition. -/// -/// The boundaries are computed for the fields referenced by the scan's projection and filter. -/// All partitions translate through the first opener's cached entry (the cache lives on the -/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges -/// tiling the file's rows exactly once. -#[derive(Debug)] -pub(crate) struct NaturalSplits { - /// Sorted row boundaries of the natural splits; split `i` covers - /// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the - /// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the - /// per-partition layout walk in `prepare`. - row_boundaries: Arc<[u64]>, - /// For each split, the byte a DataFusion byte range must contain to own it (see - /// [`split_assignment_byte`]); one entry per split, sorted because split midpoints - /// increase monotonically under the row-to-byte projection. - assignment_bytes: Box<[u64]>, -} - -impl NaturalSplits { - fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { - let row_count = row_boundaries.last().copied().unwrap_or_default(); - let assignment_bytes = if row_count == 0 { - Box::default() - } else { - row_boundaries - .windows(2) - .enumerate() - .map(|(idx, boundaries)| { - split_assignment_byte( - idx, - &(boundaries[0]..boundaries[1]), - row_count, - total_size, - ) - }) - .collect() - }; - - debug_assert!(assignment_bytes.is_sorted()); - debug_assert_eq!( - assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), - row_boundaries.len() - ); - - Self { - row_boundaries, - assignment_bytes, - } - } -} - -/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use. -fn natural_splits_for_file( - natural_splits: &DashMap>, - path: &Path, - scan_builder: &ScanBuilder, - total_size: u64, -) -> DFResult> { - if let Some(splits) = natural_splits.get(path) { - return Ok(Arc::clone(splits.value())); - } - - // Compute while holding the entry so concurrent partitions opening the same file wait - // for the winner instead of all walking the layout tree; the redundant walks contend on - // the lazily-initialized layout children and dominate the cost of the computation itself. - match natural_splits.entry(path.clone()) { - Entry::Occupied(entry) => Ok(Arc::clone(entry.get())), - Entry::Vacant(entry) => { - let splits = compute_natural_splits(scan_builder, total_size)?; - entry.insert(Arc::clone(&splits)); - Ok(splits) - } - } -} - -/// Walk the layout tree to compute the file's full natural split boundaries for the fields -/// referenced by the scan's projection and filter. -fn compute_natural_splits( - scan_builder: &ScanBuilder, - total_size: u64, -) -> DFResult> { - let row_boundaries = scan_builder - .full_file_splits() - .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?; - - Ok(Arc::new(NaturalSplits::new( - row_boundaries.into(), - total_size, - ))) -} - -/// Translate a DataFusion byte range to the contiguous natural split ranges it owns. -/// Most splits are assigned by midpoint, but the leading split stays with the range that owns -/// byte 0 so a tiny first byte range still claims the first rows. -fn split_aligned_row_range( - byte_range: Range, - natural_splits: &NaturalSplits, -) -> Option> { - if byte_range.start >= byte_range.end { - return None; - } - - let first_split = natural_splits - .assignment_bytes - .partition_point(|&assignment_byte| assignment_byte < byte_range.start); - let after_last_split = natural_splits - .assignment_bytes - .partition_point(|&assignment_byte| assignment_byte < byte_range.end); - if first_split == after_last_split { - return None; - } - - Some( - natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split], - ) -} - -fn split_assignment_byte( - idx: usize, - split_range: &Range, - row_count: u64, - total_size: u64, -) -> u64 { - if idx == 0 && split_range.start == 0 { - // Byte 0 is the only stable representative for the leading split. A midpoint can fall - // into the next DataFusion byte range and leave the first range with no rows to read. - 0 - } else { - split_midpoint_to_byte(split_range, row_count, total_size) - } -} - -fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64) -> u64 { - let midpoint_row = split_range.start + (split_range.end - split_range.start) / 2; - let midpoint_byte = (u128::from(midpoint_row) * u128::from(total_size)) / u128::from(row_count); - - u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") -} - -#[cfg(test)] -mod tests; diff --git a/vortex-datafusion/src/persistent/opener/tests.rs b/vortex-datafusion/src/persistent/opener/tests.rs deleted file mode 100644 index c8c8f904151..00000000000 --- a/vortex-datafusion/src/persistent/opener/tests.rs +++ /dev/null @@ -1,1546 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt; -use std::sync::Arc; -use std::sync::LazyLock; - -use arrow_array::record_batch; -use arrow_schema::Field; -use arrow_schema::Fields; -use arrow_schema::SchemaRef; -use datafusion::arrow::array::DictionaryArray; -use datafusion::arrow::array::Int32Array; -use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::array::StringArray; -use datafusion::arrow::array::StructArray; -use datafusion::arrow::datatypes::DataType; -use datafusion::arrow::datatypes::Schema; -use datafusion::arrow::datatypes::UInt32Type; -use datafusion::arrow::util::display::FormatOptions; -use datafusion::arrow::util::pretty::pretty_format_batches_with_options; -use datafusion::logical_expr::col; -use datafusion::logical_expr::lit; -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_common::tree_node::Transformed; -use datafusion_common::tree_node::TreeNode; -use datafusion_execution::cache::default_cache::DefaultCache; -use datafusion_expr::Operator; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions as df_expr; -use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; -use datafusion_physical_expr::projection::ProjectionExpr; -use insta::assert_snapshot; -use itertools::Itertools; -use object_store::ObjectStore; -use object_store::ObjectStoreExt; -use object_store::memory::InMemory; -use rstest::rstest; -use vortex::VortexSessionDefault; -use vortex::buffer::Buffer; -use vortex::file::WriteOptionsSessionExt; -use vortex::io::VortexWrite; -use vortex::io::object_store::ObjectStoreWrite; -use vortex::metrics::DefaultMetricsRegistry; -use vortex::scan::selection::Selection; -use vortex::scan::strict_sorted_buffer::StrictSortedBuffer; -use vortex::session::VortexSession; - -use super::*; -use crate::VortexAccessPlan; -use crate::convert::exprs::DefaultExpressionConvertor; -use crate::persistent::reader::DefaultVortexReaderFactory; - -static SESSION: LazyLock = LazyLock::new(VortexSession::default); - -/// Test-only expr used to test error reporting. -#[derive(Debug, Eq, Hash, PartialEq)] -struct SnapshotErrorExpr; - -impl fmt::Display for SnapshotErrorExpr { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "snapshot_error") - } -} - -impl PhysicalExpr for SnapshotErrorExpr { - fn data_type(&self, _input_schema: &Schema) -> DFResult { - Ok(DataType::Boolean) - } - - fn nullable(&self, _input_schema: &Schema) -> DFResult { - Ok(false) - } - - fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt::Display::fmt(self, f) - } - - fn evaluate(&self, _batch: &RecordBatch) -> DFResult { - Err(DataFusionError::Internal( - "intentional snapshot error".to_owned(), - )) - } - - fn children(&self) -> Vec<&PhysicalExprRef> { - Vec::new() - } - - fn with_new_children( - self: Arc, - children: Vec, - ) -> DFResult { - assert!(children.is_empty()); - Ok(self) - } - - fn snapshot(&self) -> DFResult> { - Err(DataFusionError::Internal( - "intentional snapshot error".to_owned(), - )) - } -} - -fn natural_splits(total_size: u64, split_ranges: &[Range]) -> NaturalSplits { - let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1); - if let Some(first) = split_ranges.first() { - row_boundaries.push(first.start); - row_boundaries.extend(split_ranges.iter().map(|range| range.end)); - } - NaturalSplits::new(row_boundaries.into(), total_size) -} - -#[rstest] -#[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] -#[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] -#[case(1..8, 10, vec![0..1, 1..9, 9..10], Some(1..9))] -#[case(1..4, 16, vec![0..1, 1..2, 2..3, 3..4], None)] -#[case(0..1, 10, vec![0..2, 2..10], Some(0..2))] -#[case(0..2, 2, vec![], None)] -fn test_split_aligned_row_range( - #[case] byte_range: Range, - #[case] total_size: u64, - #[case] split_ranges: Vec>, - #[case] expected: Option>, -) { - assert_eq!( - split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)), - expected - ); -} - -#[test] -fn test_split_aligned_ranges_cover_splits_exactly_once() { - let split_ranges = vec![0..1, 1..4, 4..10, 10..13]; - let byte_ranges = [0..4, 4..8, 8..12, 12..16]; - let natural_splits = natural_splits(16, &split_ranges); - - let assigned = byte_ranges - .into_iter() - .filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits)) - .collect::>(); - - assert_eq!(assigned, vec![0..4, 4..10, 10..13]); - assert_eq!( - assigned - .iter() - .map(|range| range.end - range.start) - .sum::(), - 13 - ); - - let split_starts = split_ranges - .iter() - .map(|range| range.start) - .collect::>(); - let split_ends = split_ranges - .iter() - .map(|range| range.end) - .collect::>(); - - for range in &assigned { - assert!(split_starts.contains(&range.start)); - assert!(split_ends.contains(&range.end)); - } - - for (left, right) in assigned.iter().tuple_windows() { - assert_eq!(left.end, right.start); - } -} - -#[rstest] -#[case(vec![], 10)] -#[case(vec![0], 10)] -#[case(vec![], 0)] -#[case(vec![0], 0)] -fn test_natural_splits_empty_file(#[case] row_boundaries: Vec, #[case] total_size: u64) { - let splits = NaturalSplits::new(row_boundaries.clone().into(), total_size); - - assert!(splits.assignment_bytes.is_empty()); - assert_eq!(splits.row_boundaries.as_ref(), row_boundaries.as_slice()); - assert_eq!(split_aligned_row_range(0..u64::MAX, &splits), None); -} - -#[test] -fn test_split_aligned_row_range_keeps_colliding_assignments_together() { - let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]); - - assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]); - assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2)); - assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4)); -} - -async fn write_arrow_to_vortex( - object_store: Arc, - path: &str, - rb: RecordBatch, -) -> anyhow::Result { - let schema = rb.schema(); - let array = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; - let path = Path::parse(path)?; - - let mut write = ObjectStoreWrite::new(object_store, &path).await?; - let summary = SESSION - .write_options() - .write(&mut write, array.to_array_stream()) - .await?; - write.shutdown().await?; - - Ok(summary.size()) -} - -fn make_opener( - object_store: Arc, - table_schema: TableSchema, - filter: Option, -) -> VortexOpener { - VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), - projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), - filter, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema, - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - } -} - -#[tokio::test] -async fn test_open() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "part=1/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file_schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.partition_values = vec![ScalarValue::Int32(Some(1))]; - - let table_schema = TableSchema::builder(Arc::clone(&file_schema)) - .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) - .build(); - - // filter matches partition value - let filter = col("part").eq(lit(1)); - let filter = logical2physical(&filter, table_schema.table_schema()); - - let opener = make_opener( - Arc::clone(&object_store), - table_schema.clone(), - Some(filter), - ); - let stream = opener.open(file.clone()).unwrap().await.unwrap(); - - let data = stream.try_collect::>().await?; - let num_batches = data.len(); - let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); - - assert_eq!((num_batches, num_rows), (1, 3)); - - // filter doesn't matches partition value - let filter = col("part").eq(lit(2)); - let filter = logical2physical(&filter, table_schema.table_schema()); - - let opener = make_opener( - Arc::clone(&object_store), - table_schema.clone(), - Some(filter), - ); - let stream = opener.open(file.clone()).unwrap().await.unwrap(); - - let data = stream.try_collect::>().await?; - let num_batches = data.len(); - let num_rows = data.iter().map(|rb| rb.num_rows()).sum::(); - assert_eq!((num_batches, num_rows), (0, 0)); - - 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; - let file_path = "part=1/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file_schema = Arc::new( - batch.schema().as_ref().clone().with_metadata( - [("table".to_string(), "metadata".to_string())] - .into_iter() - .collect(), - ), - ); - let table_schema = TableSchema::builder(file_schema) - .with_table_partition_cols(vec![Arc::new( - Field::new("part", DataType::Int32, false).with_metadata( - [("partition".to_string(), "metadata".to_string())] - .into_iter() - .collect(), - ), - )]) - .build(); - let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); - let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); - - assert_eq!( - expected_schema.metadata().get("table"), - Some(&"metadata".to_string()) - ); - assert_eq!( - expected_schema.field(1).metadata().get("partition"), - Some(&"metadata".to_string()) - ); - - for projection_pushdown in [false, true] { - let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); - opener.projection = projection.clone(); - opener.projection_pushdown = projection_pushdown; - - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.partition_values = vec![ScalarValue::Int32(Some(1))]; - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert!(!batches.is_empty()); - for batch in batches { - assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); - } - } - - Ok(()) -} - -#[tokio::test] -async fn test_open_all_valid_nullable_columns_with_nonnullable_table_schema() -> anyhow::Result<()> -{ - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "nullable/file.vortex"; - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), - vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], - )?; - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - let expected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let table_schema = TableSchema::from(Arc::clone(&expected_schema)); - - for projection_pushdown in [false, true] { - let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); - opener.projection_pushdown = projection_pushdown; - - let file = PartitionedFile::new(file_path.to_string(), data_size); - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert_eq!(batches.len(), 1); - assert_eq!(batches[0].schema().as_ref(), expected_schema.as_ref()); - } - - Ok(()) -} - -#[tokio::test] -async fn test_file_pruning_replaces_partition_columns_without_file_statistics() -> anyhow::Result<()> -{ - let object_store = Arc::new(InMemory::new()) as Arc; - let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - let table_schema = TableSchema::builder(Arc::clone(&file_schema)) - .with_table_partition_cols(vec![Arc::new(Field::new("part", DataType::Int32, false))]) - .build(); - - let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; - let predicate = Arc::new(df_expr::BinaryExpr::new( - Arc::clone(&partition_column), - Operator::Gt, - df_expr::lit(ScalarValue::Int32(Some(1))), - )) as PhysicalExprRef; - let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( - vec![partition_column], - predicate, - )) as PhysicalExprRef; - - let mut opener = make_opener(object_store, table_schema, None); - opener.file_pruning_predicate = Some(dynamic_predicate); - let df_metrics = opener.df_metrics.clone(); - - // The file does not exist and has no statistics. Replacing `part` with 1 - // makes the predicate false, so pruning must happen before any file I/O. - let mut file = PartitionedFile::new("missing.vortex", 1); - file.partition_values = vec![ScalarValue::Int32(Some(1))]; - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert!(batches.is_empty()); - assert_eq!( - df_metrics - .clone_inner() - .sum_by_name("num_predicate_creation_errors") - .map(|metric| metric.as_usize()), - Some(0) - ); - - Ok(()) -} - -#[tokio::test] -async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "metrics/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); - statistics.column_statistics[0].null_count = Precision::Exact(0); - let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); - - let mut opener = make_opener(object_store, TableSchema::from(batch.schema()), None); - opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); - let df_metrics = opener.df_metrics.clone(); - - let batches = opener.open(file)?.await?.try_collect::>().await?; - - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); - assert_eq!( - df_metrics - .clone_inner() - .sum_by_name("num_predicate_creation_errors") - .map(|metric| metric.as_usize()), - Some(1) - ); - - Ok(()) -} - -#[tokio::test] -async fn test_open_empty_file() -> anyhow::Result<()> { - use futures::TryStreamExt; - - let object_store = Arc::new(InMemory::new()) as Arc; - let data_batch = record_batch!(("a", Int32, Vec::::new())).unwrap(); - let file_path = "part=1/empty.vortex"; - let file_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, data_batch.clone()).await?; - - let file_schema = data_batch.schema(); - // Parallel scans may attach a byte range even for empty files; the - // opener must return early before attempting split-aligned translation. - let file = - PartitionedFile::new_with_range(file_path.to_string(), file_size, 0, file_size as i64); - - let table_schema = TableSchema::from(Arc::clone(&file_schema)); - - let opener = make_opener(object_store, table_schema, None); - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - assert_eq!(data.len(), 0); - - Ok(()) -} - -#[tokio::test] -async fn test_open_populates_file_metadata_cache() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "cached/file.vortex"; - let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file = PartitionedFile::new(file_path.to_string(), data_size); - let table_schema = TableSchema::from(batch.schema()); - - let cache: Arc = Arc::new( - DefaultCache::::new(64 * 1024 * 1024), - ); - let mut opener = make_opener(Arc::clone(&object_store), table_schema, None); - opener.file_metadata_cache = Some(Arc::clone(&cache)); - - // The first open misses the cache and must write the parsed footer back. - let stream = opener.open(file.clone())?.await?; - stream.try_collect::>().await?; - - let entry = cache - .get(file.path()) - .ok_or_else(|| anyhow::anyhow!("footer was not cached after open"))?; - assert!(entry.is_valid_for(&file.object_meta)); - assert!( - entry - .file_metadata - .as_any() - .downcast_ref::() - .is_some() - ); - - // The second open hits the cache and still returns the same data. - let stream = opener.open(file.clone())?.await?; - let data = stream.try_collect::>().await?; - assert_eq!(data.iter().map(|rb| rb.num_rows()).sum::(), 3); - - Ok(()) -} - -#[rstest] -#[tokio::test] -async fn test_open_files_different_table_schema() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - - let file1 = { - let file1_path = "/path/file1.vortex"; - let batch1 = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); - let data_size1 = - write_arrow_to_vortex(Arc::clone(&object_store), file1_path, batch1).await?; - PartitionedFile::new(file1_path.to_string(), data_size1) - }; - - let file2 = { - let file2_path = "/path/file2.vortex"; - let batch2 = record_batch!(("a", Int16, vec![Some(-1), Some(-2), Some(-3)])).unwrap(); - let data_size2 = - write_arrow_to_vortex(Arc::clone(&object_store), file2_path, batch2).await?; - PartitionedFile::new(file2_path.to_string(), data_size2) - }; - - // Table schema has can accommodate both files - let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( - "a", - DataType::Int32, - true, - )]))); - - let make_opener = |filter| VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone(&object_store))), - projection: ProjectionExprs::from_indices(&[0], table_schema.file_schema()), - filter: Some(filter), - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: table_schema.clone(), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - let filter = col("a").lt(lit(100_i32)); - let filter = logical2physical(&filter, table_schema.table_schema()); - - let opener1 = make_opener(Arc::clone(&filter)); - let stream = opener1.open(file1)?.await?; - - let format_opts = FormatOptions::new().with_types_info(true); - - let data = stream.try_collect::>().await?; - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+ - | a | - | Int32 | - +-------+ - | 1 | - | 2 | - | 3 | - +-------+ - "); - - let opener2 = make_opener(Arc::clone(&filter)); - let stream = opener2.open(file2)?.await?; - - let data = stream.try_collect::>().await?; - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+ - | a | - | Int32 | - +-------+ - | -1 | - | -2 | - | -3 | - +-------+ - "); - - Ok(()) -} - -#[tokio::test] -// This test verifies that files with different column order than the -// table schema can be opened without errors. The fix ensures that the -// schema mapper is only used for type casting, not for reordering, -// since the vortex projection already handles reordering. -async fn test_schema_different_column_order() -> anyhow::Result<()> { - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - // File has columns in order: c, b, a - let batch = record_batch!( - ("c", Int32, vec![Some(300), Some(301), Some(302)]), - ("b", Int32, vec![Some(200), Some(201), Some(202)]), - ("a", Int32, vec![Some(100), Some(101), Some(102)]) - ) - .unwrap(); - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - let file = PartitionedFile::new(file_path.to_string(), data_size); - - // Table schema has columns in different order: a, b, c - let table_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - ])); - - let opener = VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), - projection: ProjectionExprs::from_indices(&[0, 1, 2], &table_schema), - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: TableSchema::from(Arc::clone(&table_schema)), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - let stream = opener.open(file)?.await?; - - let format_opts = FormatOptions::new().with_types_info(true); - let data = stream.try_collect::>().await?; - - // Verify the output has columns in table schema order (a, b, c) - // not file order (c, b, a) - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+-------+-------+ - | a | b | c | - | Int32 | Int32 | Int32 | - +-------+-------+-------+ - | 100 | 200 | 300 | - | 101 | 201 | 301 | - | 102 | 202 | 302 | - +-------+-------+-------+ - "); - - Ok(()) -} - -#[tokio::test] -// This test verifies that expression rewriting doesn't fail when there is -// a nested schema mismatch between the physical file schema and logical -// table schema. -async fn test_adapter_logical_physical_struct_mismatch() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - let file_struct_fields = Fields::from(vec![ - Field::new("field1", DataType::Utf8, true), - Field::new("field2", DataType::Utf8, true), - ]); - let struct_array = StructArray::new( - file_struct_fields.clone(), - vec![ - Arc::new(StringArray::from(vec!["value1", "value2", "value3"])), - Arc::new(StringArray::from(vec!["a", "b", "c"])), - ], - None, - ); - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![Field::new( - "my_struct", - DataType::Struct(file_struct_fields), - true, - )])), - vec![Arc::new(struct_array)], - )?; - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - // Table schema has an extra utf8 field. - let table_schema = TableSchema::from(Arc::new(Schema::new(vec![Field::new( - "my_struct", - DataType::Struct(Fields::from(vec![ - Field::new( - "field1", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - Field::new( - "field2", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - Field::new("field3", DataType::Utf8, true), - ])), - true, - )]))); - - let opener = make_opener( - Arc::clone(&object_store), - table_schema.clone(), - // expression references my_struct column which has different fields in each - // field. - Some(logical2physical( - &col("my_struct").is_not_null(), - table_schema.table_schema(), - )), - ); - - // The opener should be able to open the file with a filter on the - // struct column. - let data = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? - .await? - .try_collect::>() - .await?; - - assert_eq!(data.len(), 1); - assert_eq!(data[0].num_rows(), 3); - - Ok(()) -} - -#[tokio::test] -// Minimal reproducing test for the schema projection bug. -// Before the fix, this would fail with a cast error when the file schema -// and table schema have different field orders and we project a subset of columns. -async fn test_projection_bug_minimal_repro() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - // File has columns in order: a, b, c with simple types - let batch = record_batch!( - ("a", Int32, vec![Some(1)]), - ("b", Utf8, vec![Some("test")]), - ("c", Int32, vec![Some(2)]) - ) - .unwrap(); - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - // Table schema has columns in DIFFERENT order: c, a, b - // and different types that require casting (Utf8 -> Dictionary) - let table_schema = TableSchema::from(Arc::new(Schema::new(vec![ - Field::new("c", DataType::Int32, true), - Field::new("a", DataType::Int32, true), - Field::new( - "b", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - ), - ]))); - - // Project columns [0, 2] from table schema, which should give us: c, b - // Before the fix, the schema adapter would get confused about which fields - // to select from the file, causing incorrect type mappings. - let projection = vec![0, 2]; - - let opener = VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone(&object_store))), - projection: ProjectionExprs::from_indices(projection.as_ref(), table_schema.file_schema()), - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: table_schema.clone(), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - // This should succeed and return the correctly projected and cast data - let data = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? - .await? - .try_collect::>() - .await?; - - // Verify the columns are in the right order and have the right values - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - let format_opts = FormatOptions::new().with_types_info(true); - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+--------------------------+ - | c | b | - | Int32 | Dictionary(UInt32, Utf8) | - +-------+--------------------------+ - | 2 | test | - +-------+--------------------------+ - "); - - Ok(()) -} - -fn make_test_batch_with_10_rows() -> RecordBatch { - record_batch!( - ("a", Int32, (0..=9).map(Some).collect::>()), - ( - "b", - Utf8, - (0..=9).map(|i| Some(format!("r{}", i))).collect::>() - ) - ) - .unwrap() -} - -fn make_test_opener( - object_store: Arc, - schema: SchemaRef, - projection: ProjectionExprs, -) -> VortexOpener { - VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(object_store)), - projection, - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema: TableSchema::from(schema), - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - } -} - -#[tokio::test] -// Test that Selection::IncludeByIndex filters to specific row indices. -async fn test_selection_include_by_index() -> anyhow::Result<()> { - use datafusion::arrow::util::pretty::pretty_format_batches_with_options; - - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.extensions.insert( - VortexAccessPlan::default().with_selection(Selection::IncludeByIndex( - StrictSortedBuffer::try_new(Buffer::from_iter(vec![1, 3, 5, 7]))?, - )), - ); - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0, 1], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - let format_opts = FormatOptions::new().with_types_info(true); - - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+------+ - | a | b | - | Int32 | Utf8 | - +-------+------+ - | 1 | r1 | - | 3 | r3 | - | 5 | r5 | - | 7 | r7 | - +-------+------+ - "); - - Ok(()) -} - -#[tokio::test] -// Test that Selection::ExcludeByIndex excludes specific row indices. -async fn test_selection_exclude_by_index() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.extensions.insert( - VortexAccessPlan::default().with_selection(Selection::ExcludeByIndex( - StrictSortedBuffer::try_new(Buffer::from_iter(vec![0, 2, 4, 6, 8]))?, - )), - ); - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0, 1], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - let format_opts = FormatOptions::new().with_types_info(true); - - assert_snapshot!(pretty_format_batches_with_options(&data, &format_opts)?.to_string(), @r" - +-------+------+ - | a | b | - | Int32 | Utf8 | - +-------+------+ - | 1 | r1 | - | 3 | r3 | - | 5 | r5 | - | 7 | r7 | - | 9 | r9 | - +-------+------+ - "); - - Ok(()) -} - -#[tokio::test] -// Test that Selection::All returns all rows. -async fn test_selection_all() -> anyhow::Result<()> { - use vortex::scan::selection::Selection; - - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let mut file = PartitionedFile::new(file_path.to_string(), data_size); - file.extensions - .insert(VortexAccessPlan::default().with_selection(Selection::All)); - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); - assert_eq!(total_rows, 10); - - Ok(()) -} - -#[tokio::test] -// Test that when no extensions are provided, all rows are returned (backward compatibility). -async fn test_selection_no_extensions() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = make_test_batch_with_10_rows(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let schema = batch.schema(); - let file = PartitionedFile::new(file_path.to_string(), data_size); - // file.extensions is None by default - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0], &schema), - ); - - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - let total_rows: usize = data.iter().map(|rb| rb.num_rows()).sum(); - assert_eq!(total_rows, 10); - - Ok(()) -} - -#[tokio::test] -async fn test_projection_expr_pushdown() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - let file_path = "/path/file.vortex"; - - let batch = record_batch!( - ("a", Int32, vec![Some(1), Some(2), Some(3)]), - ("b", Int32, vec![Some(10), Some(20), Some(30)]) - ) - .unwrap(); - let data_size = - write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; - - let file_schema = batch.schema(); - let table_schema = TableSchema::from(Arc::clone(&file_schema)); - - // Create a projection that includes an arithmetic expression: a + b * 2 - let col_a = df_expr::col("a", &file_schema)?; - let col_b = df_expr::col("b", &file_schema)?; - let two = df_expr::lit(ScalarValue::Int32(Some(2))); - - // b * 2 - let b_times_2 = df_expr::binary(col_b, Operator::Multiply, two, &file_schema)?; - // a + (b * 2) - let a_plus_b_times_2 = df_expr::binary(col_a, Operator::Plus, b_times_2, &file_schema)?; - - let projection = ProjectionExprs::new(vec![ProjectionExpr::new( - a_plus_b_times_2, - "result".to_string(), - )]); - - let opener = VortexOpener { - partition: 1, - session: SESSION.clone(), - vortex_reader_factory: Arc::new(DefaultVortexReaderFactory::new(Arc::clone(&object_store))), - projection, - filter: None, - file_pruning_predicate: None, - expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), - table_schema, - limit: None, - metrics_registry: Arc::new(DefaultMetricsRegistry::default()), - df_metrics: ExecutionPlanMetricsSet::new(), - layout_readers: Default::default(), - natural_splits: Default::default(), - has_output_ordering: false, - expression_convertor: Arc::new(DefaultExpressionConvertor::default()), - file_metadata_cache: None, - projection_pushdown: false, - scan_concurrency: None, - }; - - let file = PartitionedFile::new(file_path.to_string(), data_size); - let stream = opener.open(file)?.await?; - let data = stream.try_collect::>().await?; - - // Expected: a + b * 2 - // row 0: 1 + 10 * 2 = 21 - // row 1: 2 + 20 * 2 = 42 - // row 2: 3 + 30 * 2 = 63 - assert_snapshot!(pretty_format_batches_with_options(&data, &FormatOptions::new().with_types_info(true))?.to_string(), @r" - +--------+ - | result | - | Int32 | - +--------+ - | 21 | - | 42 | - | 63 | - +--------+ - "); - - Ok(()) -} - -/// When a Struct contains Dictionary fields, writing to vortex and reading back -/// should preserve the Dictionary type. -#[tokio::test] -async fn test_struct_with_dictionary_roundtrip() -> anyhow::Result<()> { - let object_store = Arc::new(InMemory::new()) as Arc; - - let struct_fields = Fields::from(vec![ - Field::new_dictionary("a", DataType::UInt32, DataType::Utf8, true), - Field::new_dictionary("b", DataType::UInt32, DataType::Utf8, true), - ]); - let struct_array = StructArray::new( - struct_fields.clone(), - vec![ - Arc::new(DictionaryArray::::from_iter(["x", "y", "x"])), - Arc::new(DictionaryArray::::from_iter(["p", "p", "q"])), - ], - None, - ); - - let schema = Arc::new(Schema::new(vec![Field::new( - "labels", - DataType::Struct(struct_fields.clone()), - false, - )])); - let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(struct_array)])?; - - let file_path = "/test.vortex"; - let data_size = write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch).await?; - - let opener = make_test_opener( - Arc::clone(&object_store), - Arc::clone(&schema), - ProjectionExprs::from_indices(&[0], &schema), - ); - let data: Vec<_> = opener - .open(PartitionedFile::new(file_path.to_string(), data_size))? - .await? - .try_collect() - .await?; - - assert_eq!( - data[0].schema().field(0).data_type(), - &DataType::Struct(struct_fields), - "Struct(Dictionary) type should be preserved" - ); - Ok(()) -} - -#[derive(Debug)] -struct ModuloAdapterFactory; - -#[derive(Debug)] -struct ModuloAdapter(Arc); - -impl PhysicalExprAdapterFactory for ModuloAdapterFactory { - fn create( - &self, - logical: SchemaRef, - physical: SchemaRef, - ) -> DFResult> { - Ok(Arc::new(ModuloAdapter( - DefaultPhysicalExprAdapterFactory.create(logical, physical)?, - ))) - } -} - -impl datafusion_physical_expr_adapter::PhysicalExprAdapter for ModuloAdapter { - fn rewrite(&self, expr: PhysicalExprRef) -> DFResult { - self.0 - .rewrite(expr)? - .transform_up(|expr| { - if expr - .downcast_ref::() - .is_some_and(|c| c.name() == "b") - { - Ok(Transformed::yes(Arc::new(df_expr::BinaryExpr::new( - expr, - Operator::Modulo, - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(2)))), - )) as PhysicalExprRef)) - } else { - Ok(Transformed::no(expr)) - } - }) - .map(|result| result.data) - } -} - -struct DelegatingConvertor(DefaultExpressionConvertor); - -impl ExpressionConvertor for DelegatingConvertor { - fn try_convert( - &self, - expr: &PhysicalExprRef, - schema: &Schema, - ) -> DFResult> { - self.0.try_convert(expr, schema) - } -} - -#[rstest] -#[tokio::test] -async fn test_adapted_filter_fallback_and_limit( - #[values(false, true)] projection_pushdown: bool, - #[values(false, true)] custom_convertor: bool, - #[values(false, true)] zero_columns: bool, -) -> anyhow::Result<()> { - let ctx = crate::common_tests::TestSessionContext::new(projection_pushdown); - let batch = record_batch!( - ("a", Int32, vec![10, 20, 30, 40]), - ("b", Int32, vec![Some(1), None, Some(2), Some(4)]) - )?; - ctx.write_arrow_batch("adapted.vortex", &batch).await?; - let metadata = ctx.store.head(&Path::from("adapted.vortex")).await?; - let mut source = crate::VortexSource::new(TableSchema::from(batch.schema()), SESSION.clone()) - .with_projection_pushdown(projection_pushdown); - if custom_convertor { - source = source.with_expression_convertor(Arc::new(DelegatingConvertor( - DefaultExpressionConvertor::default(), - ))); - } - let filter: PhysicalExprRef = Arc::new(df_expr::BinaryExpr::new( - Arc::new(df_expr::Column::new("b", 1)), - Operator::Eq, - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(0)))), - )); - let accepted = datafusion_datasource::file::FileSource::try_pushdown_filters( - &source, - vec![filter], - &datafusion_common::config::ConfigOptions::new(), - )?; - assert!(matches!( - accepted.filters.as_slice(), - [datafusion_physical_plan::filter_pushdown::PushedDown::Yes] - )); - let mut source = accepted - .updated_node - .ok_or_else(|| anyhow::anyhow!("Expected updated source"))?; - // Alias a computed output to b, which is also an unprojected residual input. - let projection = if zero_columns { - ProjectionExprs::from(Vec::::new()) - } else { - vec![ProjectionExpr { - expr: Arc::new(df_expr::BinaryExpr::new( - Arc::new(df_expr::Column::new("a", 0)), - Operator::Plus, - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), - )), - alias: "b".into(), - }] - .into() - }; - source = source - .try_pushdown_projection(&projection)? - .ok_or_else(|| anyhow::anyhow!("Expected projected source"))?; - let config = datafusion_datasource::file_scan_config::FileScanConfigBuilder::new( - datafusion_execution::object_store::ObjectStoreUrl::local_filesystem(), - source, - ) - .with_expr_adapter(Some(Arc::new(ModuloAdapterFactory))) - .with_limit(Some(1)) - .with_file(PartitionedFile::new("adapted.vortex", metadata.size)) - .build(); - let plan: Arc = Arc::new( - datafusion_datasource::source::DataSourceExec::new(Arc::new(config)), - ); - let batches = datafusion_physical_plan::collect(plan, ctx.session.task_ctx()).await?; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - if zero_columns { - assert!(batches.iter().all(|batch| batch.num_columns() == 0)); - } else { - assert_batches_eq!(["+----+", "| b |", "+----+", "| 31 |", "+----+"], &batches); - } - 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!( - DefaultExpressionConvertor::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_convertor = Arc::new(ResidualConvertor); - 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 ResidualConvertor; -impl ExpressionConvertor for ResidualConvertor { - fn try_convert( - &self, - _expr: &PhysicalExprRef, - _schema: &Schema, - ) -> DFResult> { - Ok(None) - } -} - -#[rstest] -#[tokio::test] -async fn test_physical_in_list( - #[values(false, true)] negated: bool, - #[values(0, 1, 2)] list_kind: usize, -) -> anyhow::Result<()> { - let store: Arc = Arc::new(InMemory::new()); - let batch = record_batch!( - ("id", Int32, vec![0, 1, 2]), - ("a", Int32, vec![Some(1), Some(2), None]), - ("b", Int32, vec![Some(1), None, Some(2)]) - )?; - let size = write_arrow_to_vortex(Arc::clone(&store), "in-list.vortex", batch.clone()).await?; - let list: Vec = match list_kind { - 0 => vec![], - 1 => vec![Arc::new(df_expr::Column::new("b", 2))], - _ => vec![ - Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), - Arc::new(df_expr::Literal::new(ScalarValue::Int32(None))), - ], - }; - let filter: PhysicalExprRef = Arc::new(df_expr::InListExpr::try_new( - Arc::new(df_expr::Column::new("a", 1)), - list, - negated, - &batch.schema(), - )?); - assert_eq!( - DefaultExpressionConvertor::default() - .try_convert(&filter, &batch.schema())? - .is_some(), - list_kind == 2, - ); - 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("in-list.vortex", size))? - .await? - .try_collect::>() - .await?; - let actual = concat_batches(&expected.schema(), &actual)?; - assert_eq!(actual, expected); - Ok(()) -} - -#[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_convertor = Arc::new(ResidualConvertor); - 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_convertor = Arc::new(ResidualConvertor); - 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 f8ed18f3d10..9e2312b263c 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -5,10 +5,8 @@ use std::fmt::Formatter; use std::sync::Arc; use std::sync::Weak; -use arrow_schema::DataType; use datafusion_common::Result as DFResult; use datafusion_common::config::ConfigOptions; -use datafusion_common::exec_datafusion_err; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; @@ -475,9 +473,6 @@ impl FileSource for VortexSource { let supported_filters = filters .into_iter() .map(|expr| { - if expr.data_type(self.table_schema.table_schema())? != DataType::Boolean { - return Err(exec_datafusion_err!("Filter must be Boolean: {expr}")); - } if self .expression_convertor .try_convert(&expr, self.table_schema.table_schema())? diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 60d5a85b869..66d9bf88eb6 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -646,12 +646,11 @@ async fn arrow_uuid_extension_roundtrip_nested_struct() -> anyhow::Result<()> { async fn test_predicate_memtable_oracle( #[case] projection: &str, #[case] predicate: &str, - #[values(false, true)] projection_pushdown: bool, - #[values(false, true)] predicate_pushdown: bool, + #[values(false, true)] pushdown: bool, ) -> anyhow::Result<()> { let options = crate::VortexTableOptions { - projection_pushdown, - predicate_pushdown, + projection_pushdown: pushdown, + predicate_pushdown: pushdown, ..Default::default() }; let ctx = TestSessionContext::new_with_factory(Arc::new( From d37c4e1454eb210024f3c9a240c012c6fc981ceb Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 13:16:29 +0100 Subject: [PATCH 09/10] Fix nullable cast Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/exprs/mod.rs | 12 ++++++----- vortex-datafusion/src/convert/exprs/tests.rs | 21 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index b05cba79b1f..46d5174a1b8 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -341,7 +341,12 @@ impl DefaultExpressionConvertor { return Err(Unconverted::Unsupported); } let child = self.convert_expr(cast_expr.expr(), schema, input_dtype)?; - let cast_dtype = self.arrow_dtype(cast_expr.target_field())?; + // 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()) @@ -595,10 +600,7 @@ fn converted_dtype(expr: &Expression, input_dtype: &DType) -> Conversion 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 - || (cast.expr().nullable(schema)? && !cast.target_field().is_nullable()) - { + if options.safe || options.format_options != DEFAULT_FORMAT_OPTIONS { return Ok(false); } let source = cast.expr().data_type(schema)?; diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index e7e0e3bfc93..1ed2f3e79df 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -620,6 +620,27 @@ fn test_expr_from_df_like(#[case] negated: bool, #[case] case_insensitive: bool) 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; From 26fd045aa0ca4afdcd93c6e7c1be9049cf007ae9 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 9 Sep 2026 16:21:10 +0100 Subject: [PATCH 10/10] Rename Signed-off-by: Adam Gutglick --- .../integrations/datafusion.md | 4 +- vortex-datafusion/src/convert/exprs/mod.rs | 32 ++--- vortex-datafusion/src/convert/exprs/tests.rs | 34 +++--- vortex-datafusion/src/convert/mod.rs | 6 +- vortex-datafusion/src/persistent/format.rs | 114 +++++++++--------- vortex-datafusion/src/persistent/opener.rs | 38 +++--- vortex-datafusion/src/persistent/source.rs | 44 +++---- vortex-datafusion/src/v2/source.rs | 14 +-- 8 files changed, 143 insertions(+), 143 deletions(-) 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-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 46d5174a1b8..b358c5f3ed9 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -63,12 +63,12 @@ pub struct ProcessedProjection { /// Trait for converting DataFusion expressions to Vortex ones. /// -/// Custom convertors implement a single schema-aware decision. Conversion should preserve -/// DataFusion values, nulls, and evaluation errors; see [`DefaultExpressionConvertor`] for +/// 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 convertor +/// # Implementing a custom converter /// /// ``` /// use std::sync::Arc; @@ -77,12 +77,12 @@ pub struct ProcessedProjection { /// use datafusion_common::Result as DFResult; /// use datafusion_physical_expr::PhysicalExpr; /// use vortex::expr::Expression; -/// use vortex_datafusion::convert::DefaultExpressionConvertor; -/// use vortex_datafusion::convert::ExpressionConvertor; +/// use vortex_datafusion::convert::DefaultExpressionConverter; +/// use vortex_datafusion::convert::ExpressionConverter; /// -/// struct CustomExpressionConvertor(DefaultExpressionConvertor); +/// struct CustomExpressionConverter(DefaultExpressionConverter); /// -/// impl ExpressionConvertor for CustomExpressionConvertor { +/// impl ExpressionConverter for CustomExpressionConverter { /// fn try_convert( /// &self, /// expr: &Arc, @@ -92,11 +92,11 @@ pub struct ProcessedProjection { /// } /// } /// -/// let _convertor: Arc = Arc::new(CustomExpressionConvertor( -/// DefaultExpressionConvertor::default(), +/// let _converter: Arc = Arc::new(CustomExpressionConverter( +/// DefaultExpressionConverter::default(), /// )); /// ``` -pub trait ExpressionConvertor: Send + Sync { +pub trait ExpressionConverter: Send + Sync { /// Convert an expression for native evaluation against this schema. /// /// Returns None for valid but unsupported expressions. Malformed expressions and @@ -193,12 +193,12 @@ impl From for Unconverted { /// Conversion result where `?` propagates both unsupported expressions and errors. type Conversion = Result; -/// The default [`ExpressionConvertor`] implementation. +/// 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 DefaultExpressionConvertor { +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, @@ -206,7 +206,7 @@ pub struct DefaultExpressionConvertor { session: VortexSession, } -impl Default for DefaultExpressionConvertor { +impl Default for DefaultExpressionConverter { fn default() -> Self { Self { session: VortexSession::default(), @@ -214,8 +214,8 @@ impl Default for DefaultExpressionConvertor { } } -impl DefaultExpressionConvertor { - /// Create a convertor that resolves Arrow extension types using `session`'s +impl DefaultExpressionConverter { + /// Create a converter that resolves Arrow extension types using `session`'s /// dtype registry. pub fn new(session: VortexSession) -> Self { Self { session } @@ -567,7 +567,7 @@ impl DefaultExpressionConvertor { } } -impl ExpressionConvertor for DefaultExpressionConvertor { +impl ExpressionConverter for DefaultExpressionConverter { fn try_convert( &self, expr: &Arc, diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 1ed2f3e79df..120cafde885 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -82,16 +82,16 @@ fn array_length_expr(args: Vec>, schema: &Schema) -> Arc, schema: &Schema) -> DFResult { - Ok(DefaultExpressionConvertor::default() + Ok(DefaultExpressionConverter::default() .try_convert(expr, schema)? .is_some()) } -/// Convert `expr` natively, failing the test if the convertor declines it. +/// Convert `expr` natively, failing the test if the converter declines it. fn convert(expr: Arc, schema: &Schema) -> DFResult { - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, schema)? .ok_or_else(|| exec_datafusion_err!("Expected native conversion for {expr}")) } @@ -115,9 +115,9 @@ fn int_literals(values: impl IntoIterator>) -> Vec, @@ -147,7 +147,7 @@ fn test_duplicate_aliases_fall_back_before_conversion() -> DFResult<()> { ]); let output_schema = projection.project_schema(&schema)?; let processed = - FailingConvertor.split_projection(projection.clone(), &schema, &output_schema)?; + FailingConverter.split_projection(projection.clone(), &schema, &output_schema)?; assert_eq!( processed.scan_projection, pack( @@ -395,7 +395,7 @@ fn test_in_list_malformed_literal() -> DFResult<()> { ))), ])?; assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &schema) .is_err() ); @@ -806,7 +806,7 @@ fn test_projection_ignores_unreferenced_unsupported_field() -> DFResult<()> { "sum", )]); let output_schema = projection.project_schema(&schema)?; - let processed = DefaultExpressionConvertor::default().split_projection( + let processed = DefaultExpressionConverter::default().split_projection( projection, &schema, &output_schema, @@ -871,7 +871,7 @@ 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!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&col_expr, &test_schema) .is_err() ); @@ -1065,7 +1065,7 @@ fn test_length_function_invalid_arity(#[case] function: ScalarUDF, #[case] arity Arc::new(ConfigOptions::new()), )); assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &Schema::empty()) .is_err() ); @@ -1162,7 +1162,7 @@ fn test_case_when_datafusion_vortex_equivalence() -> anyhow::Result<()> { fn assert_native_matches(expr: Arc, batch: RecordBatch) -> anyhow::Result<()> { let session = VortexSession::default(); - let converted = DefaultExpressionConvertor::new(session.clone()) + 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] { @@ -1493,7 +1493,7 @@ fn test_malformed_get_field_returns_error( Arc::new(ConfigOptions::new()), )); assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &schema) .is_err() ); @@ -1508,7 +1508,7 @@ fn test_column_identity_validation_skips_dynamic_filters() -> DFResult<()> { ]); let expr: Arc = Arc::new(df_expr::Column::new("a", 1)); assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &schema) .is_err() ); @@ -1606,7 +1606,7 @@ fn test_case_non_boolean_condition_returns_error() -> DFResult<()> { None, )?); assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &Schema::empty()) .is_err() ); @@ -1618,7 +1618,7 @@ 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!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &Schema::empty()) .is_err() ); @@ -1630,7 +1630,7 @@ fn test_literal_invalid_row_count(#[values(0, 2)] len: usize) { let expr: Arc = Arc::new(df_expr::Literal::new(ScalarValue::Struct(Arc::new(array)))); assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&expr, &Schema::empty()) .is_err() ); 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 2e9afbb173e..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,37 +827,37 @@ mod tests { } #[derive(Default)] - struct ExpressionConvertorCalls { + struct ExpressionConverterCalls { try_convert: AtomicBool, } - impl ExpressionConvertorCalls { + impl ExpressionConverterCalls { fn reset(&self) { 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 { + impl ExpressionConverter for TestExpressionConverter { fn try_convert( &self, expr: &Arc, @@ -881,24 +881,24 @@ 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(), )?; @@ -979,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 diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 267b58ae9fd..8d7696d6c2b 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -58,7 +58,7 @@ 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::raw_projection; use crate::convert::schema::calculate_physical_schema; @@ -103,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, @@ -144,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) @@ -300,7 +300,7 @@ impl FileOpener for VortexOpener { )); } for expr in split_conjunction(filter) { - let converted = expr_convertor + let converted = expr_converter .try_convert(expr, &this_file_schema) .map_err(|e| { exec_datafusion_err!( @@ -340,13 +340,13 @@ impl FileOpener for VortexOpener { leftover_projection: projection, } } else if projection_pushdown { - expr_convertor.split_projection( + expr_converter.split_projection( projection, &this_file_schema, output_schema.as_ref(), )? } else { - expr_convertor.no_pushdown_projection(projection, &this_file_schema)? + expr_converter.no_pushdown_projection(projection, &this_file_schema)? }; // The schema of the stream returned from the vortex scan. @@ -705,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); @@ -886,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, @@ -1256,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, @@ -1343,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, @@ -1497,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, @@ -1557,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, @@ -1764,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, @@ -1883,7 +1883,7 @@ mod tests { ))), )); assert!( - DefaultExpressionConvertor::default() + DefaultExpressionConverter::default() .try_convert(&filter, &logical)? .is_some() ); @@ -1914,7 +1914,7 @@ mod tests { ); opener.projection = Vec::::new().into(); // Force a supported physical expression to use the residual path. - opener.expression_convertor = Arc::new(ResidualConvertor); + opener.expression_converter = Arc::new(ResidualConverter); let batches = opener .open(PartitionedFile::new("literal.vortex", size))? .await? @@ -1928,8 +1928,8 @@ mod tests { Ok(()) } - struct ResidualConvertor; - impl ExpressionConvertor for ResidualConvertor { + struct ResidualConverter; + impl ExpressionConverter for ResidualConverter { fn try_convert( &self, _expr: &PhysicalExprRef, @@ -1995,7 +1995,7 @@ mod tests { Arc::new(df_expr::Literal::new(ScalarValue::Int32(Some(1)))), )); let mut opener = make_opener(store, TableSchema::from(batch.schema()), Some(filter)); - opener.expression_convertor = Arc::new(ResidualConvertor); + opener.expression_converter = Arc::new(ResidualConverter); let result = opener .open(PartitionedFile::new("filter-error.vortex", size))? .await? @@ -2024,7 +2024,7 @@ mod tests { false, ))))), ); - opener.expression_convertor = Arc::new(ResidualConvertor); + opener.expression_converter = Arc::new(ResidualConverter); opener.projection = vec![ProjectionExpr { expr: Arc::new(SnapshotErrorExpr), alias: "failure".into(), diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 9e2312b263c..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; @@ -137,7 +137,7 @@ use crate::persistent::reader::VortexReaderFactory; /// /// - when disabled, `VortexSource` still prunes unreferenced top-level columns, /// but DataFusion applies the full projection after the scan, -/// - when enabled, the default convertor evaluates fully supported projections +/// - 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`]: @@ -199,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, @@ -221,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, @@ -232,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, @@ -260,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 } @@ -357,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, @@ -474,7 +474,7 @@ impl FileSource for VortexSource { .into_iter() .map(|expr| { if self - .expression_convertor + .expression_converter .try_convert(&expr, self.table_schema.table_schema())? .is_some() { @@ -570,11 +570,11 @@ mod tests { use super::*; use crate::convert::exprs::ProcessedProjection; - struct TrackingExpressionConvertor { - inner: DefaultExpressionConvertor, + struct TrackingExpressionConverter { + inner: DefaultExpressionConverter, } - impl ExpressionConvertor for TrackingExpressionConvertor { + impl ExpressionConverter for TrackingExpressionConverter { fn try_convert( &self, expr: &PhysicalExprRef, @@ -664,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(), @@ -686,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/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 258a97fc591..706e3bbc7ef 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -118,8 +118,8 @@ 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::stats::stats_set_to_df; @@ -544,18 +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,7 +602,7 @@ impl DataSource for VortexDataSource { )); } - let convertor = DefaultExpressionConvertor::default(); + let converter = DefaultExpressionConverter::default(); let filters = filters .into_iter() .map(|filter| { @@ -618,7 +618,7 @@ impl DataSource for VortexDataSource { // so we can safely claim PushedDown::Yes for them. let converted = filters .iter() - .map(|expr| convertor.try_convert(expr, &self.projected_schema)) + .map(|expr| converter.try_convert(expr, &self.projected_schema)) .collect::>>()?; let pushdown_results: Vec = converted .iter()