Repository: tobymao/sqlglot

Issue:
Decimal-Casts are wrongly removed
Optimizing SQL removes casts, resulting in wrong / failing queries. Casting a decimal to a smaller decimal will usually result in an exception, or null depending on the platform. Running the following query on Databricks will write null, if the value is larger than a `decimal(4,2)`.
```sql
insert into tab2 (val1, val2)
select cast(x as decimal(4,2)), try_cast(x as decimal(4,2))
from tab
```

When optimizing the query, all casts are removed. This causes the query to always fail on Databricks (`[CAST_OVERFLOW_IN_TABLE_INSERT] Fail to assign a value of "DECIMAL(10,6)" type to the "DECIMAL(4,2)" type column or variable ... due to an overflow.`)
```sql
SELECT
  `tab`.`x` AS `x`,
  `tab`.`x` AS `x`
FROM `tab` AS `tab`
```

```python
import sqlglot
from sqlglot.optimizer import optimize

print(
    optimize(
        sqlglot.parse_one(
            """
insert into tab2 (val1, val2)
select cast(x as decimal(4,2)), try_cast(x as decimal(4,2))
from tab
        """,
            read="databricks",
        ),
        schema={"tab": {"x": "decimal(10,6)"}},
    ).sql(pretty=True, dialect="databricks")
)
```

Discussion:

Solution patch:
diff --git a/sqlglot/dialects/hive.py b/sqlglot/dialects/hive.py
index 86c9e06e..6c2446d0 100644
--- a/sqlglot/dialects/hive.py
+++ b/sqlglot/dialects/hive.py
@@ -610,6 +610,8 @@ class Hive(Dialect):
             e: f"({self.expressions(e, 'this', indent=False)})",
             exp.NotForReplicationColumnConstraint: lambda *_: "",
             exp.OnProperty: lambda *_: "",
+            exp.PartitionedByBucket: lambda self, e: self.func("BUCKET", e.expression, e.this),
+            exp.PartitionByTruncate: lambda self, e: self.func("TRUNCATE", e.expression, e.this),
             exp.PrimaryKeyColumnConstraint: lambda *_: "PRIMARY KEY",
             exp.WeekOfYear: rename_func("WEEKOFYEAR"),
             exp.DayOfMonth: rename_func("DAYOFMONTH"),
diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py
index 45d64278..496d53e0 100644
--- a/sqlglot/dialects/presto.py
+++ b/sqlglot/dialects/presto.py
@@ -57,8 +57,12 @@ def _no_sort_array(self: Presto.Generator, expression: exp.SortArray) -> str:
 
 def _schema_sql(self: Presto.Generator, expression: exp.Schema) -> str:
     if isinstance(expression.parent, exp.PartitionedByProperty):
-        columns = ", ".join(f"'{c.name}'" for c in expression.expressions)
-        return f"ARRAY[{columns}]"
+        partition_exprs = [
+            self.sql(c) if isinstance(c, (exp.Func, exp.Property)) else self.sql(c, "this")
+            for c in expression.expressions
+        ]
+
+        return self.sql(exp.Array(expressions=[exp.Literal.string(c) for c in partition_exprs]))
 
     if expression.parent:
         for schema in expression.parent.find_all(exp.Schema):
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 9faea823..e08c6098 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2917,6 +2917,14 @@ class PartitionedByProperty(Property):
     arg_types = {"this": True}
 
 
+class PartitionedByBucket(Property):
+    arg_types = {"this": True, "expression": True}
+
+
+class PartitionByTruncate(Property):
+    arg_types = {"this": True, "expression": True}
+
+
 # https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/CREATE_TABLE/
 class PartitionByRangeProperty(Property):
     arg_types = {"partition_expressions": True, "create_expressions": True}
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index d1d14519..6664c7a2 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -170,6 +170,8 @@ class Generator(metaclass=_Generator):
         exp.Operator: lambda self, e: self.binary(e, ""),  # The operator is produced in `binary`
         exp.OutputModelProperty: lambda self, e: f"OUTPUT{self.sql(e, 'this')}",
         exp.PathColumnConstraint: lambda self, e: f"PATH {self.sql(e, 'this')}",
+        exp.PartitionedByBucket: lambda self, e: self.func("BUCKET", e.this, e.expression),
+        exp.PartitionByTruncate: lambda self, e: self.func("TRUNCATE", e.this, e.expression),
         exp.PivotAny: lambda self, e: f"ANY{self.sql(e, 'this')}",
         exp.ProjectionPolicyColumnConstraint: lambda self,
         e: f"PROJECTION POLICY {self.sql(e, 'this')}",
diff --git a/sqlglot/optimizer/canonicalize.py b/sqlglot/optimizer/canonicalize.py
index e5632dab..d654867d 100644
--- a/sqlglot/optimizer/canonicalize.py
+++ b/sqlglot/optimizer/canonicalize.py
@@ -93,15 +93,18 @@ def remove_redundant_casts(expression: exp.Expression) -> exp.Expression:
     if (
         isinstance(expression, exp.Cast)
         and expression.this.type
-        and expression.to.this == expression.this.type.this
+        and expression.to == expression.this.type
     ):
         return expression.this
+
     if (
         isinstance(expression, (exp.Date, exp.TsOrDsToDate))
         and expression.this.type
         and expression.this.type.this == exp.DataType.Type.DATE
+        and not expression.this.type.expressions
     ):
         return expression.this
+
     return expression
 
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 94543b75..edd4f8b1 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -1100,8 +1100,32 @@ class Parser(metaclass=_Parser):
         "WITH": lambda self: self.expression(
             exp.Properties, expressions=self._parse_wrapped_properties()
         ),
+        "BUCKET": lambda self: self._parse_partitioned_by_bucket_or_truncate(),
+        "TRUNCATE": lambda self: self._parse_partitioned_by_bucket_or_truncate(),
     }
 
+    def _parse_partitioned_by_bucket_or_truncate(self) -> exp.Expression:
+        klass = (
+            exp.PartitionedByBucket
+            if self._prev.text.upper() == "BUCKET"
+            else exp.PartitionByTruncate
+        )
+
+        args = self._parse_wrapped_csv(lambda: self._parse_primary() or self._parse_column())
+        this, expression = seq_get(args, 0), seq_get(args, 1)
+
+        if isinstance(this, exp.Literal):
+            # Check for Iceberg partition transforms (bucket / truncate) and ensure their arguments are in the right order
+            #  - For Hive, it's `bucket(<num buckets>, <col name>)` or `truncate(<num_chars>, <col_name>)`
+            #  - For Trino, it's reversed - `bucket(<col name>, <num buckets>)` or `truncate(<col_name>, <num_chars>)`
+            # Both variants are canonicalized in the latter i.e `bucket(<col name>, <num buckets>)`
+            #
+            # Hive ref: https://docs.aws.amazon.com/athena/latest/ug/querying-iceberg-creating-tables.html#querying-iceberg-partitioning
+            # Trino ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html#ctas-table-properties
+            this, expression = expression, this
+
+        return self.expression(klass, this=this, expression=expression)
+
     ALTER_PARSERS = {
         "ADD": lambda self: self._parse_alter_table_add(),
         "AS": lambda self: self._parse_select(),
@@ -1132,6 +1156,8 @@ class Parser(metaclass=_Parser):
         "PRIMARY KEY",
         "UNIQUE",
         "WATERMARK",
+        "BUCKET",
+        "TRUNCATE",
     }
 
     NO_PAREN_FUNCTION_PARSERS = {


Test patch:
diff --git a/tests/dialects/test_athena.py b/tests/dialects/test_athena.py
index a2c798c2..2cae90c9 100644
--- a/tests/dialects/test_athena.py
+++ b/tests/dialects/test_athena.py
@@ -202,6 +202,67 @@ class TestAthena(Validator):
             identify=True,
         )
 
+    def test_create_table(self):
+        # There are two CREATE TABLE syntaxes
+        # Both hit Athena's Hive engine but creating an Iceberg table is different from creating a normal Hive table
+
+        table_schema = exp.Schema(
+            this=exp.to_table("foo.bar"),
+            expressions=[
+                exp.ColumnDef(this=exp.to_identifier("a"), kind=exp.DataType.build("int")),
+                exp.ColumnDef(this=exp.to_identifier("b"), kind=exp.DataType.build("varchar")),
+            ],
+        )
+
+        # Hive tables - CREATE EXTERNAL TABLE
+        ct_hive = exp.Create(
+            this=table_schema,
+            kind="TABLE",
+            properties=exp.Properties(
+                expressions=[
+                    exp.ExternalProperty(),
+                    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")])
+                    ),
+                ]
+            ),
+        )
+        self.assertEqual(
+            ct_hive.sql(dialect=self.dialect, identify=True),
+            "CREATE EXTERNAL TABLE `foo`.`bar` (`a` INT, `b` STRING) STORED AS PARQUET LOCATION 's3://foo' PARTITIONED BY (`partition_col`)",
+        )
+
+        # Iceberg tables - CREATE TABLE... TBLPROPERTIES ('table_type'='iceberg')
+        # no EXTERNAL keyword and the 'table_type=iceberg' property must be set
+        # ref: https://docs.aws.amazon.com/athena/latest/ug/querying-iceberg-creating-tables.html#querying-iceberg-partitioning
+        ct_iceberg = exp.Create(
+            this=table_schema,
+            kind="TABLE",
+            properties=exp.Properties(
+                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"),
+                                exp.PartitionedByBucket(
+                                    this=exp.to_column("a"), expression=exp.Literal.number(4)
+                                ),
+                            ]
+                        )
+                    ),
+                    exp.Property(this=exp.var("table_type"), value=exp.Literal.string("iceberg")),
+                ]
+            ),
+        )
+        self.assertEqual(
+            ct_iceberg.sql(dialect=self.dialect, identify=True),
+            "CREATE TABLE `foo`.`bar` (`a` INT, `b` STRING) STORED AS PARQUET LOCATION 's3://foo' PARTITIONED BY (`partition_col`, BUCKET(4, `a`)) TBLPROPERTIES ('table_type'='iceberg')",
+        )
+
     def test_ctas(self):
         # Hive tables use 'external_location' to specify the table location, Iceberg tables use 'location' to specify the table location
         # In addition, Hive tables used 'partitioned_by' to specify the partition fields and Iceberg tables use 'partitioning' to specify the partition fields
@@ -223,7 +284,11 @@ class TestAthena(Validator):
         )
         self.assertEqual(
             ctas_hive.sql(dialect=self.dialect, identify=True),
-            "CREATE TABLE \"foo\".\"bar\" WITH (format='parquet', external_location='s3://foo', partitioned_by=ARRAY['partition_col']) AS SELECT 1",
+            "CREATE TABLE \"foo\".\"bar\" WITH (format='parquet', external_location='s3://foo', partitioned_by=ARRAY['\"partition_col\"']) AS SELECT 1",
+        )
+        self.assertEqual(
+            ctas_hive.sql(dialect=self.dialect, identify=False),
+            "CREATE TABLE foo.bar WITH (format='parquet', external_location='s3://foo', partitioned_by=ARRAY['partition_col']) AS SELECT 1",
         )
 
         ctas_iceberg = exp.Create(
@@ -234,7 +299,14 @@ class TestAthena(Validator):
                     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")])
+                        this=exp.Schema(
+                            expressions=[
+                                exp.to_column("partition_col"),
+                                exp.PartitionedByBucket(
+                                    this=exp.to_column("a"), expression=exp.Literal.number(4)
+                                ),
+                            ]
+                        )
                     ),
                 ]
             ),
@@ -242,5 +314,9 @@ class TestAthena(Validator):
         )
         self.assertEqual(
             ctas_iceberg.sql(dialect=self.dialect, identify=True),
-            "CREATE TABLE \"foo\".\"bar\" WITH (table_type='iceberg', location='s3://foo', partitioning=ARRAY['partition_col']) AS SELECT 1",
+            "CREATE TABLE \"foo\".\"bar\" WITH (table_type='iceberg', location='s3://foo', partitioning=ARRAY['\"partition_col\"', 'BUCKET(\"a\", 4)']) AS SELECT 1",
+        )
+        self.assertEqual(
+            ctas_iceberg.sql(dialect=self.dialect, identify=False),
+            "CREATE TABLE foo.bar WITH (table_type='iceberg', location='s3://foo', partitioning=ARRAY['partition_col', 'BUCKET(a, 4)']) AS SELECT 1",
         )
diff --git a/tests/dialects/test_spark.py b/tests/dialects/test_spark.py
index b570aea2..6fb5ae8c 100644
--- a/tests/dialects/test_spark.py
+++ b/tests/dialects/test_spark.py
@@ -56,7 +56,7 @@ class TestSpark(Validator):
             "CREATE TABLE x USING ICEBERG PARTITIONED BY (MONTHS(y)) LOCATION 's3://z'",
             write={
                 "duckdb": "CREATE TABLE x",
-                "presto": "CREATE TABLE x WITH (FORMAT='ICEBERG', PARTITIONED_BY=ARRAY['MONTHS'])",
+                "presto": "CREATE TABLE x WITH (FORMAT='ICEBERG', PARTITIONED_BY=ARRAY['MONTHS(y)'])",
                 "hive": "CREATE TABLE x STORED AS ICEBERG PARTITIONED BY (MONTHS(y)) LOCATION 's3://z'",
                 "spark": "CREATE TABLE x USING ICEBERG PARTITIONED BY (MONTHS(y)) LOCATION 's3://z'",
             },
diff --git a/tests/dialects/test_trino.py b/tests/dialects/test_trino.py
index b4c6bcba..c90c462a 100644
--- a/tests/dialects/test_trino.py
+++ b/tests/dialects/test_trino.py
@@ -93,6 +93,16 @@ class TestTrino(Validator):
             "CREATE TABLE foo.bar WITH (LOCATION='s3://bucket/foo/bar') AS SELECT 1"
         )
 
+        # Hive connector syntax (partitioned_by)
+        self.validate_identity(
+            "CREATE TABLE foo (a VARCHAR, b INTEGER, c DATE) WITH (PARTITIONED_BY=ARRAY['a', 'b'])"
+        )
+
+        # Iceberg connector syntax (partitioning, can contain Iceberg transform expressions)
+        self.validate_identity(
+            "CREATE TABLE foo (a VARCHAR, b INTEGER, c DATE) WITH (PARTITIONING=ARRAY['a', 'bucket(4, b)', 'month(c)'])",
+        )
+
     def test_analyze(self):
         self.validate_identity("ANALYZE tbl")
         self.validate_identity("ANALYZE tbl WITH (prop1=val1, prop2=val2)")
diff --git a/tests/fixtures/optimizer/canonicalize.sql b/tests/fixtures/optimizer/canonicalize.sql
index 3610f162..63c24892 100644
--- a/tests/fixtures/optimizer/canonicalize.sql
+++ b/tests/fixtures/optimizer/canonicalize.sql
@@ -124,6 +124,12 @@ SELECT CAST(CAST(`t`.`some_col` AS DATE) AS DATETIME) < CAST(CAST(`t`.`other_col
 --------------------------------------
 -- Remove redundant casts
 --------------------------------------
+CAST(CAST("foo" AS DECIMAL(4, 2)) AS DECIMAL(8, 4)) AS "x";
+CAST(CAST("foo" AS DECIMAL(4, 2)) AS DECIMAL(8, 4)) AS "x";
+
+CAST(CAST("foo" AS DECIMAL(4, 2)) AS DECIMAL(4, 2)) AS "x";
+CAST("foo" AS DECIMAL(4, 2)) AS "x";
+
 CAST(CAST('2023-01-01' AS DATE) AS DATE);
 CAST('2023-01-01' AS DATE);