Merge pull request #1088 from zhang2014/constant_folding

Constant folding

common/functions/src/scalars/function.rs

@@ -26,4 +26,9 @@ pub trait Function: fmt::Display + Sync + Send + DynClone {
     fn return_type(&self, args: &[DataType]) -> Result<DataType>;
     fn nullable(&self, _input_schema: &DataSchema) -> Result<bool>;
     fn eval(&self, columns: &[DataColumn], _input_rows: usize) -> Result<DataColumn>;
+
+    // If function returns the same result when same arguments, it is deterministic function.
+    fn is_deterministic(&self) -> bool {
+        true
+    }
 }

common/functions/src/scalars/udfs/crash_me.rs

@@ -44,6 +44,10 @@ impl Function for CrashMeFunction {
     fn eval(&self, _columns: &[DataColumn], _input_rows: usize) -> Result<DataColumn> {
         panic!("crash me function");
     }
+
+    fn is_deterministic(&self) -> bool {
+        false
+    }
 }
 
 impl fmt::Display for CrashMeFunction {

common/functions/src/scalars/udfs/database.rs

@@ -41,6 +41,10 @@ impl Function for DatabaseFunction {
     fn num_arguments(&self) -> usize {
         1
     }
+
+    fn is_deterministic(&self) -> bool {
+        false
+    }
 }
 
 impl fmt::Display for DatabaseFunction {

common/functions/src/scalars/udfs/exists.rs

@@ -64,6 +64,10 @@ impl Function for ExistsFunction {
     fn num_arguments(&self) -> usize {
         1
     }
+
+    fn is_deterministic(&self) -> bool {
+        false
+    }
 }
 
 impl fmt::Display for ExistsFunction {

common/functions/src/scalars/udfs/sleep.rs

@@ -91,6 +91,10 @@ impl Function for SleepFunction {
             }
         }
     }
+
+    fn is_deterministic(&self) -> bool {
+        false
+    }
 }
 
 impl fmt::Display for SleepFunction {

common/functions/src/scalars/udfs/version.rs

@@ -44,6 +44,10 @@ impl Function for VersionFunction {
     fn num_arguments(&self) -> usize {
         1
     }
+
+    fn is_deterministic(&self) -> bool {
+        false
+    }
 }
 
 impl fmt::Display for VersionFunction {

common/planners/src/plan_expression.rs

@@ -48,7 +48,11 @@ pub enum Expression {
     /// Column name.
     Column(String),
     /// Constant value.
-    Literal(DataValue),
+    /// Note: When literal represents a column, its column_name will not be None
+    Literal {
+        value: DataValue,
+        column_name: Option<String>,
+    },
     /// A unary expression such as "NOT foo"
     UnaryExpression { op: String, expr: Box<Expression> },
 
@@ -101,15 +105,50 @@ pub enum Expression {
 }
 
 impl Expression {
+    pub fn create_literal(value: DataValue) -> Expression {
+        Expression::Literal {
+            value,
+            column_name: None,
+        }
+    }
+
     pub fn column_name(&self) -> String {
         match self {
             Expression::Alias(name, _expr) => name.clone(),
-            Expression::ScalarFunction { op, .. } => {
+            Expression::Column(name) => name.clone(),
+            Expression::Literal {
+                column_name: Some(name),
+                ..
+            } => name.clone(),
+            Expression::UnaryExpression { op, expr } => {
+                format!("({} {})", op, expr.column_name())
+            }
+            Expression::BinaryExpression { op, left, right } => {
+                format!("({} {} {})", left.column_name(), op, right.column_name())
+            }
+            Expression::ScalarFunction { op, args } => {
                 match OP_SET.get(&op.to_lowercase().as_ref()) {
                     Some(_) => format!("{}()", op),
-                    None => format!("{:?}", self),
+                    None => {
+                        let args_column_name =
+                            args.iter().map(Expression::column_name).collect::<Vec<_>>();
+
+                        format!("{}({})", op, args_column_name.join(", "))
+                    }
+                }
+            }
+            Expression::AggregateFunction { op, distinct, args } => {
+                let args_column_name = args.iter().map(Expression::column_name).collect::<Vec<_>>();
+
+                match distinct {
+                    true => format!("{}(distinct {})", op, args_column_name.join(", ")),
+                    false => format!("{}({})", op, args_column_name.join(", ")),
                 }
             }
+            Expression::Sort { expr, .. } => expr.column_name(),
+            Expression::Cast { expr, data_type } => {
+                format!("cast({} as {:?})", expr.column_name(), data_type)
+            }
             Expression::Subquery { name, .. } => name.clone(),
             Expression::ScalarSubquery { name, .. } => name.clone(),
             _ => format!("{:?}", self),
@@ -164,7 +203,7 @@ impl Expression {
         match self {
             Expression::Alias(_, expr) => expr.to_data_type(input_schema),
             Expression::Column(s) => Ok(input_schema.field_with_name(s)?.data_type().clone()),
-            Expression::Literal(v) => Ok(v.data_type()),
+            Expression::Literal { value, .. } => Ok(value.data_type()),
             Expression::Subquery { query_plan, .. } => Ok(Self::to_subquery_type(query_plan)),
             Expression::ScalarSubquery { query_plan, .. } => {
                 Ok(Self::to_scalar_subquery_type(query_plan))
@@ -246,7 +285,7 @@ impl fmt::Debug for Expression {
         match self {
             Expression::Alias(alias, v) => write!(f, "{:?} as {:#}", v, alias),
             Expression::Column(ref v) => write!(f, "{:#}", v),
-            Expression::Literal(ref v) => write!(f, "{:#}", v),
+            Expression::Literal { ref value, .. } => write!(f, "{:#}", value),
             Expression::Subquery { name, .. } => write!(f, "subquery({})", name),
             Expression::ScalarSubquery { name, .. } => write!(f, "scalar subquery({})", name),
             Expression::BinaryExpression { op, left, right } => {

common/planners/src/plan_expression_chain.rs

@@ -56,10 +56,10 @@ impl ExpressionChain {
                 };
                 self.actions.push(ExpressionAction::Input(input));
             }
-            Expression::Literal(l) => {
+            Expression::Literal { value, .. } => {
                 let value = ActionConstant {
                     name: expr.column_name(),
-                    value: l.clone(),
+                    value: value.clone(),
                 };
 
                 self.actions.push(ExpressionAction::Constant(value));

common/planners/src/plan_expression_common.rs

@@ -295,7 +295,7 @@ where F: Fn(&Expression) -> Result<Option<Expression>> {
             }),
 
             Expression::Column(_)
-            | Expression::Literal(_)
+            | Expression::Literal { .. }
             | Expression::Subquery { .. }
             | Expression::ScalarSubquery { .. } => Ok(expr.clone()),
         },

common/planners/src/plan_expression_literal.rs

@@ -12,13 +12,13 @@ pub trait Literal {
 
 impl Literal for &str {
     fn to_literal(&self) -> Expression {
-        Expression::Literal(DataValue::Utf8(Some(self.to_string())))
+        Expression::create_literal(DataValue::Utf8(Some(self.to_string())))
     }
 }
 
 impl Literal for String {
     fn to_literal(&self) -> Expression {
-        Expression::Literal(DataValue::Utf8(Some(self.clone())))
+        Expression::create_literal(DataValue::Utf8(Some(self.clone())))
     }
 }
 
@@ -27,7 +27,7 @@ macro_rules! make_literal {
         #[allow(missing_docs)]
         impl Literal for $TYPE {
             fn to_literal(&self) -> Expression {
-                Expression::Literal(DataValue::$SCALAR(Some(self.clone())))
+                Expression::create_literal(DataValue::$SCALAR(Some(self.clone())))
             }
         }
     };

common/planners/src/plan_rewriter.rs

@@ -141,7 +141,10 @@ pub trait PlanRewriter {
             }),
             Expression::Wildcard => Ok(Expression::Wildcard),
             Expression::Column(column_name) => Ok(Expression::Column(column_name.clone())),
-            Expression::Literal(value) => Ok(Expression::Literal(value.clone())),
+            Expression::Literal { value, column_name } => Ok(Expression::Literal {
+                value: value.clone(),
+                column_name: column_name.clone(),
+            }),
             Expression::Subquery { name, query_plan } => {
                 let new_subquery = self.rewrite_subquery_plan(query_plan)?;
                 Ok(Expression::Subquery {
@@ -480,7 +483,7 @@ impl RewriteHelper {
                 })
             }
             Expression::Wildcard
-            | Expression::Literal(_)
+            | Expression::Literal { .. }
             | Expression::Subquery { .. }
             | Expression::ScalarSubquery { .. }
             | Expression::Sort { .. } => Ok(expr.clone()),
@@ -533,7 +536,7 @@ impl RewriteHelper {
         Ok(match expr {
             Expression::Alias(_, expr) => vec![expr.as_ref().clone()],
             Expression::Column(_) => vec![],
-            Expression::Literal(_) => vec![],
+            Expression::Literal { .. } => vec![],
             Expression::Subquery { .. } => vec![],
             Expression::ScalarSubquery { .. } => vec![],
             Expression::UnaryExpression { expr, .. } => {
@@ -555,7 +558,7 @@ impl RewriteHelper {
         Ok(match expr {
             Expression::Alias(_, expr) => Self::expression_plan_columns(expr)?,
             Expression::Column(_) => vec![expr.clone()],
-            Expression::Literal(_) => vec![],
+            Expression::Literal { .. } => vec![],
             Expression::Subquery { .. } => vec![],
             Expression::ScalarSubquery { .. } => vec![],
             Expression::UnaryExpression { expr, .. } => Self::expression_plan_columns(expr)?,
@@ -628,7 +631,7 @@ impl RewriteHelper {
                 Expression::Alias(alias.clone(), Box::from(expressions[0].clone()))
             }
             Expression::Column(_) => expr.clone(),
-            Expression::Literal(_) => expr.clone(),
+            Expression::Literal { .. } => expr.clone(),
             Expression::BinaryExpression { op, .. } => Expression::BinaryExpression {
                 left: Box::new(expressions[0].clone()),
                 op: op.clone(),

fusequery/query/src/api/rpc/flight_actions_test.rs

@@ -21,7 +21,7 @@ async fn test_shuffle_action_try_into() -> Result<()> {
         stage_id: String::from("stage_id"),
         plan: parse_query("SELECT number FROM numbers(5)")?,
         sinks: vec![String::from("stream_id")],
-        scatters_expression: Expression::Literal(DataValue::UInt64(Some(1))),
+        scatters_expression: Expression::create_literal(DataValue::UInt64(Some(1))),
     };
 
     let from_action = FlightAction::PrepareShuffleAction(shuffle_action);
@@ -36,7 +36,7 @@ async fn test_shuffle_action_try_into() -> Result<()> {
             assert_eq!(action.sinks, vec![String::from("stream_id")]);
             assert_eq!(
                 action.scatters_expression,
-                Expression::Literal(DataValue::UInt64(Some(1)))
+                Expression::create_literal(DataValue::UInt64(Some(1)))
             );
         }
     }

fusequery/query/src/api/rpc/flight_dispatcher_test.rs

@@ -51,7 +51,7 @@ async fn test_run_shuffle_action_with_no_scatters() -> Result<()> {
                 stage_id: stage_id.clone(),
                 plan: parse_query("SELECT number FROM numbers(5)")?,
                 sinks: vec![stream_id.clone()],
-                scatters_expression: Expression::Literal(DataValue::UInt64(Some(1))),
+                scatters_expression: Expression::create_literal(DataValue::UInt64(Some(1))),
             }),
         )?;
 

fusequery/query/src/api/rpc/flight_scatter_hash.rs

@@ -84,7 +84,7 @@ impl HashFlightScatter {
                     expr: Box::new(expr),
                     data_type: DataType::UInt64,
                 },
-                Expression::Literal(DataValue::UInt64(Some(num as u64))),
+                Expression::create_literal(DataValue::UInt64(Some(num as u64))),
             ],
         }
     }

fusequery/query/src/api/rpc/flight_service_test.rs

@@ -159,7 +159,7 @@ fn do_action_request(query_id: &str, stage_id: &str) -> Result<Request<Action>>
         stage_id: String::from(stage_id),
         plan: parse_query("SELECT number FROM numbers(5)")?,
         sinks: vec![String::from("stream_id")],
-        scatters_expression: Expression::Literal(DataValue::UInt64(Some(1))),
+        scatters_expression: Expression::create_literal(DataValue::UInt64(Some(1))),
     });
 
     Ok(Request::new(flight_action.try_into()?))

fusequery/query/src/datasources/system/numbers_table.rs

@@ -78,8 +78,8 @@ impl Table for NumbersTable {
     ) -> Result<ReadDataSourcePlan> {
         let mut total = None;
         let ScanPlan { table_args, .. } = scan.clone();
-        if let Some(Expression::Literal(v)) = table_args {
-            total = Some(v.as_u64()?);
+        if let Some(Expression::Literal { value, .. }) = table_args {
+            total = Some(value.as_u64()?);
         }
 
         let total = total.ok_or_else(|| {

fusequery/query/src/datasources/system/numbers_table_test.rs

@@ -19,7 +19,7 @@ async fn test_number_table() -> Result<()> {
     let scan = &ScanPlan {
         schema_name: "scan_test".to_string(),
         table_schema: DataSchemaRefExt::create(vec![]),
-        table_args: Some(Expression::Literal(DataValue::UInt64(Some(8)))),
+        table_args: Some(Expression::create_literal(DataValue::UInt64(Some(8)))),
         projected_schema: DataSchemaRefExt::create(vec![DataField::new(
             "number",
             DataType::UInt64,

fusequery/query/src/functions/context_function.rs

@@ -26,10 +26,10 @@ impl ContextFunction {
         }
 
         Ok(match name.to_lowercase().as_str() {
-            "database" => vec![Expression::Literal(DataValue::Utf8(Some(
+            "database" => vec![Expression::create_literal(DataValue::Utf8(Some(
                 ctx.get_current_database(),
             )))],
-            "version" => vec![Expression::Literal(DataValue::Utf8(Some(
+            "version" => vec![Expression::create_literal(DataValue::Utf8(Some(
                 ctx.get_fuse_version(),
             )))],
             _ => vec![],