Repository: tobymao/sqlglot

Issue:
Distinct before expression in ArgMax / ArgMin not supported
queries such as
`SELECT MAX_BY(DISTINCT selected_col, filtered_col, 10) FROM table`

are supported in Snowflake, Trino, Clickhouse (using `argMaxDistinct` or `argMax(DISTINCT col, ...)`), probably more dialects.

The following throws an error:
```Python
from sqlglot import parse_one
parse_one("SELECT MAX_BY(DISTINCT selected_col, filtered_col) FROM table")
```

```
ParseError: Required keyword: 'expression' missing for <class 'sqlglot.expressions.ArgMax'>. Line 1, Col: 54.
  SELECT MAX_BY(DISTINCT selected_col, filtered_col) FROM table
```

This impacts ArgMax, ArgMin, maybe more functions?

Discussion:

Solution patch:
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 9e99b473..46080674 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -1184,6 +1184,12 @@ class Parser(metaclass=_Parser):
     KEY_VALUE_DEFINITIONS = (exp.Alias, exp.EQ, exp.PropertyEQ, exp.Slice)
 
     FUNCTION_PARSERS = {
+        **{
+            name: lambda self: self._parse_max_min_by(exp.ArgMax) for name in exp.ArgMax.sql_names()
+        },
+        **{
+            name: lambda self: self._parse_max_min_by(exp.ArgMin) for name in exp.ArgMin.sql_names()
+        },
         "CAST": lambda self: self._parse_cast(self.STRICT_CAST),
         "CEIL": lambda self: self._parse_ceil_floor(exp.Ceil),
         "CONVERT": lambda self: self._parse_convert(self.STRICT_CAST),
@@ -8224,3 +8230,16 @@ class Parser(metaclass=_Parser):
             this=exp.var("FORMAT_NAME"),
             value=self._parse_string() or self._parse_table_parts(),
         )
+
+    def _parse_max_min_by(self, expr_type: t.Type[exp.AggFunc]) -> exp.AggFunc:
+        args: t.List[exp.Expression] = []
+
+        if self._match(TokenType.DISTINCT):
+            args.append(self.expression(exp.Distinct, expressions=[self._parse_assignment()]))
+            self._match(TokenType.COMMA)
+
+        args.extend(self._parse_csv(self._parse_assignment))
+
+        return self.expression(
+            expr_type, this=seq_get(args, 0), expression=seq_get(args, 1), count=seq_get(args, 2)
+        )


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index ed7f4e0a..c7a9f9d1 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -2578,3 +2578,11 @@ SINGLE = TRUE""",
         self.assertEqual(expr.find(exp.Placeholder), exp.Placeholder(this="1"))
         self.validate_identity("SELECT :1, :2")
         self.validate_identity("SELECT :1 + :2")
+
+    def test_max_by_min_by(self):
+        max_by = self.validate_identity("MAX_BY(DISTINCT selected_col, filtered_col)")
+        min_by = self.validate_identity("MIN_BY(DISTINCT selected_col, filtered_col)")
+
+        for node in (max_by, min_by):
+            self.assertEqual(len(node.this.expressions), 1)
+            self.assertIsInstance(node.expression, exp.Column)
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
index ce1a32f1..7ccb2f4f 100644
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -889,3 +889,5 @@ CAST(x AS INT128)
 CAST(x AS UINT128)
 CAST(x AS UINT256)
 SELECT export
+SELECT ARG_MAX(DISTINCT selected_col, filtered_col) FROM table
+SELECT ARG_MIN(DISTINCT selected_col, filtered_col) FROM table
\ No newline at end of file