Fix(athena): Generate PartitionedByProperty correctly on CTAS for an Iceberg table (#4654)

Author: erindru

sqlglot/dialects/athena.py

@@ -32,6 +32,18 @@ def _generate_as_hive(expression: exp.Expression) -> bool:
     return False
 
 
+def _is_iceberg_table(properties: exp.Properties) -> bool:
+    table_type_property = next(
+        (
+            p
+            for p in properties.expressions
+            if isinstance(p, exp.Property) and p.name == "table_type"
+        ),
+        None,
+    )
+    return bool(table_type_property and table_type_property.text("value").lower() == "iceberg")
+
+
 def _location_property_sql(self: Athena.Generator, e: exp.LocationProperty):
     # If table_type='iceberg', the LocationProperty is called 'location'
     # Otherwise, it's called 'external_location'
@@ -40,20 +52,25 @@ def _location_property_sql(self: Athena.Generator, e: exp.LocationProperty):
     prop_name = "external_location"
 
     if isinstance(e.parent, exp.Properties):
-        table_type_property = next(
-            (
-                p
-                for p in e.parent.expressions
-                if isinstance(p, exp.Property) and p.name == "table_type"
-            ),
-            None,
-        )
-        if table_type_property and table_type_property.text("value").lower() == "iceberg":
+        if _is_iceberg_table(e.parent):
             prop_name = "location"
 
     return f"{prop_name}={self.sql(e, 'this')}"
 
 
+def _partitioned_by_property_sql(self: Athena.Generator, e: exp.PartitionedByProperty):
+    # If table_type='iceberg' then the table property for partitioning is called 'partitioning'
+    # If table_type='hive' it's called 'partitioned_by'
+    # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html#ctas-table-properties
+
+    prop_name = "partitioned_by"
+    if isinstance(e.parent, exp.Properties):
+        if _is_iceberg_table(e.parent):
+            prop_name = "partitioning"
+
+    return f"{prop_name}={self.sql(e, 'this')}"
+
+
 class Athena(Trino):
     """
     Over the years, it looks like AWS has taken various execution engines, bolted on AWS-specific modifications and then
@@ -132,7 +149,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.PartitionedByProperty: _partitioned_by_property_sql,
             exp.LocationProperty: _location_property_sql,
         }
 

tests/dialects/test_athena.py

@@ -198,7 +198,8 @@ def test_dml_quoting(self):
 
     def test_ctas(self):
         # Hive tables use 'external_location' to specify the table location, Iceberg tables use 'location' to specify the table location
-        # The 'table_type' property is used to determine if it's a Hive or an Iceberg table
+        # In addition, Hive tables used 'partitioned_by' to specify the partition fields and Iceberg tables use 'partitioning' to specify the partition fields
+        # The 'table_type' property is used to determine if it's a Hive or an Iceberg table. If it's omitted, it defaults to Hive
         # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html#ctas-table-properties
         ctas_hive = exp.Create(
             this=exp.to_table("foo.bar"),
@@ -207,13 +208,16 @@ def test_ctas(self):
                 expressions=[
                     exp.FileFormatProperty(this=exp.Literal.string("parquet")),
                     exp.LocationProperty(this=exp.Literal.string("s3://foo")),
+                    exp.PartitionedByProperty(
+                        this=exp.Schema(expressions=[exp.to_column("partition_col")])
+                    ),
                 ]
             ),
             expression=exp.select("1"),
         )
         self.assertEqual(
             ctas_hive.sql(dialect=self.dialect, identify=True),
-            "CREATE TABLE \"foo\".\"bar\" WITH (format='parquet', external_location='s3://foo') AS SELECT 1",
+            "CREATE TABLE \"foo\".\"bar\" WITH (format='parquet', external_location='s3://foo', partitioned_by=ARRAY['partition_col']) AS SELECT 1",
         )
 
         ctas_iceberg = exp.Create(
@@ -223,11 +227,14 @@ def test_ctas(self):
                 expressions=[
                     exp.Property(this=exp.var("table_type"), value=exp.Literal.string("iceberg")),
                     exp.LocationProperty(this=exp.Literal.string("s3://foo")),
+                    exp.PartitionedByProperty(
+                        this=exp.Schema(expressions=[exp.to_column("partition_col")])
+                    ),
                 ]
             ),
             expression=exp.select("1"),
         )
         self.assertEqual(
             ctas_iceberg.sql(dialect=self.dialect, identify=True),
-            "CREATE TABLE \"foo\".\"bar\" WITH (table_type='iceberg', location='s3://foo') AS SELECT 1",
+            "CREATE TABLE \"foo\".\"bar\" WITH (table_type='iceberg', location='s3://foo', partitioning=ARRAY['partition_col']) AS SELECT 1",
         )