[query-engine] Expand expressions to support user-defined functions (#1478)

Relates to #1479

## Changes

* Make it possible to declare and invoke user-defined functions in query
expression tree

## Details

Implementation and KQL parsing will come as follow-ups.

Author: CodeBlanch

rust/experimental/query_engine/engine-recordset/src/scalars/scalar_expressions.rs

@@ -460,6 +460,8 @@ where
                 }
             }
         }
+        ScalarExpression::Argument(_) => todo!(),
+        ScalarExpression::InvokeFunction(_) => todo!(),
     };
 
     execution_context.add_diagnostic_if_enabled(

rust/experimental/query_engine/expressions/src/pipeline_expression.rs

@@ -8,6 +8,7 @@ pub struct PipelineExpression {
     query: Box<str>,
     query_location: QueryLocation,
     constants: Vec<StaticScalarExpression>,
+    functions: Vec<PipelineFunction>,
     initializations: Vec<PipelineInitialization>,
     expressions: Vec<DataExpression>,
 }
@@ -18,6 +19,7 @@ impl PipelineExpression {
             query: query.into(),
             query_location: QueryLocation::new(0, query.len(), 1, 1).unwrap(),
             constants: Vec::new(),
+            functions: Vec::new(),
             initializations: Vec::new(),
             expressions: Vec::new(),
         }
@@ -38,6 +40,11 @@ impl PipelineExpression {
         self.constants.len() - 1
     }
 
+    pub(crate) fn push_function(&mut self, value: PipelineFunction) -> usize {
+        self.functions.push(value);
+        self.functions.len() - 1
+    }
+
     pub fn get_constants(&self) -> &[StaticScalarExpression] {
         &self.constants
     }
@@ -46,6 +53,14 @@ impl PipelineExpression {
         self.constants.get(constant_id)
     }
 
+    pub fn get_functions(&self) -> &[PipelineFunction] {
+        &self.functions
+    }
+
+    pub fn get_function(&self, function_id: usize) -> Option<&PipelineFunction> {
+        self.functions.get(function_id)
+    }
+
     pub fn get_expressions(&self) -> &[DataExpression] {
         &self.expressions
     }
@@ -69,12 +84,14 @@ impl PipelineExpression {
     pub fn get_resolution_scope(&self) -> PipelineResolutionScope<'_> {
         PipelineResolutionScope {
             constants: &self.constants,
+            functions: &self.functions,
         }
     }
 
     pub(crate) fn optimize(&mut self) -> Result<(), Vec<ExpressionError>> {
         let scope = PipelineResolutionScope {
             constants: &self.constants,
+            functions: &self.functions,
         };
 
         let mut errors = Vec::new();
@@ -125,6 +142,22 @@ impl std::fmt::Display for PipelineExpression {
             }
         }
 
+        if self.functions.is_empty() {
+            writeln!(f, "├── Functions: []")?;
+        } else {
+            writeln!(f, "├── Functions:")?;
+            let last_idx = self.functions.len() - 1;
+            for (i, func) in self.functions.iter().enumerate() {
+                if i == last_idx {
+                    writeln!(f, "│   └── {i}")?;
+                    func.fmt_with_indent(f, "│       ")?;
+                } else {
+                    writeln!(f, "│   ├── {i}")?;
+                    func.fmt_with_indent(f, "│   │   ")?;
+                }
+            }
+        }
+
         if self.initializations.is_empty() {
             writeln!(f, "├── Initializations: []")?;
         } else {
@@ -163,12 +196,17 @@ impl std::fmt::Display for PipelineExpression {
 
 pub struct PipelineResolutionScope<'a> {
     constants: &'a Vec<StaticScalarExpression>,
+    functions: &'a Vec<PipelineFunction>,
 }
 
 impl<'a> PipelineResolutionScope<'a> {
     pub fn get_constant(&self, constant_id: usize) -> Option<&'a StaticScalarExpression> {
         self.constants.get(constant_id)
     }
+
+    pub fn get_function(&self, function_id: usize) -> Option<&'a PipelineFunction> {
+        self.functions.get(function_id)
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -220,6 +258,14 @@ impl PipelineExpressionBuilder {
         self
     }
 
+    pub fn with_functions(mut self, functions: Vec<PipelineFunction>) -> PipelineExpressionBuilder {
+        for f in functions {
+            self.push_function(f);
+        }
+
+        self
+    }
+
     pub fn with_global_variables(
         mut self,
         variables: Vec<(&str, ScalarExpression)>,
@@ -246,6 +292,10 @@ impl PipelineExpressionBuilder {
         self.pipeline.push_constant(value)
     }
 
+    pub fn push_function(&mut self, value: PipelineFunction) -> usize {
+        self.pipeline.push_function(value)
+    }
+
     pub fn push_global_variable(&mut self, name: &str, value: ScalarExpression) {
         self.pipeline.push_global_variable(name, value)
     }
@@ -267,6 +317,210 @@ impl AsRef<PipelineExpression> for PipelineExpressionBuilder {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub enum PipelineFunctionExpression {
+    Transform(TransformExpression),
+    Return(ScalarExpression),
+}
+
+impl PipelineFunctionExpression {
+    pub(crate) fn fmt_with_indent(
+        &self,
+        f: &mut std::fmt::Formatter<'_>,
+        indent: &str,
+    ) -> std::fmt::Result {
+        match self {
+            PipelineFunctionExpression::Transform(t) => {
+                write!(f, "Transform: ")?;
+                t.fmt_with_indent(f, format!("{indent}           ").as_str())
+            }
+            PipelineFunctionExpression::Return(s) => {
+                write!(f, "Return: ")?;
+                s.fmt_with_indent(f, format!("{indent}        ").as_str())
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct PipelineFunction {
+    query_location: QueryLocation,
+    parameters: Vec<PipelineFunctionParameter>,
+    return_value_type: Option<ValueType>,
+    implementation: PipelineFunctionImplementation,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum PipelineFunctionImplementation {
+    Expressions(Vec<PipelineFunctionExpression>),
+    External(Box<str>),
+}
+
+impl PipelineFunctionImplementation {
+    pub(crate) fn fmt_with_indent(
+        &self,
+        f: &mut std::fmt::Formatter<'_>,
+        indent: &str,
+    ) -> std::fmt::Result {
+        match self {
+            PipelineFunctionImplementation::Expressions(e) => {
+                if e.is_empty() {
+                    writeln!(f, "(Expressions): None")?;
+                } else {
+                    writeln!(f, "(Expressions):")?;
+                    let last_idx = e.len() - 1;
+                    for (i, p) in e.iter().enumerate() {
+                        if i == last_idx {
+                            write!(f, "{indent}    └── ")?;
+                            p.fmt_with_indent(f, format!("{indent}        ").as_str())?;
+                        } else {
+                            write!(f, "{indent}    ├── ")?;
+                            p.fmt_with_indent(f, format!("{indent}    │   ").as_str())?;
+                        }
+                    }
+                }
+            }
+            PipelineFunctionImplementation::External(id) => {
+                writeln!(f, "(External): {id}")?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl PipelineFunction {
+    pub fn new_with_expressions(
+        query_location: QueryLocation,
+        parameters: Vec<PipelineFunctionParameter>,
+        return_value_type: Option<ValueType>,
+        expressions: Vec<PipelineFunctionExpression>,
+    ) -> PipelineFunction {
+        Self {
+            query_location,
+            parameters,
+            return_value_type,
+            implementation: PipelineFunctionImplementation::Expressions(expressions),
+        }
+    }
+
+    pub fn new_external(
+        name: &str,
+        parameters: Vec<PipelineFunctionParameter>,
+        return_value_type: Option<ValueType>,
+    ) -> PipelineFunction {
+        Self {
+            query_location: QueryLocation::new(0, 0, 1, 1).unwrap(),
+            parameters,
+            return_value_type,
+            implementation: PipelineFunctionImplementation::External(name.into()),
+        }
+    }
+
+    pub fn get_parameters(&self) -> &[PipelineFunctionParameter] {
+        &self.parameters
+    }
+
+    pub fn get_return_value_type(&self) -> Option<ValueType> {
+        self.return_value_type.clone()
+    }
+
+    pub fn get_implementation(&self) -> &PipelineFunctionImplementation {
+        &self.implementation
+    }
+
+    pub(crate) fn fmt_with_indent(
+        &self,
+        f: &mut std::fmt::Formatter<'_>,
+        indent: &str,
+    ) -> std::fmt::Result {
+        if self.parameters.is_empty() {
+            writeln!(f, "{indent}├── Parameters: None")?;
+        } else {
+            writeln!(f, "{indent}├── Parameters: ")?;
+            let last_idx = self.parameters.len() - 1;
+            for (i, p) in self.parameters.iter().enumerate() {
+                if i == last_idx {
+                    p.fmt_with_indent(f, format!("{indent}│   └── ").as_str())?;
+                } else {
+                    p.fmt_with_indent(f, format!("{indent}│   ├── ").as_str())?;
+                }
+            }
+        }
+
+        match &self.return_value_type {
+            Some(t) => writeln!(f, "{indent}├── ReturnType: {t}")?,
+            None => writeln!(f, "{indent}├── ReturnType: None")?,
+        }
+
+        write!(f, "{indent}└── Implementation")?;
+        self.implementation.fmt_with_indent(f, indent)?;
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct PipelineFunctionParameter {
+    query_location: QueryLocation,
+    name: Box<str>,
+    parameter_type: PipelineFunctionParameterType,
+}
+
+impl PipelineFunctionParameter {
+    pub fn new(
+        query_location: QueryLocation,
+        name: &str,
+        parameter_type: PipelineFunctionParameterType,
+    ) -> PipelineFunctionParameter {
+        Self {
+            query_location,
+            name: name.into(),
+            parameter_type,
+        }
+    }
+
+    pub fn get_name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn get_parameter_type(&self) -> PipelineFunctionParameterType {
+        self.parameter_type.clone()
+    }
+
+    pub(crate) fn fmt_with_indent(
+        &self,
+        f: &mut std::fmt::Formatter<'_>,
+        indent: &str,
+    ) -> std::fmt::Result {
+        write!(f, "{indent}{} = ", &self.name)?;
+
+        let value_type = match &self.parameter_type {
+            PipelineFunctionParameterType::Scalar(v) => {
+                write!(f, "Scalar(")?;
+                v
+            }
+            PipelineFunctionParameterType::MutableValue(v) => {
+                write!(f, "MutableValue(")?;
+                v
+            }
+        };
+
+        match value_type {
+            Some(v) => writeln!(f, "{v:?})")?,
+            None => writeln!(f, "Any)")?,
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum PipelineFunctionParameterType {
+    Scalar(Option<ValueType>),
+    MutableValue(Option<ValueType>),
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;