Remove input schema from PhysicalExpr, move the validation logic to physical expression planner (#6122)

mingmwang · web-flow · commit 569f6fe87ec0 · 2023-04-26T06:45:29.000-04:00
* Remove input schema from PhysicalExpr, move the validation logic to phyiscal expression planner

* fix fmt
diff --git a/datafusion/physical-expr/src/expressions/datetime.rs b/datafusion/physical-expr/src/expressions/datetime.rs
@@ -37,45 +37,16 @@ pub struct DateTimeIntervalExpr {
     lhs: Arc<dyn PhysicalExpr>,
     op: Operator,
     rhs: Arc<dyn PhysicalExpr>,
-    // TODO: move type checking to the planning phase and not in the physical expr
-    // so we can remove this
-    input_schema: Schema,
 }
 
 impl DateTimeIntervalExpr {
     /// Create a new instance of DateIntervalExpr
-    pub fn try_new(
+    pub fn new(
         lhs: Arc<dyn PhysicalExpr>,
         op: Operator,
         rhs: Arc<dyn PhysicalExpr>,
-        input_schema: &Schema,
-    ) -> Result<Self> {
-        match (
-            lhs.data_type(input_schema)?,
-            op,
-            rhs.data_type(input_schema)?,
-        ) {
-            (
-                DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _),
-                Operator::Plus | Operator::Minus,
-                DataType::Interval(_),
-            )
-            | (DataType::Timestamp(_, _), Operator::Minus, DataType::Timestamp(_, _))
-            | (DataType::Interval(_), Operator::Plus, DataType::Timestamp(_, _))
-            | (
-                DataType::Interval(_),
-                Operator::Plus | Operator::Minus,
-                DataType::Interval(_),
-            ) => Ok(Self {
-                lhs,
-                op,
-                rhs,
-                input_schema: input_schema.clone(),
-            }),
-            (lhs, _, rhs) => Err(DataFusionError::Execution(format!(
-                "Invalid operation {op} between '{lhs}' and '{rhs}' for DateIntervalExpr"
-            ))),
-        }
+    ) -> Self {
+        Self { lhs, op, rhs }
     }
 
     /// Get the left-hand side expression
@@ -202,12 +173,11 @@ impl PhysicalExpr for DateTimeIntervalExpr {
         self: Arc<Self>,
         children: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        Ok(Arc::new(DateTimeIntervalExpr::try_new(
+        Ok(Arc::new(DateTimeIntervalExpr::new(
             children[0].clone(),
             self.op,
             children[1].clone(),
-            &self.input_schema,
-        )?))
+        )))
     }
 }
 
@@ -220,6 +190,36 @@ impl PartialEq<dyn Any> for DateTimeIntervalExpr {
     }
 }
 
+/// create a DateIntervalExpr
+pub fn date_time_interval_expr(
+    lhs: Arc<dyn PhysicalExpr>,
+    op: Operator,
+    rhs: Arc<dyn PhysicalExpr>,
+    input_schema: &Schema,
+) -> Result<Arc<dyn PhysicalExpr>> {
+    match (
+        lhs.data_type(input_schema)?,
+        op,
+        rhs.data_type(input_schema)?,
+    ) {
+        (
+            DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _),
+            Operator::Plus | Operator::Minus,
+            DataType::Interval(_),
+        )
+        | (DataType::Timestamp(_, _), Operator::Minus, DataType::Timestamp(_, _))
+        | (DataType::Interval(_), Operator::Plus, DataType::Timestamp(_, _))
+        | (
+            DataType::Interval(_),
+            Operator::Plus | Operator::Minus,
+            DataType::Interval(_),
+        ) => Ok(Arc::new(DateTimeIntervalExpr::new(lhs, op, rhs))),
+        (lhs, _, rhs) => Err(DataFusionError::Execution(format!(
+            "Invalid operation {op} between '{lhs}' and '{rhs}' for DateIntervalExpr"
+        ))),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -535,7 +535,7 @@ mod tests {
         let lhs = create_physical_expr(&dt, &dfs, &schema, &props)?;
         let rhs = create_physical_expr(&interval, &dfs, &schema, &props)?;
 
-        let cut = DateTimeIntervalExpr::try_new(lhs, op, rhs, &schema)?;
+        let cut = date_time_interval_expr(lhs, op, rhs, &schema)?;
         let res = cut.evaluate(&batch)?;
 
         let mut builder = Date32Builder::with_capacity(8);
@@ -613,7 +613,7 @@ mod tests {
         let lhs_str = format!("{lhs}");
         let rhs_str = format!("{rhs}");
 
-        let cut = DateTimeIntervalExpr::try_new(lhs, op, rhs, &schema)?;
+        let cut = DateTimeIntervalExpr::new(lhs, op, rhs);
 
         assert_eq!(lhs_str, format!("{}", cut.lhs()));
         assert_eq!(op, cut.op().clone());
diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs
@@ -47,8 +47,7 @@ pub struct InListExpr {
     expr: Arc<dyn PhysicalExpr>,
     list: Vec<Arc<dyn PhysicalExpr>>,
     negated: bool,
-    static_filter: Option<Box<dyn Set>>,
-    input_schema: Schema,
+    static_filter: Option<Arc<dyn Set>>,
 }
 
 impl Debug for InListExpr {
@@ -62,7 +61,7 @@ impl Debug for InListExpr {
 }
 
 /// A type-erased container of array elements
-trait Set: Send + Sync {
+pub trait Set: Send + Sync {
     fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray>;
 }
 
@@ -172,36 +171,36 @@ where
 }
 
 /// Creates a `Box<dyn Set>` for the given list of `IN` expressions and `batch`
-fn make_set(array: &dyn Array) -> Result<Box<dyn Set>> {
+fn make_set(array: &dyn Array) -> Result<Arc<dyn Set>> {
     Ok(downcast_primitive_array! {
-        array => Box::new(ArraySet::new(array, make_hash_set(array))),
+        array => Arc::new(ArraySet::new(array, make_hash_set(array))),
         DataType::Boolean => {
             let array = as_boolean_array(array)?;
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         },
         DataType::Decimal128(_, _) => {
             let array = as_primitive_array::<Decimal128Type>(array)?;
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         }
         DataType::Decimal256(_, _) => {
             let array = as_primitive_array::<Decimal256Type>(array)?;
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         }
         DataType::Utf8 => {
             let array = as_string_array(array)?;
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         }
         DataType::LargeUtf8 => {
             let array = as_largestring_array(array);
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         }
         DataType::Binary => {
             let array = as_generic_binary_array::<i32>(array)?;
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         }
         DataType::LargeBinary => {
             let array = as_generic_binary_array::<i64>(array)?;
-            Box::new(ArraySet::new(array, make_hash_set(array)))
+            Arc::new(ArraySet::new(array, make_hash_set(array)))
         }
         DataType::Dictionary(_, _) => unreachable!("dictionary should have been flattened"),
         d => return Err(DataFusionError::NotImplemented(format!("DataType::{d} not supported in InList")))
@@ -233,7 +232,7 @@ fn evaluate_list(
 fn try_cast_static_filter_to_set(
     list: &[Arc<dyn PhysicalExpr>],
     schema: &Schema,
-) -> Result<Box<dyn Set>> {
+) -> Result<Arc<dyn Set>> {
     let batch = RecordBatch::new_empty(Arc::new(schema.clone()));
     make_set(evaluate_list(list, &batch)?.as_ref())
 }
@@ -244,15 +243,13 @@ impl InListExpr {
         expr: Arc<dyn PhysicalExpr>,
         list: Vec<Arc<dyn PhysicalExpr>>,
         negated: bool,
-        schema: &Schema,
+        static_filter: Option<Arc<dyn Set>>,
     ) -> Self {
-        let static_filter = try_cast_static_filter_to_set(&list, schema).ok();
         Self {
             expr,
             list,
             negated,
             static_filter,
-            input_schema: schema.clone(),
         }
     }
 
@@ -325,12 +322,13 @@ impl PhysicalExpr for InListExpr {
         self: Arc<Self>,
         children: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        in_list(
+        // assume the static_filter will not change during the rewrite process
+        Ok(Arc::new(InListExpr::new(
             children[0].clone(),
             children[1..].to_vec(),
-            &self.negated,
-            &self.input_schema,
-        )
+            self.negated,
+            self.static_filter.clone(),
+        )))
     }
 }
 
@@ -364,7 +362,13 @@ pub fn in_list(
             )));
         }
     }
-    Ok(Arc::new(InListExpr::new(expr, list, *negated, schema)))
+    let static_filter = try_cast_static_filter_to_set(&list, schema).ok();
+    Ok(Arc::new(InListExpr::new(
+        expr,
+        list,
+        *negated,
+        static_filter,
+    )))
 }
 
 #[cfg(test)]
diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs
@@ -78,7 +78,7 @@ pub use cast::{
     cast, cast_column, cast_with_options, CastExpr, DEFAULT_DATAFUSION_CAST_OPTIONS,
 };
 pub use column::{col, Column, UnKnownColumn};
-pub use datetime::DateTimeIntervalExpr;
+pub use datetime::{date_time_interval_expr, DateTimeIntervalExpr};
 pub use get_indexed_field::GetIndexedFieldExpr;
 pub use in_list::{in_list, InListExpr};
 pub use is_not_null::{is_not_null, IsNotNullExpr};
diff --git a/datafusion/physical-expr/src/intervals/test_utils.rs b/datafusion/physical-expr/src/intervals/test_utils.rs
@@ -19,7 +19,7 @@
 
 use std::sync::Arc;
 
-use crate::expressions::{BinaryExpr, DateTimeIntervalExpr, Literal};
+use crate::expressions::{date_time_interval_expr, BinaryExpr, Literal};
 use crate::PhysicalExpr;
 use arrow_schema::Schema;
 use datafusion_common::{DataFusionError, ScalarValue};
@@ -78,30 +78,22 @@ pub fn gen_conjunctive_temporal_expr(
     d: ScalarValue,
     schema: &Schema,
 ) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
-    let left_and_1 = Arc::new(DateTimeIntervalExpr::try_new(
+    let left_and_1 = date_time_interval_expr(
         left_col.clone(),
         op_1,
         Arc::new(Literal::new(a)),
         schema,
-    )?);
-    let left_and_2 = Arc::new(DateTimeIntervalExpr::try_new(
+    )?;
+    let left_and_2 = date_time_interval_expr(
         right_col.clone(),
         op_2,
         Arc::new(Literal::new(b)),
         schema,
-    )?);
-    let right_and_1 = Arc::new(DateTimeIntervalExpr::try_new(
-        left_col,
-        op_3,
-        Arc::new(Literal::new(c)),
-        schema,
-    )?);
-    let right_and_2 = Arc::new(DateTimeIntervalExpr::try_new(
-        right_col,
-        op_4,
-        Arc::new(Literal::new(d)),
-        schema,
-    )?);
+    )?;
+    let right_and_1 =
+        date_time_interval_expr(left_col, op_3, Arc::new(Literal::new(c)), schema)?;
+    let right_and_2 =
+        date_time_interval_expr(right_col, op_4, Arc::new(Literal::new(d)), schema)?;
     let left_expr = Arc::new(BinaryExpr::new(left_and_1, Operator::Gt, left_and_2));
     let right_expr = Arc::new(BinaryExpr::new(right_and_1, Operator::Lt, right_and_2));
     Ok(Arc::new(BinaryExpr::new(
diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs
@@ -19,7 +19,7 @@ use crate::var_provider::is_system_variables;
 use crate::{
     execution_props::ExecutionProps,
     expressions::{
-        self, binary, like, Column, DateTimeIntervalExpr, GetIndexedFieldExpr, Literal,
+        self, binary, date_time_interval_expr, like, Column, GetIndexedFieldExpr, Literal,
     },
     functions, udf,
     var_provider::VarType,
@@ -195,42 +195,22 @@ pub fn create_physical_expr(
                     DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _),
                     Operator::Plus | Operator::Minus,
                     DataType::Interval(_),
-                ) => Ok(Arc::new(DateTimeIntervalExpr::try_new(
-                    lhs,
-                    *op,
-                    rhs,
-                    input_schema,
-                )?)),
+                ) => Ok(date_time_interval_expr(lhs, *op, rhs, input_schema)?),
                 (
                     DataType::Interval(_),
                     Operator::Plus | Operator::Minus,
                     DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _),
-                ) => Ok(Arc::new(DateTimeIntervalExpr::try_new(
-                    rhs,
-                    *op,
-                    lhs,
-                    input_schema,
-                )?)),
+                ) => Ok(date_time_interval_expr(rhs, *op, lhs, input_schema)?),
                 (
                     DataType::Timestamp(_, _),
                     Operator::Minus,
                     DataType::Timestamp(_, _),
-                ) => Ok(Arc::new(DateTimeIntervalExpr::try_new(
-                    lhs,
-                    *op,
-                    rhs,
-                    input_schema,
-                )?)),
+                ) => Ok(date_time_interval_expr(lhs, *op, rhs, input_schema)?),
                 (
                     DataType::Interval(_),
                     Operator::Plus | Operator::Minus,
                     DataType::Interval(_),
-                ) => Ok(Arc::new(DateTimeIntervalExpr::try_new(
-                    lhs,
-                    *op,
-                    rhs,
-                    input_schema,
-                )?)),
+                ) => Ok(date_time_interval_expr(lhs, *op, rhs, input_schema)?),
                 _ => {
                     // Note that the logical planner is responsible
                     // for type coercion on the arguments (e.g. if one
diff --git a/datafusion/physical-expr/src/utils.rs b/datafusion/physical-expr/src/utils.rs
@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::equivalence::EquivalentClass;
-use crate::expressions::{BinaryExpr, Column, InListExpr, UnKnownColumn};
+use crate::expressions::{BinaryExpr, Column, UnKnownColumn};
 use crate::{
     EquivalenceProperties, PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement,
 };
@@ -586,28 +586,7 @@ pub fn reassign_predicate_columns(
                 column.name(),
                 index,
             ))));
-        } else if let Some(in_list) = expr_any.downcast_ref::<InListExpr>() {
-            // transform child first
-            let expr = reassign_predicate_columns(
-                in_list.expr().clone(),
-                schema,
-                ignore_not_found,
-            )?;
-            let list = in_list
-                .list()
-                .iter()
-                .map(|expr| {
-                    reassign_predicate_columns(expr.clone(), schema, ignore_not_found)
-                })
-                .collect::<Result<Vec<_>>>()?;
-            return Ok(Transformed::Yes(Arc::new(InListExpr::new(
-                expr,
-                list,
-                in_list.negated(),
-                schema.as_ref(),
-            ))));
         }
-
         Ok(Transformed::No(expr))
     })
 }
diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
