SNOW-1652349: Add support for iceberg to write_pandas (#2056)

Author: sfc-gh-jrose

DESCRIPTION.md

@@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added a <19.0.0 pin to pyarrow as a workaround to a bug affecting Azure Batch.
   - Optimized distribution package lookup to speed up import.
   - Fixed a bug where privatelink OCSP Cache url could not be determined if privatelink account name was specified in uppercase
+  - Added support for iceberg tables to `write_pandas`
 
 - v3.13.2(January 29, 2025)
   - Changed not to use scoped temporary objects.

src/snowflake/connector/pandas_tools.py

@@ -215,6 +215,42 @@ def _create_temp_file_format(
     return file_format_location
 
 
+def _convert_value_to_sql_option(value: Union[str, bool, int, float]) -> str:
+    if isinstance(value, str):
+        if len(value) > 1 and value.startswith("'") and value.endswith("'"):
+            return value
+        else:
+            value = value.replace(
+                "'", "''"
+            )  # escape single quotes before adding a pair of quotes
+            return f"'{value}'"
+    else:
+        return str(value)
+
+
+def _iceberg_config_statement_helper(iceberg_config: dict[str, str]) -> str:
+    ALLOWED_CONFIGS = {
+        "EXTERNAL_VOLUME",
+        "CATALOG",
+        "BASE_LOCATION",
+        "CATALOG_SYNC",
+        "STORAGE_SERIALIZATION_POLICY",
+    }
+
+    normalized = {
+        k.upper(): _convert_value_to_sql_option(v)
+        for k, v in iceberg_config.items()
+        if v is not None
+    }
+
+    if invalid_configs := set(normalized.keys()) - ALLOWED_CONFIGS:
+        raise ProgrammingError(
+            f"Invalid iceberg configurations option(s) provided {', '.join(sorted(invalid_configs))}"
+        )
+
+    return " ".join(f"{k}={v}" for k, v in normalized.items())
+
+
 def write_pandas(
     conn: SnowflakeConnection,
     df: pandas.DataFrame,
@@ -231,6 +267,7 @@ def write_pandas(
     overwrite: bool = False,
     table_type: Literal["", "temp", "temporary", "transient"] = "",
     use_logical_type: bool | None = None,
+    iceberg_config: dict[str, str] | None = None,
     **kwargs: Any,
 ) -> tuple[
     bool,
@@ -295,6 +332,14 @@ def write_pandas(
             Snowflake can interpret Parquet logical types during data loading. To enable Parquet logical types,
             set use_logical_type as True. Set to None to use Snowflakes default. For more information, see:
             https://docs.snowflake.com/en/sql-reference/sql/create-file-format
+        iceberg_config: A dictionary that can contain the following iceberg configuration values:
+                * external_volume: specifies the identifier for the external volume where
+                    the Iceberg table stores its metadata files and data in Parquet format
+                * catalog: specifies either Snowflake or a catalog integration to use for this table
+                * base_location: the base directory that snowflake can write iceberg metadata and files to
+                * catalog_sync: optionally sets the catalog integration configured for Polaris Catalog
+                * storage_serialization_policy: specifies the storage serialization policy for the table
+
 
 
     Returns:
@@ -479,9 +524,14 @@ def drop_object(name: str, object_type: str) -> None:
             quote_identifiers,
         )
 
+        iceberg = "ICEBERG " if iceberg_config else ""
+        iceberg_config_statement = _iceberg_config_statement_helper(
+            iceberg_config or {}
+        )
+
         create_table_sql = (
-            f"CREATE {table_type.upper()} TABLE IF NOT EXISTS identifier(?) "
-            f"({create_table_columns})"
+            f"CREATE {table_type.upper()} {iceberg}TABLE IF NOT EXISTS identifier(?) "
+            f"({create_table_columns}) {iceberg_config_statement}"
             f" /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
         )
         params = (target_table_location,)

test/integ/conftest.py

@@ -246,6 +246,11 @@ def conn_cnx() -> Callable[..., ContextManager[SnowflakeConnection]]:
     return db
 
 
+@pytest.fixture(scope="module")
+def module_conn_cnx() -> Callable[..., ContextManager[SnowflakeConnection]]:
+    return db
+
+
 @pytest.fixture()
 def negative_conn_cnx() -> Callable[..., ContextManager[SnowflakeConnection]]:
     """Use this if an incident is expected and we don't want GS to create a dump file about the incident."""

test/integ/pandas/test_pandas_tools.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import math
+import re
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Callable, Generator
 from unittest import mock
@@ -26,10 +27,14 @@
 
 try:
     from snowflake.connector.options import pandas
-    from snowflake.connector.pandas_tools import write_pandas
+    from snowflake.connector.pandas_tools import (
+        _iceberg_config_statement_helper,
+        write_pandas,
+    )
 except ImportError:
     pandas = None
     write_pandas = None
+    _iceberg_config_statement_helper = None
 
 if TYPE_CHECKING:
     from snowflake.connector import SnowflakeConnection
@@ -1029,6 +1034,34 @@ def mock_execute(*args, **kwargs):
             cnx.execute_string(f"drop schema if exists {target_schema}")
 
 
+def test__iceberg_config_statement_helper():
+    config = {
+        "EXTERNAL_VOLUME": "vol",
+        "CATALOG": "'SNOWFLAKE'",
+        "BASE_LOCATION": "/root",
+        "CATALOG_SYNC": "foo",
+        "STORAGE_SERIALIZATION_POLICY": "bar",
+    }
+    assert (
+        _iceberg_config_statement_helper(config)
+        == "EXTERNAL_VOLUME='vol' CATALOG='SNOWFLAKE' BASE_LOCATION='/root' CATALOG_SYNC='foo' STORAGE_SERIALIZATION_POLICY='bar'"
+    )
+
+    config["STORAGE_SERIALIZATION_POLICY"] = None
+    assert (
+        _iceberg_config_statement_helper(config)
+        == "EXTERNAL_VOLUME='vol' CATALOG='SNOWFLAKE' BASE_LOCATION='/root' CATALOG_SYNC='foo'"
+    )
+
+    config["foo"] = True
+    config["bar"] = True
+    with pytest.raises(
+        ProgrammingError,
+        match=re.escape("Invalid iceberg configurations option(s) provided BAR, FOO"),
+    ):
+        _iceberg_config_statement_helper(config)
+
+
 def test_write_pandas_with_on_error(
     conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
 ):