address comments

Author: Yingjian Wu

tests/catalog/test_base.py

@@ -766,26 +766,3 @@ def test_table_properties_raise_for_none_value(catalog: InMemoryCatalog) -> None
     with pytest.raises(ValidationError) as exc_info:
         _ = given_catalog_has_a_table(catalog, properties=property_with_none)
     assert "None type is not a supported value in properties: property_name" in str(exc_info.value)
-
-
-def test_abort_table_transaction_on_exception(catalog: InMemoryCatalog) -> None:
-    tbl = given_catalog_has_a_table(catalog)
-    # Populate some initial data
-    data = pa.Table.from_pylist(
-        [{"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}],
-        schema=TEST_TABLE_SCHEMA.as_arrow(),
-    )
-    tbl.append(data)
-
-    # Data to overwrite
-    data = pa.Table.from_pylist(
-        [{"x": 7, "y": 8, "z": 9}, {"x": 7, "y": 8, "z": 9}, {"x": 7, "y": 8, "z": 9}],
-        schema=TEST_TABLE_SCHEMA.as_arrow(),
-    )
-
-    with pytest.raises(ValueError):
-        with tbl.transaction() as txn:
-            txn.overwrite(data)
-            raise ValueError
-
-    assert len(tbl.scan().to_pandas()) == 2  # type: ignore

tests/integration/test_writes/test_writes.py

@@ -1448,3 +1448,26 @@ def test_rewrite_manifest_after_partition_evolution(session_catalog: Catalog) ->
             EqualTo("category", "A"),
         ),
     )
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("format_version", [1, 2])
+def test_abort_table_transaction_on_exception(
+    spark: SparkSession, session_catalog: Catalog, arrow_table_with_null: pa.Table, format_version: int
+) -> None:
+    identifier = "default.table_test_abort_table_transaction_on_exception"
+    tbl = _create_table(session_catalog, identifier, properties={"format-version": format_version})
+
+    # Pre-populate some data
+    tbl.append(arrow_table_with_null)
+    assert len(tbl.scan().to_pandas()) == 3
+
+    # try to commit a transaction that raises exception at the middle
+    with pytest.raises(ValueError):
+        with tbl.transaction() as txn:
+            txn.append(arrow_table_with_null)
+            raise ValueError
+            txn.append(arrow_table_with_null)  # type: ignore
+
+    # Validate the transaction is aborted
+    assert len(tbl.scan().to_pandas()) == 3  # type: ignore