Repository: tobymao/sqlglot

Issue:
Generated SQL for Array Aggregations with Order By doesn't work in BigQuery
Hi there, I ran into an [issue](https://github.com/ibis-project/ibis/issues/10236) using Ibis and I was able to trace the problem back `sqlglot`.

A minimal (ish) reproduction is

```
>>> from sqlglot import select, expressions as exp
>>>
>>> query = (
...     select(
...         exp.Column(this="color"),
...         exp.ArrayAgg(
...             this=exp.Order(
...                 this=exp.Column(this="id"),
...                 expressions=[
...                     exp.Ordered(
...                         this=exp.Column(this="id"),
...                     )
...                 ]
...             ),
...         ).as_("ids"),
...     )
...     .from_("colors")
...     .group_by(exp.Literal.number(1))
... )
>>>
>>> print(query.sql(dialect="bigquery"))
SELECT color, ARRAY_AGG(id ORDER BY id NULLS LAST) AS ids FROM colors GROUP BY 1
```

which I ran against main. I also made a table in BigQuery 

```
CREATE TABLE `my-project.my_dataset.colors` AS (
  SELECT 1 AS id, 'red' AS color
  UNION ALL
  SELECT 2 AS id, 'red' AS color
  UNION ALL
  SELECT 3 AS id, 'blue' AS color
)
```

So I could actually run the query. I get

```
BadRequest: 400 POST https://bigquery.googleapis.com/bigquery/v2/projects/my-project/queries?prettyPrint=false: NULLS LAST not supported with ascending sort order in aggregate functions.
```

I fudged with the table names (project and dataset) but that's it. I tried to find some documentation but I couldn't.

edit - Relevant documentation for ARRAY_AGG can be found [here](https://cloud.google.com/bigquery/docs/reference/standard-sql/aggregate_functions#array_agg)

Discussion:
@yjabri: Would it be possible to have this released in `25.20.x`? The current version of Ibis (9.5.0) has `sqlglot = ">=23.4,<25.21"`@georgesittas: This has already been released with https://github.com/tobymao/sqlglot/blob/main/CHANGELOG.md#v25241---2024-10-01@yjabri: Looks like it was in v25.24.1. Since Ibis currently has a strict upper limit of 25.21 I was wondering if it would be possible to backport the patch on top of the latest 25.20.x release 25.20.2.

edit - looks like the latest tag in 25.20.x is 25.20.2 not 25.20.1. I only ask because I noticed that [25.20.2](https://pypi.org/project/sqlglot/#history) was released on Sept 24, after the release of 25.22.0.@tobymao: sorry we don’t offer that@yjabri: All good, thanks for the quick responses.

Solution patch:
diff --git a/sqlglot/dialects/athena.py b/sqlglot/dialects/athena.py
index b1ea3efb64..108c7643f3 100644
--- a/sqlglot/dialects/athena.py
+++ b/sqlglot/dialects/athena.py
@@ -48,7 +48,7 @@ def _location_property_sql(self: Athena.Generator, e: exp.LocationProperty):
             ),
             None,
         )
-        if table_type_property and table_type_property.text("value") == "iceberg":
+        if table_type_property and table_type_property.text("value").lower() == "iceberg":
             prop_name = "location"
 
     return f"{prop_name}={self.sql(e, 'this')}"
@@ -132,6 +132,7 @@ class Generator(Trino.Generator):
         TRANSFORMS = {
             **Trino.Generator.TRANSFORMS,
             exp.FileFormatProperty: lambda self, e: f"format={self.sql(e, 'this')}",
+            exp.PartitionedByProperty: lambda self, e: f"partitioned_by={self.sql(e, 'this')}",
             exp.LocationProperty: _location_property_sql,
         }
 
diff --git a/sqlglot/dialects/spark2.py b/sqlglot/dialects/spark2.py
index 0c62a7b2d5..15c2d2d231 100644
--- a/sqlglot/dialects/spark2.py
+++ b/sqlglot/dialects/spark2.py
@@ -13,7 +13,7 @@
     unit_to_str,
 )
 from sqlglot.dialects.hive import Hive
-from sqlglot.helper import seq_get
+from sqlglot.helper import seq_get, ensure_list
 from sqlglot.transforms import (
     preprocess,
     remove_unique_constraints,
@@ -21,6 +21,11 @@
     move_schema_columns_to_partitioned_by,
 )
 
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+
+    from sqlglot.optimizer.annotate_types import TypeAnnotator
+
 
 def _map_sql(self: Spark2.Generator, expression: exp.Map) -> str:
     keys = expression.args.get("keys")
@@ -110,10 +115,48 @@ def temporary_storage_provider(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def _annotate_by_similar_args(
+    self: TypeAnnotator, expression: E, *args: str, target_type: exp.DataType | exp.DataType.Type
+) -> E:
+    """
+    Infers the type of the expression according to the following rules:
+    - If all args are of the same type OR any arg is of target_type, the expr is inferred as such
+    - If any arg is of UNKNOWN type and none of target_type, the expr is inferred as UNKNOWN
+    """
+    self._annotate_args(expression)
+
+    expressions: t.List[exp.Expression] = []
+    for arg in args:
+        arg_expr = expression.args.get(arg)
+        expressions.extend(expr for expr in ensure_list(arg_expr) if expr)
+
+    last_datatype = None
+
+    has_unknown = False
+    for expr in expressions:
+        if expr.is_type(exp.DataType.Type.UNKNOWN):
+            has_unknown = True
+        elif expr.is_type(target_type):
+            has_unknown = False
+            last_datatype = target_type
+            break
+        else:
+            last_datatype = expr.type
+
+    self._set_type(expression, exp.DataType.Type.UNKNOWN if has_unknown else last_datatype)
+    return expression
+
+
 class Spark2(Hive):
     ANNOTATORS = {
         **Hive.ANNOTATORS,
         exp.Substring: lambda self, e: self._annotate_by_args(e, "this"),
+        exp.Concat: lambda self, e: _annotate_by_similar_args(
+            self, e, "expressions", target_type=exp.DataType.Type.TEXT
+        ),
+        exp.Pad: lambda self, e: _annotate_by_similar_args(
+            self, e, "this", "fill_pattern", target_type=exp.DataType.Type.TEXT
+        ),
     }
 
     class Parser(Hive.Parser):
diff --git a/sqlglot/dialects/sqlite.py b/sqlglot/dialects/sqlite.py
index 6737c1c112..aa8df3e535 100644
--- a/sqlglot/dialects/sqlite.py
+++ b/sqlglot/dialects/sqlite.py
@@ -251,7 +251,7 @@ def datediff_sql(self, expression: exp.DateDiff) -> str:
             elif unit == "NANOSECOND":
                 sql = f"{sql} * 8640000000000.0"
             else:
-                self.unsupported("DATEDIFF unsupported for '{unit}'.")
+                self.unsupported(f"DATEDIFF unsupported for '{unit}'.")
 
             return f"CAST({sql} AS INTEGER)"
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 1f85126b7f..bce7c0b514 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -203,7 +203,8 @@ class Generator(metaclass=_Generator):
     }
 
     # Whether null ordering is supported in order by
-    # True: Full Support, None: No support, False: No support in window specifications
+    # True: Full Support, None: No support, False: No support for certain cases
+    # such as window specifications, aggregate functions etc
     NULL_ORDERING_SUPPORTED: t.Optional[bool] = True
 
     # Whether ignore nulls is inside the agg or outside.
@@ -2345,6 +2346,18 @@ def ordered_sql(self, expression: exp.Ordered) -> str:
                     f"'{nulls_sort_change.strip()}' translation not supported in window functions"
                 )
                 nulls_sort_change = ""
+            elif (
+                self.NULL_ORDERING_SUPPORTED is False
+                and (isinstance(expression.find_ancestor(exp.AggFunc, exp.Select), exp.AggFunc))
+                and (
+                    (asc and nulls_sort_change == " NULLS LAST")
+                    or (desc and nulls_sort_change == " NULLS FIRST")
+                )
+            ):
+                self.unsupported(
+                    f"'{nulls_sort_change.strip()}' translation not supported for aggregate functions with {sort_order} sort order"
+                )
+                nulls_sort_change = ""
             elif self.NULL_ORDERING_SUPPORTED is None:
                 if expression.this.is_int:
                     self.unsupported(


Test patch:
diff --git a/tests/dialects/test_athena.py b/tests/dialects/test_athena.py
index ca91d4a1c1..ef96938b10 100644
--- a/tests/dialects/test_athena.py
+++ b/tests/dialects/test_athena.py
@@ -62,8 +62,12 @@ def test_ddl(self):
 
         # CTAS goes to the Trino engine, where the table properties cant be encased in single quotes like they can for Hive
         # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html#ctas-table-properties
+        # They're also case sensitive and need to be lowercase, otherwise you get eg "Table properties [FORMAT] are not supported."
         self.validate_identity(
-            "CREATE TABLE foo WITH (table_type='ICEBERG', external_location='s3://foo/') AS SELECT * FROM a"
+            "CREATE TABLE foo WITH (table_type='ICEBERG', location='s3://foo/', format='orc', partitioning=ARRAY['bucket(id, 5)']) AS SELECT * FROM a"
+        )
+        self.validate_identity(
+            "CREATE TABLE foo WITH (table_type='HIVE', external_location='s3://foo/', format='parquet', partitioned_by=ARRAY['ds']) AS SELECT * FROM a"
         )
         self.validate_identity(
             "CREATE TABLE foo AS WITH foo AS (SELECT a, b FROM bar) SELECT * FROM foo"
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index e2adfeabce..d854165f2c 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -1985,3 +1985,17 @@ def test_range_type(self):
         self.validate_identity(
             "SELECT RANGE(CAST('2022-10-01 14:53:27 America/Los_Angeles' AS TIMESTAMP), CAST('2022-10-01 16:00:00 America/Los_Angeles' AS TIMESTAMP))"
         )
+
+    def test_null_ordering(self):
+        # Aggregate functions allow "NULLS FIRST" only with ascending order and
+        # "NULLS LAST" only with descending
+        for sort_order, null_order in (("ASC", "NULLS LAST"), ("DESC", "NULLS FIRST")):
+            self.validate_all(
+                f"SELECT color, ARRAY_AGG(id ORDER BY id {sort_order}) AS ids FROM colors GROUP BY 1",
+                read={
+                    "": f"SELECT color, ARRAY_AGG(id ORDER BY id {sort_order} {null_order}) AS ids FROM colors GROUP BY 1"
+                },
+                write={
+                    "bigquery": f"SELECT color, ARRAY_AGG(id ORDER BY id {sort_order}) AS ids FROM colors GROUP BY 1",
+                },
+            )
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 857ba1aab5..30093bb194 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1378,25 +1378,41 @@ def gen_expr(depth: int) -> exp.Expression:
         self.assertEqual(18, normalization_distance(gen_expr(3), max_=100))
         self.assertEqual(110, normalization_distance(gen_expr(10), max_=100))
 
-    def test_custom_annotators(self):
+    def test_spark_annotators(self):
+        """Test Spark annotators, mainly built-in string/binary functions"""
+
+        spark_schema = {"tbl": {"bin_col": "BINARY", "str_col": "STRING"}}
+
+        from sqlglot.dialects import Dialect
+
+        def _assert_func_return_type(func: str, dialect: str, target_type: str):
+            ast = parse_one(f"SELECT {func} FROM tbl", read=dialect)
+            annotators = Dialect.get_or_raise(dialect).ANNOTATORS
+            annotated = annotate_types(ast, annotators=annotators, schema=spark_schema)
+
+            self.assertEqual(
+                annotated.expressions[0].type.sql(dialect),
+                exp.DataType.build(target_type).sql(dialect),
+            )
+
+        str_col, bin_col = "tbl.str_col", "tbl.bin_col"
+
         # In Spark hierarchy, SUBSTRING result type is dependent on input expr type
         for dialect in ("spark2", "spark", "databricks"):
-            for expr_type_pair in (
-                ("col", "STRING"),
-                ("col", "BINARY"),
-                ("'str_literal'", "STRING"),
-                ("CAST('str_literal' AS BINARY)", "BINARY"),
-            ):
-                with self.subTest(
-                    f"Testing {dialect}'s SUBSTRING() result type for {expr_type_pair}"
-                ):
-                    expr, type = expr_type_pair
-                    ast = parse_one(f"SELECT substring({expr}, 2, 3) AS x FROM tbl", read=dialect)
-
-                    subst_type = (
-                        optimizer.optimize(ast, schema={"tbl": {"col": type}}, dialect=dialect)
-                        .expressions[0]
-                        .type
-                    )
-
-                    self.assertEqual(subst_type.sql(dialect), exp.DataType.build(type).sql(dialect))
+            _assert_func_return_type(f"SUBSTRING({str_col}, 0, 0)", dialect, "STRING")
+            _assert_func_return_type(f"SUBSTRING({bin_col}, 0, 0)", dialect, "BINARY")
+
+            _assert_func_return_type(f"CONCAT({bin_col}, {bin_col})", dialect, "BINARY")
+            _assert_func_return_type(f"CONCAT({bin_col}, {str_col})", dialect, "STRING")
+            _assert_func_return_type(f"CONCAT({str_col}, {bin_col})", dialect, "STRING")
+            _assert_func_return_type(f"CONCAT({str_col}, {str_col})", dialect, "STRING")
+
+            _assert_func_return_type(f"CONCAT({str_col}, foo)", dialect, "STRING")
+            _assert_func_return_type(f"CONCAT({bin_col}, bar)", dialect, "UNKNOWN")
+            _assert_func_return_type("CONCAT(foo, bar)", dialect, "UNKNOWN")
+
+            for func in ("LPAD", "RPAD"):
+                _assert_func_return_type(f"{func}({bin_col}, 1, {bin_col})", dialect, "BINARY")
+                _assert_func_return_type(f"{func}({bin_col}, 1, {str_col})", dialect, "STRING")
+                _assert_func_return_type(f"{func}({str_col}, 1, {bin_col})", dialect, "STRING")
+                _assert_func_return_type(f"{func}({str_col}, 1, {str_col})", dialect, "STRING")