Repository: cloudspannerecosystem/memefish (go)

Issue:
Support lambda expressions
Lambda expressions have not yet been supported.

Syntax

https://cloud.google.com/spanner/docs/reference/standard-sql/functions-reference#lambdas

```
(arg[, ...]) -> body_expression
```

```
arg -> body_expression
```

Examples

https://cloud.google.com/spanner/docs/reference/standard-sql/array_functions#array_filter

```
SELECT
  ARRAY_FILTER([1 ,2, 3], e -> e > 1) AS a1,
  ARRAY_FILTER([0, 2, 3], (e, i) -> e > i) AS a2;

/*-------+-------*
 | a1    | a2    |
 +-------+-------+
 | [2,3] | [2,3] |
 *-------+-------*/
```

https://cloud.google.com/spanner/docs/reference/standard-sql/array_functions#array_includes

```
SELECT
  ARRAY_INCLUDES([1, 2, 3], e -> e > 3) AS a1,
  ARRAY_INCLUDES([1, 2, 3], e -> e > 0) AS a2;

/*-------+------*
 | a1    | a2   |
 +-------+------+
 | false | true |
 *-------+------*/
```

https://cloud.google.com/spanner/docs/reference/standard-sql/array_functions#array_transform

```
SELECT
  ARRAY_TRANSFORM([1, 2, 3], e -> e + 1) AS a1,
  ARRAY_TRANSFORM([1, 2, 3], (e, i) -> e + i) AS a2;

/*---------+---------*
 | a1      | a2      |
 +---------+---------+
 | [2,3,4] | [1,3,5] |
 *---------+---------*/
```

Reference

https://cloud.google.com/spanner/docs/release-notes#March_21_2023
> The following functions and expressions have been added to the GoogleSQL dialect:
> * [`ARRAY_FILTER` function](https://cloud.google.com/spanner/docs/reference/standard-sql/array_functions.md#array_filter)
> * [`ARRAY_TRANSFORM` function](https://cloud.google.com/spanner/docs/reference/standard-sql/array_functions.md#array_transform)
> * [Lambda expressions](https://cloud.google.com/spanner/docs/reference/standard-sql/functions-reference#lambdas)

FYI

* In spansql, they are developed, but not merged.
  * https://github.com/googleapis/google-cloud-go/issues/8768
  * https://github.com/googleapis/google-cloud-go/pull/8769

Pull request description:
Implement lambda expressions
This PR implements lambda expressions.
Note: `ast.LambdaArg` is a one of `ast.Arg`. It is a same with ZetaSQL AST node.

- close #153 

Patch:
diff --git a/ast/ast.go b/ast/ast.go
index 1de8aac..e13a979 100644
--- a/ast/ast.go
+++ b/ast/ast.go
@@ -172,6 +172,7 @@ func (CallExpr) isExpr()              {}
 func (CountStarExpr) isExpr()         {}
 func (CastExpr) isExpr()              {}
 func (ExtractExpr) isExpr()           {}
+func (ReplaceFieldsExpr) isExpr()     {}
 func (CaseExpr) isExpr()              {}
 func (IfExpr) isExpr()                {}
 func (ParenExpr) isExpr()             {}
@@ -217,6 +218,7 @@ type Arg interface {
 func (ExprArg) isArg()     {}
 func (IntervalArg) isArg() {}
 func (SequenceArg) isArg() {}
+func (LambdaArg) isArg()   {}
 
 // NullHandlingModifier represents IGNORE/RESPECT NULLS of aggregate function calls
 type NullHandlingModifier interface {
@@ -1192,6 +1194,21 @@ type SequenceArg struct {
 	Expr Expr
 }
 
+// LambdaArg is lambda expression argument of the generic function call.
+//
+//	{{if .Lparen.Invalid}}{{.Args | sqlJoin ", "}}{{else}}({{.Args | sqlJoin ", "}}) -> {{.Expr | sql}}
+//
+// Note: Args won't be empty. If Lparen is not appeared, Args have exactly one element.
+type LambdaArg struct {
+	// pos = Lparen || Args[0].pos
+	// end = Expr.end
+
+	Lparen token.Pos // optional
+
+	Args []*Ident // if Lparen.Invalid() then len(Args) = 1 else len(Args) > 0
+	Expr Expr
+}
+
 // NamedArg represents a name and value pair in named arguments
 //
 //	{{.Name | sql}} => {{.Value | sql}}
@@ -1273,6 +1290,31 @@ type ExtractExpr struct {
 	AtTimeZone *AtTimeZone // optional
 }
 
+// ReplaceFieldsArg is value AS field_path node in ReplaceFieldsExpr.
+//
+//	{{.Expr | sql}} AS {{.Field | sql}}
+type ReplaceFieldsArg struct {
+	// pos = Expr.pos
+	// end = Field.end
+
+	Expr  Expr
+	Field *Path
+}
+
+// ReplaceFieldsExpr is REPLACE_FIELDS call expression node.
+//
+//	REPLACE_FIELDS({{.Expr.| sql}}, {{.Fields | sqlJoin ", "}})
+type ReplaceFieldsExpr struct {
+	// pos = ReplaceFields
+	// end = Rparen + 1
+
+	ReplaceFields token.Pos // position of "REPLACE_FIELDS" keyword
+	Rparen        token.Pos // position of ")"
+
+	Expr   Expr
+	Fields []*ReplaceFieldsArg
+}
+
 // AtTimeZone is AT TIME ZONE clause in EXTRACT call.
 //
 //	AT TIME ZONE {{.Expr | sql}}
diff --git a/ast/pos.go b/ast/pos.go
index dd0cd88..c6cc1f4 100644
--- a/ast/pos.go
+++ b/ast/pos.go
@@ -422,6 +422,14 @@ func (s *SequenceArg) End() token.Pos {
 	return nodeEnd(wrapNode(s.Expr))
 }
 
+func (l *LambdaArg) Pos() token.Pos {
+	return posChoice(l.Lparen, nodePos(nodeSliceIndex(l.Args, 0)))
+}
+
+func (l *LambdaArg) End() token.Pos {
+	return nodeEnd(wrapNode(l.Expr))
+}
+
 func (n *NamedArg) Pos() token.Pos {
 	return nodePos(wrapNode(n.Name))
 }
@@ -478,6 +486,22 @@ func (e *ExtractExpr) End() token.Pos {
 	return posAdd(e.Rparen, 1)
 }
 
+func (r *ReplaceFieldsArg) Pos() token.Pos {
+	return nodePos(wrapNode(r.Expr))
+}
+
+func (r *ReplaceFieldsArg) End() token.Pos {
+	return nodeEnd(wrapNode(r.Field))
+}
+
+func (r *ReplaceFieldsExpr) Pos() token.Pos {
+	return r.ReplaceFields
+}
+
+func (r *ReplaceFieldsExpr) End() token.Pos {
+	return posAdd(r.Rparen, 1)
+}
+
 func (a *AtTimeZone) Pos() token.Pos {
 	return a.At
 }
diff --git a/ast/sql.go b/ast/sql.go
index 12f864d..fa8caaf 100644
--- a/ast/sql.go
+++ b/ast/sql.go
@@ -506,6 +506,15 @@ func (c *CallExpr) SQL() string {
 		")"
 }
 
+func (l *LambdaArg) SQL() string {
+	// This implementation is not exactly matched with the doc comment for simplicity.
+	return strOpt(!l.Lparen.Invalid(), "(") +
+		sqlJoin(l.Args, ", ") +
+		strOpt(!l.Lparen.Invalid(), ")") +
+		" -> " +
+		l.Expr.SQL()
+}
+
 func (n *NamedArg) SQL() string { return n.Name.SQL() + " => " + n.Value.SQL() }
 
 func (i *IgnoreNulls) SQL() string { return "IGNORE NULLS" }
@@ -549,6 +558,12 @@ func (a *AtTimeZone) SQL() string {
 	return "AT TIME ZONE " + a.Expr.SQL()
 }
 
+func (r *ReplaceFieldsArg) SQL() string { return r.Expr.SQL() + " AS " + r.Field.SQL() }
+
+func (r *ReplaceFieldsExpr) SQL() string {
+	return "REPLACE_FIELDS(" + r.Expr.SQL() + ", " + sqlJoin(r.Fields, ", ") + ")"
+}
+
 func (c *CastExpr) SQL() string {
 	return strOpt(c.Safe, "SAFE_") + "CAST(" + c.Expr.SQL() + " AS " + c.Type.SQL() + ")"
 }
diff --git a/parser.go b/parser.go
index 0655fc5..f7c30dd 100644
--- a/parser.go
+++ b/parser.go
@@ -1469,6 +1469,8 @@ func (p *Parser) parseLit() ast.Expr {
 		switch {
 		case id.IsKeywordLike("SAFE_CAST"):
 			return p.parseCastExpr()
+		case id.IsKeywordLike("REPLACE_FIELDS"):
+			return p.parseReplaceFieldsExpr()
 		}
 		p.nextToken()
 		switch p.Token.Kind {
@@ -1588,6 +1590,42 @@ func (p *Parser) tryParseNamedArg() *ast.NamedArg {
 	}
 }
 
+func (p *Parser) lookaheadLambdaArg() bool {
+	lexer := p.Lexer.Clone()
+	defer func() {
+		p.Lexer = lexer
+	}()
+
+	if p.Token.Kind != "(" && p.Token.Kind != token.TokenIdent {
+		return false
+	}
+
+	// Note: all lambda patterns can be parsed as expr -> expr.
+	p.parseExpr()
+	return p.Token.Kind == "->"
+}
+
+func (p *Parser) parseLambdaArg() *ast.LambdaArg {
+	lparen := token.InvalidPos
+	var args []*ast.Ident
+	if p.Token.Kind == "(" {
+		lparen = p.expect("(").Pos
+		args = parseCommaSeparatedList(p, p.parseIdent)
+		p.expect(")")
+	} else {
+		args = []*ast.Ident{p.parseIdent()}
+	}
+
+	p.expect("->")
+	expr := p.parseExpr()
+
+	return &ast.LambdaArg{
+		Lparen: lparen,
+		Args:   args,
+		Expr:   expr,
+	}
+}
+
 func (p *Parser) parseArg() ast.Arg {
 	if i := p.tryParseIntervalArg(); i != nil {
 		return i
@@ -1595,6 +1633,9 @@ func (p *Parser) parseArg() ast.Arg {
 	if s := p.tryParseSequenceArg(); s != nil {
 		return s
 	}
+	if p.lookaheadLambdaArg() {
+		return p.parseLambdaArg()
+	}
 	return p.parseExprArg()
 }
 
@@ -1754,6 +1795,33 @@ func (p *Parser) parseIfExpr() *ast.IfExpr {
 	}
 }
 
+func (p *Parser) parseReplaceFieldsArg() *ast.ReplaceFieldsArg {
+	expr := p.parseExpr()
+	p.expect("AS")
+	field := p.parsePath()
+
+	return &ast.ReplaceFieldsArg{
+		Expr:  expr,
+		Field: field,
+	}
+}
+
+func (p *Parser) parseReplaceFieldsExpr() *ast.ReplaceFieldsExpr {
+	replaceFields := p.expectKeywordLike("REPLACE_FIELDS").Pos
+	p.expect("(")
+	expr := p.parseExpr()
+	p.expect(",")
+	fields := parseCommaSeparatedList(p, p.parseReplaceFieldsArg)
+	rparen := p.expect(")").Pos
+
+	return &ast.ReplaceFieldsExpr{
+		ReplaceFields: replaceFields,
+		Rparen:        rparen,
+		Expr:          expr,
+		Fields:        fields,
+	}
+}
+
 func (p *Parser) parseCastExpr() *ast.CastExpr {
 	if p.Token.Kind != "CAST" && !p.Token.IsKeywordLike("SAFE_CAST") {
 		panic(p.errorfAtToken(&p.Token, `expected CAST keyword or SAFE_CAST pseudo keyword, but: %v`, p.Token.Kind))


Test patch:
diff --git a/testdata/input/expr/array_functions_array_filter_parenless_lambda.sql b/testdata/input/expr/array_functions_array_filter_parenless_lambda.sql
new file mode 100644
index 0000000..e00d2ea
--- /dev/null
+++ b/testdata/input/expr/array_functions_array_filter_parenless_lambda.sql
@@ -0,0 +1,1 @@
+ARRAY_FILTER([1 ,2, 3], e -> e > 1)
\ No newline at end of file
diff --git a/testdata/input/expr/array_functions_array_filter_two_args_lambda.sql b/testdata/input/expr/array_functions_array_filter_two_args_lambda.sql
new file mode 100644
index 0000000..205e21f
--- /dev/null
+++ b/testdata/input/expr/array_functions_array_filter_two_args_lambda.sql
@@ -0,0 +1,1 @@
+ARRAY_FILTER([0, 2, 3], (e, i) -> e > i)
\ No newline at end of file
diff --git a/testdata/input/expr/protocol_buffer_functions_replace_fields.sql b/testdata/input/expr/protocol_buffer_functions_replace_fields.sql
new file mode 100644
index 0000000..64caec0
--- /dev/null
+++ b/testdata/input/expr/protocol_buffer_functions_replace_fields.sql
@@ -0,0 +1,6 @@
+REPLACE_FIELDS(
+  NEW Book(
+    "The Hummingbird" AS title,
+    NEW BookDetails(10 AS chapters) AS details),
+  "The Hummingbird II" AS title,
+  11 AS details.chapters)
\ No newline at end of file
diff --git a/testdata/result/expr/array_functions_array_filter_parenless_lambda.sql.txt b/testdata/result/expr/array_functions_array_filter_parenless_lambda.sql.txt
new file mode 100644
index 0000000..982a7ba
--- /dev/null
+++ b/testdata/result/expr/array_functions_array_filter_parenless_lambda.sql.txt
@@ -0,0 +1,72 @@
+--- array_functions_array_filter_parenless_lambda.sql
+ARRAY_FILTER([1 ,2, 3], e -> e > 1)
+--- AST
+&ast.CallExpr{
+  Rparen: 34,
+  Func:   &ast.Ident{
+    NamePos: 0,
+    NameEnd: 12,
+    Name:    "ARRAY_FILTER",
+  },
+  Distinct: false,
+  Args:     []ast.Arg{
+    &ast.ExprArg{
+      Expr: &ast.ArrayLiteral{
+        Array:  -1,
+        Lbrack: 13,
+        Rbrack: 21,
+        Type:   nil,
+        Values: []ast.Expr{
+          &ast.IntLiteral{
+            ValuePos: 14,
+            ValueEnd: 15,
+            Base:     10,
+            Value:    "1",
+          },
+          &ast.IntLiteral{
+            ValuePos: 17,
+            ValueEnd: 18,
+            Base:     10,
+            Value:    "2",
+          },
+          &ast.IntLiteral{
+            ValuePos: 20,
+            ValueEnd: 21,
+            Base:     10,
+            Value:    "3",
+          },
+        },
+      },
+    },
+    &ast.LambdaArg{
+      Lparen: -1,
+      Args:   []*ast.Ident{
+        &ast.Ident{
+          NamePos: 24,
+          NameEnd: 25,
+          Name:    "e",
+        },
+      },
+      Expr: &ast.BinaryExpr{
+        Op:   ">",
+        Left: &ast.Ident{
+          NamePos: 29,
+          NameEnd: 30,
+          Name:    "e",
+        },
+        Right: &ast.IntLiteral{
+          ValuePos: 33,
+          ValueEnd: 34,
+          Base:     10,
+          Value:    "1",
+        },
+      },
+    },
+  },
+  NamedArgs:    []*ast.NamedArg(nil),
+  NullHandling: nil,
+  Having:       nil,
+}
+
+--- SQL
+ARRAY_FILTER([1, 2, 3], e -> e > 1)
diff --git a/testdata/result/expr/array_functions_array_filter_two_args_lambda.sql.txt b/testdata/result/expr/array_functions_array_filter_two_args_lambda.sql.txt
new file mode 100644
index 0000000..7d83c24
--- /dev/null
+++ b/testdata/result/expr/array_functions_array_filter_two_args_lambda.sql.txt
@@ -0,0 +1,76 @@
+--- array_functions_array_filter_two_args_lambda.sql
+ARRAY_FILTER([0, 2, 3], (e, i) -> e > i)
+--- AST
+&ast.CallExpr{
+  Rparen: 39,
+  Func:   &ast.Ident{
+    NamePos: 0,
+    NameEnd: 12,
+    Name:    "ARRAY_FILTER",
+  },
+  Distinct: false,
+  Args:     []ast.Arg{
+    &ast.ExprArg{
+      Expr: &ast.ArrayLiteral{
+        Array:  -1,
+        Lbrack: 13,
+        Rbrack: 21,
+        Type:   nil,
+        Values: []ast.Expr{
+          &ast.IntLiteral{
+            ValuePos: 14,
+            ValueEnd: 15,
+            Base:     10,
+            Value:    "0",
+          },
+          &ast.IntLiteral{
+            ValuePos: 17,
+            ValueEnd: 18,
+            Base:     10,
+            Value:    "2",
+          },
+          &ast.IntLiteral{
+            ValuePos: 20,
+            ValueEnd: 21,
+            Base:     10,
+            Value:    "3",
+          },
+        },
+      },
+    },
+    &ast.LambdaArg{
+      Lparen: 24,
+      Args:   []*ast.Ident{
+        &ast.Ident{
+          NamePos: 25,
+          NameEnd: 26,
+          Name:    "e",
+        },
+        &ast.Ident{
+          NamePos: 28,
+          NameEnd: 29,
+          Name:    "i",
+        },
+      },
+      Expr: &ast.BinaryExpr{
+        Op:   ">",
+        Left: &ast.Ident{
+          NamePos: 34,
+          NameEnd: 35,
+          Name:    "e",
+        },
+        Right: &ast.Ident{
+          NamePos: 38,
+          NameEnd: 39,
+          Name:    "i",
+        },
+      },
+    },
+  },
+  NamedArgs:    []*ast.NamedArg(nil),
+  NullHandling: nil,
+  Having:       nil,
+}
+
+--- SQL
+ARRAY_FILTER([0, 2, 3], (e, i) -> e > i)
diff --git a/testdata/result/expr/protocol_buffer_functions_replace_fields.sql.txt b/testdata/result/expr/protocol_buffer_functions_replace_fields.sql.txt
new file mode 100644
index 0000000..5587e16
--- /dev/null
+++ b/testdata/result/expr/protocol_buffer_functions_replace_fields.sql.txt
@@ -0,0 +1,126 @@
+--- protocol_buffer_functions_replace_fields.sql
+REPLACE_FIELDS(
+  NEW Book(
+    "The Hummingbird" AS title,
+    NEW BookDetails(10 AS chapters) AS details),
+  "The Hummingbird II" AS title,
+  11 AS details.chapters)
+--- AST
+&ast.ReplaceFieldsExpr{
+  ReplaceFields: 0,
+  Rparen:        166,
+  Expr:          &ast.NewConstructor{
+    New:  18,
+    Type: &ast.NamedType{
+      Path: []*ast.Ident{
+        &ast.Ident{
+          NamePos: 22,
+          NameEnd: 26,
+          Name:    "Book",
+        },
+      },
+    },
+    Args: []ast.NewConstructorArg{
+      &ast.Alias{
+        Expr: &ast.StringLiteral{
+          ValuePos: 32,
+          ValueEnd: 49,
+          Value:    "The Hummingbird",
+        },
+        As: &ast.AsAlias{
+          As:    50,
+          Alias: &ast.Ident{
+            NamePos: 53,
+            NameEnd: 58,
+            Name:    "title",
+          },
+        },
+      },
+      &ast.Alias{
+        Expr: &ast.NewConstructor{
+          New:  64,
+          Type: &ast.NamedType{
+            Path: []*ast.Ident{
+              &ast.Ident{
+                NamePos: 68,
+                NameEnd: 79,
+                Name:    "BookDetails",
+              },
+            },
+          },
+          Args: []ast.NewConstructorArg{
+            &ast.Alias{
+              Expr: &ast.IntLiteral{
+                ValuePos: 80,
+                ValueEnd: 82,
+                Base:     10,
+                Value:    "10",
+              },
+              As: &ast.AsAlias{
+                As:    83,
+                Alias: &ast.Ident{
+                  NamePos: 86,
+                  NameEnd: 94,
+                  Name:    "chapters",
+                },
+              },
+            },
+          },
+          Rparen: 94,
+        },
+        As: &ast.AsAlias{
+          As:    96,
+          Alias: &ast.Ident{
+            NamePos: 99,
+            NameEnd: 106,
+            Name:    "details",
+          },
+        },
+      },
+    },
+    Rparen: 106,
+  },
+  Fields: []*ast.ReplaceFieldsArg{
+    &ast.ReplaceFieldsArg{
+      Expr: &ast.StringLiteral{
+        ValuePos: 111,
+        ValueEnd: 131,
+        Value:    "The Hummingbird II",
+      },
+      Field: &ast.Path{
+        Idents: []*ast.Ident{
+          &ast.Ident{
+            NamePos: 135,
+            NameEnd: 140,
+            Name:    "title",
+          },
+        },
+      },
+    },
+    &ast.ReplaceFieldsArg{
+      Expr: &ast.IntLiteral{
+        ValuePos: 144,
+        ValueEnd: 146,
+        Base:     10,
+        Value:    "11",
+      },
+      Field: &ast.Path{
+        Idents: []*ast.Ident{
+          &ast.Ident{
+            NamePos: 150,
+            NameEnd: 157,
+            Name:    "details",
+          },
+          &ast.Ident{
+            NamePos: 158,
+            NameEnd: 166,
+            Name:    "chapters",
+          },
+        },
+      },
+    },
+  },
+}
+
+--- SQL
+REPLACE_FIELDS(NEW Book("The Hummingbird" AS title, NEW BookDetails(10 AS chapters) AS details), "The Hummingbird II" AS title, 11 AS details.chapters)