Repository: tobymao/sqlglot

Issue:
Clickhouse/Trino/Duckdb max_by() doesn't transpile
Trino and Duckdb have an aggregate function `max_by` which returns the value corresponding to the maximum of the 2nd argument.  Clickhouse calls their function `argMax`.

sqlglot currently ignores this function instead of renaming it.

```python
In [120]: import sqlglot

In [121]: sql = "SELECT max_by(a.id, a.timestamp) FROM a"

In [122]: sqlglot.transpile(sql, read="trino", write="clickhouse")[0]
Out[122]: 'SELECT max_by(a.id, a.timestamp) FROM a'

In [123]: sqlglot.transpile(sql, read="duckdb", write="clickhouse")[0]
Out[123]: 'SELECT max_by(a.id, a.timestamp) FROM a'

In [124]: sqlglot.transpile(sql, read="trino", write="duckdb")[0]
Out[124]: 'SELECT MAX_BY(a.id, a.timestamp) FROM a'

# NOTE this is valid because duckdb has an argmax alias for arg_max (also max_by).
In [125]: sqlglot.transpile("SELECT argMax(a.id, a.timestamp) FROM a", read="clickhouse", write="duckdb")[0]
Out[125]: 'SELECT ARGMAX(a.id, a.timestamp) FROM a'

In [126]: sqlglot.transpile("SELECT argMax(a.id, a.timestamp) FROM a", read="clickhouse", write="trino")[0]
Out[126]: 'SELECT ARGMAX(a.id, a.timestamp) FROM a'

```

**Official Documentation**
Please include links to official SQL documentation related to your issue.

https://clickhouse.com/docs/en/sql-reference/aggregate-functions/reference/argmax
https://trino.io/docs/current/functions/aggregate.html?highlight=max_by#max_by
https://duckdb.org/docs/sql/aggregates.html


Discussion:

Patch:
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
index 1abbef18..a424ea45 100644
--- a/sqlglot/dataframe/sql/functions.py
+++ b/sqlglot/dataframe/sql/functions.py
@@ -84,11 +84,11 @@ def min(col: ColumnOrName) -> Column:
 
 
 def max_by(col: ColumnOrName, ord: ColumnOrName) -> Column:
-    return Column.invoke_anonymous_function(col, "MAX_BY", ord)
+    return Column.invoke_expression_over_column(col, expression.ArgMax, expression=ord)
 
 
 def min_by(col: ColumnOrName, ord: ColumnOrName) -> Column:
-    return Column.invoke_anonymous_function(col, "MIN_BY", ord)
+    return Column.invoke_expression_over_column(col, expression.ArgMin, expression=ord)
 
 
 def count(col: ColumnOrName) -> Column:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 8cb8698b..51baba23 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -8,6 +8,7 @@ from sqlglot import exp, generator, parser, tokens, transforms
 from sqlglot._typing import E
 from sqlglot.dialects.dialect import (
     Dialect,
+    arg_max_or_min_no_count,
     binary_from_function,
     date_add_interval_sql,
     datestrtodate_sql,
@@ -434,6 +435,8 @@ class BigQuery(Dialect):
         TRANSFORMS = {
             **generator.Generator.TRANSFORMS,
             exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
+            exp.ArgMax: arg_max_or_min_no_count("MAX_BY"),
+            exp.ArgMin: arg_max_or_min_no_count("MIN_BY"),
             exp.ArraySize: rename_func("ARRAY_LENGTH"),
             exp.Cast: transforms.preprocess([transforms.remove_precision_parameterized_types]),
             exp.CollateProperty: lambda self, e: f"DEFAULT COLLATE {self.sql(e, 'this')}"
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
index a935cce3..30f728c7 100644
--- a/sqlglot/dialects/clickhouse.py
+++ b/sqlglot/dialects/clickhouse.py
@@ -5,6 +5,7 @@ import typing as t
 from sqlglot import exp, generator, parser, tokens, transforms
 from sqlglot.dialects.dialect import (
     Dialect,
+    arg_max_or_min_no_count,
     inline_array_sql,
     no_pivot_sql,
     rename_func,
@@ -373,6 +374,8 @@ class ClickHouse(Dialect):
             exp.Select: transforms.preprocess([transforms.eliminate_qualify]),
             exp.AnyValue: rename_func("any"),
             exp.ApproxDistinct: rename_func("uniq"),
+            exp.ArgMax: arg_max_or_min_no_count("argMax"),
+            exp.ArgMin: arg_max_or_min_no_count("argMin"),
             exp.Array: inline_array_sql,
             exp.CastToStrType: rename_func("CAST"),
             exp.CurrentDate: lambda self, e: self.func("CURRENT_DATE"),
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index ac6b42cb..739e8d72 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -822,3 +822,13 @@ def generatedasidentitycolumnconstraint_sql(
     start = self.sql(expression, "start") or "1"
     increment = self.sql(expression, "increment") or "1"
     return f"IDENTITY({start}, {increment})"
+
+
+def arg_max_or_min_no_count(name: str) -> t.Callable[[Generator, exp.ArgMax | exp.ArgMin], str]:
+    def _arg_max_or_min_sql(self: Generator, expression: exp.ArgMax | exp.ArgMin) -> str:
+        if expression.args.get("count"):
+            self.unsupported(f"Only two arguments are supported in function {name}.")
+
+        return self.func(name, expression.this, expression.expression)
+
+    return _arg_max_or_min_sql
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
index 7bbe4eec..287e03af 100644
--- a/sqlglot/dialects/duckdb.py
+++ b/sqlglot/dialects/duckdb.py
@@ -6,6 +6,7 @@ from sqlglot import exp, generator, parser, tokens
 from sqlglot.dialects.dialect import (
     Dialect,
     approx_count_distinct_sql,
+    arg_max_or_min_no_count,
     arrow_json_extract_scalar_sql,
     arrow_json_extract_sql,
     binary_from_function,
@@ -249,6 +250,8 @@ class DuckDB(Dialect):
             if e.expressions and e.expressions[0].find(exp.Select)
             else inline_array_sql(self, e),
             exp.ArraySize: rename_func("ARRAY_LENGTH"),
+            exp.ArgMax: arg_max_or_min_no_count("ARG_MAX"),
+            exp.ArgMin: arg_max_or_min_no_count("ARG_MIN"),
             exp.ArraySort: _array_sort_sql,
             exp.ArraySum: rename_func("LIST_SUM"),
             exp.BitwiseXor: rename_func("XOR"),
diff --git a/sqlglot/dialects/hive.py b/sqlglot/dialects/hive.py
index 23a6f3e2..7bff5533 100644
--- a/sqlglot/dialects/hive.py
+++ b/sqlglot/dialects/hive.py
@@ -6,6 +6,7 @@ from sqlglot import exp, generator, parser, tokens, transforms
 from sqlglot.dialects.dialect import (
     Dialect,
     approx_count_distinct_sql,
+    arg_max_or_min_no_count,
     create_with_partitions_sql,
     format_time_lambda,
     if_sql,
@@ -431,6 +432,8 @@ class Hive(Dialect):
             exp.Property: _property_sql,
             exp.AnyValue: rename_func("FIRST"),
             exp.ApproxDistinct: approx_count_distinct_sql,
+            exp.ArgMax: arg_max_or_min_no_count("MAX_BY"),
+            exp.ArgMin: arg_max_or_min_no_count("MIN_BY"),
             exp.ArrayConcat: rename_func("CONCAT"),
             exp.ArrayJoin: lambda self, e: self.func("CONCAT_WS", e.expression, e.this),
             exp.ArraySize: rename_func("SIZE"),
diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py
index 791e82dc..aac368cb 100644
--- a/sqlglot/dialects/presto.py
+++ b/sqlglot/dialects/presto.py
@@ -280,6 +280,8 @@ class Presto(Dialect):
             exp.AnyValue: rename_func("ARBITRARY"),
             exp.ApproxDistinct: _approx_distinct_sql,
             exp.ApproxQuantile: rename_func("APPROX_PERCENTILE"),
+            exp.ArgMax: rename_func("MAX_BY"),
+            exp.ArgMin: rename_func("MIN_BY"),
             exp.Array: lambda self, e: f"ARRAY[{self.expressions(e, flat=True)}]",
             exp.ArrayConcat: rename_func("CONCAT"),
             exp.ArrayContains: rename_func("CONTAINS"),
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 3ad170f2..07be65b9 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -460,6 +460,8 @@ class Snowflake(Dialect):
 
         TRANSFORMS = {
             **generator.Generator.TRANSFORMS,
+            exp.ArgMax: rename_func("MAX_BY"),
+            exp.ArgMin: rename_func("MIN_BY"),
             exp.Array: inline_array_sql,
             exp.ArrayConcat: rename_func("ARRAY_CAT"),
             exp.ArrayJoin: rename_func("ARRAY_TO_STRING"),
diff --git a/sqlglot/dialects/teradata.py b/sqlglot/dialects/teradata.py
index 0aeb0d69..152afa6a 100644
--- a/sqlglot/dialects/teradata.py
+++ b/sqlglot/dialects/teradata.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import typing as t
 
 from sqlglot import exp, generator, parser, tokens, transforms
-from sqlglot.dialects.dialect import Dialect, max_or_greatest, min_or_least
+from sqlglot.dialects.dialect import Dialect, max_or_greatest, min_or_least, rename_func
 from sqlglot.tokens import TokenType
 
 
@@ -169,6 +169,8 @@ class Teradata(Dialect):
 
         TRANSFORMS = {
             **generator.Generator.TRANSFORMS,
+            exp.ArgMax: rename_func("MAX_BY"),
+            exp.ArgMin: rename_func("MIN_BY"),
             exp.Max: max_or_greatest,
             exp.Min: min_or_least,
             exp.Select: transforms.preprocess(
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 2d1ebc2d..5b012b15 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -4189,6 +4189,16 @@ class Abs(Func):
     pass
 
 
+class ArgMax(AggFunc):
+    arg_types = {"this": True, "expression": True, "count": False}
+    _sql_names = ["ARG_MAX", "ARGMAX", "MAX_BY"]
+
+
+class ArgMin(AggFunc):
+    arg_types = {"this": True, "expression": True, "count": False}
+    _sql_names = ["ARG_MIN", "ARGMIN", "MIN_BY"]
+
+
 class ApproxTopK(AggFunc):
     arg_types = {"this": True, "expression": False, "counters": False}
 
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index 4eeada19..58c7ac2d 100644
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -112,12 +112,18 @@ def lineage(
                 column
                 if isinstance(column, int)
                 else next(
-                    i
-                    for i, select in enumerate(scope.expression.selects)
-                    if select.alias_or_name == column
+                    (
+                        i
+                        for i, select in enumerate(scope.expression.selects)
+                        if select.alias_or_name == column or select.is_star
+                    ),
+                    -1,  # mypy will not allow a None here, but a negative index should never be returned
                 )
             )
 
+            if index == -1:
+                raise ValueError(f"Could not find {column} in {scope.expression}")
+
             for s in scope.union_scopes:
                 to_node(index, scope=s, upstream=upstream)
 


Test patch:
diff --git a/tests/dialects/test_presto.py b/tests/dialects/test_presto.py
index 96a25128..ed734b6d 100644
--- a/tests/dialects/test_presto.py
+++ b/tests/dialects/test_presto.py
@@ -558,6 +558,37 @@ class TestPresto(Validator):
             "SELECT SPLIT_TO_MAP('a:1;b:2;a:3', ';', ':', (k, v1, v2) -> CONCAT(v1, v2))"
         )
 
+        self.validate_all(
+            "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+            read={
+                "bigquery": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "clickhouse": "SELECT argMax(a.id, a.timestamp) FROM a",
+                "duckdb": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "snowflake": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "spark": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "teradata": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+            },
+            write={
+                "bigquery": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "clickhouse": "SELECT argMax(a.id, a.timestamp) FROM a",
+                "duckdb": "SELECT ARG_MAX(a.id, a.timestamp) FROM a",
+                "presto": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "snowflake": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "spark": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+                "teradata": "SELECT MAX_BY(a.id, a.timestamp) FROM a",
+            },
+        )
+        self.validate_all(
+            "SELECT MIN_BY(a.id, a.timestamp, 3) FROM a",
+            write={
+                "clickhouse": "SELECT argMin(a.id, a.timestamp) FROM a",
+                "duckdb": "SELECT ARG_MIN(a.id, a.timestamp) FROM a",
+                "presto": "SELECT MIN_BY(a.id, a.timestamp, 3) FROM a",
+                "snowflake": "SELECT MIN_BY(a.id, a.timestamp, 3) FROM a",
+                "spark": "SELECT MIN_BY(a.id, a.timestamp) FROM a",
+                "teradata": "SELECT MIN_BY(a.id, a.timestamp, 3) FROM a",
+            },
+        )
         self.validate_all(
             """JSON '"foo"'""",
             write={
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index 07ede15c..10b56155 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -235,3 +235,28 @@ class TestLineage(unittest.TestCase):
         node = node.downstream[0]
         self.assertEqual(node.name, "t3.my_column")
         self.assertEqual(node.source.sql(), "foo AS t3")
+
+    def test_lineage_cte_union(self) -> None:
+        query = """
+        WITH dataset AS (
+            SELECT *
+            FROM catalog.db.table_a
+
+            UNION
+
+            SELECT *
+            FROM catalog.db.table_b
+        )
+
+        SELECT x, created_at FROM dataset;
+        """
+        node = lineage("x", query)
+
+        self.assertEqual(node.name, "x")
+
+        downstream_a = node.downstream[0]
+        self.assertEqual(downstream_a.name, "0")
+        self.assertEqual(downstream_a.source.sql(), "SELECT * FROM catalog.db.table_a AS table_a")
+        downstream_b = node.downstream[1]
+        self.assertEqual(downstream_b.name, "0")
+        self.assertEqual(downstream_b.source.sql(), "SELECT * FROM catalog.db.table_b AS table_b")