Creating catalog dynamically

Author: datapythonista

pandas/io/iceberg.py

@@ -63,7 +63,7 @@ def read_iceberg(
     ...     catalog_properties={"s3.secret-access-key": "my-secret"},
     ...     row_filter="trip_distance >= 10.0",
     ...     selected_fields=("VendorID", "tpep_pickup_datetime"),
-    ... )
+    ... )  # doctest: +SKIP
     """
     pyiceberg_catalog = import_optional_dependency("pyiceberg.catalog")
     pyiceberg_expressions = import_optional_dependency("pyiceberg.expressions")

pandas/tests/io/data/iceberg/catalog.sqlite

@@ -26,7 +26,9 @@
 ```
 """
 
+from contextlib import contextmanager
 import pathlib
+import tempfile
 
 import pytest
 
@@ -36,94 +38,119 @@
 from pandas.io.iceberg import read_iceberg
 
 pyiceberg = pytest.importorskip("pyiceberg")
+pyiceberg_catalog = pytest.importorskip("pyiceberg.catalog")
+pq = pytest.importorskip("pyarrow.parquet")
+
+
+@contextmanager
+def create_catalog():
+    # the catalog stores the full path of data files, so the catalog needs to be
+    # created dynamically, and not saved in pandas/tests/io/data as other formats
+    with tempfile.TemporaryDirectory("pandas-iceberg.tmp") as catalog_path:
+        uri = f"sqlite:///{catalog_path}/catalog.sqlite"
+        catalog = pyiceberg_catalog.load_catalog(
+            "default",
+            type="sql",
+            uri=uri,
+            warehouse=f"file://{catalog_path}",
+        )
+        catalog.create_namespace("default")
 
+        df = pq.read_table(
+            pathlib.Path(__file__).parent / "data" / "parquet" / "simple.parquet"
+        )
+        table = catalog.create_table("default.simple", schema=df.schema)
+        table.append(df)
 
-@pytest.fixture
-def catalog_uri(datapath):
-    path = datapath("io", "data", "iceberg", "catalog.sqlite")
-    return f"sqlite:///{path}"
+        yield uri
 
 
 class TestIceberg:
-    def test_read(self, catalog_uri):
+    def test_read(self):
         expected = pd.DataFrame(
             {
                 "A": [1, 2, 3],
                 "B": ["foo", "foo", "foo"],
             }
         )
-        result = read_iceberg(
-            "default.simple",
-            catalog_properties={"uri": catalog_uri},
-        )
+        with create_catalog() as catalog_uri:
+            result = read_iceberg(
+                "default.simple",
+                catalog_properties={"uri": catalog_uri},
+            )
         tm.assert_frame_equal(result, expected)
 
-    def test_read_by_catalog_name(self, catalog_uri):
+    def test_read_by_catalog_name(self):
         config_path = pathlib.Path.home() / ".pyiceberg.yaml"
-        with open(config_path, "w") as f:
-            f.write(f"""\
-catalog:
-  pandas_tests_catalog:
-    uri: {catalog_uri}""")
-        expected = pd.DataFrame(
-            {
-                "A": [1, 2, 3],
-                "B": ["foo", "foo", "foo"],
-            }
-        )
-        result = read_iceberg(
-            "default.simple",
-            catalog_name="pandas_tests_catalog",
-        )
+        with create_catalog() as catalog_uri:
+            with open(config_path, "w") as f:
+                f.write(f"""\
+    catalog:
+      pandas_tests_catalog:
+        uri: {catalog_uri}""")
+            expected = pd.DataFrame(
+                {
+                    "A": [1, 2, 3],
+                    "B": ["foo", "foo", "foo"],
+                }
+            )
+            result = read_iceberg(
+                "default.simple",
+                catalog_name="pandas_tests_catalog",
+            )
         tm.assert_frame_equal(result, expected)
         # config_path.unlink()
 
-    def test_read_with_row_filter(self, catalog_uri):
+    def test_read_with_row_filter(self):
         expected = pd.DataFrame(
             {
                 "A": [2, 3],
                 "B": ["foo", "foo"],
             }
         )
-        result = read_iceberg(
-            "default.simple",
-            catalog_properties={"uri": catalog_uri},
-            row_filter="A > 1",
-        )
+        with create_catalog() as catalog_uri:
+            result = read_iceberg(
+                "default.simple",
+                catalog_properties={"uri": catalog_uri},
+                row_filter="A > 1",
+            )
         tm.assert_frame_equal(result, expected)
 
-    def test_read_with_case_sensitive(self, catalog_uri):
+    def test_read_with_case_sensitive(self):
         expected = pd.DataFrame(
             {
                 "A": [1, 2, 3],
             }
         )
-        result = read_iceberg(
-            "default.simple",
-            catalog_properties={"uri": catalog_uri},
-            selected_fields=["a"],
-            case_sensitive=False,
-        )
-        tm.assert_frame_equal(result, expected)
-
-        with pytest.raises(ValueError, match="^Could not find column"):
-            read_iceberg(
+        with create_catalog() as catalog_uri:
+            result = read_iceberg(
                 "default.simple",
                 catalog_properties={"uri": catalog_uri},
                 selected_fields=["a"],
-                case_sensitive=True,
+                case_sensitive=False,
             )
+        tm.assert_frame_equal(result, expected)
 
-    def test_read_with_limit(self, catalog_uri):
+        with create_catalog() as catalog_uri:
+            with pytest.raises(ValueError, match="^Could not find column"):
+                read_iceberg(
+                    "default.simple",
+                    catalog_properties={"uri": catalog_uri},
+                    selected_fields=["a"],
+                    case_sensitive=True,
+                )
+
+    def test_read_with_limit(self):
         expected = pd.DataFrame(
             {
                 "A": [1, 2],
                 "B": ["foo", "foo"],
             }
         )
-        result = read_iceberg(
-            "default.simple",
-            catalog_properties={"uri": catalog_uri},
-            limit=2,
-        )
+        with create_catalog() as catalog_uri:
+            result = read_iceberg(
+                "default.simple",
+                catalog_properties={"uri": catalog_uri},
+                limit=2,
+            )
         tm.assert_frame_equal(result, expected)