fix(bigquery): correct percentile and distinct-array transpilation

sqlglot/generators/bigquery.py

@@ -206,11 +206,13 @@ def _array_agg_fold_sql(
     self: BigQueryGenerator, expression: exp.ArrayDistinct | exp.ArraySort
 ) -> str:
     distinct = False
+    ignore_nulls = False
     sort: exp.ArraySort | None = None
     agg: exp.Expr = expression
     while isinstance(agg, (exp.ArrayDistinct, exp.ArraySort)) and not agg.expression:
         if isinstance(agg, exp.ArrayDistinct):
             distinct = True
+            ignore_nulls |= agg.args.get("check_null") is False
         elif sort is None:
             sort = agg
         agg = agg.this
@@ -238,7 +240,7 @@ def _array_agg_fold_sql(
         keys = [exp.Ordered(this=value.copy(), desc=not asc or None, nulls_first=asc)]
 
         # Aggregate ORDER BY only puts NULLs first when ascending, so other placements need a key
-        if nulls_first is not None and nulls_first != asc:
+        if not ignore_nulls and nulls_first is not None and nulls_first != asc:
             if distinct:
                 self.unsupported("ARRAY_AGG(DISTINCT) only puts NULLs first when ascending")
             else:
@@ -251,6 +253,13 @@ def _array_agg_fold_sql(
                 )
         folded = exp.Order(this=folded, expressions=keys)
     agg.set("this", folded)
+    if ignore_nulls:
+        # All-NULL input yields [], but an empty input still yields NULL.
+        agg = exp.If(
+            this=exp.Count(this=exp.Star()).eq(0),
+            true=exp.null(),
+            false=exp.func("COALESCE", exp.IgnoreNulls(this=agg), exp.array()),
+        )
     return self.sql(agg)
 
 
@@ -689,6 +698,13 @@ def frombase_sql(self, expression: exp.FromBase) -> str:
         return self.sql(exp.cast(prefixed, exp.DType.BIGINT, copy=False))
 
     def approxquantile_sql(self, expression: exp.ApproxQuantile) -> str:
+        node: exp.Expr = expression
+        while node.parent and not isinstance(node.parent, (exp.Window, exp.AggFunc, exp.Select)):
+            node = node.parent
+        if isinstance(node.parent, exp.Window) and node.arg_key == "this":
+            self.unsupported("APPROX_QUANTILES does not support window functions")
+            return self.function_fallback_sql(expression)
+
         # The k/n quantile is APPROX_QUANTILES(x, n)[k], e.g. 0.95 -> n = 20, k = 19. n is at most
         # 100,000, so the quantile has at most 5 decimal places, and Decimal comparisons are exact
         quantile = expression.args["quantile"]
@@ -698,7 +714,7 @@ def approxquantile_sql(self, expression: exp.ApproxQuantile) -> str:
             value is None
             or not 0 <= value <= 1
             or round(value, 5) != value
-            or (weight and weight.name != "1")
+            or (weight and (not weight.is_number or weight.to_py() != 1))
         ):
             self.unsupported("APPROX_QUANTILE needs a literal quantile and no weight")
             return self.function_fallback_sql(expression)

sqlglot/parsers/duckdb.py

@@ -138,6 +138,10 @@ class DuckDBParser(parser.Parser):
 
     FUNCTIONS = {
         **{k: v for k, v in parser.Parser.FUNCTIONS.items() if k not in ("DATE_SUB", "GLOB")},
+        **dict.fromkeys(
+            ("ARRAY_DISTINCT", "LIST_DISTINCT"),
+            lambda args: exp.ArrayDistinct(this=seq_get(args, 0), check_null=False),
+        ),
         "ANY_VALUE": lambda args: exp.IgnoreNulls(this=exp.AnyValue.from_arg_list(args)),
         "ARRAY_PREPEND": _build_array_prepend,
         "ARRAY_REVERSE_SORT": _build_list_sort("ARRAY_REVERSE_SORT", desc=True),
@@ -179,7 +183,6 @@ class DuckDBParser(parser.Parser):
         "JSON_EXTRACT_PATH": parser.build_extract_json_with_path(exp.JSONExtract),
         "JSON_EXTRACT_STRING": parser.build_extract_json_with_path(exp.JSONExtractScalar),
         "LIST": exp.ArrayAgg.from_arg_list,
-        "LIST_DISTINCT": exp.ArrayDistinct.from_arg_list,
         "LIST_APPEND": exp.ArrayAppend.from_arg_list,
         "LIST_CONCAT": parser.build_array_concat,
         "LIST_CONTAINS": exp.ArrayContains.from_arg_list,

sqlglot/parsers/trino.py

@@ -5,7 +5,7 @@
 
 from sqlglot import exp, parser
 from sqlglot.helper import ensure_list
-from sqlglot.parsers.presto import PrestoParser
+from sqlglot.parsers.presto import PrestoParser, _build_approx_percentile
 from sqlglot.tokens import TokenType
 
 FRACTIONAL_SECONDS_RE = re.compile(r":\d+\.(\d+)")
@@ -31,6 +31,11 @@ class TrinoParser(PrestoParser):
 
     FUNCTIONS = {
         **PrestoParser.FUNCTIONS,
+        "APPROX_PERCENTILE": lambda args: (
+            exp.ApproxQuantile(this=args[0], weight=args[1], quantile=args[2])
+            if len(args) == 3
+            else _build_approx_percentile(args)
+        ),
         "VERSION": exp.CurrentVersion.from_arg_list,
     }
 

tests/dialects/test_bigquery.py

@@ -3543,6 +3543,38 @@ def test_array_agg(self):
                 },
             )
 
+    def test_array_distinct_fold(self):
+        for source, order in (
+            ("LIST_DISTINCT(LIST(x))", ""),
+            ("ARRAY_DISTINCT(ARRAY_AGG(x))", ""),
+            ("LIST_SORT(LIST_DISTINCT(LIST(x)))", " ORDER BY x"),
+            ("LIST_DISTINCT(LIST_SORT(LIST(x), 'DESC', 'NULLS FIRST'))", " ORDER BY x DESC"),
+        ):
+            with self.subTest(source=source):
+                self.validate_all(
+                    f"IF(COUNT(*) = 0, NULL, COALESCE(ARRAY_AGG(DISTINCT x IGNORE NULLS{order}), []))",
+                    read={"duckdb": source},
+                )
+
+        for values, where in (
+            ("1, NULL, 2, 1", ""),
+            ("NULL", ""),
+            ("1", " WHERE FALSE"),
+            ("'a', NULL, 'b', 'a'", ""),
+            ("CAST(NULL AS STRING)", ""),
+            ("'a'", " WHERE FALSE"),
+        ):
+            with self.subTest(values=values, where=where):
+                self.validate_all(
+                    "SELECT ARRAY_LENGTH(IF(COUNT(*) = 0, NULL, "
+                    "COALESCE(ARRAY_AGG(DISTINCT x IGNORE NULLS), []))) "
+                    f"FROM UNNEST([{values}]) AS x{where}",
+                    read={
+                        "duckdb": "SELECT ARRAY_LENGTH(LIST_DISTINCT(LIST(x))) "
+                        f"FROM UNNEST([{values}]) AS t(x){where}"
+                    },
+                )
+
     def test_array_concat(self):
         self.validate_all(
             "WITH x AS ( SELECT 1 AS id), test_cte AS ( SELECT ARRAY_CONCAT(( SELECT id FROM x WHERE FALSE)) AS result ) SELECT * FROM test_cte;",
@@ -4194,6 +4226,46 @@ def test_approx_quantiles(self):
         self.validate_identity("APPROX_QUANTILES(x, 2 IGNORE NULLS)")
         self.validate_identity("APPROX_QUANTILES(DISTINCT x, 2 RESPECT NULLS)")
 
+    def test_approx_quantile_window(self):
+        for source, target in (
+            ("ORDER BY APPROX_PERCENTILE(x, 0.5)", "ORDER BY APPROX_QUANTILES(x, 2)[1] NULLS LAST"),
+            (
+                "PARTITION BY APPROX_PERCENTILE(x, 0.5) ORDER BY y",
+                "PARTITION BY APPROX_QUANTILES(x, 2)[1] ORDER BY y NULLS LAST",
+            ),
+            (
+                "ORDER BY APPROX_PERCENTILE(x, 0.5) FILTER(WHERE x > 0)",
+                "ORDER BY APPROX_QUANTILES(IF(x > 0, x, NULL), 2)[1] NULLS LAST",
+            ),
+        ):
+            with self.subTest(source=source):
+                self.validate_all(
+                    f"SELECT ROW_NUMBER() OVER ({target}) FROM t GROUP BY y",
+                    read={"trino": f"SELECT ROW_NUMBER() OVER ({source}) FROM t GROUP BY y"},
+                )
+
+        for sql in (
+            "APPROX_PERCENTILE(x, 0.5) OVER ()",
+            "APPROX_PERCENTILE(x, 0.5) OVER (PARTITION BY y)",
+            "APPROX_PERCENTILE(x, 0.5) FILTER(WHERE x > 0) OVER ()",
+        ):
+            with (
+                self.subTest(sql=sql),
+                self.assertRaisesRegex(
+                    UnsupportedError, "APPROX_QUANTILES does not support window functions"
+                ),
+            ):
+                parse_one(sql, read="trino").sql("bigquery", unsupported_level=ErrorLevel.RAISE)
+
+        self.validate_all(
+            "SELECT SUM(APPROX_QUANTILES(x, 2)[1]) OVER () FROM t",
+            read={"trino": "SELECT SUM(APPROX_PERCENTILE(x, 0.5)) OVER () FROM t"},
+        )
+        self.validate_all(
+            "SELECT SUM((SELECT APPROX_QUANTILES(x, 2)[1] FROM t)) OVER () FROM u",
+            read={"trino": "SELECT SUM((SELECT APPROX_PERCENTILE(x, 0.5) FROM t)) OVER () FROM u"},
+        )
+
     def test_approx_quantiles_to_duckdb(self):
         self.validate_all(
             "APPROX_QUANTILES(x, 1)",

tests/dialects/test_presto.py

@@ -7,6 +7,16 @@
 class TestPresto(Validator):
     dialect = "presto"
 
+    def test_approx_percentile(self):
+        expression = self.validate_identity("APPROX_PERCENTILE(x, 0.5, 0.01)")
+        self.assertEqual(expression.args["quantile"], exp.Literal.number("0.5"))
+        self.assertEqual(expression.args["accuracy"], exp.Literal.number("0.01"))
+        self.assertIsNone(expression.args.get("weight"))
+        self.validate_all(
+            "APPROX_PERCENTILE(x, 0.5, 0.01)",
+            write={"bigquery": "APPROX_QUANTILES(x, 2)[1]"},
+        )
+
     def test_temporal_precision(self):
         message = "Presto only supports millisecond precision for temporal types"
         for kind, value in (("TIMESTAMP", "2020-01-01 00:00:00"), ("TIME", "12:00:00")):

tests/dialects/test_trino.py

@@ -40,6 +40,7 @@ def test_bigquery_functions(self):
             write={"bigquery": "CAST(REGEXP_REPLACE('-ff', '^([+-]?)', '\\\\10x') AS INT64)"},
         )
         for sql, expected in (
+            ("APPROX_PERCENTILE(col, 1, 0.5)", "APPROX_QUANTILES(col, 2)[1]"),
             ("APPROX_PERCENTILE(col, 1, 0.5, 0.001)", "APPROX_QUANTILES(col, 2)[1]"),
             # APPROX_QUANTILES takes at most 100,000 buckets
             ("APPROX_PERCENTILE(col, 0.00001)", "APPROX_QUANTILES(col, 100000)[1]"),
@@ -55,6 +56,21 @@ def test_bigquery_functions(self):
 
         for sql, expected, warning in (
             ("FROM_BASE(s, 2)", "FROM_BASE(s, 2)", "FROM_BASE is only supported for base 16"),
+            (
+                "APPROX_PERCENTILE(col, 2, 0.5)",
+                "APPROX_QUANTILE(col, 0.5, 2)",
+                "APPROX_QUANTILE needs a literal quantile and no weight",
+            ),
+            (
+                "APPROX_PERCENTILE(col, w, 0.5)",
+                "APPROX_QUANTILE(col, 0.5, w)",
+                "APPROX_QUANTILE needs a literal quantile and no weight",
+            ),
+            (
+                'APPROX_PERCENTILE(col, "1", 0.5)',
+                "APPROX_QUANTILE(col, 0.5, `1`)",
+                "APPROX_QUANTILE needs a literal quantile and no weight",
+            ),
             (
                 "APPROX_PERCENTILE(col, 2, 0.5, 0.01)",
                 "APPROX_QUANTILE(col, 0.5, 0.01, 2)",
@@ -80,6 +96,19 @@ def test_bigquery_functions(self):
             with self.subTest(sql), self.assertLogs(generator_logger, level="WARNING") as cm:
                 self.assertEqual(parse_one(sql, read="trino").sql("bigquery"), expected)
             self.assertIn(warning, cm.output[0])
+            with self.assertRaises(UnsupportedError):
+                parse_one(sql, read="trino").sql("bigquery", unsupported_level=ErrorLevel.RAISE)
+
+    def test_approx_percentile(self):
+        for weight in ("1", "w", "w + 1"):
+            for quantile in ("0.5", "ARRAY[0.25, 0.5]"):
+                with self.subTest(weight=weight, quantile=quantile):
+                    expression = self.validate_identity(
+                        f"APPROX_PERCENTILE(x, {weight}, {quantile})"
+                    )
+                    self.assertEqual(expression.args["weight"], self.parse_one(weight))
+                    self.assertEqual(expression.args["quantile"], self.parse_one(quantile))
+                    self.assertIsNone(expression.args.get("accuracy"))
 
     def test_decimal_beyond_numeric_range(self):
         # BigQuery's NUMERIC holds at most 29 integer digits and 9 fractional digits