Assign configurable prefix to partition field names to prevent conflict with schema field names (#241)

* Updating how partitions are handled / calculated to satisfy constraints introduced in pyiceberg 0.10.0 which require partition spec field names to be different from any existing schema field names when a transform is being applied.

* Pass env into docker compose

* Correctly handle when a field isn't found in the schema

* Update partition-related fixtures to create partition columns with internal utility used to prevent partition name / column name clashes

* Add more tests to cover new functionality

* Fix tests for updates to how partition names are calculated

* Fix tests to reflect partition naming with calculated suffixes

* Absolute paths prevent cache_dir from being set in different dev environments that might not have the same directories

* Trying to speed up test initialization / teardown.

No need to recreate containers. Provide option to opt out of teardown to enable faster repeated test execution

* Switching to assigning a configurable prefix to the partition field name to make it unique.

This is more predictable and configurable than the previous implementation that added transform-specific suffixes.

* Formatting

* changelog item

* Remove unnecessary instance attribute

* Formatting

* fixing some typing errors

* Handle situations that I don't think actually happen in the wild but are needed to make the type-checking gods happy

* Update tests to account for partition_field_name_prefix

* Fix wrong partition field naming

* formatting

* Require non-empty partition_field_name_prefix

* Fix docstring

* Add migration note

* Include PR number

Author: zyd14

libraries/dagster-iceberg/CHANGELOG.md

@@ -10,3 +10,4 @@
 ### Changed
 
 - Rename I/O managers from `<Storage><Engine>IOManager` to `<Engine><Storage>IOManager`.
+- Support for pyiceberg 0.10.0 (#241). Change behavior for how partition field names are calculated to address validation introduced in [pyiceberg #2505](https://github.com/apache/iceberg-python/pull/2305). Partition field names calculated for Dagster asset partitions now include a configurable prefix before the column name they reference. **Migration note**: When updating partition specs on existing tables, new partition field names will be generated using the configured prefix (default: `part_`). For example, a partition field previously named `timestamp` will become `part_timestamp`. Existing data remains accessible, but queries referencing partition fields will need to be updated to use the new naming convention. To maintain backward compatibility, existing tables with unchanged partition specs will retain their original field names until their partition specs are updated. 

libraries/dagster-iceberg/pyproject.toml

@@ -109,7 +109,7 @@ ignore = [
 "kitchen-sink/*" = ["INP001"]
 
 [tool.pytest.ini_options]
-cache_dir = "/home/vscode/workspace/.cache/pytest"
+cache_dir = "./.cache/pytest"
 
 [tool.pyright]
 exclude = [".venv", ".github", "docs", "tests"]

libraries/dagster-iceberg/src/dagster_iceberg/_utils/__init__.py

@@ -1,7 +1,15 @@
-from dagster import __version__
-from packaging import version
-
-from dagster_iceberg._utils.io import table_writer as table_writer
+from dagster_iceberg._utils.io import (
+    DEFAULT_PARTITION_FIELD_NAME_PREFIX as DEFAULT_PARTITION_FIELD_NAME_PREFIX,
+)
+from dagster_iceberg._utils.io import (
+    DEFAULT_WRITE_MODE as DEFAULT_WRITE_MODE,
+)
+from dagster_iceberg._utils.io import (
+    WriteMode as WriteMode,
+)
+from dagster_iceberg._utils.io import (
+    table_writer as table_writer,
+)
 from dagster_iceberg._utils.partitions import (
     DagsterPartitionToDaftSqlPredicateMapper as DagsterPartitionToDaftSqlPredicateMapper,
 )
@@ -11,19 +19,4 @@
 from dagster_iceberg._utils.partitions import (
     DagsterPartitionToPolarsSqlPredicateMapper as DagsterPartitionToPolarsSqlPredicateMapper,
 )
-
-
-def preview(wrapped=None):
-    if version.parse(__version__) >= version.parse("1.10.0"):
-        from dagster._annotations import (
-            preview as decorator,  # pyright: ignore[reportAttributeAccessIssue]
-        )
-    else:
-        from dagster._annotations import (
-            experimental as decorator,  # pyright: ignore[reportAttributeAccessIssue]
-        )
-
-    if wrapped is not None:
-        return decorator(wrapped)
-
-    return decorator
+from dagster_iceberg._utils.warnings import preview as preview

libraries/dagster-iceberg/src/dagster_iceberg/_utils/io.py

@@ -35,6 +35,7 @@ class WriteMode(Enum):
 
 
 DEFAULT_WRITE_MODE: Final[WriteMode] = WriteMode.overwrite
+DEFAULT_PARTITION_FIELD_NAME_PREFIX: Final[str] = "part"
 
 
 def table_writer(
@@ -44,6 +45,7 @@ def table_writer(
     schema_update_mode: str,
     partition_spec_update_mode: str,
     dagster_run_id: str,
+    partition_field_name_prefix: str = DEFAULT_PARTITION_FIELD_NAME_PREFIX,
     dagster_partition_key: str | None = None,
     table_properties: dict[str, str] | None = None,
     write_mode: WriteMode = DEFAULT_WRITE_MODE,
@@ -113,6 +115,7 @@ def table_writer(
                 table=table.refresh(),
                 table_slice=table_slice,
                 partition_spec_update_mode=partition_spec_update_mode,
+                partition_field_name_prefix=partition_field_name_prefix,
             )
         if table_properties is not None:
             update_table_properties(
@@ -135,6 +138,7 @@ def table_writer(
                 # When creating new tables with dagster partitions, we always update
                 # the partition spec
                 partition_spec_update_mode="update",
+                partition_field_name_prefix=partition_field_name_prefix,
             )
 
     row_filter: E.BooleanExpression

libraries/dagster-iceberg/src/dagster_iceberg/_utils/partitions.py

@@ -28,6 +28,50 @@
 K = TypeVar("K")
 
 
+def partition_field_name_for(column: str, prefix: str) -> str:
+    """Generate a partition field name with configurable prefix to avoid column name conflicts.
+
+    This function creates stable, compliant spec field names that avoid conflicts
+    with table column names when using transforms in pyiceberg 0.10.0+.
+
+    Args:
+        column: The source column name
+        prefix: The prefix to apply to the column name
+
+    Returns:
+        A unique partition field name with prefix applied to column name
+    """
+    return f"{prefix}_{column}"
+
+
+def _get_partition_field_by_source_column(
+    schema: Schema, spec: PartitionSpec, column_name: str
+) -> PartitionField | None:
+    """Find a partition field in the spec by its source column name.
+
+    This function looks up partition fields by matching the source column ID
+    rather than the partition field name, which is more reliable for change detection.
+
+    Args:
+        schema: The table schema
+        spec: The partition specification
+        column_name: The source column name to find
+
+    Returns:
+        The matching PartitionField if found, None otherwise
+    """
+    try:
+        field = schema.find_field(column_name)
+    except ValueError:
+        # Column doesn't exist in schema
+        return None
+
+    for pf in spec.fields:
+        if pf.source_id == field.field_id:
+            return pf
+    return None
+
+
 class DagsterPartitionToPredicateMapper(Generic[K]):
     def __init__(
         self,
@@ -211,6 +255,7 @@ def update_table_partition_spec(
     table: Table,
     table_slice: TableSlice,
     partition_spec_update_mode: str,
+    partition_field_name_prefix: str = "part",
 ):
     partition_dimensions = cast(
         "Sequence[TablePartitionDimension]",
@@ -222,11 +267,17 @@ def update_table_partition_spec(
         exception_types=ValueError,
         table_slice=table_slice,
         partition_spec_update_mode=partition_spec_update_mode,
+        partition_field_name_prefix=partition_field_name_prefix,
     )
 
 
 class PyIcebergPartitionSpecUpdaterWithRetry(IcebergOperationWithRetry):
-    def operation(self, table_slice: TableSlice, partition_spec_update_mode: str):
+    def operation(
+        self,
+        table_slice: TableSlice,
+        partition_spec_update_mode: str,
+        partition_field_name_prefix: str,
+    ):
         self.logger.debug("Updating table partition spec")
         IcebergTableSpecUpdater(
             partition_mapping=PartitionMapper(
@@ -235,6 +286,7 @@ def operation(self, table_slice: TableSlice, partition_spec_update_mode: str):
                 iceberg_partition_spec=self.table.spec(),
             ),
             partition_spec_update_mode=partition_spec_update_mode,
+            partition_field_name_prefix=partition_field_name_prefix,
         ).update_table_spec(table=self.table)
 
 
@@ -383,9 +435,11 @@ def updated_dagster_time_partition_field(self) -> str | None:
                 time_partition_partitions.start,
                 time_partition_partitions.end,
             )
-            # Check if field is present in iceberg partition spec
-            current_time_partition_field = self.get_iceberg_partition_field_by_name(
-                time_partition.partition_expr,
+            # Check if field is present in iceberg partition spec by source column
+            current_time_partition_field = _get_partition_field_by_source_column(
+                schema=self.iceberg_table_schema,
+                spec=self.iceberg_partition_spec,
+                column_name=time_partition.partition_expr,
             )
             if current_time_partition_field is not None and (
                 time_partition_transformation != current_time_partition_field.transform
@@ -415,21 +469,43 @@ def updated(self) -> list[TablePartitionDimension]:
 
     def deleted(self) -> list[PartitionField]:
         """Retrieve partition fields need to be removed from the iceberg table."""
-        return [
-            p
-            for p in self.iceberg_partition_spec.fields
-            if p.name in self.deleted_partition_field_names
-        ]
+        # Get current dagster partition column names
+        current_partition_columns = set(
+            self.get_dagster_partition_dimension_names(
+                allow_empty_dagster_partitions=True
+            )
+        )
+
+        # Find partition fields whose source columns are no longer in dagster partitions
+        deleted_fields = []
+        for partition_field in self.iceberg_partition_spec.fields:
+            # Find the source column name for this partition field
+            source_column_name = None
+            for column in self.iceberg_table_schema.fields:
+                if column.field_id == partition_field.source_id:
+                    source_column_name = column.name
+                    break
+
+            # If source column is not in current dagster partitions, mark for deletion
+            if (
+                source_column_name is not None
+                and source_column_name not in current_partition_columns
+            ):
+                deleted_fields.append(partition_field)
+
+        return deleted_fields
 
 
 class IcebergTableSpecUpdater:
     def __init__(
         self,
         partition_mapping: PartitionMapper,
         partition_spec_update_mode: str,
+        partition_field_name_prefix: str,
     ):
         self.partition_spec_update_mode = partition_spec_update_mode
         self.partition_mapping = partition_mapping
+        self.partition_field_name_prefix = partition_field_name_prefix
         self.logger = logging.getLogger(
             "dagster_iceberg._utils.partitions.IcebergTableSpecUpdater",
         )
@@ -444,7 +520,17 @@ def _changes(
         }
 
     def _spec_update(self, update: UpdateSpec, partition: TablePartitionDimension):
-        self._spec_delete(update=update, partition_name=partition.partition_expr)
+        # Find the existing partition field by source column to get its current name
+        existing_field = _get_partition_field_by_source_column(
+            schema=self.partition_mapping.iceberg_table_schema,
+            spec=self.partition_mapping.iceberg_partition_spec,
+            column_name=partition.partition_expr,
+        )
+        if existing_field is not None:
+            # Delete the existing field by its current partition field name
+            self._spec_delete(update=update, partition_name=existing_field.name)
+
+        # Add the new field with updated transform and new name
         self._spec_new(update=update, partition=partition)
 
     def _spec_delete(self, update: UpdateSpec, partition_name: str):
@@ -456,16 +542,20 @@ def _spec_new(self, update: UpdateSpec, partition: TablePartitionDimension):
             transform = diff_to_transformation(*partition.partitions)
         else:
             transform = IdentityTransform()
+
+        # Generate a unique partition field name that avoids conflicts with column names
+        # when using transforms (required for pyiceberg 0.10.0+ compatibility)
+        partition_field_name = partition_field_name_for(
+            partition.partition_expr, prefix=self.partition_field_name_prefix
+        )
+
         self.logger.debug("Setting new partition column: %s", partition.partition_expr)
         self.logger.debug("Using transform: %s", transform)
+        self.logger.debug("Using partition field name: %s", partition_field_name)
         update.add_field(
             source_column_name=partition.partition_expr,
             transform=transform,
-            # Name the partition field spec the same as the column name.
-            #  We rely on this throughout this codebase because it makes
-            #  it a lot easier to make the mapping between dagster partitions
-            #  and Iceberg partition fields.
-            partition_field_name=partition.partition_expr,
+            partition_field_name=partition_field_name,
         )
 
     @property

libraries/dagster-iceberg/src/dagster_iceberg/_utils/warnings.py

@@ -0,0 +1,18 @@
+from dagster import __version__
+from packaging import version
+
+
+def preview(wrapped=None):
+    if version.parse(__version__) >= version.parse("1.10.0"):
+        from dagster._annotations import (
+            preview as decorator,  # pyright: ignore[reportAttributeAccessIssue]
+        )
+    else:
+        from dagster._annotations import (
+            experimental as decorator,  # pyright: ignore[reportAttributeAccessIssue]
+        )
+
+    if wrapped is not None:
+        return decorator(wrapped)
+
+    return decorator

libraries/dagster-iceberg/src/dagster_iceberg/config.py

@@ -2,8 +2,9 @@
 
 from dagster import Config
 from dagster._annotations import public
+from pydantic import Field, model_validator
 
-from dagster_iceberg._utils import preview
+from dagster_iceberg._utils import DEFAULT_PARTITION_FIELD_NAME_PREFIX, preview
 
 
 @public
@@ -45,3 +46,13 @@ class IcebergCatalogConfig(Config):
     """
 
     properties: dict[str, Any]
+    partition_field_name_prefix: str = Field(
+        default=DEFAULT_PARTITION_FIELD_NAME_PREFIX,
+        description="Prefix to apply to the partition field names. This is required to avoid conflicts with schema field names when defining partitions using non-identity transforms in pyiceberg 0.10.0+. Defaults to 'part'.",
+    )
+
+    @model_validator(mode="after")
+    def validate_partition_field_name_prefix(self) -> "IcebergCatalogConfig":
+        if self.partition_field_name_prefix == "":
+            raise ValueError("partition_field_name_prefix cannot be an empty string")
+        return self

libraries/dagster-iceberg/src/dagster_iceberg/handler.py

@@ -14,8 +14,14 @@
 from pyiceberg import table as ibt
 from pyiceberg.catalog import Catalog
 
-from dagster_iceberg._utils import preview, table_writer
-from dagster_iceberg._utils.io import DEFAULT_WRITE_MODE, WriteMode
+from dagster_iceberg._utils import (
+    DEFAULT_PARTITION_FIELD_NAME_PREFIX,
+    DEFAULT_WRITE_MODE,
+    WriteMode,
+    preview,
+    table_writer,
+)
+from dagster_iceberg.config import IcebergCatalogConfig
 
 if TYPE_CHECKING:
     from pyiceberg.table.snapshots import Snapshot
@@ -57,6 +63,7 @@ def handle_output(
         partition_spec_update_mode = metadata.get("partition_spec_update_mode", "error")
         schema_update_mode = metadata.get("schema_update_mode", "error")
 
+        partition_field_name_prefix = self._get_partition_field_name_prefix(context)
         write_mode_with_output_override = self._get_write_mode(context)
 
         table_writer(
@@ -71,6 +78,7 @@ def handle_output(
             ),
             table_properties=table_properties_usr,
             write_mode=write_mode_with_output_override,
+            partition_field_name_prefix=partition_field_name_prefix,
         )
 
         table_ = connection.load_table(f"{table_slice.schema}.{table_slice.table}")
@@ -91,7 +99,33 @@ def handle_output(
             },
         )
 
+    def _get_partition_field_name_prefix(self, context: OutputContext) -> str:
+        """Get partition_field_name_prefix from asset definition metadata if available, otherwise fall back to IO manager config."""
+        if (
+            context.resource_config is None
+        ):  # This doesn't seem to ever actually happen, but that's the way OutputContext is typed
+            raise ValueError(
+                "Resource config is required to get partition_field_name_prefix. Unexpected value None found for resource_config."
+            )
+
+        config = context.resource_config.get("config", {})
+        if isinstance(config, dict):
+            partition_field_name_prefix = config.get(
+                "partition_field_name_prefix", DEFAULT_PARTITION_FIELD_NAME_PREFIX
+            )
+        elif isinstance(config, IcebergCatalogConfig):
+            partition_field_name_prefix = config.partition_field_name_prefix
+        else:
+            raise ValueError(
+                f"Unable to retrieve partition_field_name_prefix from `config` attribute of resource_config with unexpected type {type(config)}"
+            )
+
+        return context.definition_metadata.get(
+            "partition_field_name_prefix", partition_field_name_prefix
+        )
+
     def _get_write_mode(self, context: OutputContext) -> WriteMode:
+        """Get write mode from asset definition metadata if available, otherwise from output metadata"""
         try:
             definition_write_mode = WriteMode(
                 context.definition_metadata.get("write_mode", DEFAULT_WRITE_MODE)