[arrow-datafusion] branch master updated: Expose and document a simpler public API for simplify expressions (#3719)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new fef45e74d Expose and document a simpler public API for simplify expressions (#3719)
fef45e74d is described below

commit fef45e74d6772cf1f4aa8b32338ac4509fa24ab4
Author: ygf11 <ya...@gmail.com>
AuthorDate: Fri Oct 7 22:52:29 2022 +0800

    Expose and document a simpler public API for simplify expressions (#3719)
    
    * Expose and document a simpler public API for simplify expressions
    
    * Fix minor comment
    
    * Fixed signature of schema
---
 datafusion/optimizer/src/simplify_expressions.rs | 42 +++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/datafusion/optimizer/src/simplify_expressions.rs b/datafusion/optimizer/src/simplify_expressions.rs
index 2901bfa8a..bb17a9925 100644
--- a/datafusion/optimizer/src/simplify_expressions.rs
+++ b/datafusion/optimizer/src/simplify_expressions.rs
@@ -950,12 +950,30 @@ macro_rules! assert_contains {
     };
 }
 
+/// Apply simplification and constant propagation to ([Expr]).
+///
+/// # Arguments
+///
+/// * `expr` - The logical expression
+/// * `schema` - The DataFusion schema for the expr, used to resolve `Column` references
+///                      to qualified or unqualified fields by name.
+/// * `props` - The Arrow schema for the input, used for determining expression data types
+///                    when performing type coercion.
+pub fn simplify_expr(
+    expr: Expr,
+    schema: &DFSchemaRef,
+    props: &ExecutionProps,
+) -> Result<Expr> {
+    let info = SimplifyContext::new(vec![schema], props);
+    expr.simplify(&info)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use arrow::array::{ArrayRef, Int32Array};
     use chrono::{DateTime, TimeZone, Utc};
-    use datafusion_common::DFField;
+    use datafusion_common::{DFField, ToDFSchema};
     use datafusion_expr::logical_plan::table_scan;
     use datafusion_expr::{
         and, binary_expr, call_fn, col, create_udf, lit, lit_timestamp_nano,
@@ -2553,4 +2571,26 @@ mod tests {
 
         assert_optimized_plan_eq(&plan, expected);
     }
+
+    #[test]
+    fn simplify_expr_api_test() {
+        let schema = Schema::new(vec![Field::new("x", DataType::Int32, false)])
+            .to_dfschema_ref()
+            .unwrap();
+        let props = ExecutionProps::new();
+
+        // x + (1 + 3) -> x + 4
+        {
+            let expr = col("x") + (lit(1) + lit(3));
+            let simplifed_expr = simplify_expr(expr, &schema, &props).unwrap();
+            assert_eq!(simplifed_expr, col("x") + lit(4));
+        }
+
+        // x * 1 -> x
+        {
+            let expr = col("x") * lit(1);
+            let simplifed_expr = simplify_expr(expr, &schema, &props).unwrap();
+            assert_eq!(simplifed_expr, col("x"));
+        }
+    }
 }