Merge pull request #9030 from andylokandy/fnctx

refactor(expr): rename FunctionContext

Author: Xuanwo

src/query/codegen/src/writes/register.rs

@@ -35,7 +35,7 @@ pub fn codegen_register() {
             use std::sync::Arc;
             
             use crate::Function;
-            use crate::FunctionContext;
+            use crate::EvalContext;
             use crate::FunctionDomain;
             use crate::FunctionRegistry;
             use crate::FunctionSignature;
@@ -82,7 +82,7 @@ pub fn codegen_register() {
                     func: G,
                 ) where
                     F: Fn({arg_f_closure_sig}) -> FunctionDomain<O> + 'static + Clone + Copy + Send + Sync,
-                    G: Fn({arg_g_closure_sig} FunctionContext) -> O::Scalar + 'static + Clone + Copy + Send + Sync,
+                    G: Fn({arg_g_closure_sig} EvalContext) -> O::Scalar + 'static + Clone + Copy + Send + Sync,
                 {{
                     self.register_passthrough_nullable_{n_args}_arg::<{arg_generics} O, _, _>(
                         name,
@@ -154,7 +154,7 @@ pub fn codegen_register() {
                     func: G,
                 ) where
                     F: Fn({arg_f_closure_sig}) -> FunctionDomain<O> + 'static + Clone + Copy + Send + Sync,
-                    G: for<'a> Fn({arg_g_closure_sig} FunctionContext) -> Result<Value<O>, String> + 'static + Clone + Copy + Send + Sync,
+                    G: for<'a> Fn({arg_g_closure_sig} EvalContext) -> Result<Value<O>, String> + 'static + Clone + Copy + Send + Sync,
                 {{
                     let has_nullable = &[{arg_sig_type} O::data_type()]
                         .iter()
@@ -257,7 +257,7 @@ pub fn codegen_register() {
                     func: G,
                 ) where
                     F: Fn({arg_f_closure_sig}) -> FunctionDomain<NullableType<O>> + 'static + Clone + Copy + Send + Sync,
-                    G: for<'a> Fn({arg_g_closure_sig} FunctionContext) -> Result<Value<NullableType<O>>, String> + 'static + Clone + Copy + Send + Sync,
+                    G: for<'a> Fn({arg_g_closure_sig} EvalContext) -> Result<Value<NullableType<O>>, String> + 'static + Clone + Copy + Send + Sync,
                 {{
                     let has_nullable = &[{arg_sig_type} O::data_type()]
                         .iter()
@@ -340,7 +340,7 @@ pub fn codegen_register() {
                     func: G,
                 ) where
                     F: Fn({arg_f_closure_sig}) -> FunctionDomain<O> + 'static + Clone + Copy + Send + Sync,
-                    G: for <'a> Fn({arg_g_closure_sig} FunctionContext) -> Result<Value<O>, String> + 'static + Clone + Copy + Send + Sync,
+                    G: for <'a> Fn({arg_g_closure_sig} EvalContext) -> Result<Value<O>, String> + 'static + Clone + Copy + Send + Sync,
                 {{
                     self.funcs
                         .entry(name.to_string())
@@ -441,8 +441,8 @@ pub fn codegen_register() {
             source,
             "
                 pub fn vectorize_{n_args}_arg<{arg_generics_bound} O: ArgType>(
-                    func: impl Fn({arg_input_closure_sig} FunctionContext) -> O::Scalar + Copy + Send + Sync,
-                ) -> impl Fn({arg_output_closure_sig} FunctionContext) -> Result<Value<O>, String> + Copy + Send + Sync {{
+                    func: impl Fn({arg_input_closure_sig} EvalContext) -> O::Scalar + Copy + Send + Sync,
+                ) -> impl Fn({arg_output_closure_sig} EvalContext) -> Result<Value<O>, String> + Copy + Send + Sync {{
                     move |{func_args} ctx| match ({args_tuple}) {{
                         ({arg_scalar}) => Ok(Value::Scalar(func({func_args} ctx))),
                         {match_arms}
@@ -535,8 +535,8 @@ pub fn codegen_register() {
             source,
             "
                 pub fn vectorize_with_builder_{n_args}_arg<{arg_generics_bound} O: ArgType>(
-                    func: impl Fn({arg_input_closure_sig} &mut O::ColumnBuilder, FunctionContext) -> Result<(), String> + Copy + Send + Sync,
-                ) -> impl Fn({arg_output_closure_sig} FunctionContext) -> Result<Value<O>, String> + Copy + Send + Sync {{
+                    func: impl Fn({arg_input_closure_sig} &mut O::ColumnBuilder, EvalContext) -> Result<(), String> + Copy + Send + Sync,
+                ) -> impl Fn({arg_output_closure_sig} EvalContext) -> Result<Value<O>, String> + Copy + Send + Sync {{
                     move |{func_args} ctx| match ({args_tuple}) {{
                         ({arg_scalar}) => {{
                             let mut builder = O::create_builder(1, ctx.generics);
@@ -640,8 +640,8 @@ pub fn codegen_register() {
             source,
             "
                 pub fn passthrough_nullable_{n_args}_arg<{arg_generics_bound} O: ArgType>(
-                    func: impl for <'a> Fn({arg_input_closure_sig} FunctionContext) -> Result<Value<O>, String> + Copy + Send + Sync,
-                ) -> impl for <'a> Fn({arg_output_closure_sig} FunctionContext) -> Result<Value<NullableType<O>>, String> + Copy + Send + Sync {{
+                    func: impl for <'a> Fn({arg_input_closure_sig} EvalContext) -> Result<Value<O>, String> + Copy + Send + Sync,
+                ) -> impl for <'a> Fn({arg_output_closure_sig} EvalContext) -> Result<Value<NullableType<O>>, String> + Copy + Send + Sync {{
                     move |{closure_args} ctx| match ({args_tuple}) {{
                         {scalar_nones_pats} => Ok(Value::Scalar(None)),
                         ({arg_scalar}) => Ok(Value::Scalar(Some(
@@ -747,8 +747,8 @@ pub fn codegen_register() {
             source,
             "
                 pub fn combine_nullable_{n_args}_arg<{arg_generics_bound} O: ArgType>(
-                    func: impl for <'a> Fn({arg_input_closure_sig} FunctionContext) -> Result<Value<NullableType<O>>, String> + Copy + Send + Sync,
-                ) -> impl for <'a> Fn({arg_output_closure_sig} FunctionContext) -> Result<Value<NullableType<O>>, String> + Copy + Send + Sync {{
+                    func: impl for <'a> Fn({arg_input_closure_sig} EvalContext) -> Result<Value<NullableType<O>>, String> + Copy + Send + Sync,
+                ) -> impl for <'a> Fn({arg_output_closure_sig} EvalContext) -> Result<Value<NullableType<O>>, String> + Copy + Send + Sync {{
                     move |{closure_args} ctx| match ({args_tuple}) {{
                         {scalar_nones_pats} => Ok(Value::Scalar(None)),
                         ({arg_scalar}) => Ok(Value::Scalar(
@@ -825,8 +825,8 @@ pub fn codegen_register() {
             source,
             "
                 fn erase_function_generic_{n_args}_arg<{arg_generics_bound} O: ArgType>(
-                    func: impl for <'a> Fn({arg_g_closure_sig} FunctionContext) -> Result<Value<O>, String>,
-                ) -> impl Fn(&[ValueRef<AnyType>], FunctionContext) -> Result<Value<AnyType>, String> {{
+                    func: impl for <'a> Fn({arg_g_closure_sig} EvalContext) -> Result<Value<O>, String>,
+                ) -> impl Fn(&[ValueRef<AnyType>], EvalContext) -> Result<Value<AnyType>, String> {{
                     move |args, ctx| {{
                         {let_args}
                         func({func_args} ctx).map(Value::upcast)

src/query/expression/src/evaluator.rs

@@ -16,14 +16,13 @@ use std::collections::HashMap;
 #[cfg(debug_assertions)]
 use std::sync::Mutex;
 
-use chrono_tz::Tz;
 use common_arrow::arrow::bitmap;
 use itertools::Itertools;
 
 use crate::chunk::Chunk;
 use crate::expression::Expr;
 use crate::expression::Span;
-use crate::function::FunctionContext;
+use crate::function::EvalContext;
 use crate::property::Domain;
 use crate::type_check::check_simple_cast;
 use crate::types::any::AnyType;
@@ -39,21 +38,26 @@ use crate::values::ColumnBuilder;
 use crate::values::Scalar;
 use crate::values::Value;
 use crate::ColumnIndex;
+use crate::FunctionContext;
 use crate::FunctionDomain;
 use crate::FunctionRegistry;
 use crate::Result;
 
 pub struct Evaluator<'a> {
     input_columns: &'a Chunk,
-    tz: Tz,
+    fn_ctx: FunctionContext,
     fn_registry: &'a FunctionRegistry,
 }
 
 impl<'a> Evaluator<'a> {
-    pub fn new(input_columns: &'a Chunk, tz: Tz, fn_registry: &'a FunctionRegistry) -> Self {
+    pub fn new(
+        input_columns: &'a Chunk,
+        fn_ctx: FunctionContext,
+        fn_registry: &'a FunctionRegistry,
+    ) -> Self {
         Evaluator {
             input_columns,
-            tz,
+            fn_ctx,
             fn_registry,
         }
     }
@@ -82,10 +86,10 @@ impl<'a> Evaluator<'a> {
                         .all_equal()
                 );
                 let cols_ref = cols.iter().map(Value::as_ref).collect::<Vec<_>>();
-                let ctx = FunctionContext {
+                let ctx = EvalContext {
                     generics,
                     num_rows: self.input_columns.num_rows(),
-                    tz: self.tz,
+                    tz: self.fn_ctx.tz,
                 };
                 (function.eval)(cols_ref.as_slice(), ctx).map_err(|msg| (span.clone(), msg))
             }
@@ -110,9 +114,13 @@ impl<'a> Evaluator<'a> {
             if !*RECURSING.lock().unwrap() {
                 *RECURSING.lock().unwrap() = true;
                 assert_eq!(
-                    ConstantFolder::new(self.input_columns.domains(), self.tz, self.fn_registry)
-                        .fold(expr)
-                        .1,
+                    ConstantFolder::new(
+                        self.input_columns.domains(),
+                        self.fn_ctx,
+                        self.fn_registry
+                    )
+                    .fold(expr)
+                    .1,
                     None,
                     "domain calculation should not return any domain for expressions that are possible to fail"
                 );
@@ -401,7 +409,7 @@ impl<'a> Evaluator<'a> {
             span,
             cast_fn,
             [(value, src_type.clone())],
-            self.tz,
+            self.fn_ctx,
             num_rows,
             self.fn_registry,
         )?;
@@ -412,19 +420,19 @@ impl<'a> Evaluator<'a> {
 
 pub struct ConstantFolder<'a, Index: ColumnIndex> {
     input_domains: HashMap<Index, Domain>,
-    tz: Tz,
+    fn_ctx: FunctionContext,
     fn_registry: &'a FunctionRegistry,
 }
 
 impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> {
     pub fn new(
         input_domains: HashMap<Index, Domain>,
-        tz: Tz,
+        fn_ctx: FunctionContext,
         fn_registry: &'a FunctionRegistry,
     ) -> Self {
         ConstantFolder {
             input_domains,
-            tz,
+            fn_ctx,
             fn_registry,
         }
     }
@@ -485,7 +493,7 @@ impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> {
 
                 if inner_expr.as_constant().is_some() {
                     let chunk = Chunk::empty();
-                    let evaluator = Evaluator::new(&chunk, self.tz, self.fn_registry);
+                    let evaluator = Evaluator::new(&chunk, self.fn_ctx, self.fn_registry);
                     // Since we know the expression is constant, it'll be safe to change its column index type.
                     let cast_expr = cast_expr.project_column_ref(|_| unreachable!());
                     if let Ok(Value::Scalar(scalar)) = evaluator.run(&cast_expr) {
@@ -561,7 +569,7 @@ impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> {
 
                 if all_args_is_scalar {
                     let chunk = Chunk::empty();
-                    let evaluator = Evaluator::new(&chunk, self.tz, self.fn_registry);
+                    let evaluator = Evaluator::new(&chunk, self.fn_ctx, self.fn_registry);
                     // Since we know the expression is constant, it'll be safe to change its column index type.
                     let func_expr = func_expr.project_column_ref(|_| unreachable!());
                     if let Ok(Value::Scalar(scalar)) = evaluator.run(&func_expr) {
@@ -753,7 +761,7 @@ impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> {
             span,
             cast_fn,
             [(domain.clone(), src_type.clone())],
-            self.tz,
+            self.fn_ctx,
             self.fn_registry,
         )?;
         assert_eq!(&ty, dest_type);

src/query/expression/src/function.rs

@@ -42,7 +42,12 @@ pub struct FunctionSignature {
 }
 
 #[derive(Clone, Copy)]
-pub struct FunctionContext<'a> {
+pub struct FunctionContext {
+    pub tz: Tz,
+}
+
+#[derive(Clone, Copy)]
+pub struct EvalContext<'a> {
     pub generics: &'a GenericMap,
     pub num_rows: usize,
     pub tz: Tz,
@@ -70,9 +75,7 @@ pub struct Function {
     pub calc_domain: Box<dyn Fn(&[Domain]) -> FunctionDomain<AnyType> + Send + Sync>,
     #[allow(clippy::type_complexity)]
     pub eval: Box<
-        dyn Fn(&[ValueRef<AnyType>], FunctionContext) -> Result<Value<AnyType>, String>
-            + Send
-            + Sync,
+        dyn Fn(&[ValueRef<AnyType>], EvalContext) -> Result<Value<AnyType>, String> + Send + Sync,
     >,
 }
 
@@ -204,8 +207,8 @@ impl FunctionRegistry {
 
 pub fn wrap_nullable<F>(
     f: F,
-) -> impl Fn(&[ValueRef<AnyType>], FunctionContext) -> Result<Value<AnyType>, String> + Copy
-where F: Fn(&[ValueRef<AnyType>], FunctionContext) -> Result<Value<AnyType>, String> + Copy {
+) -> impl Fn(&[ValueRef<AnyType>], EvalContext) -> Result<Value<AnyType>, String> + Copy
+where F: Fn(&[ValueRef<AnyType>], EvalContext) -> Result<Value<AnyType>, String> + Copy {
     move |args, ctx| {
         type T = NullableType<AnyType>;
         type Result = AnyType;