Push down comparisons through REE-to-dictionary casts

vortex-datafusion/src/convert/exprs/mod.rs

@@ -342,24 +342,29 @@ impl DefaultExpressionConverter {
         }
 
         if let Some(cast_expr) = expr.downcast_ref::<df_expr::CastExpr>() {
-            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
             // to the target field's metadata.
             let cast_dtype = self
                 .arrow_dtype(cast_expr.target_field())?
                 .with_nullability(Nullability::from(cast_expr.nullable(schema)?));
-            // Matching Arrow storage types do not imply matching extension semantics.
             let child_dtype = converted_dtype(&child, input_dtype)?;
-            if (child_dtype.is_extension() || cast_dtype.is_extension())
-                && !child_dtype.eq_ignore_nullability(&cast_dtype)
-            {
-                return Err(Unconverted::Unsupported);
-            }
-            return Ok(cast(child, cast_dtype));
+            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::<df_expr::NotExpr>() {
@@ -600,25 +605,41 @@ fn converted_dtype(expr: &Expression, input_dtype: &DType) -> Conversion<DType>
         .map_err(|_| Unconverted::Unsupported)
 }
 
-fn supported_cast(cast: &df_expr::CastExpr, schema: &Schema) -> DFResult<bool> {
+/// 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<CastPushdown> {
     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<Operator> {

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<dyn PhysicalExpr> = 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::<Cast>() || 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<Arc<dyn PhysicalExpr>> {
+    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<Dict<Utf8>>` columns to `Dict<Utf8>` 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<dyn PhysicalExpr> = 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<dyn PhysicalExpr> = 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::<UInt32Type>::try_new(
+        UInt32Array::from(vec![Some(0), None, Some(1), Some(0)]),
+        Arc::new(StringArray::from(vec!["x", "y", "unused"])),
+    )?;
+    let ree = RunArray::<Int32Type>::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(