Introduce a helper ExpressionEvaluator to simplify expression evaluation

query/src/common/expression_evaluator.rs

@@ -0,0 +1,168 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_datablocks::DataBlock;
+use common_datavalues::ColumnRef;
+use common_datavalues::ColumnWithField;
+use common_datavalues::DataField;
+use common_datavalues::DataType;
+use common_datavalues::DataTypeImpl;
+use common_exception::ErrorCode;
+use common_exception::Result;
+use common_functions::scalars::CastFunction;
+use common_functions::scalars::FunctionContext;
+use common_functions::scalars::FunctionFactory;
+use common_planners::Expression;
+
+pub struct ExpressionEvaluator;
+
+impl ExpressionEvaluator {
+    pub fn eval(
+        func_ctx: FunctionContext,
+        expression: &Expression,
+        block: &DataBlock,
+    ) -> Result<ColumnRef> {
+        match expression {
+            Expression::Column(name) => block.try_column_by_name(name.as_str()).cloned(),
+
+            Expression::Alias(_, _) => Err(ErrorCode::UnImplement(
+                "Unsupported Alias scalar expression",
+            )),
+
+            Expression::QualifiedColumn(_) => Err(ErrorCode::LogicalError(
+                "Unsupported QualifiedColumn scalar expression",
+            )),
+
+            Expression::Literal {
+                value, data_type, ..
+            } => data_type.create_constant_column(value, block.num_rows()),
+
+            Expression::UnaryExpression { op, expr } => {
+                let result = Self::eval(func_ctx.clone(), expr, block)?;
+                let arg_types: Vec<DataTypeImpl> = vec![result.data_type()];
+                let arg_types2: Vec<&DataTypeImpl> = arg_types.iter().collect();
+                let func = FunctionFactory::instance().get(op, &arg_types2)?;
+
+                let columns = [ColumnWithField::new(
+                    result.clone(),
+                    DataField::new("", result.data_type()),
+                )];
+                func.eval(func_ctx, &columns, block.num_rows())
+            }
+
+            Expression::BinaryExpression { left, op, right } => {
+                let left_result = Self::eval(func_ctx.clone(), left.as_ref(), block)?;
+                let right_result = Self::eval(func_ctx.clone(), right.as_ref(), block)?;
+                let arg_types: Vec<DataTypeImpl> =
+                    vec![left_result.data_type(), right_result.data_type()];
+                let arg_types2: Vec<&DataTypeImpl> = arg_types.iter().collect();
+                let func = FunctionFactory::instance().get(op, &arg_types2)?;
+
+                let columns = [
+                    ColumnWithField::new(
+                        left_result.clone(),
+                        DataField::new("", left_result.data_type()),
+                    ),
+                    ColumnWithField::new(
+                        right_result.clone(),
+                        DataField::new("", right_result.data_type()),
+                    ),
+                ];
+                func.eval(func_ctx, &columns, block.num_rows())
+            }
+
+            Expression::ScalarFunction { op, args } => {
+                let results = args
+                    .iter()
+                    .map(|expr| Self::eval(func_ctx.clone(), expr, block))
+                    .collect::<Result<Vec<ColumnRef>>>()?;
+                let arg_types: Vec<DataTypeImpl> =
+                    results.iter().map(|col| col.data_type()).collect();
+                let arg_types2: Vec<&DataTypeImpl> = arg_types.iter().collect();
+
+                let func = FunctionFactory::instance().get(op, &arg_types2)?;
+                let columns: Vec<ColumnWithField> = results
+                    .into_iter()
+                    .map(|col| {
+                        ColumnWithField::new(col.clone(), DataField::new("", col.data_type()))
+                    })
+                    .collect();
+                func.eval(func_ctx, &columns, block.num_rows())
+            }
+
+            Expression::AggregateFunction { .. } => Err(ErrorCode::LogicalError(
+                "Unsupported AggregateFunction scalar expression",
+            )),
+
+            Expression::Sort { .. } => Err(ErrorCode::LogicalError(
+                "Unsupported Sort scalar expression",
+            )),
+
+            Expression::Cast {
+                expr, data_type, ..
+            } => {
+                let result = Self::eval(func_ctx.clone(), expr, block)?;
+                let func_name = "cast".to_string();
+                let type_name = data_type.name();
+
+                let func = if data_type.is_nullable() {
+                    CastFunction::create_try(&func_name, &type_name)
+                } else {
+                    CastFunction::create(&func_name, &type_name)
+                }?;
+
+                let columns = [ColumnWithField::new(
+                    result.clone(),
+                    DataField::new("", result.data_type()),
+                )];
+
+                func.eval(func_ctx, &columns, block.num_rows())
+            }
+
+            Expression::ScalarSubquery { name, .. } => {
+                // Scalar subquery results are ready in the expression input
+                let expr = Expression::Column(name.clone());
+                Self::eval(func_ctx, &expr, block)
+            }
+
+            Expression::Subquery { name, .. } => {
+                // Subquery results are ready in the expression input
+                let expr = Expression::Column(name.clone());
+                Self::eval(func_ctx, &expr, block)
+            }
+
+            Expression::MapAccess { args, .. } => {
+                let results = args
+                    .iter()
+                    .map(|expr| Self::eval(func_ctx.clone(), expr, block))
+                    .collect::<Result<Vec<ColumnRef>>>()?;
+                let arg_types: Vec<DataTypeImpl> =
+                    results.iter().map(|col| col.data_type()).collect();
+                let arg_types2: Vec<&DataTypeImpl> = arg_types.iter().collect();
+
+                let func_name = "get_path";
+                let func = FunctionFactory::instance().get(func_name, &arg_types2)?;
+                let columns: Vec<ColumnWithField> = results
+                    .into_iter()
+                    .map(|col| {
+                        ColumnWithField::new(col.clone(), DataField::new("", col.data_type()))
+                    })
+                    .collect();
+                func.eval(func_ctx, &columns, block.num_rows())
+            }
+
+            Expression::Wildcard => Err(ErrorCode::LogicalError("Unsupported wildcard expression")),
+        }
+    }
+}

query/src/common/mod.rs

@@ -13,9 +13,11 @@
 // limitations under the License.
 
 pub mod context_function;
+mod expression_evaluator;
 mod hashtable;
 mod meta;
 pub mod service;
 
+pub use expression_evaluator::ExpressionEvaluator;
 pub use hashtable::*;
 pub use meta::MetaClientProvider;

query/src/sessions/query_ctx.rs

@@ -29,6 +29,7 @@ use common_contexts::DalMetrics;
 use common_datablocks::DataBlock;
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_functions::scalars::FunctionContext;
 use common_infallible::Mutex;
 use common_infallible::RwLock;
 use common_io::prelude::FormatSettings;
@@ -394,6 +395,14 @@ impl QueryContext {
         blocks.clear();
         result
     }
+
+    pub fn try_get_function_context(&self) -> Result<FunctionContext> {
+        Ok(FunctionContext {
+            tz: String::from_utf8(self.get_settings().get_timezone()?).map_err(|_| {
+                ErrorCode::LogicalError("Timezone has been checked and should be valid.")
+            })?,
+        })
+    }
 }
 
 impl TrySpawn for QueryContext {

query/tests/it/common/expression_evaluator.rs

@@ -0,0 +1,91 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_base::tokio;
+use common_datablocks::DataBlock;
+use common_datavalues::Column;
+use common_datavalues::DataField;
+use common_datavalues::DataValue;
+use common_datavalues::Int32Type;
+use common_datavalues::Int64Type;
+use common_datavalues::PrimitiveColumn;
+use common_datavalues::StringType;
+use common_exception::Result;
+use common_planners::Expression;
+use databend_query::common::ExpressionEvaluator;
+
+use crate::tests::create_query_context;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_expression_evaluator() -> Result<()> {
+    // A expression which contains Unary, Binary, ScalarFunction, Column, Literal and Cast:
+    //
+    //  CAST(-POW(a + b, 2) AS STRING)
+    let expr = Expression::Cast {
+        expr: Box::new(Expression::UnaryExpression {
+            op: "negate".to_string(),
+            expr: Box::new(Expression::ScalarFunction {
+                op: "pow".to_string(),
+                args: vec![
+                    Expression::BinaryExpression {
+                        left: Box::new(Expression::Column("a".to_string())),
+                        op: "+".to_string(),
+                        right: Box::new(Expression::Column("b".to_string())),
+                    },
+                    Expression::Literal {
+                        value: DataValue::Int64(2),
+                        column_name: None,
+                        data_type: Int64Type::arc(),
+                    },
+                ],
+            }),
+        }),
+        data_type: StringType::arc(),
+        pg_style: false,
+    };
+
+    // Block layout:
+    //
+    //  a  b
+    // ------
+    //  1  1
+    //  2  2
+    //  3  3
+    let mut block = DataBlock::empty();
+    block = block.add_column(
+        PrimitiveColumn::<i32>::new_from_vec(vec![1, 2, 3]).arc(),
+        DataField::new("a", Int32Type::arc()),
+    )?;
+    block = block.add_column(
+        PrimitiveColumn::<i32>::new_from_vec(vec![1, 2, 3]).arc(),
+        DataField::new("b", Int32Type::arc()),
+    )?;
+
+    let func_ctx = create_query_context().await?.try_get_function_context()?;
+    let result = ExpressionEvaluator::eval(func_ctx, &expr, &block)?;
+
+    assert_eq!(
+        result.get(0),
+        DataValue::String("-4.0".to_string().into_bytes())
+    );
+    assert_eq!(
+        result.get(1),
+        DataValue::String("-16.0".to_string().into_bytes())
+    );
+    assert_eq!(
+        result.get(2),
+        DataValue::String("-36.0".to_string().into_bytes())
+    );
+    Ok(())
+}

query/tests/it/common/mod.rs

@@ -1,4 +1,4 @@
-// Copyright 2021 Datafuse Labs.
+// Copyright 2022 Datafuse Labs.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,4 +12,5 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod expression_evaluator;
 mod hashtable;