feat(query): Support array lambda function

src/query/ast/src/ast/expr.rs

@@ -180,6 +180,7 @@ pub enum Expr {
         args: Vec<Expr>,
         params: Vec<Literal>,
         window: Option<Window>,
+        lambda: Option<Lambda>,
     },
     /// `CASE ... WHEN ... ELSE ...` expression
     Case {
@@ -376,6 +377,12 @@ pub enum WindowFrameBound {
     Following(Option<Box<Expr>>),
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct Lambda {
+    pub params: Vec<Identifier>,
+    pub expr: Box<Expr>,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum BinaryOperator {
     Plus,
@@ -882,6 +889,21 @@ impl Display for WindowSpec {
     }
 }
 
+impl Display for Lambda {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        if self.params.len() == 1 {
+            write!(f, "{}", self.params[0])?;
+        } else {
+            write!(f, "(")?;
+            write_comma_separated_list(f, self.params.clone())?;
+            write!(f, ")")?;
+        }
+        write!(f, " -> {}", self.expr)?;
+
+        Ok(())
+    }
+}
+
 impl Display for Expr {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -1043,6 +1065,7 @@ impl Display for Expr {
                 args,
                 params,
                 window,
+                lambda,
                 ..
             } => {
                 write!(f, "{name}")?;
@@ -1056,6 +1079,9 @@ impl Display for Expr {
                     write!(f, "DISTINCT ")?;
                 }
                 write_comma_separated_list(f, args)?;
+                if let Some(lambda) = lambda {
+                    write!(f, ", {lambda}")?;
+                }
                 write!(f, ")")?;
 
                 if let Some(window) = window {

src/query/ast/src/ast/format/ast_format.rs

@@ -431,6 +431,7 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         args: &'ast [Expr],
         _params: &'ast [Literal],
         _over: &'ast Option<Window>,
+        _lambda: &'ast Option<Lambda>,
     ) {
         let mut children = Vec::with_capacity(args.len());
         for arg in args.iter() {

src/query/ast/src/parser/expr.rs

@@ -291,8 +291,9 @@ pub enum ExprElement {
         distinct: bool,
         name: Identifier,
         args: Vec<Expr>,
-        window: Option<Window>,
         params: Vec<Literal>,
+        window: Option<Window>,
+        lambda: Option<Lambda>,
     },
     /// `CASE ... WHEN ... ELSE ...` expression
     Case {
@@ -491,13 +492,15 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
                 args,
                 params,
                 window,
+                lambda,
             } => Expr::FunctionCall {
                 span: transform_span(elem.span.0),
                 distinct,
                 name,
                 args,
                 params,
                 window,
+                lambda,
             },
             ExprElement::Case {
                 operand,
@@ -830,6 +833,25 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             args: opt_args.unwrap_or_default(),
             params: vec![],
             window: None,
+            lambda: None,
+        },
+    );
+
+    let function_call_with_lambda = map(
+        rule! {
+            #function_name
+            ~ "(" ~ #subexpr(0) ~ "," ~ #ident ~ RArrow ~ #subexpr(0) ~ ")"
+        },
+        |(name, _, arg, _, param, _, expr, _)| ExprElement::FunctionCall {
+            distinct: false,
+            name,
+            args: vec![arg],
+            params: vec![],
+            window: None,
+            lambda: Some(Lambda {
+                params: vec![param],
+                expr: Box::new(expr),
+            }),
         },
     );
 
@@ -845,6 +867,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             args: opt_args.unwrap_or_default(),
             params: vec![],
             window: Some(window.1),
+            lambda: None,
         },
     );
 
@@ -860,6 +883,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             args: opt_args.unwrap_or_default(),
             params: params.map(|x| x.1).unwrap_or_default(),
             window: None,
+            lambda: None,
         },
     );
 
@@ -1028,6 +1052,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             | #trim_from : "`TRIM([(BOTH | LEADEING | TRAILING) ... FROM ...)`"
             | #is_distinct_from: "`... IS [NOT] DISTINCT FROM ...`"
             | #count_all_with_window : "`COUNT(*) OVER ...`"
+            | #function_call_with_lambda : "<function>"
             | #function_call_with_window : "<function>"
             | #function_call_with_params : "<function>"
             | #function_call : "<function>"

src/query/ast/src/visitors/visitor.rs

@@ -222,6 +222,7 @@ pub trait Visitor<'ast>: Sized {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn visit_function_call(
         &mut self,
         _span: Span,
@@ -230,6 +231,7 @@ pub trait Visitor<'ast>: Sized {
         args: &'ast [Expr],
         _params: &'ast [Literal],
         over: &'ast Option<Window>,
+        lambda: &'ast Option<Lambda>,
     ) {
         for arg in args {
             walk_expr(self, arg);
@@ -238,6 +240,9 @@ pub trait Visitor<'ast>: Sized {
         if let Some(over) = over {
             self.visit_window(over);
         }
+        if let Some(lambda) = lambda {
+            walk_expr(self, &lambda.expr)
+        }
     }
 
     fn visit_window(&mut self, window: &'ast Window) {

src/query/ast/src/visitors/visitor_mut.rs

@@ -236,6 +236,7 @@ pub trait VisitorMut: Sized {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn visit_function_call(
         &mut self,
         _span: Span,
@@ -244,6 +245,7 @@ pub trait VisitorMut: Sized {
         args: &mut [Expr],
         _params: &mut [Literal],
         over: &mut Option<Window>,
+        lambda: &mut Option<Lambda>,
     ) {
         for arg in args.iter_mut() {
             walk_expr_mut(self, arg);
@@ -269,6 +271,9 @@ pub trait VisitorMut: Sized {
                 }
             }
         }
+        if let Some(lambda) = lambda {
+            walk_expr_mut(self, &mut lambda.expr)
+        }
     }
 
     fn visit_frame_bound(&mut self, bound: &mut WindowFrameBound) {

src/query/ast/src/visitors/walk.rs

@@ -94,7 +94,8 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expr: &'a Expr) {
             args,
             params,
             window,
-        } => visitor.visit_function_call(*span, *distinct, name, args, params, window),
+            lambda,
+        } => visitor.visit_function_call(*span, *distinct, name, args, params, window, lambda),
         Expr::Case {
             span,
             operand,

src/query/ast/src/visitors/walk_mut.rs

@@ -94,7 +94,8 @@ pub fn walk_expr_mut<V: VisitorMut>(visitor: &mut V, expr: &mut Expr) {
             args,
             params,
             window,
-        } => visitor.visit_function_call(*span, *distinct, name, args, params, window),
+            lambda,
+        } => visitor.visit_function_call(*span, *distinct, name, args, params, window, lambda),
         Expr::Case {
             span,
             operand,

src/query/ast/tests/it/parser.rs

@@ -662,6 +662,8 @@ fn test_expr() {
         r#"COUNT() OVER (ORDER BY hire_date ROWS UNBOUNDED PRECEDING)"#,
         r#"COUNT() OVER (ORDER BY hire_date ROWS CURRENT ROW)"#,
         r#"COUNT() OVER (ORDER BY hire_date ROWS 3 PRECEDING)"#,
+        r#"ARRAY_APPLY([1,2,3], x -> x + 1)"#,
+        r#"ARRAY_FILTER(col, y -> y % 2 = 0)"#,
     ];
 
     for case in cases {