Repository: TobikoData/sqlmesh
Pull request #4441: Add iceberg support to table_diff
Author: login: williebsweet
Created: 2025-05-16T21:53:49Z | Merged: 2025-05-26T15:07:00Z

Resolves #4208

PR contains the following:
- Added a check of both source and target table formats 
- Added an error, should someone try to diff an Iceberg table and a Hive table
- Set the table format for a Iceberg/Iceberg diff to Iceberg
- Tests for various combinations of source/target/temp table formats

Linked issue #4208: `table_diff` on Athena fails for Iceberg tables
Based on issue a user reported in Slack: https://tobiko-data.slack.com/archives/C044BRE5W4S/p1744125717966919?thread_ts=1742833178.418699&cid=C044BRE5W4S

The table_diff [creates a temporary table](https://github.com/TobikoData/sqlmesh/blob/9952bbba55b53c5f8dff06fa8d2cc82636406e92/sqlmesh/core/table_diff.py#L425) in order to hold the diff results.

In the case of Athena, this temporary table will be created as Hive by default. However, if the source table is Iceberg, the temporary table also  needs to be created as Iceberg because some data types are not compatible between the two.

Timestamp fields are a good example:
```
pyathena.error.OperationalError: NOT_SUPPORTED: Unsupported Hive type: timestamp(6) with time zone. You may need to manually clean the data at location 's3://athena/sqlmesh/tables/1543f306-59e7-4310-b2b0-f75e0558761a' before retrying. Athena will not delete data in your account.
```
Linked issue comments: user: login: williebsweet
created_at: 2025-05-19T13:25:14Z
body: @erindru Created a PR to address this. I installed the edited version locally and it resolved my issue.
author_association: CONTRIBUTOR

Diff:
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 7b59acb912..8ea7f1e9cc 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -8,6 +8,7 @@
 
 from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.mixins import RowDiffMixin
+from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter
 from sqlglot import exp, parse_one
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
@@ -15,6 +16,7 @@
 from sqlglot.optimizer.scope import find_all_in_scope
 
 from sqlmesh.utils.pydantic import PydanticModel
+from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
@@ -431,7 +433,26 @@ def name(e: exp.Expression) -> str:
             schema = to_schema(temp_schema, dialect=self.dialect)
             temp_table = exp.table_("diff", db=schema.db, catalog=schema.catalog, quoted=True)
 
-            with self.adapter.temp_table(query, name=temp_table) as table:
+            temp_table_kwargs = {}
+            if isinstance(self.adapter, AthenaEngineAdapter):
+                # Athena has two table formats: Hive (the default) and Iceberg. TableDiff requires that
+                # the formats be the same for the source, target, and temp tables.
+                source_table_type = self.adapter._query_table_type(self.source_table)
+                target_table_type = self.adapter._query_table_type(self.target_table)
+
+                if source_table_type == "iceberg" and target_table_type == "iceberg":
+                    temp_table_kwargs["table_format"] = "iceberg"
+                # Sets the temp table's format to Iceberg.
+                # If neither source nor target table is Iceberg, it defaults to Hive (Athena's default).
+                elif source_table_type == "iceberg" or target_table_type == "iceberg":
+                    raise SQLMeshError(
+                        f"Source table '{self.source}' format '{source_table_type}' and target table '{self.target}' format '{target_table_type}' "
+                        f"do not match for Athena. Diffing between different table formats is not supported."
+                    )
+
+            with self.adapter.temp_table(
+                query, name=temp_table, columns_to_types=None, **temp_table_kwargs
+            ) as table:
                 summary_sums = [
                     exp.func("SUM", "s_exists").as_("s_count"),
                     exp.func("SUM", "t_exists").as_("t_count"),
diff --git a/tests/core/engine_adapter/test_athena.py b/tests/core/engine_adapter/test_athena.py
index f96f5656a2..2fd45ac3df 100644
--- a/tests/core/engine_adapter/test_athena.py
+++ b/tests/core/engine_adapter/test_athena.py
@@ -10,6 +10,7 @@
 from sqlmesh.core.model import load_sql_based_model
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.core.table_diff import TableDiff
 
 from tests.core.engine_adapter import to_sql_calls
 
@@ -21,6 +22,16 @@ def adapter(make_mocked_engine_adapter: t.Callable) -> AthenaEngineAdapter:
     return make_mocked_engine_adapter(AthenaEngineAdapter)
 
 
+@pytest.fixture
+def table_diff(adapter: AthenaEngineAdapter) -> TableDiff:
+    return TableDiff(
+        adapter=adapter,
+        source="source_table",
+        target="target_table",
+        on=["id"],
+    )
+
+
 @pytest.mark.parametrize(
     "config_s3_warehouse_location,table_properties,table,expected_location",
     [
@@ -483,3 +494,74 @@ def test_iceberg_partition_transforms(adapter: AthenaEngineAdapter):
         # Trino syntax - CTAS
         """CREATE TABLE IF NOT EXISTS "test_table" WITH (table_type='iceberg', partitioning=ARRAY['MONTH(business_date)', 'BUCKET(colb, 4)', 'colc'], location='s3://bucket/prefix/test_table/', is_external=false) AS SELECT CAST("business_date" AS TIMESTAMP) AS "business_date", CAST("colb" AS VARCHAR) AS "colb", CAST("colc" AS VARCHAR) AS "colc" FROM (SELECT CAST(1 AS TIMESTAMP) AS "business_date", CAST(2 AS VARCHAR) AS "colb", 'foo' AS "colc" LIMIT 0) AS "_subquery\"""",
     ]
+
+
+@pytest.mark.parametrize(
+    "source_format, target_format, expected_temp_format, expect_error",
+    [
+        ("hive", "hive", None, False),
+        ("iceberg", "hive", None, True),  # Expect error for mismatched formats
+        ("hive", "iceberg", None, True),  # Expect error for mismatched formats
+        ("iceberg", "iceberg", "iceberg", False),
+        (None, "iceberg", None, True),  # Source doesn't exist or type unknown, target is iceberg
+        (
+            "iceberg",
+            None,
+            "iceberg",
+            True,
+        ),  # Target doesn't exist or type unknown, source is iceberg
+        (None, "hive", None, False),  # Source doesn't exist or type unknown, target is hive
+        ("hive", None, None, False),  # Target doesn't exist or type unknown, source is hive
+        (None, None, None, False),  # Both don't exist or types unknown
+    ],
+)
+def test_table_diff_temp_table_format(
+    table_diff: TableDiff,
+    mocker: MockerFixture,
+    source_format: t.Optional[str],
+    target_format: t.Optional[str],
+    expected_temp_format: t.Optional[str],
+    expect_error: bool,
+):
+    adapter = t.cast(AthenaEngineAdapter, table_diff.adapter)
+
+    # Mock _query_table_type to return specified formats
+    def mock_query_table_type(table_name: exp.Table) -> t.Optional[str]:
+        if table_name.name == "source_table":
+            return source_format
+        if table_name.name == "target_table":
+            return target_format
+        return "hive"  # Default for other tables if any
+
+    mocker.patch.object(adapter, "_query_table_type", side_effect=mock_query_table_type)
+
+    # Mock temp_table to capture kwargs
+    mock_temp_table = mocker.patch.object(adapter, "temp_table", autospec=True)
+    mock_temp_table.return_value.__enter__.return_value = exp.to_table("diff_table")
+
+    # Mock fetchdf and other calls made within row_diff to avoid actual DB interaction
+    mocker.patch.object(adapter, "fetchdf", return_value=pd.DataFrame())
+    mocker.patch.object(adapter, "get_data_objects", return_value=[])
+    mocker.patch.object(adapter, "columns", return_value={"id": exp.DataType.build("int")})
+
+    if expect_error:
+        with pytest.raises(
+            SQLMeshError,
+            match="do not match for Athena. Diffing between different table formats is not supported.",
+        ):
+            table_diff.row_diff()
+        mock_temp_table.assert_not_called()  # temp_table should not be called if formats mismatch
+        return
+
+    try:
+        table_diff.row_diff()
+    except Exception:
+        pass  # We only care about the temp_table call args for non-error cases
+
+    mock_temp_table.assert_called_once()
+    _, called_kwargs = mock_temp_table.call_args
+
+    if expected_temp_format:
+        assert called_kwargs.get("table_format") == expected_temp_format
+    else:
+        assert "table_format" not in called_kwargs


Reviews:
user: login: izeigerman
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-19T17:09:31Z
user: login: erindru
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-19T22:19:25Z
user: login: williebsweet
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T12:19:02Z
user: login: erindru
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-21T20:10:34Z
user: login: erindru
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-25T22:56:37Z
user: login: erindru
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-25T22:57:43Z

Review comments:
diff_hunk: @@ -2203,6 +2203,19 @@ def temp_table(
             query_or_df, columns_to_types=columns_to_types, target_table=name
         )
 
+        if hasattr(self, "_query_table_type_or_raise") and hasattr(self, "source_table"):
path: sqlmesh/core/engine_adapter/base.py
user: login: izeigerman
body: Should this logic be in `athena.py` where these attributes are available? 
created_at: 2025-05-19T17:09:31Z
author_association: COLLABORATOR
diff_hunk: @@ -2203,6 +2203,19 @@ def temp_table(
             query_or_df, columns_to_types=columns_to_types, target_table=name
         )
 
+        if hasattr(self, "_query_table_type_or_raise") and hasattr(self, "source_table"):
path: sqlmesh/core/engine_adapter/base.py
user: login: erindru
body: I... dont actually see how this works. Where does `self.source_table` come from? This isn't an engine adapter property.

table_diff [calls](https://github.com/TobikoData/sqlmesh/blob/529e52fa9d890dae3e56bb958e49f5788676e5dc/sqlmesh/core/table_diff.py#L434) `EngineAdapter.temp_table(...)` and passes the query that it's trying to store the results of, which looks something like:
```
with source as (select ... from source_table),
target as (select ... from target_table),
joined as (select ... from source full outer join target)
select * from joined
```

I actually think it's table_diff that should be supplying the table format, not the engine adapter introspecting the query to try and look up all the tables referenced in the query to see if any of them are Iceberg.


created_at: 2025-05-19T22:19:25Z
author_association: COLLABORATOR
diff_hunk: @@ -2203,6 +2203,19 @@ def temp_table(
             query_or_df, columns_to_types=columns_to_types, target_table=name
         )
 
+        if hasattr(self, "_query_table_type_or_raise") and hasattr(self, "source_table"):
path: sqlmesh/core/engine_adapter/base.py
user: login: williebsweet
body: @erindru Something more like [this previous commit](https://github.com/TobikoData/sqlmesh/pull/4441/commits/d45c73e487a5fc6c84f1cf7d6ea02b37822ead9a)?
created_at: 2025-05-21T12:19:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2203,6 +2203,19 @@ def temp_table(
             query_or_df, columns_to_types=columns_to_types, target_table=name
         )
 
+        if hasattr(self, "_query_table_type_or_raise") and hasattr(self, "source_table"):
path: sqlmesh/core/engine_adapter/base.py
user: login: erindru
body: Yes, exactly. With perhaps a comment to explain why there is an Athena-specific check and the addition of some tests
created_at: 2025-05-21T20:10:33Z
author_association: COLLABORATOR
diff_hunk: @@ -431,7 +433,26 @@ def name(e: exp.Expression) -> str:
             schema = to_schema(temp_schema, dialect=self.dialect)
             temp_table = exp.table_("diff", db=schema.db, catalog=schema.catalog, quoted=True)
 
-            with self.adapter.temp_table(query, name=temp_table) as table:
+            temp_table_kwargs = {}
+            if isinstance(self.adapter, AthenaEngineAdapter):
+                # Athena has two table formats: Hive (the default) and Iceberg. TableDiff requires that
+                # the formats be the same for the source, target, and temp tables.
+                source_table_type = self.adapter._query_table_type(self.source_table)
+                target_table_type = self.adapter._query_table_type(self.target_table)
+
+                if source_table_type == "iceberg" and target_table_type == "iceberg":
+                    temp_table_kwargs["table_format"] = "iceberg"
+                # Sets the temp table's format to Iceberg.
+                # If neither source nor target table is Iceberg, it defaults to Hive (Athena's default).
+                elif source_table_type == "iceberg" or target_table_type == "iceberg":
+                    raise SQLMeshError(
path: sqlmesh/core/table_diff.py
user: login: erindru
body: I _think_ this is right. Table diff relies on copying a sample of data from each table into a single table, which on most engines is fine because they have a consistent table format.

I know that timestamp types in particular are handled differently between Athena/Iceberg and Athena/Hive so I dont think we can say "if Iceberg is detected - make the diff table Iceberg" because copying Hive data into it verbatim can still cause an error
created_at: 2025-05-25T22:56:36Z
author_association: COLLABORATOR

PR comments:
user: login: CLAassistant
created_at: 2025-05-16T21:53:54Z
body: [![CLA assistant check](https://cla-assistant.io/pull/badge/signed)](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=4441) <br/>All committers have signed the CLA.
author_association: NONE
user: login: izeigerman
created_at: 2025-05-19T17:09:49Z
body: @williebsweet thanks for the submission! Some tests will be helpful
author_association: COLLABORATOR
user: login: williebsweet
created_at: 2025-05-23T12:49:16Z
body: @erindru I moved the logic back into `table_diff`. I did the following:
- Added a check of both source and target table formats 
- Added an error, should someone try to diff an Iceberg table and a Hive table
- Set the table format for a Iceberg/Iceberg diff to Iceberg

@izeigerman I also added tests for the scenarios above.
author_association: CONTRIBUTOR
user: login: williebsweet
created_at: 2025-05-23T13:39:24Z
body: Forgot to mention - I also tested this with my company's SQLMesh project (installing the `--editable` package) and I was able to run a diff. 🎆 
author_association: CONTRIBUTOR
user: login: erindru
created_at: 2025-05-25T22:57:15Z
body: Thanks @williebsweet ! This looks much better
author_association: COLLABORATOR