Repository: tobymao/sqlglot
Pull request: Add support for clickhouse's parametric function syntax
Phew... this was more tricky than expected.

Fixes #882 

Problem statement:
ClickHouse quantile/quantileIf functions fail to parse
ClickHouse has a somewhat unique way of spelling these:

```
9bf006fcf914 :) select quantile(0.5)(a) from (select 1 a union all select 2);

SELECT quantile(0.5)(a)
FROM
(
    SELECT 1 AS a
    UNION ALL
    SELECT 2
)

Query id: fe37056a-efe0-45bf-9d34-95de2916ecef

┌─quantile(0.5)(a)─┐
│              1.5 │
└──────────────────┘

1 row in set. Elapsed: 0.002 sec.

9bf006fcf914 :) select quantileIf(0.5)(a, true) from (select 1 a union all select 2);

SELECT quantileIf(0.5)(a, true)
FROM
(
    SELECT 1 AS a
    UNION ALL
    SELECT 2
)

Query id: d3b2e971-c9af-4a6c-8ba3-93caacc246bc

┌─quantileIf(0.5)(a, true)─┐
│                      1.5 │
└──────────────────────────┘

1 row in set. Elapsed: 0.002 sec.

9bf006fcf914 :) select quantileIf(0.5)(a, false) from (select 1 a union all select 2);

SELECT quantileIf(0.5)(a, false)
FROM
(
    SELECT 1 AS a
    UNION ALL
    SELECT 2
)

Query id: 27307fa4-2b25-4382-8e6c-ef3d4eb1b306

┌─quantileIf(0.5)(a, false)─┐
│                       nan │
└───────────────────────────┘

1 row in set. Elapsed: 0.002 sec.
```

When parsed with `sqlglot` I get a parse error:

```
In [1]: import sqlglot as sg

In [2]: sg.__version__
Out[2]: '10.3.2'

In [3]: sg.parse_one("select quantile(0.5)(a) from (select 1 a union all select 2)", read="clickhouse")
...
ParseError: Required keyword: 'quantile' missing for <class 'sqlglot.expressions.Quantile'>. Line 1, Col: 20.
  select quantile(0.5)(a) from (select 1 a union all select 2)
```

Hints from issue discussion:
I'll take a look at this soon. Besides the `quantile` argument missing in the `Quantile` expression, I don't think SQLGlot can handle the syntax `foo(a)(b)` yet.

Commit message: Add support for clickhouse's parametric function syntax

Patch to review:
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
index 7136340228..6e655c2118 100644
--- a/sqlglot/dialects/clickhouse.py
+++ b/sqlglot/dialects/clickhouse.py
@@ -11,6 +11,13 @@ def _lower_func(sql):
     return sql[:index].lower() + sql[index:]
 
 
+def _quantile_sql(self, expression: exp.Quantile) -> str:
+    if_suffix = "If" if len(expression.this) == 2 else ""
+    args = self.format_args(self.expressions(expression, "quantile"))
+    params = self.format_args(self.expressions(expression, "this"))
+    return f"quantile{if_suffix}({args})({params})"
+
+
 class ClickHouse(Dialect):
     normalize_functions = None
     null_ordering = "nulls_are_last"
@@ -37,6 +44,8 @@ class Parser(parser.Parser):
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,  # type: ignore
             "MAP": parse_var_map,
+            "QUANTILE": lambda args, params: exp.Quantile(this=params, quantile=args),
+            "QUANTILEIF": lambda args, params: exp.Quantile(this=params, quantile=args),
         }
 
         JOIN_KINDS = {*parser.Parser.JOIN_KINDS, TokenType.ANY, TokenType.ASOF}  # type: ignore
@@ -76,6 +85,7 @@ class Generator(generator.Generator):
             exp.Final: lambda self, e: f"{self.sql(e, 'this')} FINAL",
             exp.Map: lambda self, e: _lower_func(var_map_sql(self, e)),
             exp.VarMap: lambda self, e: _lower_func(var_map_sql(self, e)),
+            exp.Quantile: _quantile_sql,
         }
 
         EXPLICIT_UNION = True
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 308f36385e..6280c5ee73 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -2049,7 +2049,23 @@ def _parse_function(self, functions=None):
             args = self._parse_csv(self._parse_lambda)
 
             if function:
-                this = function(args)
+                params = None
+
+                # Clickhouse supports function calls like foo(x, y)(z), so for these we need to also parse the
+                # second parameter list (i.e. "(z)") and the corresponding function will receive both arg lists.
+                if function.__code__.co_name == "<lambda>" and function.__code__.co_argcount == 2:
+                    index = self._index
+                    self._match_r_paren()
+
+                    if self._match(TokenType.L_PAREN):
+                        params = self._parse_csv(self._parse_lambda)
+                    else:
+                        self._retreat(index)
+
+                    this = function(args, params)
+                else:
+                    this = function(args)
+
                 self.validate_expression(this, args)
             else:
                 this = self.expression(exp.Anonymous, this=this, expressions=args)
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
index c95c967b19..801e1af6dc 100644
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -14,6 +14,7 @@ def test_clickhouse(self):
         self.validate_identity("SELECT * FROM foo LEFT ASOF JOIN bla")
         self.validate_identity("SELECT * FROM foo ASOF JOIN bla")
         self.validate_identity("SELECT * FROM foo ANY JOIN bla")
+        self.validate_identity("SELECT quantile(0.5)(a)")
 
         self.validate_all(
             "SELECT fname, lname, age FROM person ORDER BY age DESC NULLS FIRST, fname ASC NULLS LAST, lname",
@@ -38,3 +39,9 @@ def test_clickhouse(self):
             "SELECT x #! comment",
             write={"": "SELECT x /* comment */"},
         )
+        self.validate_all(
+            "SELECT quantileIf(0.5)(a, true)",
+            write={
+                "clickhouse": "SELECT quantileIf(0.5)(a, TRUE)",
+            },
+        )


Code review comments:
diff_hunk: @@ -2049,7 +2049,23 @@ def _parse_function(self, functions=None):
             args = self._parse_csv(self._parse_lambda)
 
             if function:
-                this = function(args)
+                params = None
+
+                # Clickhouse supports function calls like foo(x, y)(z), so for these we need to also parse the
+                # second parameter list (i.e. "(z)") and the corresponding function will receive both arg lists.
+                if function.__code__.co_name == "<lambda>" and function.__code__.co_argcount == 2:
+                    index = self._index
+                    self._match_r_paren()
path: sqlglot/parser.py
text: @user1:
i don't like doing the retreat here, can we do this another way?

@author:
fixed.

diff_hunk: @@ -2049,7 +2049,23 @@ def _parse_function(self, functions=None):
             args = self._parse_csv(self._parse_lambda)
 
             if function:
-                this = function(args)
+                params = None
+
+                # Clickhouse supports function calls like foo(x, y)(z), so for these we need to also parse the
+                # second parameter list (i.e. "(z)") and the corresponding function will receive both arg lists.
+                if function.__code__.co_name == "<lambda>" and function.__code__.co_argcount == 2:
path: sqlglot/parser.py
text: @author:
The `<lambda>` check was added because without it we would also match functions like `format_time_lambda` which expects two arguments (but the second one is a helper we hard code into the calls, if I'm not mistaken).

The contract here is that we're basically expecting functions like `quantile` in clickhouse to have their own entry in the `FUNCTIONS` map, where the value is a lambda that expects two arguments: one for each argument list.


@user1:
this doesn't seem correct to me, format_time_lambda should return a lambda that only has one argument

@author:
fixed.