feat: extend interval arithmetic to support scalar functions (UDFs)

datafusion/functions/src/math/ceil.rs

@@ -191,11 +191,70 @@ impl ScalarUDFImpl for CeilFunc {
     }
 
     fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> {
-        let data_type = inputs[0].data_type();
-        Interval::make_unbounded(&data_type)
+        let [input] = inputs else {
+            let data_type = inputs
+                .first()
+                .map(|i| i.data_type())
+                .unwrap_or(DataType::Float64);
+            return Interval::make_unbounded(&data_type);
+        };
+        let data_type = input.data_type();
+        match (ceil_scalar(input.lower()), ceil_scalar(input.upper())) {
+            (Some(lo), Some(hi)) => Interval::try_new(lo, hi)
+                .or_else(|_| Interval::make_unbounded(&data_type)),
+            _ => Interval::make_unbounded(&data_type),
+        }
+    }
+
+    fn propagate_constraints(
+        &self,
+        interval: &Interval,
+        inputs: &[&Interval],
+    ) -> Result<Option<Vec<Interval>>> {
+        let [input_interval] = inputs else {
+            return Ok(Some(inputs.iter().map(|i| (*i).clone()).collect()));
+        };
+        // ceil(x) ∈ [N, M] → x ∈ (N−1, M] — conservative closed: [N−1, M]
+        let lo = match interval.lower() {
+            ScalarValue::Float64(Some(n)) if n.is_finite() => {
+                Some(ScalarValue::Float64(Some(n - 1.0)))
+            }
+            ScalarValue::Float32(Some(n)) if n.is_finite() => {
+                Some(ScalarValue::Float32(Some(n - 1.0)))
+            }
+            _ => None,
+        };
+        let hi = match interval.upper() {
+            ScalarValue::Float64(Some(n)) if n.is_finite() => {
+                Some(ScalarValue::Float64(Some(*n)))
+            }
+            ScalarValue::Float32(Some(n)) if n.is_finite() => {
+                Some(ScalarValue::Float32(Some(*n)))
+            }
+            _ => None,
+        };
+        match (lo, hi) {
+            (Some(lo), Some(hi)) => {
+                let constraint = Interval::try_new(lo, hi)?;
+                Ok(input_interval.intersect(constraint)?.map(|r| vec![r]))
+            }
+            _ => Ok(Some(vec![(*input_interval).clone()])),
+        }
     }
 
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
 }
+
+fn ceil_scalar(v: &ScalarValue) -> Option<ScalarValue> {
+    match v {
+        ScalarValue::Float64(Some(f)) if f.is_finite() => {
+            Some(ScalarValue::Float64(Some(f.ceil())))
+        }
+        ScalarValue::Float32(Some(f)) if f.is_finite() => {
+            Some(ScalarValue::Float32(Some(f.ceil())))
+        }
+        _ => None,
+    }
+}

datafusion/physical-expr/src/intervals/utils.rs

@@ -22,6 +22,7 @@ use std::sync::Arc;
 use crate::{
     PhysicalExpr,
     expressions::{BinaryExpr, CastExpr, Column, Literal, NegativeExpr},
+    scalar_function::ScalarFunctionExpr,
 };
 
 use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano};
@@ -56,6 +57,12 @@ pub fn check_support(expr: &Arc<dyn PhysicalExpr>, schema: &SchemaRef) -> bool {
         check_support(cast.expr(), schema)
     } else if let Some(negative) = expr.downcast_ref::<NegativeExpr>() {
         check_support(negative.arg(), schema)
+    } else if let Some(scalar_fn) = expr.downcast_ref::<ScalarFunctionExpr>() {
+        is_datatype_supported(scalar_fn.return_type())
+            && scalar_fn
+                .args()
+                .iter()
+                .all(|arg| check_support(arg, schema))
     } else {
         false
     }
@@ -193,3 +200,146 @@ fn interval_dt_to_duration_ms(dt: &IntervalDayTime) -> Result<i64> {
         )
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::expressions::{Column, Literal};
+    use crate::scalar_function::ScalarFunctionExpr;
+    use arrow::datatypes::{Field, Schema};
+    use datafusion_common::config::ConfigOptions;
+    use datafusion_common::{Result, ScalarValue};
+    use datafusion_expr::{
+        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+    };
+
+    fn f64_schema() -> SchemaRef {
+        Arc::new(Schema::new(vec![Field::new("x", DataType::Float64, false)]))
+    }
+
+    fn utf8_schema() -> SchemaRef {
+        Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]))
+    }
+
+    fn col_x() -> Arc<dyn PhysicalExpr> {
+        Arc::new(Column::new("x", 0))
+    }
+
+    fn lit_f64(v: f64) -> Arc<dyn PhysicalExpr> {
+        Arc::new(Literal::new(ScalarValue::Float64(Some(v))))
+    }
+
+    fn scalar_fn_expr(
+        udf: Arc<datafusion_expr::ScalarUDF>,
+        args: Vec<Arc<dyn PhysicalExpr>>,
+        return_type: DataType,
+    ) -> Arc<dyn PhysicalExpr> {
+        let name = udf.name().to_string();
+        Arc::new(ScalarFunctionExpr::new(
+            &name,
+            udf,
+            args,
+            Field::new("result", return_type, true).into(),
+            Arc::new(ConfigOptions::default()),
+        ))
+    }
+
+    /// A minimal UDF whose declared return type is Utf8, used to test that
+    /// check_support rejects functions with unsupported return types without
+    /// relying on an invalid ceil-returns-Utf8 combination.
+    #[derive(Debug, PartialEq, Eq, Hash)]
+    struct Utf8UDF {
+        signature: Signature,
+    }
+
+    impl Utf8UDF {
+        fn new() -> Self {
+            Self {
+                signature: Signature::uniform(
+                    1,
+                    vec![DataType::Float64],
+                    Volatility::Immutable,
+                ),
+            }
+        }
+    }
+
+    impl ScalarUDFImpl for Utf8UDF {
+        fn as_any(&self) -> &dyn std::any::Any {
+            self
+        }
+        fn name(&self) -> &str {
+            "utf8_udf"
+        }
+        fn signature(&self) -> &Signature {
+            &self.signature
+        }
+        fn return_type(&self, _: &[DataType]) -> Result<DataType> {
+            Ok(DataType::Utf8)
+        }
+        fn invoke_with_args(&self, _: ScalarFunctionArgs) -> Result<ColumnarValue> {
+            Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
+        }
+    }
+
+    fn utf8_udf() -> Arc<datafusion_expr::ScalarUDF> {
+        Arc::new(datafusion_expr::ScalarUDF::from(Utf8UDF::new()))
+    }
+
+    #[test]
+    fn test_check_support_scalar_fn_supported_return_type() {
+        // ceil(x) returns Float64 — both return type and child are supported
+        let schema = f64_schema();
+        let expr = scalar_fn_expr(
+            datafusion_functions::math::ceil(),
+            vec![col_x()],
+            DataType::Float64,
+        );
+        assert!(check_support(&expr, &schema));
+    }
+
+    #[test]
+    fn test_check_support_scalar_fn_unsupported_return_type() {
+        // utf8_udf(x) returns Utf8 — not in is_datatype_supported
+        let schema = f64_schema();
+        let expr = scalar_fn_expr(utf8_udf(), vec![col_x()], DataType::Utf8);
+        assert!(!check_support(&expr, &schema));
+    }
+
+    #[test]
+    fn test_check_support_scalar_fn_unsupported_child() {
+        // ceil applied to a Utf8 column — child fails is_datatype_supported
+        let schema = utf8_schema();
+        let col_s = Arc::new(Column::new("s", 0)) as Arc<dyn PhysicalExpr>;
+        let expr = scalar_fn_expr(
+            datafusion_functions::math::ceil(),
+            vec![col_s],
+            DataType::Float64,
+        );
+        assert!(!check_support(&expr, &schema));
+    }
+
+    #[test]
+    fn test_check_support_scalar_fn_in_binary_expr() {
+        // ceil(x) > 5.0 — the main use case: ScalarFunctionExpr inside a BinaryExpr
+        let schema = f64_schema();
+        let ceil_x = scalar_fn_expr(
+            datafusion_functions::math::ceil(),
+            vec![col_x()],
+            DataType::Float64,
+        );
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(ceil_x, Operator::Gt, lit_f64(5.0)));
+        assert!(check_support(&expr, &schema));
+    }
+
+    #[test]
+    fn test_check_support_scalar_fn_in_binary_expr_unsupported_return() {
+        // utf8_udf(x) > 5.0 where f returns Utf8 — should be false
+        let schema = f64_schema();
+        let fn_expr = scalar_fn_expr(utf8_udf(), vec![col_x()], DataType::Utf8);
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(fn_expr, Operator::Gt, lit_f64(5.0)));
+        assert!(!check_support(&expr, &schema));
+    }
+}

datafusion/physical-plan/src/filter.rs

@@ -1209,6 +1209,66 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_filter_statistics_ceil_scalar_fn() -> Result<()> {
+        // Table: x Float64, min=8.0, max=16.0, 100 rows.
+        // Filter: ceil(x) > 12.0
+        //
+        // The range [8.0, 16.0] lies within a single IEEE-754 binade so
+        // Float64 cardinality is proportional to the value range, making
+        // the selectivity estimate predictable.
+        //
+        // With check_support recognising ScalarFunctionExpr and CeilFunc
+        // implementing evaluate_bounds/propagate_constraints the solver
+        // narrows x to roughly [11.0, 16.0]:
+        //   ceil(x) > 12 → x ∈ (11, 16] → conservative [11, 16]
+        //   selectivity ≈ (16−11)/(16−8) = 5/8 = 0.625 → ~62 rows
+        //
+        // Without the fix the estimate stays at 100 (no interval analysis).
+        let schema = Schema::new(vec![Field::new("x", DataType::Float64, false)]);
+        let input = Arc::new(StatisticsExec::new(
+            Statistics {
+                num_rows: Precision::Inexact(100),
+                total_byte_size: Precision::Absent,
+                column_statistics: vec![ColumnStatistics {
+                    min_value: Precision::Inexact(ScalarValue::Float64(Some(8.0))),
+                    max_value: Precision::Inexact(ScalarValue::Float64(Some(16.0))),
+                    ..Default::default()
+                }],
+            },
+            schema.clone(),
+        ));
+
+        let x = col("x", &schema)?;
+        let ceil_udf = datafusion_functions::math::ceil();
+        let config = Arc::new(ConfigOptions::new());
+        let ceil_x: Arc<dyn PhysicalExpr> =
+            Arc::new(datafusion_physical_expr::ScalarFunctionExpr::try_new(
+                Arc::clone(&ceil_udf),
+                vec![x],
+                &schema,
+                config,
+            )?);
+        let predicate = binary(ceil_x, Operator::Gt, lit(12.0f64), &schema)?;
+
+        let filter = Arc::new(FilterExec::try_new(predicate, input)?);
+        let statistics = filter.partition_statistics(None)?;
+
+        let num_rows = statistics.num_rows.get_value().copied().unwrap_or(100);
+        // Interval analysis must narrow the estimate below the full 100-row input.
+        assert!(
+            num_rows < 100,
+            "expected interval analysis to narrow row estimate, got {num_rows}"
+        );
+        // The conservative bound is x ∈ [11, 16] out of [8, 16] → ~62 rows.
+        // Allow a generous range to be robust to float-cardinality rounding.
+        assert!(
+            num_rows >= 50,
+            "expected at least 50 rows after ceil(x) > 12.0 on [8,16], got {num_rows}"
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_filter_statistics_basic_expr() -> Result<()> {
         // Table: