Python: Use Pydantic for (de)serialization

python/pyproject.toml

@@ -41,6 +41,8 @@ packages = [
 python = "^3.8"
 mmh3 = "^3.0.0"
 
+pydantic = "^1.9.1"
+
 pyarrow = { version = "^8.0.0", optional = true }
 
 [tool.poetry.dev-dependencies]

python/src/iceberg/schema.py

@@ -28,6 +28,8 @@
     TypeVar,
 )
 
+from pydantic import Field, PrivateAttr
+
 from iceberg.files import StructProtocol
 from iceberg.types import (
     IcebergType,
@@ -37,27 +39,38 @@
     PrimitiveType,
     StructType,
 )
+from iceberg.utils.iceberg_base_model import IcebergBaseModel
 
 T = TypeVar("T")
 
 
-class Schema:
+class Schema(IcebergBaseModel):
     """A table Schema
 
     Example:
         >>> from iceberg import schema
         >>> from iceberg import types
     """
 
-    def __init__(self, *columns: NestedField, schema_id: int, identifier_field_ids: list[int] | None = None):
-        self._struct = StructType(*columns)
-        self._schema_id = schema_id
-        self._identifier_field_ids = identifier_field_ids or []
-        self._name_to_id: dict[str, int] = index_by_name(self)
-        self._name_to_id_lower: dict[str, int] = {}  # Should be accessed through self._lazy_name_to_id_lower()
-        self._id_to_field: dict[int, NestedField] = {}  # Should be accessed through self._lazy_id_to_field()
-        self._id_to_name: dict[int, str] = {}  # Should be accessed through self._lazy_id_to_name()
-        self._id_to_accessor: dict[int, Accessor] = {}  # Should be accessed through self._lazy_id_to_accessor()
+    fields: tuple[NestedField, ...] = Field()
+    schema_id: int = Field(alias="schema-id")
+    identifier_field_ids: list[int] = Field(alias="identifier-field-ids", default_factory=list)
+
+    _name_to_id: dict[str, int] = PrivateAttr()
+    # Should be accessed through self._lazy_name_to_id_lower()
+    _name_to_id_lower: dict[str, int] = PrivateAttr(default_factory=dict)
+    # Should be accessed through self._lazy_id_to_field()
+    _id_to_field: dict[int, NestedField] = PrivateAttr(default_factory=dict)
+    # Should be accessed through self._lazy_id_to_name()
+    _id_to_name: dict[int, str] = PrivateAttr(default_factory=dict)
+    # Should be accessed through self._lazy_id_to_accessor()
+    _id_to_accessor: dict[int, Accessor] = PrivateAttr(default_factory=dict)
+
+    def __init__(self, *fields: NestedField, **data):
+        if fields:
+            data["fields"] = fields
+        super().__init__(**data)
+        self._name_to_id = index_by_name(self)
 
     def __str__(self):
         return "table {\n" + "\n".join(["  " + str(field) for field in self.columns]) + "\n}"
@@ -85,16 +98,7 @@ def __eq__(self, other) -> bool:
     @property
     def columns(self) -> tuple[NestedField, ...]:
         """A list of the top-level fields in the underlying struct"""
-        return self._struct.fields
-
-    @property
-    def schema_id(self) -> int:
-        """The ID of this Schema"""
-        return self._schema_id
-
-    @property
-    def identifier_field_ids(self) -> list[int]:
-        return self._identifier_field_ids
+        return self.fields
 
     def _lazy_id_to_field(self) -> dict[int, NestedField]:
         """Returns an index of field ID to NestedField instance
@@ -134,7 +138,7 @@ def _lazy_id_to_accessor(self) -> dict[int, Accessor]:
 
     def as_struct(self) -> StructType:
         """Returns the underlying struct"""
-        return self._struct
+        return StructType(*self.fields)
 
     def find_field(self, name_or_id: str | int, case_sensitive: bool = True) -> NestedField:
         """Find a field using a field name or field ID
@@ -343,23 +347,23 @@ def _(obj: StructType, visitor: SchemaVisitor[T]) -> T:
 def _(obj: ListType, visitor: SchemaVisitor[T]) -> T:
     """Visit a ListType with a concrete SchemaVisitor"""
 
-    visitor.before_list_element(obj.element)
-    result = visit(obj.element.field_type, visitor)
-    visitor.after_list_element(obj.element)
+    visitor.before_list_element(obj.element_field)
+    result = visit(obj.element_type, visitor)
+    visitor.after_list_element(obj.element_field)
 
     return visitor.list(obj, result)
 
 
 @visit.register(MapType)
 def _(obj: MapType, visitor: SchemaVisitor[T]) -> T:
     """Visit a MapType with a concrete SchemaVisitor"""
-    visitor.before_map_key(obj.key)
-    key_result = visit(obj.key.field_type, visitor)
-    visitor.after_map_key(obj.key)
+    visitor.before_map_key(obj.key_field)
+    key_result = visit(obj.key_type, visitor)
+    visitor.after_map_key(obj.key_field)
 
-    visitor.before_map_value(obj.value)
-    value_result = visit(obj.value.field_type, visitor)
-    visitor.after_list_element(obj.value)
+    visitor.before_map_value(obj.value_field)
+    value_result = visit(obj.value_type, visitor)
+    visitor.after_list_element(obj.value_field)
 
     return visitor.map(obj, key_result, value_result)
 
@@ -389,13 +393,13 @@ def field(self, field: NestedField, field_result) -> dict[int, NestedField]:
 
     def list(self, list_type: ListType, element_result) -> dict[int, NestedField]:
         """Add the list element ID to the index"""
-        self._index[list_type.element.field_id] = list_type.element
+        self._index[list_type.element_field.field_id] = list_type.element_field
         return self._index
 
     def map(self, map_type: MapType, key_result, value_result) -> dict[int, NestedField]:
         """Add the key ID and value ID as individual items in the index"""
-        self._index[map_type.key.field_id] = map_type.key
-        self._index[map_type.value.field_id] = map_type.value
+        self._index[map_type.key_field.field_id] = map_type.key_field
+        self._index[map_type.value_field.field_id] = map_type.value_field
         return self._index
 
     def primitive(self, primitive) -> dict[int, NestedField]:
@@ -458,13 +462,13 @@ def field(self, field: NestedField, field_result: dict[str, int]) -> dict[str, i
 
     def list(self, list_type: ListType, element_result: dict[str, int]) -> dict[str, int]:
         """Add the list element name to the index"""
-        self._add_field(list_type.element.name, list_type.element.field_id)
+        self._add_field(list_type.element_field.name, list_type.element_field.field_id)
         return self._index
 
     def map(self, map_type: MapType, key_result: dict[str, int], value_result: dict[str, int]) -> dict[str, int]:
         """Add the key name and value name as individual items in the index"""
-        self._add_field(map_type.key.name, map_type.key.field_id)
-        self._add_field(map_type.value.name, map_type.value.field_id)
+        self._add_field(map_type.key_field.name, map_type.key_field.field_id)
+        self._add_field(map_type.value_field.name, map_type.value_field.field_id)
         return self._index
 
     def _add_field(self, name: str, field_id: int):

python/src/iceberg/serializers.py

@@ -0,0 +1,75 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import codecs
+import json
+from typing import Union
+
+from iceberg.io.base import InputFile, InputStream, OutputFile
+from iceberg.table.metadata import TableMetadata, TableMetadataV1, TableMetadataV2
+
+
+class FromByteStream:
+    """A collection of methods that deserialize dictionaries into Iceberg objects"""
+
+    @staticmethod
+    def table_metadata(byte_stream: InputStream, encoding: str = "utf-8") -> TableMetadata:
+        """Instantiate a TableMetadata object from a byte stream
+
+        Args:
+            byte_stream: A file-like byte stream object
+            encoding (default "utf-8"): The byte encoder to use for the reader
+        """
+        reader = codecs.getreader(encoding)
+        metadata = json.load(reader(byte_stream))  # type: ignore
+        return TableMetadata.parse_obj(metadata)  # type: ignore
+
+
+class FromInputFile:
+    """A collection of methods that deserialize InputFiles into Iceberg objects"""
+
+    @staticmethod
+    def table_metadata(input_file: InputFile, encoding: str = "utf-8") -> TableMetadata:
+        """Create a TableMetadata instance from an input file
+
+        Args:
+            input_file (InputFile): A custom implementation of the iceberg.io.file.InputFile abstract base class
+            encoding (str): Encoding to use when loading bytestream
+
+        Returns:
+            TableMetadata: A table metadata instance
+
+        """
+        return FromByteStream.table_metadata(byte_stream=input_file.open(), encoding=encoding)
+
+
+class ToOutputFile:
+    """A collection of methods that serialize Iceberg objects into files given an OutputFile instance"""
+
+    @staticmethod
+    def table_metadata(
+        metadata: Union[TableMetadataV1, TableMetadataV2], output_file: OutputFile, overwrite: bool = False
+    ) -> None:
+        """Write a TableMetadata instance to an output file
+
+        Args:
+            output_file (OutputFile): A custom implementation of the iceberg.io.file.OutputFile abstract base class
+            overwrite (bool): Where to overwrite the file if it already exists. Defaults to `False`.
+        """
+        f = output_file.create(overwrite=overwrite)
+        f.write(metadata.json().encode("utf-8"))
+        f.close()

python/src/iceberg/table/metadata.py

@@ -0,0 +1,151 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import List, Literal, Union
+
+from pydantic import Field
+
+from iceberg.schema import Schema
+from iceberg.utils.iceberg_base_model import IcebergBaseModel
+
+
+class TableMetadataCommonFields(IcebergBaseModel):
+    """Metadata for an Iceberg table as specified in the Apache Iceberg
+    spec (https://iceberg.apache.org/spec/#iceberg-table-spec)"""
+
+    table_uuid: str = Field(alias="table-uuid")
+    """A UUID that identifies the table, generated when the table is created.
+    Implementations must throw an exception if a table’s UUID does not match
+    the expected UUID after refreshing metadata."""
+
+    location: str
+    """The table’s base location. This is used by writers to determine where
+    to store data files, manifest files, and table metadata files."""
+
+    last_updated_ms: int = Field(alias="last-updated-ms")
+    """Timestamp in milliseconds from the unix epoch when the table
+    was last updated. Each table metadata file should update this
+    field just before writing."""
+
+    last_column_id: int = Field(alias="last-column-id")
+    """An integer; the highest assigned column ID for the table.
+    This is used to ensure columns are always assigned an unused ID
+    when evolving schemas."""
+
+    schemas: List[Schema] = Field()
+    """A list of schemas, stored as objects with schema-id."""
+
+    current_schema_id: int = Field(alias="current-schema-id")
+    """ID of the table’s current schema."""
+
+    partition_specs: list = Field(alias="partition-specs")
+    """A list of partition specs, stored as full partition spec objects."""
+
+    default_spec_id: int = Field(alias="default-spec-id")
+    """ID of the “current” spec that writers should use by default."""
+
+    last_partition_id: int = Field(alias="last-partition-id")
+    """An integer; the highest assigned partition field ID across all
+    partition specs for the table. This is used to ensure partition fields
+    are always assigned an unused ID when evolving specs."""
+
+    properties: dict
+    """	A string to string map of table properties. This is used to
+    control settings that affect reading and writing and is not intended
+    to be used for arbitrary metadata. For example, commit.retry.num-retries
+    is used to control the number of commit retries."""
+
+    current_snapshot_id: int = Field(alias="current-snapshot-id")
+    """ID of the current table snapshot."""
+
+    snapshots: list
+    """A list of valid snapshots. Valid snapshots are snapshots for which
+    all data files exist in the file system. A data file must not be
+    deleted from the file system until the last snapshot in which it was
+    listed is garbage collected."""
+
+    snapshot_log: list = Field(alias="snapshot-log")
+    """A list (optional) of timestamp and snapshot ID pairs that encodes
+    changes to the current snapshot for the table. Each time the
+    current-snapshot-id is changed, a new entry should be added with the
+    last-updated-ms and the new current-snapshot-id. When snapshots are
+    expired from the list of valid snapshots, all entries before a snapshot
+    that has expired should be removed."""
+
+    metadata_log: list = Field(alias="metadata-log")
+    """A list (optional) of timestamp and metadata file location pairs that
+    encodes changes to the previous metadata files for the table. Each time
+    a new metadata file is created, a new entry of the previous metadata
+    file location should be added to the list. Tables can be configured to
+    remove oldest metadata log entries and keep a fixed-size log of the most
+    recent entries after a commit."""
+
+    sort_orders: list = Field(alias="sort-orders")
+    """A list of sort orders, stored as full sort order objects."""
+
+    default_sort_order_id: int = Field(alias="default-sort-order-id")
+    """Default sort order id of the table. Note that this could be used by
+    writers, but is not used when reading because reads use the specs stored
+     in manifest files."""
+
+
+class TableMetadataV1(TableMetadataCommonFields, IcebergBaseModel):
+    format_version: Literal[1] = Field(alias="format-version")
+    """An integer version number for the format. Currently, this can be 1 or 2
+    based on the spec. Implementations must throw an exception if a table’s
+    version is higher than the supported version."""
+
+    schema_: Schema = Field(alias="schema")
+    """The table’s current schema. (Deprecated: use schemas and
+    current-schema-id instead)"""
+
+    partition_spec: dict = Field(alias="partition-spec")
+    """The table’s current partition spec, stored as only fields.
+    Note that this is used by writers to partition data, but is
+    not used when reading because reads use the specs stored in
+    manifest files. (Deprecated: use partition-specs and default-spec-id
+    instead)"""
+
+
+class TableMetadataV2(TableMetadataCommonFields, IcebergBaseModel):
+    format_version: Literal[2] = Field(alias="format-version")
+    """An integer version number for the format. Currently, this can be 1 or 2
+    based on the spec. Implementations must throw an exception if a table’s
+    version is higher than the supported version."""
+
+    last_sequence_number: int = Field(alias="last-sequence-number")
+    """The table’s highest assigned sequence number, a monotonically
+    increasing long that tracks the order of snapshots in a table."""
+
+
+class TableMetadata:
+    # Once this has been resolved, we can simplify this: https://github.com/samuelcolvin/pydantic/issues/3846
+    # TableMetadata = Annotated[Union[TableMetadataV1, TableMetadataV2], Field(alias="format-version", discriminator="format-version")]
+
+    @staticmethod
+    def parse_obj(data: dict) -> Union[TableMetadataV1, TableMetadataV2]:
+        if "format-version" not in data:
+            raise ValueError(f"Missing format-version in TableMetadata: {data}")
+
+        format_version = data["format-version"]
+
+        if format_version == 1:
+            return TableMetadataV1(**data)
+        elif format_version == 2:
+            return TableMetadataV2(**data)
+        else:
+            raise ValueError(f"Unknown format version: {format_version}")