diff --git a/vortex-datafusion/src/convert/exprs/mod.rs b/vortex-datafusion/src/convert/exprs/mod.rs index 2f5c4b3862b..b3b4563db09 100644 --- a/vortex-datafusion/src/convert/exprs/mod.rs +++ b/vortex-datafusion/src/convert/exprs/mod.rs @@ -342,9 +342,7 @@ impl DefaultExpressionConverter { } if let Some(cast_expr) = expr.downcast_ref::() { - if !supported_cast(cast_expr, schema)? { - return Err(Unconverted::Unsupported); - } + let pushdown = classify_cast(cast_expr, schema)?; let child = self.convert_expr(cast_expr.expr(), schema, input_dtype)?; // DataFusion casts preserve input nulls even when the target field is declared // non-nullable. Match the expression's runtime nullability rather than narrowing @@ -352,14 +350,21 @@ impl DefaultExpressionConverter { let cast_dtype = self .arrow_dtype(cast_expr.target_field())? .with_nullability(Nullability::from(cast_expr.nullable(schema)?)); - // Matching Arrow storage types do not imply matching extension semantics. let child_dtype = converted_dtype(&child, input_dtype)?; - if (child_dtype.is_extension() || cast_dtype.is_extension()) - && !child_dtype.eq_ignore_nullability(&cast_dtype) - { - return Err(Unconverted::Unsupported); - } - return Ok(cast(child, cast_dtype)); + return match pushdown { + // Matching Arrow types do not imply matching Vortex dtypes, e.g. for extensions. + CastPushdown::Drop if child_dtype.eq_ignore_nullability(&cast_dtype) => Ok(child), + CastPushdown::Drop => Err(Unconverted::Unsupported), + CastPushdown::Keep => { + // Matching Arrow storage types do not imply matching extension semantics. + if (child_dtype.is_extension() || cast_dtype.is_extension()) + && !child_dtype.eq_ignore_nullability(&cast_dtype) + { + return Err(Unconverted::Unsupported); + } + Ok(cast(child, cast_dtype)) + } + }; } if let Some(not_expr) = expr.downcast_ref::() { @@ -600,25 +605,41 @@ fn converted_dtype(expr: &Expression, input_dtype: &DType) -> Conversion .map_err(|_| Unconverted::Unsupported) } -fn supported_cast(cast: &df_expr::CastExpr, schema: &Schema) -> DFResult { +/// How a supported DataFusion cast is converted. +enum CastPushdown { + /// The Arrow cast cannot fail and keeps the Vortex dtype, so the converter omits it. + Drop, + /// The cast becomes a Vortex `cast`. + Keep, +} + +/// Classifies a cast for pushdown; unsupported casts stay in DataFusion. +fn classify_cast(cast: &df_expr::CastExpr, schema: &Schema) -> Conversion { use DataType::*; let options = cast.cast_options(); if options.safe || options.format_options != DEFAULT_FORMAT_OPTIONS { - return Ok(false); + return Err(Unconverted::Unsupported); } 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) - )) + Ok(match (&source, target) { + _ if source == *target => CastPushdown::Drop, + // Arrow expands runs by gathering with `Int32` indices, which fit Int16/Int32 run ends. + // Gathering a dictionary only moves keys and reuses the values, so nothing can overflow. + (RunEndEncoded(run_ends, values), Dictionary(..)) + if matches!(run_ends.data_type(), Int16 | Int32) && values.data_type() == target => + { + CastPushdown::Drop + } + (Int8, Int16 | Int32 | Int64) + | (Int16, Int32 | Int64) + | (Int32, Int64) + | (UInt8, UInt16 | UInt32 | UInt64) + | (UInt16, UInt32 | UInt64) + | (UInt32, UInt64) + | (Float32, Float64) => CastPushdown::Keep, + _ => return Err(Unconverted::Unsupported), + }) } fn try_operator_from_df(value: &DFOperator) -> Option { diff --git a/vortex-datafusion/src/convert/exprs/tests.rs b/vortex-datafusion/src/convert/exprs/tests.rs index 6b3f737d766..6e5a31b7a30 100644 --- a/vortex-datafusion/src/convert/exprs/tests.rs +++ b/vortex-datafusion/src/convert/exprs/tests.rs @@ -4,15 +4,28 @@ use std::sync::Arc; use arrow_array::Array; +use arrow_array::DictionaryArray; use arrow_array::FixedSizeBinaryArray; +use arrow_array::Int32Array; use arrow_array::RecordBatch; +use arrow_array::RunArray; +use arrow_array::StringArray; use arrow_array::StructArray; +use arrow_array::UInt32Array; +use arrow_array::types::Int32Type; +use arrow_array::types::UInt32Type; use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::IntervalUnit; use arrow_schema::Schema; use arrow_schema::TimeUnit as ArrowTimeUnit; use arrow_schema::extension::Uuid as ArrowUuid; +use datafusion::common::DFSchema; +use datafusion::execution::context::ExecutionProps; +use datafusion::logical_expr::simplify::SimplifyContext; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; +use datafusion::physical_expr::create_physical_expr; +use datafusion::prelude::SessionContext; use datafusion_common::ScalarValue; use datafusion_common::arrow::buffer::NullBuffer; use datafusion_common::arrow::compute::cast; @@ -39,6 +52,7 @@ use vortex::expr::between; use vortex::scalar::Scalar; use vortex::scalar_fn::fns::between::BetweenOptions; use vortex::scalar_fn::fns::between::StrictComparison; +use vortex::scalar_fn::fns::cast::Cast; use vortex::scalar_fn::fns::list_contains::ListContains; use vortex::scalar_fn::fns::literal::Literal; use vortex_array::assert_arrays_eq; @@ -377,7 +391,7 @@ fn test_contiguous_in_list_converts_to_between() -> DFResult<()> { fn test_in_list_from_array_converts_to_set_lookup(#[case] data_type: DataType) -> DFResult<()> { let schema = Schema::new(vec![Field::new("a", data_type.clone(), true)]); let values = cast( - &arrow_array::StringArray::from(vec![Some("x"), None, Some("y"), Some("x")]), + &StringArray::from(vec![Some("x"), None, Some("y"), Some("x")]), &data_type, )?; let expr = Arc::new(df_expr::InListExpr::try_new_from_array( @@ -1260,9 +1274,7 @@ fn test_case_when_datafusion_vortex_equivalence() -> anyhow::Result<()> { DataType::Int32, false, )])), - vec![Arc::new(arrow_array::Int32Array::from(vec![ - 1, 5, 10, 15, 20, - ]))], + vec![Arc::new(Int32Array::from(vec![1, 5, 10, 15, 20]))], )?; let value: Arc = Arc::new(df_expr::Column::new("value", 0)); let int32 = @@ -1457,6 +1469,131 @@ fn test_cast_options_fall_back() -> DFResult<()> { Ok(()) } +fn dict_utf8(key: DataType) -> DataType { + DataType::Dictionary(Box::new(key), Box::new(DataType::Utf8)) +} + +fn ree_dict_utf8(run_ends: DataType) -> DataType { + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", run_ends, false)), + Arc::new(Field::new("values", dict_utf8(DataType::UInt32), true)), + ) +} + +fn ree_utf8(run_ends: DataType) -> DataType { + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", run_ends, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ) +} + +fn contains_cast(expr: &Expression) -> bool { + expr.is::() || expr.children().iter().any(contains_cast) +} + +/// Plans `predicate` as DataFusion does: type coercion, simplification, then physical planning. +fn plan_predicate(predicate: &str, schema: &Schema) -> DFResult> { + let df_schema = Arc::new(DFSchema::try_from(schema.clone())?); + let simplifier = ExprSimplifier::new( + SimplifyContext::builder() + .with_schema(Arc::clone(&df_schema)) + .build(), + ); + let expr = SessionContext::new().parse_sql_expr(predicate, &df_schema)?; + let expr = simplifier.simplify(simplifier.coerce(expr, &df_schema)?)?; + create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &Default::default(), + ) +} + +/// Regression test for #10634: coercion casts `REE>` columns to `Dict` before +/// comparing them, and the converter must drop that cast rather than decline the predicate. +#[rstest] +#[case::ree_eq("ree_dict = 'x'")] +#[case::ree_or("ree_dict = 'x' OR ree_dict = 'y'")] +#[case::ree_in("ree_dict IN ('x', 'y', 'z')")] +#[case::dict_eq("dict = 'x'")] +#[case::plain_eq("plain = 'x'")] +fn test_coerced_string_comparison_drops_encoding_cast(#[case] predicate: &str) -> DFResult<()> { + let schema = Schema::new(vec![ + Field::new("ree_dict", ree_dict_utf8(DataType::Int32), true), + Field::new("dict", dict_utf8(DataType::UInt32), true), + Field::new("plain", DataType::Utf8, true), + ]); + let converted = convert(plan_predicate(predicate, &schema)?, &schema)?; + assert!(!contains_cast(&converted), "{}", converted.display_tree()); + Ok(()) +} + +#[rstest] +#[case::ree_to_dict(ree_dict_utf8(DataType::Int32), dict_utf8(DataType::UInt32))] +#[case::ree_int16_run_ends(ree_dict_utf8(DataType::Int16), dict_utf8(DataType::UInt32))] +#[case::ree_identity(ree_dict_utf8(DataType::Int64), ree_dict_utf8(DataType::Int64))] +#[case::int_identity(DataType::Int32, DataType::Int32)] +fn test_infallible_encoding_cast_is_dropped( + #[case] source: DataType, + #[case] target: DataType, +) -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", source, true)]); + let expr: Arc = Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + target, + None, + )); + let converted = convert(expr, &schema)?; + assert!(!contains_cast(&converted), "{}", converted.display_tree()); + Ok(()) +} + +// Arrow expands runs with `Int32` indices and decodes offset-based values through `take`, which +// may overflow; key changes may overflow the target keys. +#[rstest] +#[case::ree_int64_run_ends(ree_dict_utf8(DataType::Int64), dict_utf8(DataType::UInt32))] +#[case::ree_narrow_keys(ree_dict_utf8(DataType::Int32), dict_utf8(DataType::UInt8))] +#[case::ree_decode(ree_dict_utf8(DataType::Int32), DataType::Utf8)] +#[case::ree_plain_utf8(ree_utf8(DataType::Int32), DataType::Utf8)] +#[case::dict_narrow_keys(dict_utf8(DataType::UInt32), dict_utf8(DataType::UInt8))] +#[case::dict_decode(dict_utf8(DataType::UInt32), DataType::Utf8)] +#[case::dict_encode(DataType::Utf8, dict_utf8(DataType::UInt8))] +#[case::large_utf8(DataType::LargeUtf8, DataType::Utf8)] +fn test_fallible_encoding_cast_falls_back( + #[case] source: DataType, + #[case] target: DataType, +) -> DFResult<()> { + let schema = Schema::new(vec![Field::new("a", source, true)]); + let expr: Arc = Arc::new(df_expr::CastExpr::new( + Arc::new(df_expr::Column::new("a", 0)), + target, + None, + )); + assert!(!converts(&expr, &schema)?); + Ok(()) +} + +#[rstest] +#[case::eq("ree_dict = 'x'")] +#[case::not_eq("ree_dict <> 'x'")] +#[case::or("ree_dict = 'x' OR ree_dict = 'y'")] +#[case::in_list("ree_dict IN ('y', 'unused')")] +fn test_ree_dict_comparison_matches_datafusion(#[case] predicate: &str) -> anyhow::Result<()> { + // Rows: x, x, NULL, y, y, x. The dictionary also holds a value no key references. + let values = DictionaryArray::::try_new( + UInt32Array::from(vec![Some(0), None, Some(1), Some(0)]), + Arc::new(StringArray::from(vec!["x", "y", "unused"])), + )?; + let ree = RunArray::::try_new(&Int32Array::from(vec![2, 3, 5, 6]), &values)?; + let schema = Arc::new(Schema::new(vec![Field::new( + "ree_dict", + ree.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ree)])?; + assert_native_matches(plan_predicate(predicate, &schema)?, batch) +} + #[rstest] #[case::add(DFOperator::Plus)] #[case::sub(DFOperator::Minus)] @@ -1654,7 +1791,7 @@ fn test_native_nested_get_field( } let inner = Arc::new(StructArray::new( vec![Field::new("leaf", DataType::Int32, nullable_leaf)].into(), - vec![Arc::new(arrow_array::Int32Array::from(values))], + vec![Arc::new(Int32Array::from(values))], inner_validity.map(NullBuffer::from), )); let outer = Arc::new(StructArray::new( diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 1767ee0585b..30ef93453ce 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -4,9 +4,18 @@ use std::sync::Arc; use anyhow::anyhow; +use datafusion::arrow::array::Array; +use datafusion::arrow::array::DictionaryArray; use datafusion::arrow::array::Int32Array; use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::array::RunArray; +use datafusion::arrow::array::StringArray; +use datafusion::arrow::array::UInt32Array; use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::datatypes::Field; +use datafusion::arrow::datatypes::Int32Type; +use datafusion::arrow::datatypes::Schema; +use datafusion::arrow::datatypes::UInt32Type; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::datasource::provider::DefaultTableFactory; use datafusion::execution::SessionStateBuilder; @@ -14,7 +23,11 @@ use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; use datafusion_common::GetExt; use datafusion_common::arrow::compute::concat_batches; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::source::DataSourceExec; +use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::display::DisplayableExecutionPlan; +use datafusion_physical_plan::filter::FilterExec; use insta::assert_snapshot; use object_store::ObjectStore; use object_store::memory::InMemory; @@ -42,6 +55,7 @@ use vortex::layout::layouts::table::TableStrategy; use vortex::session::VortexSession; use crate::VortexFormatFactory; +use crate::VortexSource; use crate::common_tests::TestSessionContext; fn make_session( @@ -794,3 +808,83 @@ async fn test_predicate_memtable_oracle( } Ok(()) } + +/// Returns whether `plan` filters rows in DataFusion and whether a Vortex scan accepted a +/// predicate. +fn filter_placement(plan: &dyn ExecutionPlan) -> (bool, bool) { + let mut in_datafusion = plan.downcast_ref::().is_some(); + let mut in_scan = plan + .downcast_ref::() + .and_then(|exec| exec.data_source().downcast_ref::()) + .and_then(|config| config.file_source.downcast_ref::()) + .is_some_and(|source| source.predicate().is_some()); + for child in plan.children() { + let (child_in_datafusion, child_in_scan) = filter_placement(child.as_ref()); + in_datafusion |= child_in_datafusion; + in_scan |= child_in_scan; + } + (in_datafusion, in_scan) +} + +/// Regression test for #10634: string comparisons on `REE>` columns are pushed into +/// the scan and return the same rows as DataFusion. +#[rstest] +#[case::eq("ree_dict = 'x'")] +#[case::not_eq("ree_dict <> 'x'")] +#[case::or("ree_dict = 'x' OR ree_dict = 'y'")] +#[case::in_list("ree_dict IN ('y', 'unused')")] +#[case::and("ree_dict = 'x' AND id > 1")] +#[tokio::test] +async fn test_ree_dict_predicate_memtable_oracle( + #[case] predicate: &str, + #[values(false, true)] pushdown: bool, +) -> anyhow::Result<()> { + let options = crate::VortexTableOptions { + projection_pushdown: pushdown, + predicate_pushdown: pushdown, + ..Default::default() + }; + let ctx = TestSessionContext::new_with_factory(Arc::new( + VortexFormatFactory::new().with_options(options), + )); + // Rows: x, x, NULL, y, y, x. The dictionary also holds a value no key references. + let values = DictionaryArray::::try_new( + UInt32Array::from(vec![Some(0), None, Some(1), Some(0)]), + Arc::new(StringArray::from(vec!["x", "y", "unused"])), + )?; + let ree = RunArray::::try_new(&Int32Array::from(vec![2, 3, 5, 6]), &values)?; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("ree_dict", ree.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..6)), Arc::new(ree)], + )?; + ctx.write_arrow_batch("ree_dict.vortex", &batch).await?; + let actual_table = ctx + .table_provider("actual", "/ree_dict.vortex", schema.as_ref().clone()) + .await?; + ctx.session.register_table("actual", actual_table)?; + ctx.session.register_table( + "oracle", + Arc::new(datafusion::datasource::MemTable::try_new( + schema, + vec![vec![batch]], + )?), + )?; + + let query = |table: &str| format!("SELECT id FROM {table} WHERE {predicate} ORDER BY id"); + let actual = ctx.session.sql(&query("actual")).await?; + let plan = actual.clone().create_physical_plan().await?; + assert_eq!( + filter_placement(plan.as_ref()), + (!pushdown, pushdown), + "{}", + DisplayableExecutionPlan::new(plan.as_ref()).indent(true) + ); + let actual = actual.collect().await?; + let expected = ctx.session.sql(&query("oracle")).await?.collect().await?; + assert_eq!(batch_values(&actual), batch_values(&expected)); + Ok(()) +}