Merge pull request #10683 from leiysky/project-set

refactor(planner): Refactor set returning function

Author: BohuTANG

src/query/expression/src/type_check.rs

@@ -211,7 +211,7 @@ pub fn check_cast<Index: ColumnIndex>(
     }
 }
 
-fn wrap_nullable_for_try_cast(span: Span, ty: &DataType) -> Result<DataType> {
+pub fn wrap_nullable_for_try_cast(span: Span, ty: &DataType) -> Result<DataType> {
     match ty {
         DataType::Null => Err(ErrorCode::from_string_no_backtrace(
             "TRY_CAST() to NULL is not supported".to_string(),

src/query/functions/src/lib.rs

@@ -22,6 +22,7 @@ use scalars::BUILTIN_FUNCTIONS;
 
 pub mod aggregates;
 pub mod scalars;
+pub mod srfs;
 
 pub fn is_builtin_function(name: &str) -> bool {
     BUILTIN_FUNCTIONS.contains(name) || AggregateFunctionFactory::instance().contains(name)

src/query/functions/src/srfs/check.rs

@@ -0,0 +1,75 @@
+// 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_exception::ErrorCode;
+use common_exception::Result;
+use common_expression::type_check::check_cast;
+use common_expression::type_check::try_unify_signature;
+use common_expression::types::DataType;
+use common_expression::ColumnIndex;
+use common_expression::Expr;
+
+use crate::scalars::BUILTIN_FUNCTIONS;
+
+#[allow(clippy::type_complexity)]
+pub fn try_check_srf<Index: ColumnIndex>(
+    expected_arg_types: &[DataType],
+    expected_return_types: &[DataType],
+    args: &[Expr<Index>],
+) -> Result<(Vec<Expr<Index>>, Vec<DataType>, Vec<DataType>)> {
+    let subst = try_unify_signature(
+        args.iter().map(Expr::data_type),
+        expected_arg_types.iter(),
+        &[],
+    )?;
+
+    let checked_args = args
+        .iter()
+        .zip(expected_arg_types.iter())
+        .map(|(arg, sig_type)| {
+            let sig_type = subst.apply(sig_type)?;
+            let is_try = BUILTIN_FUNCTIONS.is_auto_try_cast_rule(arg.data_type(), &sig_type);
+            check_cast(
+                arg.span(),
+                is_try,
+                arg.clone(),
+                &sig_type,
+                &BUILTIN_FUNCTIONS,
+            )
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let generics = subst
+        .0
+        .keys()
+        .cloned()
+        .max()
+        .map(|max_generic_idx| {
+            (0..max_generic_idx + 1)
+                .map(|idx| match subst.0.get(&idx) {
+                    Some(ty) => Ok(ty.clone()),
+                    None => Err(ErrorCode::from_string_no_backtrace(format!(
+                        "unable to resolve generic T{idx}"
+                    ))),
+                })
+                .collect::<Result<Vec<_>>>()
+        })
+        .unwrap_or_else(|| Ok(vec![]))?;
+
+    // TODO: we only support one return type for now, so we can reuse
+    // the current implementation of `Substitution` to get the return type.
+    let return_type = subst.apply(&expected_return_types[0])?;
+
+    Ok((checked_args, vec![return_type], generics))
+}

src/query/functions/src/srfs/evaluator.rs

@@ -0,0 +1,61 @@
+// 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_exception::Result;
+use common_expression::types::AnyType;
+use common_expression::DataBlock;
+use common_expression::Evaluator;
+use common_expression::FunctionContext;
+use common_expression::FunctionRegistry;
+use common_expression::Value;
+
+use crate::srfs::SrfExpr;
+
+pub struct SrfEvaluator<'a> {
+    input_columns: &'a DataBlock,
+    func_ctx: FunctionContext,
+    fn_registry: &'a FunctionRegistry,
+}
+
+impl<'a> SrfEvaluator<'a> {
+    pub fn new(
+        input_columns: &'a DataBlock,
+        func_ctx: FunctionContext,
+        fn_registry: &'a FunctionRegistry,
+    ) -> Self {
+        SrfEvaluator {
+            input_columns,
+            func_ctx,
+            fn_registry,
+        }
+    }
+
+    /// Evaluate an SRF. This will essentially returns a `DataBlock` for each row in the input `DataBlock`.
+    pub fn run(&self, srf_expr: &SrfExpr) -> Result<Vec<(Vec<Value<AnyType>>, usize)>> {
+        let evaluator = Evaluator::new(self.input_columns, self.func_ctx, self.fn_registry);
+        let arg_results = srf_expr
+            .args
+            .iter()
+            .map(|arg| evaluator.run(arg))
+            .collect::<Result<Vec<_>>>()?;
+        let arg_result_refs = arg_results
+            .iter()
+            .map(|arg| arg.as_ref())
+            .collect::<Vec<_>>();
+
+        let results = (*srf_expr.srf.eval)(&arg_result_refs, self.input_columns.num_rows());
+
+        Ok(results)
+    }
+}

src/query/functions/src/srfs/mod.rs

@@ -0,0 +1,232 @@
+// 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.
+
+mod check;
+mod evaluator;
+mod register;
+
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+use common_exception::ErrorCode;
+use common_exception::Result;
+use common_expression::types::AnyType;
+use common_expression::types::DataType;
+use common_expression::ColumnIndex;
+use common_expression::Expr;
+use common_expression::RemoteExpr;
+use common_expression::Value;
+use common_expression::ValueRef;
+use ctor::ctor;
+pub use evaluator::SrfEvaluator;
+use serde::Deserialize;
+use serde::Serialize;
+
+use crate::scalars::BUILTIN_FUNCTIONS;
+use crate::srfs::check::try_check_srf;
+use crate::srfs::register::builtin_set_returning_functions;
+
+#[ctor]
+pub static BUILTIN_SET_RETURNING_FUNCTIONS: SetReturningFunctionRegistry =
+    builtin_set_returning_functions();
+
+pub type SetReturningFunctionID = (String, usize);
+
+#[derive(Clone)]
+pub struct SrfExpr<Index: ColumnIndex = usize> {
+    pub id: SetReturningFunctionID,
+    pub srf: Arc<SetReturningFunction>,
+    pub args: Vec<Expr<Index>>,
+    pub generic_types: Vec<DataType>,
+    pub return_types: Vec<DataType>,
+}
+
+impl SrfExpr {
+    pub fn sql_display(&self) -> String {
+        let mut args = vec![];
+        for arg in &self.args {
+            args.push(arg.sql_display());
+        }
+        format!("{}({})", self.id.0, args.join(", "))
+    }
+}
+
+impl Debug for SrfExpr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SrfExpr")
+            .field("id", &self.id)
+            .field("args", &self.args)
+            .field("generic_types", &self.generic_types)
+            .field("return_types", &self.return_types)
+            .finish()
+    }
+}
+
+impl SrfExpr {
+    pub fn into_remote_srf_expr(self) -> RemoteSrfExpr {
+        let mut args = vec![];
+        for arg in self.args {
+            args.push(arg.as_remote_expr());
+        }
+        RemoteSrfExpr {
+            id: self.id,
+            args,
+            generic_types: self.generic_types,
+            return_types: self.return_types,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct RemoteSrfExpr<Index: ColumnIndex = usize> {
+    pub id: SetReturningFunctionID,
+    pub args: Vec<RemoteExpr<Index>>,
+    pub generic_types: Vec<DataType>,
+    pub return_types: Vec<DataType>,
+}
+
+impl RemoteSrfExpr {
+    pub fn into_srf_expr(self) -> SrfExpr {
+        let mut args = vec![];
+        for arg in self.args {
+            args.push(arg.as_expr(&BUILTIN_FUNCTIONS));
+        }
+        SrfExpr {
+            id: self.id.clone(),
+            srf: BUILTIN_SET_RETURNING_FUNCTIONS.get(&self.id).unwrap(),
+            args,
+            generic_types: self.generic_types,
+            return_types: self.return_types,
+        }
+    }
+}
+
+#[allow(clippy::type_complexity)]
+pub struct SetReturningFunction {
+    pub name: String,
+    pub arg_types: Vec<DataType>,
+    pub return_types: Vec<DataType>,
+
+    pub eval:
+        Box<dyn Fn(&[ValueRef<AnyType>], usize) -> Vec<(Vec<Value<AnyType>>, usize)> + Send + Sync>,
+}
+
+#[derive(Default)]
+pub struct SetReturningFunctionRegistry {
+    pub srfs: HashMap<String, Vec<Arc<SetReturningFunction>>>,
+}
+
+impl SetReturningFunctionRegistry {
+    pub fn register<F>(
+        &mut self,
+        name: &str,
+        arg_types: &[DataType],
+        return_types: &[DataType],
+        eval: F,
+    ) where
+        F: Fn(&[ValueRef<AnyType>], usize) -> Vec<(Vec<Value<AnyType>>, usize)>
+            + Send
+            + Sync
+            + 'static,
+    {
+        if return_types.iter().any(|t| !t.is_nullable()) {
+            panic!("return type of srf must be nullable");
+        }
+
+        let name = name.to_string();
+        let srf = Arc::new(SetReturningFunction {
+            name: name.to_string(),
+            arg_types: arg_types.to_vec(),
+            return_types: return_types.to_vec(),
+            eval: Box::new(eval),
+        });
+        self.srfs.entry(name).or_insert_with(Vec::new).push(srf);
+    }
+
+    pub fn contains(&self, name: &str) -> bool {
+        self.srfs.contains_key(name)
+    }
+
+    pub fn search_candidates<Index: ColumnIndex>(
+        &self,
+        name: &str,
+        args: &[Expr<Index>],
+    ) -> Vec<(Arc<SetReturningFunction>, usize)> {
+        let mut candidates = vec![];
+        self.srfs
+            .get(name)
+            .unwrap_or(&vec![])
+            .iter()
+            .enumerate()
+            .for_each(|(index, srf)| {
+                if srf.name == name && srf.arg_types.len() == args.len() {
+                    candidates.push((srf.clone(), index));
+                }
+            });
+
+        candidates
+    }
+
+    pub fn get(&self, id: &SetReturningFunctionID) -> Option<Arc<SetReturningFunction>> {
+        self.srfs
+            .get(&id.0)
+            .and_then(|srfs| srfs.get(id.1))
+            .cloned()
+    }
+}
+
+pub fn check_srf<Index: ColumnIndex>(
+    name: &str,
+    args: &[Expr<Index>],
+    registry: &SetReturningFunctionRegistry,
+) -> Result<SrfExpr<Index>> {
+    let candidates = registry.search_candidates(name, args);
+
+    if candidates.is_empty() && !registry.contains(name) {
+        return Err(ErrorCode::UnknownFunction(format!(
+            "srf `{name}` does not exist"
+        )));
+    }
+
+    let mut fail_reasons = Vec::with_capacity(candidates.len());
+    for (srf, index) in candidates.iter() {
+        match try_check_srf(&srf.arg_types, &srf.return_types, args) {
+            Ok((checked_args, return_types, generic_types)) => {
+                return Ok(SrfExpr {
+                    id: (srf.name.clone(), *index),
+                    srf: srf.clone(),
+                    args: checked_args,
+                    generic_types,
+                    return_types,
+                });
+            }
+            Err(err) => fail_reasons.push(err),
+        }
+    }
+
+    Err(ErrorCode::SemanticError(format!(
+        "no overload satisfies {name}({}): {}",
+        args.iter()
+            .map(|x| x.data_type().to_string())
+            .collect::<Vec<_>>()
+            .join(", "),
+        fail_reasons
+            .into_iter()
+            .map(|e| e.message())
+            .collect::<Vec<_>>()
+            .join(", ")
+    )))
+}