pyiceberg/schema.py

# 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.
# pylint: disable=W0511
from __future__ import annotations

import itertools
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import cached_property, partial, singledispatch
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Generic,
    List,
    Literal,
    Optional,
    Set,
    Tuple,
    TypeVar,
    Union,
)

from pydantic import Field, PrivateAttr, model_validator

from pyiceberg.exceptions import ResolveError
from pyiceberg.typedef import EMPTY_DICT, IcebergBaseModel, StructProtocol
from pyiceberg.types import (
    BinaryType,
    BooleanType,
    DateType,
    DecimalType,
    DoubleType,
    FixedType,
    FloatType,
    IcebergType,
    IntegerType,
    ListType,
    LongType,
    MapType,
    NestedField,
    PrimitiveType,
    StringType,
    StructType,
    TimestampNanoType,
    TimestampType,
    TimestamptzNanoType,
    TimestamptzType,
    TimeType,
    UnknownType,
    UUIDType,
)

if TYPE_CHECKING:
    import pyarrow as pa

    from pyiceberg.table.name_mapping import (
        NameMapping,
    )

T = TypeVar("T")
P = TypeVar("P")

INITIAL_SCHEMA_ID = 0


class Schema(IcebergBaseModel):
    """A table Schema.

    Example:
        >>> from pyiceberg import schema
        >>> from pyiceberg import types
    """

    type: Literal["struct"] = "struct"
    fields: Tuple[NestedField, ...] = Field(default_factory=tuple)
    schema_id: int = Field(alias="schema-id", default=INITIAL_SCHEMA_ID)
    identifier_field_ids: List[int] = Field(alias="identifier-field-ids", default_factory=list)

    _name_to_id: Dict[str, int] = PrivateAttr()

    def __init__(self, *fields: NestedField, **data: Any):
        if fields:
            data["fields"] = fields
        super().__init__(**data)
        self._name_to_id = index_by_name(self)

    def __str__(self) -> str:
        """Return the string representation of the Schema class."""
        return "table {\n" + "\n".join(["  " + str(field) for field in self.columns]) + "\n}"

    def __repr__(self) -> str:
        """Return the string representation of the Schema class."""
        return f"Schema({', '.join(repr(column) for column in self.columns)}, schema_id={self.schema_id}, identifier_field_ids={self.identifier_field_ids})"

    def __len__(self) -> int:
        """Return the length of an instance of the Literal class."""
        return len(self.fields)

    def __eq__(self, other: Any) -> bool:
        """Return the equality of two instances of the Schema class."""
        if not other:
            return False

        if not isinstance(other, Schema):
            return False

        if len(self.columns) != len(other.columns):
            return False

        identifier_field_ids_is_equal = self.identifier_field_ids == other.identifier_field_ids
        schema_is_equal = all(lhs == rhs for lhs, rhs in zip(self.columns, other.columns))

        return identifier_field_ids_is_equal and schema_is_equal

    @model_validator(mode="after")
    def check_schema(self) -> Schema:
        if self.identifier_field_ids:
            for field_id in self.identifier_field_ids:
                self._validate_identifier_field(field_id)

        return self

    @property
    def columns(self) -> Tuple[NestedField, ...]:
        """A tuple of the top-level fields."""
        return self.fields

    @cached_property
    def _lazy_id_to_field(self) -> Dict[int, NestedField]:
        """Return an index of field ID to NestedField instance.

        This is calculated once when called for the first time. Subsequent calls to this method will use a cached index.
        """
        return index_by_id(self)

    @cached_property
    def _lazy_id_to_parent(self) -> Dict[int, int]:
        """Returns an index of field ID to parent field IDs.

        This is calculated once when called for the first time. Subsequent calls to this method will use a cached index.
        """
        return _index_parents(self)

    @cached_property
    def _lazy_name_to_id_lower(self) -> Dict[str, int]:
        """Return an index of lower-case field names to field IDs.

        This is calculated once when called for the first time. Subsequent calls to this method will use a cached index.
        """
        return {name.lower(): field_id for name, field_id in self._name_to_id.items()}

    @cached_property
    def _lazy_id_to_name(self) -> Dict[int, str]:
        """Return an index of field ID to full name.

        This is calculated once when called for the first time. Subsequent calls to this method will use a cached index.
        """
        return index_name_by_id(self)

    @cached_property
    def _lazy_id_to_accessor(self) -> Dict[int, Accessor]:
        """Return an index of field ID to accessor.

        This is calculated once when called for the first time. Subsequent calls to this method will use a cached index.
        """
        return build_position_accessors(self)

    def as_struct(self) -> StructType:
        """Return the schema as a struct."""
        return StructType(*self.fields)

    def as_arrow(self) -> "pa.Schema":
        """Return the schema as an Arrow schema."""
        from pyiceberg.io.pyarrow import schema_to_pyarrow

        return schema_to_pyarrow(self)

    def find_field(self, name_or_id: Union[str, int], case_sensitive: bool = True) -> NestedField:
        """Find a field using a field name or field ID.

        Args:
            name_or_id (Union[str, int]): Either a field name or a field ID.
            case_sensitive (bool, optional): Whether to perform a case-sensitive lookup using a field name. Defaults to True.

        Raises:
            ValueError: When the value cannot be found.

        Returns:
            NestedField: The matched NestedField.
        """
        if isinstance(name_or_id, int):
            if name_or_id not in self._lazy_id_to_field:
                raise ValueError(f"Could not find field with id: {name_or_id}")
            return self._lazy_id_to_field[name_or_id]

        if case_sensitive:
            field_id = self._name_to_id.get(name_or_id)
        else:
            field_id = self._lazy_name_to_id_lower.get(name_or_id.lower())

        if field_id is None:
            raise ValueError(f"Could not find field with name {name_or_id}, case_sensitive={case_sensitive}")

        return self._lazy_id_to_field[field_id]

    def find_type(self, name_or_id: Union[str, int], case_sensitive: bool = True) -> IcebergType:
        """Find a field type using a field name or field ID.

        Args:
            name_or_id (Union[str, int]): Either a field name or a field ID.
            case_sensitive (bool, optional): Whether to perform a case-sensitive lookup using a field name. Defaults to True.

        Returns:
            NestedField: The type of the matched NestedField.
        """
        field = self.find_field(name_or_id=name_or_id, case_sensitive=case_sensitive)
        if not field:
            raise ValueError(f"Could not find field with name or id {name_or_id}, case_sensitive={case_sensitive}")
        return field.field_type

    @property
    def highest_field_id(self) -> int:
        return max(self._lazy_id_to_name.keys(), default=0)

    @cached_property
    def name_mapping(self) -> NameMapping:
        from pyiceberg.table.name_mapping import create_mapping_from_schema

        return create_mapping_from_schema(self)

    def find_column_name(self, column_id: int) -> Optional[str]:
        """Find a column name given a column ID.

        Args:
            column_id (int): The ID of the column.

        Returns:
            str: The column name (or None if the column ID cannot be found).
        """
        return self._lazy_id_to_name.get(column_id)

    @property
    def column_names(self) -> List[str]:
        """
        Return a list of all the column names, including nested fields.

        Excludes short names.

        Returns:
            List[str]: The column names.
        """
        return list(self._lazy_id_to_name.values())

    def accessor_for_field(self, field_id: int) -> Accessor:
        """Find a schema position accessor given a field ID.

        Args:
            field_id (int): The ID of the field.

        Raises:
            ValueError: When the value cannot be found.

        Returns:
            Accessor: An accessor for the given field ID.
        """
        if field_id not in self._lazy_id_to_accessor:
            raise ValueError(f"Could not find accessor for field with id: {field_id}")

        return self._lazy_id_to_accessor[field_id]

    def identifier_field_names(self) -> Set[str]:
        """Return the names of the identifier fields.

        Returns:
            Set of names of the identifier fields
        """
        ids = set()
        for field_id in self.identifier_field_ids:
            column_name = self.find_column_name(field_id)
            if column_name is None:
                raise ValueError(f"Could not find identifier column id: {field_id}")
            ids.add(column_name)

        return ids

    def select(self, *names: str, case_sensitive: bool = True) -> Schema:
        """Return a new schema instance pruned to a subset of columns.

        Args:
            names (List[str]): A list of column names.
            case_sensitive (bool, optional): Whether to perform a case-sensitive lookup for each column name. Defaults to True.

        Returns:
            Schema: A new schema with pruned columns.

        Raises:
            ValueError: If a column is selected that doesn't exist.
        """
        try:
            if case_sensitive:
                ids = {self._name_to_id[name] for name in names}
            else:
                ids = {self._lazy_name_to_id_lower[name.lower()] for name in names}
        except KeyError as e:
            raise ValueError(f"Could not find column: {e}") from e

        return prune_columns(self, ids)

    @property
    def field_ids(self) -> Set[int]:
        """Return the IDs of the current schema."""
        return set(self._name_to_id.values())

    def _validate_identifier_field(self, field_id: int) -> None:
        """Validate that the field with the given ID is a valid identifier field.

        Args:
          field_id: The ID of the field to validate.

        Raises:
          ValueError: If the field is not valid.
        """
        field = self.find_field(field_id)
        if not field.field_type.is_primitive:
            raise ValueError(f"Identifier field {field_id} invalid: not a primitive type field")

        if not field.required:
            raise ValueError(f"Identifier field {field_id} invalid: not a required field")

        if isinstance(field.field_type, (DoubleType, FloatType)):
            raise ValueError(f"Identifier field {field_id} invalid: must not be float or double field")

        # Check whether the nested field is in a chain of required struct fields
        # Exploring from root for better error message for list and map types
        parent_id = self._lazy_id_to_parent.get(field.field_id)
        fields: List[int] = []
        while parent_id is not None:
            fields.append(parent_id)
            parent_id = self._lazy_id_to_parent.get(parent_id)

        while fields:
            parent = self.find_field(fields.pop())
            if not parent.field_type.is_struct:
                raise ValueError(f"Cannot add field {field.name} as an identifier field: must not be nested in {parent}")

            if not parent.required:
                raise ValueError(
                    f"Cannot add field {field.name} as an identifier field: must not be nested in an optional field {parent}"
                )

    def check_format_version_compatibility(self, format_version: int) -> None:
        """Check that the schema is compatible for the given table format version.

        Args:
          format_version: The Iceberg table format version.

        Raises:
          ValueError: If the schema is not compatible for the format version.
        """
        for field in self._lazy_id_to_field.values():
            if format_version < field.field_type.minimum_format_version():
                raise ValueError(
                    f"{field.field_type} is only supported in {field.field_type.minimum_format_version()} or higher. Current format version is: {format_version}"
                )


class SchemaVisitor(Generic[T], ABC):
    def before_field(self, field: NestedField) -> None:
        """Override this method to perform an action immediately before visiting a field."""

    def after_field(self, field: NestedField) -> None:
        """Override this method to perform an action immediately after visiting a field."""

    def before_list_element(self, element: NestedField) -> None:
        """Override this method to perform an action immediately before visiting an element within a ListType."""
        self.before_field(element)

    def after_list_element(self, element: NestedField) -> None:
        """Override this method to perform an action immediately after visiting an element within a ListType."""
        self.after_field(element)

    def before_map_key(self, key: NestedField) -> None:
        """Override this method to perform an action immediately before visiting a key within a MapType."""
        self.before_field(key)

    def after_map_key(self, key: NestedField) -> None:
        """Override this method to perform an action immediately after visiting a key within a MapType."""
        self.after_field(key)

    def before_map_value(self, value: NestedField) -> None:
        """Override this method to perform an action immediately before visiting a value within a MapType."""
        self.before_field(value)

    def after_map_value(self, value: NestedField) -> None:
        """Override this method to perform an action immediately after visiting a value within a MapType."""
        self.after_field(value)

    @abstractmethod
    def schema(self, schema: Schema, struct_result: T) -> T:
        """Visit a Schema."""

    @abstractmethod
    def struct(self, struct: StructType, field_results: List[T]) -> T:
        """Visit a StructType."""

    @abstractmethod
    def field(self, field: NestedField, field_result: T) -> T:
        """Visit a NestedField."""

    @abstractmethod
    def list(self, list_type: ListType, element_result: T) -> T:
        """Visit a ListType."""

    @abstractmethod
    def map(self, map_type: MapType, key_result: T, value_result: T) -> T:
        """Visit a MapType."""

    @abstractmethod
    def primitive(self, primitive: PrimitiveType) -> T:
        """Visit a PrimitiveType."""


class PreOrderSchemaVisitor(Generic[T], ABC):
    @abstractmethod
    def schema(self, schema: Schema, struct_result: Callable[[], T]) -> T:
        """Visit a Schema."""

    @abstractmethod
    def struct(self, struct: StructType, field_results: List[Callable[[], T]]) -> T:
        """Visit a StructType."""

    @abstractmethod
    def field(self, field: NestedField, field_result: Callable[[], T]) -> T:
        """Visit a NestedField."""

    @abstractmethod
    def list(self, list_type: ListType, element_result: Callable[[], T]) -> T:
        """Visit a ListType."""

    @abstractmethod
    def map(self, map_type: MapType, key_result: Callable[[], T], value_result: Callable[[], T]) -> T:
        """Visit a MapType."""

    @abstractmethod
    def primitive(self, primitive: PrimitiveType) -> T:
        """Visit a PrimitiveType."""


class SchemaWithPartnerVisitor(Generic[P, T], ABC):
    def before_field(self, field: NestedField, field_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately before visiting a field."""

    def after_field(self, field: NestedField, field_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately after visiting a field."""

    def before_list_element(self, element: NestedField, element_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately before visiting an element within a ListType."""
        self.before_field(element, element_partner)

    def after_list_element(self, element: NestedField, element_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately after visiting an element within a ListType."""
        self.after_field(element, element_partner)

    def before_map_key(self, key: NestedField, key_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately before visiting a key within a MapType."""
        self.before_field(key, key_partner)

    def after_map_key(self, key: NestedField, key_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately after visiting a key within a MapType."""
        self.after_field(key, key_partner)

    def before_map_value(self, value: NestedField, value_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately before visiting a value within a MapType."""
        self.before_field(value, value_partner)

    def after_map_value(self, value: NestedField, value_partner: Optional[P]) -> None:
        """Override this method to perform an action immediately after visiting a value within a MapType."""
        self.after_field(value, value_partner)

    @abstractmethod
    def schema(self, schema: Schema, schema_partner: Optional[P], struct_result: T) -> T:
        """Visit a schema with a partner."""

    @abstractmethod
    def struct(self, struct: StructType, struct_partner: Optional[P], field_results: List[T]) -> T:
        """Visit a struct type with a partner."""

    @abstractmethod
    def field(self, field: NestedField, field_partner: Optional[P], field_result: T) -> T:
        """Visit a nested field with a partner."""

    @abstractmethod
    def list(self, list_type: ListType, list_partner: Optional[P], element_result: T) -> T:
        """Visit a list type with a partner."""

    @abstractmethod
    def map(self, map_type: MapType, map_partner: Optional[P], key_result: T, value_result: T) -> T:
        """Visit a map type with a partner."""

    @abstractmethod
    def primitive(self, primitive: PrimitiveType, primitive_partner: Optional[P]) -> T:
        """Visit a primitive type with a partner."""


class PrimitiveWithPartnerVisitor(SchemaWithPartnerVisitor[P, T]):
    def primitive(self, primitive: PrimitiveType, primitive_partner: Optional[P]) -> T:
        """Visit a PrimitiveType."""
        if isinstance(primitive, BooleanType):
            return self.visit_boolean(primitive, primitive_partner)
        elif isinstance(primitive, IntegerType):
            return self.visit_integer(primitive, primitive_partner)
        elif isinstance(primitive, LongType):
            return self.visit_long(primitive, primitive_partner)
        elif isinstance(primitive, FloatType):
            return self.visit_float(primitive, primitive_partner)
        elif isinstance(primitive, DoubleType):
            return self.visit_double(primitive, primitive_partner)
        elif isinstance(primitive, DecimalType):
            return self.visit_decimal(primitive, primitive_partner)
        elif isinstance(primitive, DateType):
            return self.visit_date(primitive, primitive_partner)
        elif isinstance(primitive, TimeType):
            return self.visit_time(primitive, primitive_partner)
        elif isinstance(primitive, TimestampType):
            return self.visit_timestamp(primitive, primitive_partner)
        elif isinstance(primitive, TimestampNanoType):
            return self.visit_timestamp_ns(primitive, primitive_partner)
        elif isinstance(primitive, TimestamptzType):
            return self.visit_timestamptz(primitive, primitive_partner)
        elif isinstance(primitive, TimestamptzNanoType):
            return self.visit_timestamptz_ns(primitive, primitive_partner)
        elif isinstance(primitive, StringType):
            return self.visit_string(primitive, primitive_partner)
        elif isinstance(primitive, UUIDType):
            return self.visit_uuid(primitive, primitive_partner)
        elif isinstance(primitive, FixedType):
            return self.visit_fixed(primitive, primitive_partner)
        elif isinstance(primitive, BinaryType):
            return self.visit_binary(primitive, primitive_partner)
        elif isinstance(primitive, UnknownType):
            return self.visit_unknown(primitive, primitive_partner)
        else:
            raise ValueError(f"Type not recognized: {primitive}")

    @abstractmethod
    def visit_boolean(self, boolean_type: BooleanType, partner: Optional[P]) -> T:
        """Visit a BooleanType."""

    @abstractmethod
    def visit_integer(self, integer_type: IntegerType, partner: Optional[P]) -> T:
        """Visit a IntegerType."""

    @abstractmethod
    def visit_long(self, long_type: LongType, partner: Optional[P]) -> T:
        """Visit a LongType."""

    @abstractmethod
    def visit_float(self, float_type: FloatType, partner: Optional[P]) -> T:
        """Visit a FloatType."""

    @abstractmethod
    def visit_double(self, double_type: DoubleType, partner: Optional[P]) -> T:
        """Visit a DoubleType."""

    @abstractmethod
    def visit_decimal(self, decimal_type: DecimalType, partner: Optional[P]) -> T:
        """Visit a DecimalType."""

    @abstractmethod
    def visit_date(self, date_type: DateType, partner: Optional[P]) -> T:
        """Visit a DecimalType."""

    @abstractmethod
    def visit_time(self, time_type: TimeType, partner: Optional[P]) -> T:
        """Visit a DecimalType."""

    @abstractmethod
    def visit_timestamp(self, timestamp_type: TimestampType, partner: Optional[P]) -> T:
        """Visit a TimestampType."""

    @abstractmethod
    def visit_timestamp_ns(self, timestamp_ns_type: TimestampNanoType, partner: Optional[P]) -> T:
        """Visit a TimestampNanoType."""

    @abstractmethod
    def visit_timestamptz(self, timestamptz_type: TimestamptzType, partner: Optional[P]) -> T:
        """Visit a TimestamptzType."""

    @abstractmethod
    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType, partner: Optional[P]) -> T:
        """Visit a TimestamptzNanoType."""

    @abstractmethod
    def visit_string(self, string_type: StringType, partner: Optional[P]) -> T:
        """Visit a StringType."""

    @abstractmethod
    def visit_uuid(self, uuid_type: UUIDType, partner: Optional[P]) -> T:
        """Visit a UUIDType."""

    @abstractmethod
    def visit_fixed(self, fixed_type: FixedType, partner: Optional[P]) -> T:
        """Visit a FixedType."""

    @abstractmethod
    def visit_binary(self, binary_type: BinaryType, partner: Optional[P]) -> T:
        """Visit a BinaryType."""

    @abstractmethod
    def visit_unknown(self, unknown_type: UnknownType, partner: Optional[P]) -> T:
        """Visit a UnknownType."""


class PartnerAccessor(Generic[P], ABC):
    @abstractmethod
    def schema_partner(self, partner: Optional[P]) -> Optional[P]:
        """Return the equivalent of the schema as a struct."""

    @abstractmethod
    def field_partner(self, partner_struct: Optional[P], field_id: int, field_name: str) -> Optional[P]:
        """Return the equivalent struct field by name or id in the partner struct."""

    @abstractmethod
    def list_element_partner(self, partner_list: Optional[P]) -> Optional[P]:
        """Return the equivalent list element in the partner list."""

    @abstractmethod
    def map_key_partner(self, partner_map: Optional[P]) -> Optional[P]:
        """Return the equivalent map key in the partner map."""

    @abstractmethod
    def map_value_partner(self, partner_map: Optional[P]) -> Optional[P]:
        """Return the equivalent map value in the partner map."""


@singledispatch
def visit_with_partner(
    schema_or_type: Union[Schema, IcebergType], partner: P, visitor: SchemaWithPartnerVisitor[T, P], accessor: PartnerAccessor[P]
) -> T:
    raise ValueError(f"Unsupported type: {schema_or_type}")


@visit_with_partner.register(Schema)
def _(schema: Schema, partner: P, visitor: SchemaWithPartnerVisitor[P, T], accessor: PartnerAccessor[P]) -> T:
    struct_partner = accessor.schema_partner(partner)
    return visitor.schema(schema, partner, visit_with_partner(schema.as_struct(), struct_partner, visitor, accessor))  # type: ignore


@visit_with_partner.register(StructType)
def _(struct: StructType, partner: P, visitor: SchemaWithPartnerVisitor[P, T], accessor: PartnerAccessor[P]) -> T:
    field_results = []
    for field in struct.fields:
        field_partner = accessor.field_partner(partner, field.field_id, field.name)
        visitor.before_field(field, field_partner)
        try:
            field_result = visit_with_partner(field.field_type, field_partner, visitor, accessor)  # type: ignore
            field_results.append(visitor.field(field, field_partner, field_result))
        finally:
            visitor.after_field(field, field_partner)

    return visitor.struct(struct, partner, field_results)


@visit_with_partner.register(ListType)
def _(list_type: ListType, partner: P, visitor: SchemaWithPartnerVisitor[P, T], accessor: PartnerAccessor[P]) -> T:
    element_partner = accessor.list_element_partner(partner)
    visitor.before_list_element(list_type.element_field, element_partner)
    try:
        element_result = visit_with_partner(list_type.element_type, element_partner, visitor, accessor)  # type: ignore
    finally:
        visitor.after_list_element(list_type.element_field, element_partner)

    return visitor.list(list_type, partner, element_result)


@visit_with_partner.register(MapType)
def _(map_type: MapType, partner: P, visitor: SchemaWithPartnerVisitor[P, T], accessor: PartnerAccessor[P]) -> T:
    key_partner = accessor.map_key_partner(partner)
    visitor.before_map_key(map_type.key_field, key_partner)
    try:
        key_result = visit_with_partner(map_type.key_type, key_partner, visitor, accessor)  # type: ignore
    finally:
        visitor.after_map_key(map_type.key_field, key_partner)

    value_partner = accessor.map_value_partner(partner)
    visitor.before_map_value(map_type.value_field, value_partner)
    try:
        value_result = visit_with_partner(map_type.value_type, value_partner, visitor, accessor)  # type: ignore
    finally:
        visitor.after_map_value(map_type.value_field, value_partner)
    return visitor.map(map_type, partner, key_result, value_result)


@visit_with_partner.register(PrimitiveType)
def _(primitive: PrimitiveType, partner: P, visitor: SchemaWithPartnerVisitor[P, T], _: PartnerAccessor[P]) -> T:
    return visitor.primitive(primitive, partner)


class SchemaVisitorPerPrimitiveType(SchemaVisitor[T], ABC):
    def primitive(self, primitive: PrimitiveType) -> T:
        """Visit a PrimitiveType."""
        if isinstance(primitive, FixedType):
            return self.visit_fixed(primitive)
        elif isinstance(primitive, DecimalType):
            return self.visit_decimal(primitive)
        elif isinstance(primitive, BooleanType):
            return self.visit_boolean(primitive)
        elif isinstance(primitive, IntegerType):
            return self.visit_integer(primitive)
        elif isinstance(primitive, LongType):
            return self.visit_long(primitive)
        elif isinstance(primitive, FloatType):
            return self.visit_float(primitive)
        elif isinstance(primitive, DoubleType):
            return self.visit_double(primitive)
        elif isinstance(primitive, DateType):
            return self.visit_date(primitive)
        elif isinstance(primitive, TimeType):
            return self.visit_time(primitive)
        elif isinstance(primitive, TimestampType):
            return self.visit_timestamp(primitive)
        elif isinstance(primitive, TimestampNanoType):
            return self.visit_timestamp_ns(primitive)
        elif isinstance(primitive, TimestamptzType):
            return self.visit_timestamptz(primitive)
        elif isinstance(primitive, TimestamptzNanoType):
            return self.visit_timestamptz_ns(primitive)
        elif isinstance(primitive, StringType):
            return self.visit_string(primitive)
        elif isinstance(primitive, UUIDType):
            return self.visit_uuid(primitive)
        elif isinstance(primitive, BinaryType):
            return self.visit_binary(primitive)
        elif isinstance(primitive, UnknownType):
            return self.visit_unknown(primitive)
        else:
            raise ValueError(f"Type not recognized: {primitive}")

    @abstractmethod
    def visit_fixed(self, fixed_type: FixedType) -> T:
        """Visit a FixedType."""

    @abstractmethod
    def visit_decimal(self, decimal_type: DecimalType) -> T:
        """Visit a DecimalType."""

    @abstractmethod
    def visit_boolean(self, boolean_type: BooleanType) -> T:
        """Visit a BooleanType."""

    @abstractmethod
    def visit_integer(self, integer_type: IntegerType) -> T:
        """Visit a IntegerType."""

    @abstractmethod
    def visit_long(self, long_type: LongType) -> T:
        """Visit a LongType."""

    @abstractmethod
    def visit_float(self, float_type: FloatType) -> T:
        """Visit a FloatType."""

    @abstractmethod
    def visit_double(self, double_type: DoubleType) -> T:
        """Visit a DoubleType."""

    @abstractmethod
    def visit_date(self, date_type: DateType) -> T:
        """Visit a DecimalType."""

    @abstractmethod
    def visit_time(self, time_type: TimeType) -> T:
        """Visit a DecimalType."""

    @abstractmethod
    def visit_timestamp(self, timestamp_type: TimestampType) -> T:
        """Visit a TimestampType."""

    @abstractmethod
    def visit_timestamp_ns(self, timestamp_type: TimestampNanoType) -> T:
        """Visit a TimestampNanoType."""

    @abstractmethod
    def visit_timestamptz(self, timestamptz_type: TimestamptzType) -> T:
        """Visit a TimestamptzType."""

    @abstractmethod
    def visit_timestamptz_ns(self, timestamptz_ns_type: TimestamptzNanoType) -> T:
        """Visit a TimestamptzNanoType."""

    @abstractmethod
    def visit_string(self, string_type: StringType) -> T:
        """Visit a StringType."""

    @abstractmethod
    def visit_uuid(self, uuid_type: UUIDType) -> T:
        """Visit a UUIDType."""

    @abstractmethod
    def visit_binary(self, binary_type: BinaryType) -> T:
        """Visit a BinaryType."""

    @abstractmethod
    def visit_unknown(self, unknown_type: UnknownType) -> T:
        """Visit a UnknownType."""


@dataclass(init=True, eq=True, frozen=True)
class Accessor:
    """An accessor for a specific position in a container that implements the StructProtocol."""

    position: int
    inner: Optional[Accessor] = None

    def __str__(self) -> str:
        """Return the string representation of the Accessor class."""
        return f"Accessor(position={self.position},inner={self.inner})"

    def __repr__(self) -> str:
        """Return the string representation of the Accessor class."""
        return self.__str__()

    def get(self, container: StructProtocol) -> Any:
        """Return the value at self.position in `container`.

        Args:
            container (StructProtocol): A container to access at position `self.position`.

        Returns:
            Any: The value at position `self.position` in the container.
        """
        pos = self.position
        val = container[pos]
        inner = self
        while inner.inner:
            inner = inner.inner
            val = val[inner.position]

        return val


@singledispatch
def visit(obj: Union[Schema, IcebergType], visitor: SchemaVisitor[T]) -> T:
    """Apply a schema visitor to any point within a schema.

    The function traverses the schema in post-order fashion.

    Args:
        obj (Union[Schema, IcebergType]): An instance of a Schema or an IcebergType.
        visitor (SchemaVisitor[T]): An instance of an implementation of the generic SchemaVisitor base class.

    Raises:
        NotImplementedError: If attempting to visit an unrecognized object type.
    """
    raise NotImplementedError(f"Cannot visit non-type: {obj}")


@visit.register(Schema)
def _(obj: Schema, visitor: SchemaVisitor[T]) -> T:
    """Visit a Schema with a concrete SchemaVisitor."""
    return visitor.schema(obj, visit(obj.as_struct(), visitor))


@visit.register(StructType)
def _(obj: StructType, visitor: SchemaVisitor[T]) -> T:
    """Visit a StructType with a concrete SchemaVisitor."""
    results = []

    for field in obj.fields:
        visitor.before_field(field)
        result = visit(field.field_type, visitor)
        visitor.after_field(field)
        results.append(visitor.field(field, result))

    return visitor.struct(obj, results)


@visit.register(ListType)
def _(obj: ListType, visitor: SchemaVisitor[T]) -> T:
    """Visit a ListType with a concrete SchemaVisitor."""
    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_field)
    key_result = visit(obj.key_type, visitor)
    visitor.after_map_key(obj.key_field)

    visitor.before_map_value(obj.value_field)
    value_result = visit(obj.value_type, visitor)
    visitor.after_map_value(obj.value_field)

    return visitor.map(obj, key_result, value_result)


@visit.register(PrimitiveType)
def _(obj: PrimitiveType, visitor: SchemaVisitor[T]) -> T:
    """Visit a PrimitiveType with a concrete SchemaVisitor."""
    return visitor.primitive(obj)


@singledispatch
def pre_order_visit(obj: Union[Schema, IcebergType], visitor: PreOrderSchemaVisitor[T]) -> T:
    """Apply a schema visitor to any point within a schema.

    The function traverses the schema in pre-order fashion. This is a slimmed down version
    compared to the post-order traversal (missing before and after methods), mostly
    because we don't use the pre-order traversal much.

    Args:
        obj (Union[Schema, IcebergType]): An instance of a Schema or an IcebergType.
        visitor (PreOrderSchemaVisitor[T]): An instance of an implementation of the generic PreOrderSchemaVisitor base class.

    Raises:
        NotImplementedError: If attempting to visit an unrecognized object type.
    """
    raise NotImplementedError(f"Cannot visit non-type: {obj}")


@pre_order_visit.register(Schema)
def _(obj: Schema, visitor: PreOrderSchemaVisitor[T]) -> T:
    """Visit a Schema with a concrete PreOrderSchemaVisitor."""
    return visitor.schema(obj, lambda: pre_order_visit(obj.as_struct(), visitor))


@pre_order_visit.register(StructType)
def _(obj: StructType, visitor: PreOrderSchemaVisitor[T]) -> T:
    """Visit a StructType with a concrete PreOrderSchemaVisitor."""
    return visitor.struct(
        obj,
        [
            partial(
                lambda field: visitor.field(field, partial(lambda field: pre_order_visit(field.field_type, visitor), field)),
                field,
            )
            for field in obj.fields
        ],
    )


@pre_order_visit.register(ListType)
def _(obj: ListType, visitor: PreOrderSchemaVisitor[T]) -> T:
    """Visit a ListType with a concrete PreOrderSchemaVisitor."""
    return visitor.list(obj, lambda: pre_order_visit(obj.element_type, visitor))


@pre_order_visit.register(MapType)
def _(obj: MapType, visitor: PreOrderSchemaVisitor[T]) -> T:
    """Visit a MapType with a concrete PreOrderSchemaVisitor."""
    return visitor.map(obj, lambda: pre_order_visit(obj.key_type, visitor), lambda: pre_order_visit(obj.value_type, visitor))


@pre_order_visit.register(PrimitiveType)
def _(obj: PrimitiveType, visitor: PreOrderSchemaVisitor[T]) -> T:
    """Visit a PrimitiveType with a concrete PreOrderSchemaVisitor."""
    return visitor.primitive(obj)


class _IndexById(SchemaVisitor[Dict[int, NestedField]]):
    """A schema visitor for generating a field ID to NestedField index."""

    def __init__(self) -> None:
        self._index: Dict[int, NestedField] = {}

    def schema(self, schema: Schema, struct_result: Dict[int, NestedField]) -> Dict[int, NestedField]:
        return self._index

    def struct(self, struct: StructType, field_results: List[Dict[int, NestedField]]) -> Dict[int, NestedField]:
        return self._index

    def field(self, field: NestedField, field_result: Dict[int, NestedField]) -> Dict[int, NestedField]:
        """Add the field ID to the index."""
        self._index[field.field_id] = field
        return self._index

    def list(self, list_type: ListType, element_result: Dict[int, NestedField]) -> Dict[int, NestedField]:
        """Add the list element ID to the index."""
        self._index[list_type.element_field.field_id] = list_type.element_field
        return self._index

    def map(
        self, map_type: MapType, key_result: Dict[int, NestedField], value_result: Dict[int, NestedField]
    ) -> Dict[int, NestedField]:
        """Add the key ID and value ID as individual items in the index."""
        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: PrimitiveType) -> Dict[int, NestedField]:
        return self._index


def index_by_id(schema_or_type: Union[Schema, IcebergType]) -> Dict[int, NestedField]:
    """Generate an index of field IDs to NestedField instances.

    Args:
        schema_or_type (Union[Schema, IcebergType]): A schema or type to index.

    Returns:
        Dict[int, NestedField]: An index of field IDs to NestedField instances.
    """
    return visit(schema_or_type, _IndexById())


class _IndexParents(SchemaVisitor[Dict[int, int]]):
    def __init__(self) -> None:
        self.id_to_parent: Dict[int, int] = {}
        self.id_stack: List[int] = []

    def before_field(self, field: NestedField) -> None:
        self.id_stack.append(field.field_id)

    def after_field(self, field: NestedField) -> None:
        self.id_stack.pop()

    def schema(self, schema: Schema, struct_result: Dict[int, int]) -> Dict[int, int]:
        return self.id_to_parent

    def struct(self, struct: StructType, field_results: List[Dict[int, int]]) -> Dict[int, int]:
        for field in struct.fields:
            parent_id = self.id_stack[-1] if self.id_stack else None
            if parent_id is not None:
                # fields in the root struct are not added
                self.id_to_parent[field.field_id] = parent_id

        return self.id_to_parent

    def field(self, field: NestedField, field_result: Dict[int, int]) -> Dict[int, int]:
        return self.id_to_parent

    def list(self, list_type: ListType, element_result: Dict[int, int]) -> Dict[int, int]:
        self.id_to_parent[list_type.element_id] = self.id_stack[-1]
        return self.id_to_parent

    def map(self, map_type: MapType, key_result: Dict[int, int], value_result: Dict[int, int]) -> Dict[int, int]:
        self.id_to_parent[map_type.key_id] = self.id_stack[-1]
        self.id_to_parent[map_type.value_id] = self.id_stack[-1]
        return self.id_to_parent

    def primitive(self, primitive: PrimitiveType) -> Dict[int, int]:
        return self.id_to_parent


def _index_parents(schema_or_type: Union[Schema, IcebergType]) -> Dict[int, int]:
    """Generate an index of field IDs to their parent field IDs.

    Args:
        schema_or_type (Union[Schema, IcebergType]): A schema or type to index.

    Returns:
        Dict[int, int]: An index of field IDs to their parent field IDs.
    """
    return visit(schema_or_type, _IndexParents())


class _IndexByName(SchemaVisitor[Dict[str, int]]):
    """A schema visitor for generating a field name to field ID index."""

    def __init__(self) -> None:
        self._index: Dict[str, int] = {}
        self._short_name_to_id: Dict[str, int] = {}
        self._combined_index: Dict[str, int] = {}
        self._field_names: List[str] = []
        self._short_field_names: List[str] = []

    def before_map_value(self, value: NestedField) -> None:
        if not isinstance(value.field_type, StructType):
            self._short_field_names.append(value.name)
        self._field_names.append(value.name)

    def after_map_value(self, value: NestedField) -> None:
        if not isinstance(value.field_type, StructType):
            self._short_field_names.pop()
        self._field_names.pop()

    def before_list_element(self, element: NestedField) -> None:
        """Short field names omit element when the element is a StructType."""
        if not isinstance(element.field_type, StructType):
            self._short_field_names.append(element.name)
        self._field_names.append(element.name)

    def after_list_element(self, element: NestedField) -> None:
        if not isinstance(element.field_type, StructType):
            self._short_field_names.pop()
        self._field_names.pop()

    def before_field(self, field: NestedField) -> None:
        """Store the field name."""
        self._field_names.append(field.name)
        self._short_field_names.append(field.name)

    def after_field(self, field: NestedField) -> None:
        """Remove the last field name stored."""
        self._field_names.pop()
        self._short_field_names.pop()

    def schema(self, schema: Schema, struct_result: Dict[str, int]) -> Dict[str, int]:
        return self._index

    def struct(self, struct: StructType, field_results: List[Dict[str, int]]) -> Dict[str, int]:
        return self._index

    def field(self, field: NestedField, field_result: Dict[str, int]) -> Dict[str, int]:
        """Add the field name to the index."""
        self._add_field(field.name, field.field_id)
        return self._index

    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_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_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) -> None:
        """Add a field name to the index, mapping its full name to its field ID.

        Args:
            name (str): The field name.
            field_id (int): The field ID.

        Raises:
            ValueError: If the field name is already contained in the index.
        """
        full_name = name

        if self._field_names:
            full_name = ".".join([".".join(self._field_names), name])

        if full_name in self._index:
            raise ValueError(f"Invalid schema, multiple fields for name {full_name}: {self._index[full_name]} and {field_id}")
        self._index[full_name] = field_id

        if self._short_field_names:
            short_name = ".".join([".".join(self._short_field_names), name])
            self._short_name_to_id[short_name] = field_id

    def primitive(self, primitive: PrimitiveType) -> Dict[str, int]:
        return self._index

    def by_name(self) -> Dict[str, int]:
        """Return an index of combined full and short names.

        Note: Only short names that do not conflict with full names are included.
        """
        combined_index = self._short_name_to_id.copy()
        combined_index.update(self._index)
        return combined_index

    def by_id(self) -> Dict[int, str]:
        """Return an index of ID to full names."""
        id_to_full_name = {value: key for key, value in self._index.items()}
        return id_to_full_name


def index_by_name(schema_or_type: Union[Schema, IcebergType]) -> Dict[str, int]:
    """Generate an index of field names to field IDs.

    Args:
        schema_or_type (Union[Schema, IcebergType]): A schema or type to index.

    Returns:
        Dict[str, int]: An index of field names to field IDs.
    """
    if len(schema_or_type.fields) > 0:
        indexer = _IndexByName()
        visit(schema_or_type, indexer)
        return indexer.by_name()
    else:
        return EMPTY_DICT


def index_name_by_id(schema_or_type: Union[Schema, IcebergType]) -> Dict[int, str]:
    """Generate an index of field IDs full field names.

    Args:
        schema_or_type (Union[Schema, IcebergType]): A schema or type to index.

    Returns:
        Dict[str, int]: An index of field IDs to full names.
    """
    indexer = _IndexByName()
    visit(schema_or_type, indexer)
    return indexer.by_id()


Position = int


class _BuildPositionAccessors(SchemaVisitor[Dict[Position, Accessor]]):
    """A schema visitor for generating a field ID to accessor index.

    Example:
        >>> from pyiceberg.schema import Schema
        >>> from pyiceberg.types import *
        >>> schema = Schema(
        ...     NestedField(field_id=2, name="id", field_type=IntegerType(), required=False),
        ...     NestedField(field_id=1, name="data", field_type=StringType(), required=True),
        ...     NestedField(
        ...         field_id=3,
        ...         name="location",
        ...         field_type=StructType(
        ...             NestedField(field_id=5, name="latitude", field_type=FloatType(), required=False),
        ...             NestedField(field_id=6, name="longitude", field_type=FloatType(), required=False),
        ...         ),
        ...         required=True,
        ...     ),
        ...     schema_id=1,
        ...     identifier_field_ids=[1],
        ... )
        >>> result = build_position_accessors(schema)
        >>> expected = {
        ...     2: Accessor(position=0, inner=None),
        ...     1: Accessor(position=1, inner=None),
        ...     5: Accessor(position=2, inner=Accessor(position=0, inner=None)),
        ...     6: Accessor(position=2, inner=Accessor(position=1, inner=None))
        ...     3: Accessor(position=2, inner=None),
        ... }
        >>> result == expected
        True
    """

    def schema(self, schema: Schema, struct_result: Dict[Position, Accessor]) -> Dict[Position, Accessor]:
        return struct_result

    def struct(self, struct: StructType, field_results: List[Dict[Position, Accessor]]) -> Dict[Position, Accessor]:
        result = {}

        for position, field in enumerate(struct.fields):
            if field_results[position]:
                for inner_field_id, acc in field_results[position].items():
                    result[inner_field_id] = Accessor(position, inner=acc)
            result[field.field_id] = Accessor(position)

        return result

    def field(self, field: NestedField, field_result: Dict[Position, Accessor]) -> Dict[Position, Accessor]:
        return field_result

    def list(self, list_type: ListType, element_result: Dict[Position, Accessor]) -> Dict[Position, Accessor]:
        return {}

    def map(
        self, map_type: MapType, key_result: Dict[Position, Accessor], value_result: Dict[Position, Accessor]
    ) -> Dict[Position, Accessor]:
        return {}

    def primitive(self, primitive: PrimitiveType) -> Dict[Position, Accessor]:
        return {}


def build_position_accessors(schema_or_type: Union[Schema, IcebergType]) -> Dict[int, Accessor]:
    """Generate an index of field IDs to schema position accessors.

    Args:
        schema_or_type (Union[Schema, IcebergType]): A schema or type to index.

    Returns:
        Dict[int, Accessor]: An index of field IDs to accessors.
    """
    return visit(schema_or_type, _BuildPositionAccessors())


def assign_fresh_schema_ids(schema_or_type: Union[Schema, IcebergType], next_id: Optional[Callable[[], int]] = None) -> Schema:
    """Traverses the schema, and sets new IDs."""
    return pre_order_visit(schema_or_type, _SetFreshIDs(next_id_func=next_id))


class _SetFreshIDs(PreOrderSchemaVisitor[IcebergType]):
    """Traverses the schema and assigns monotonically increasing ids."""

    old_id_to_new_id: Dict[int, int]

    def __init__(self, next_id_func: Optional[Callable[[], int]] = None) -> None:
        self.old_id_to_new_id = {}
        counter = itertools.count(1)
        self.next_id_func = next_id_func if next_id_func is not None else lambda: next(counter)

    def _get_and_increment(self, current_id: int) -> int:
        new_id = self.next_id_func()
        self.old_id_to_new_id[current_id] = new_id
        return new_id

    def schema(self, schema: Schema, struct_result: Callable[[], StructType]) -> Schema:
        return Schema(
            *struct_result().fields,
            identifier_field_ids=[self.old_id_to_new_id[field_id] for field_id in schema.identifier_field_ids],
        )

    def struct(self, struct: StructType, field_results: List[Callable[[], IcebergType]]) -> StructType:
        new_ids = [self._get_and_increment(field.field_id) for field in struct.fields]
        new_fields = []
        for field_id, field, field_type in zip(new_ids, struct.fields, field_results):
            new_fields.append(
                NestedField(
                    field_id=field_id,
                    name=field.name,
                    field_type=field_type(),
                    required=field.required,
                    doc=field.doc,
                )
            )
        return StructType(*new_fields)

    def field(self, field: NestedField, field_result: Callable[[], IcebergType]) -> IcebergType:
        return field_result()

    def list(self, list_type: ListType, element_result: Callable[[], IcebergType]) -> ListType:
        element_id = self._get_and_increment(list_type.element_id)
        return ListType(
            element_id=element_id,
            element=element_result(),
            element_required=list_type.element_required,
        )

    def map(self, map_type: MapType, key_result: Callable[[], IcebergType], value_result: Callable[[], IcebergType]) -> MapType:
        key_id = self._get_and_increment(map_type.key_id)
        value_id = self._get_and_increment(map_type.value_id)
        return MapType(
            key_id=key_id,
            key_type=key_result(),
            value_id=value_id,
            value_type=value_result(),
            value_required=map_type.value_required,
        )

    def primitive(self, primitive: PrimitiveType) -> PrimitiveType:
        return primitive


# Implementation copied from Apache Iceberg repo.
def make_compatible_name(name: str) -> str:
    if not _valid_avro_name(name):
        return _sanitize_name(name)
    return name


def _valid_avro_name(name: str) -> bool:
    length = len(name)
    assert length > 0, ValueError("Can not validate empty avro name")
    first = name[0]
    if not (first.isalpha() or first == "_"):
        return False

    for character in name[1:]:
        if not (character.isalnum() or character == "_"):
            return False
    return True


def _sanitize_name(name: str) -> str:
    sb = []
    first = name[0]
    if not (first.isalpha() or first == "_"):
        sb.append(_sanitize_char(first))
    else:
        sb.append(first)

    for character in name[1:]:
        if not (character.isalnum() or character == "_"):
            sb.append(_sanitize_char(character))
        else:
            sb.append(character)
    return "".join(sb)


def _sanitize_char(character: str) -> str:
    return "_" + character if character.isdigit() else "_x" + hex(ord(character))[2:].upper()


def sanitize_column_names(schema: Schema) -> Schema:
    """Sanitize column names to make them compatible with Avro.

    The column name should be starting with '_' or digit followed by a string only contains '_', digit or alphabet,
    otherwise it will be sanitized to conform the avro naming convention.

    Args:
        schema: The schema to be sanitized.

    Returns:
        The sanitized schema.
    """
    result = visit(schema.as_struct(), _SanitizeColumnsVisitor())
    return Schema(
        *(result or StructType()).fields,
        schema_id=schema.schema_id,
        identifier_field_ids=schema.identifier_field_ids,
    )


class _SanitizeColumnsVisitor(SchemaVisitor[Optional[IcebergType]]):
    def schema(self, schema: Schema, struct_result: Optional[IcebergType]) -> Optional[IcebergType]:
        return struct_result

    def field(self, field: NestedField, field_result: Optional[IcebergType]) -> Optional[IcebergType]:
        return NestedField(
            field_id=field.field_id,
            name=make_compatible_name(field.name),
            field_type=field_result,
            doc=field.doc,
            required=field.required,
        )

    def struct(self, struct: StructType, field_results: List[Optional[IcebergType]]) -> Optional[IcebergType]:
        return StructType(*[field for field in field_results if field is not None])

    def list(self, list_type: ListType, element_result: Optional[IcebergType]) -> Optional[IcebergType]:
        return ListType(element_id=list_type.element_id, element_type=element_result, element_required=list_type.element_required)

    def map(
        self, map_type: MapType, key_result: Optional[IcebergType], value_result: Optional[IcebergType]
    ) -> Optional[IcebergType]:
        return MapType(
            key_id=map_type.key_id,
            value_id=map_type.value_id,
            key_type=key_result,
            value_type=value_result,
            value_required=map_type.value_required,
        )

    def primitive(self, primitive: PrimitiveType) -> Optional[IcebergType]:
        return primitive


def prune_columns(schema: Schema, selected: Set[int], select_full_types: bool = True) -> Schema:
    """Prunes a column by only selecting a set of field-ids.

    Args:
        schema: The schema to be pruned.
        selected: The field-ids to be included.
        select_full_types: Return the full struct when a subset is recorded

    Returns:
        The pruned schema.
    """
    result = visit(schema.as_struct(), _PruneColumnsVisitor(selected, select_full_types))
    return Schema(
        *(result or StructType()).fields,
        schema_id=schema.schema_id,
        identifier_field_ids=list(selected.intersection(schema.identifier_field_ids)),
    )


class _PruneColumnsVisitor(SchemaVisitor[Optional[IcebergType]]):
    selected: Set[int]
    select_full_types: bool

    def __init__(self, selected: Set[int], select_full_types: bool):
        self.selected = selected
        self.select_full_types = select_full_types

    def schema(self, schema: Schema, struct_result: Optional[IcebergType]) -> Optional[IcebergType]:
        return struct_result

    def struct(self, struct: StructType, field_results: List[Optional[IcebergType]]) -> Optional[IcebergType]:
        fields = struct.fields
        selected_fields = []
        same_type = True

        for idx, projected_type in enumerate(field_results):
            field = fields[idx]
            if field.field_type == projected_type:
                selected_fields.append(field)
            elif projected_type is not None:
                same_type = False
                # Type has changed, create a new field with the projected type
                selected_fields.append(
                    NestedField(
                        field_id=field.field_id,
                        name=field.name,
                        field_type=projected_type,
                        doc=field.doc,
                        required=field.required,
                    )
                )

        if selected_fields:
            if len(selected_fields) == len(fields) and same_type is True:
                # Nothing has changed, and we can return the original struct
                return struct
            else:
                return StructType(*selected_fields)
        return None

    def field(self, field: NestedField, field_result: Optional[IcebergType]) -> Optional[IcebergType]:
        if field.field_id in self.selected:
            if self.select_full_types:
                return field.field_type
            elif field.field_type.is_struct:
                return self._project_selected_struct(field_result)
            else:
                if not field.field_type.is_primitive:
                    raise ValueError(
                        f"Cannot explicitly project List or Map types, {field.field_id}:{field.name} of type {field.field_type} was selected"
                    )
                # Selected non-struct field
                return field.field_type
        elif field_result is not None:
            # This field wasn't selected but a subfield was so include that
            return field_result
        else:
            return None

    def list(self, list_type: ListType, element_result: Optional[IcebergType]) -> Optional[IcebergType]:
        if list_type.element_id in self.selected:
            if self.select_full_types:
                return list_type
            elif list_type.element_type and list_type.element_type.is_struct:
                projected_struct = self._project_selected_struct(element_result)
                return self._project_list(list_type, projected_struct)
            else:
                if not list_type.element_type.is_primitive:
                    raise ValueError(
                        f"Cannot explicitly project List or Map types, {list_type.element_id} of type {list_type.element_type} was selected"
                    )
                return list_type
        elif element_result is not None:
            return self._project_list(list_type, element_result)
        else:
            return None

    def map(
        self, map_type: MapType, key_result: Optional[IcebergType], value_result: Optional[IcebergType]
    ) -> Optional[IcebergType]:
        if map_type.value_id in self.selected:
            if self.select_full_types:
                return map_type
            elif map_type.value_type and map_type.value_type.is_struct:
                projected_struct = self._project_selected_struct(value_result)
                return self._project_map(map_type, projected_struct)
            if not map_type.value_type.is_primitive:
                raise ValueError(
                    f"Cannot explicitly project List or Map types, Map value {map_type.value_id} of type {map_type.value_type} was selected"
                )
            return map_type
        elif value_result is not None:
            return self._project_map(map_type, value_result)
        elif map_type.key_id in self.selected:
            return map_type
        return None

    def primitive(self, primitive: PrimitiveType) -> Optional[IcebergType]:
        return None

    @staticmethod
    def _project_selected_struct(projected_field: Optional[IcebergType]) -> StructType:
        if projected_field and not isinstance(projected_field, StructType):
            raise ValueError("Expected a struct")

        if projected_field is None:
            return StructType()
        else:
            return projected_field

    @staticmethod
    def _project_list(list_type: ListType, element_result: IcebergType) -> ListType:
        if list_type.element_type == element_result:
            return list_type
        else:
            return ListType(
                element_id=list_type.element_id, element_type=element_result, element_required=list_type.element_required
            )

    @staticmethod
    def _project_map(map_type: MapType, value_result: IcebergType) -> MapType:
        if map_type.value_type == value_result:
            return map_type
        else:
            return MapType(
                key_id=map_type.key_id,
                value_id=map_type.value_id,
                key_type=map_type.key_type,
                value_type=value_result,
                value_required=map_type.value_required,
            )


@singledispatch
def promote(file_type: IcebergType, read_type: IcebergType) -> IcebergType:
    """Promotes reading a file type to a read type.

    Args:
        file_type (IcebergType): The type of the Avro file.
        read_type (IcebergType): The requested read type.

    Raises:
        ResolveError: If attempting to resolve an unrecognized object type.
    """
    if file_type == read_type:
        return file_type
    else:
        raise ResolveError(f"Cannot promote {file_type} to {read_type}")


@promote.register(IntegerType)
def _(file_type: IntegerType, read_type: IcebergType) -> IcebergType:
    if isinstance(read_type, LongType):
        # Ints/Longs are binary compatible in Avro, so this is okay
        return read_type
    else:
        raise ResolveError(f"Cannot promote an int to {read_type}")


@promote.register(FloatType)
def _(file_type: FloatType, read_type: IcebergType) -> IcebergType:
    if isinstance(read_type, DoubleType):
        # A double type is wider
        return read_type
    else:
        raise ResolveError(f"Cannot promote an float to {read_type}")


@promote.register(StringType)
def _(file_type: StringType, read_type: IcebergType) -> IcebergType:
    if isinstance(read_type, BinaryType):
        return read_type
    else:
        raise ResolveError(f"Cannot promote an string to {read_type}")


@promote.register(BinaryType)
def _(file_type: BinaryType, read_type: IcebergType) -> IcebergType:
    if isinstance(read_type, StringType):
        return read_type
    else:
        raise ResolveError(f"Cannot promote an binary to {read_type}")


@promote.register(DecimalType)
def _(file_type: DecimalType, read_type: IcebergType) -> IcebergType:
    if isinstance(read_type, DecimalType):
        if file_type.precision <= read_type.precision and file_type.scale == file_type.scale:
            return read_type
        else:
            raise ResolveError(f"Cannot reduce precision from {file_type} to {read_type}")
    else:
        raise ResolveError(f"Cannot promote an decimal to {read_type}")


@promote.register(FixedType)
def _(file_type: FixedType, read_type: IcebergType) -> IcebergType:
    if isinstance(read_type, UUIDType) and len(file_type) == 16:
        # Since pyarrow reads parquet UUID as fixed 16-byte binary, the promotion is needed to ensure read compatibility
        return read_type
    else:
        raise ResolveError(f"Cannot promote {file_type} to {read_type}")


def _check_schema_compatible(requested_schema: Schema, provided_schema: Schema) -> None:
    """
    Check if the `provided_schema` is compatible with `requested_schema`.

    Both Schemas must have valid IDs and share the same ID for the same field names.

    Two schemas are considered compatible when:
    1. All `required` fields in `requested_schema` are present and are also `required` in the `provided_schema`
    2. Field Types are consistent for fields that are present in both schemas. I.e. the field type
       in the `provided_schema` can be promoted to the field type of the same field ID in `requested_schema`

    Raises:
        ValueError: If the schemas are not compatible.
    """
    pre_order_visit(requested_schema, _SchemaCompatibilityVisitor(provided_schema))


class _SchemaCompatibilityVisitor(PreOrderSchemaVisitor[bool]):
    provided_schema: Schema

    def __init__(self, provided_schema: Schema):
        from rich.console import Console
        from rich.table import Table as RichTable

        self.provided_schema = provided_schema
        self.rich_table = RichTable(show_header=True, header_style="bold")
        self.rich_table.add_column("")
        self.rich_table.add_column("Table field")
        self.rich_table.add_column("Dataframe field")
        self.console = Console(record=True)

    def _is_field_compatible(self, lhs: NestedField) -> bool:
        # Validate nullability first.
        # An optional field can be missing in the provided schema
        # But a required field must exist as a required field
        try:
            rhs = self.provided_schema.find_field(lhs.field_id)
        except ValueError:
            if lhs.required:
                self.rich_table.add_row("❌", str(lhs), "Missing")
                return False
            else:
                self.rich_table.add_row("✅", str(lhs), "Missing")
                return True

        if lhs.required and not rhs.required:
            self.rich_table.add_row("❌", str(lhs), str(rhs))
            return False

        # Check type compatibility
        if lhs.field_type == rhs.field_type:
            self.rich_table.add_row("✅", str(lhs), str(rhs))
            return True
        # We only check that the parent node is also of the same type.
        # We check the type of the child nodes when we traverse them later.
        elif any(
            (isinstance(lhs.field_type, container_type) and isinstance(rhs.field_type, container_type))
            for container_type in {StructType, MapType, ListType}
        ):
            self.rich_table.add_row("✅", str(lhs), str(rhs))
            return True
        else:
            try:
                # If type can be promoted to the requested schema
                # it is considered compatible
                promote(rhs.field_type, lhs.field_type)
                self.rich_table.add_row("✅", str(lhs), str(rhs))
                return True
            except ResolveError:
                self.rich_table.add_row("❌", str(lhs), str(rhs))
                return False

    def schema(self, schema: Schema, struct_result: Callable[[], bool]) -> bool:
        if not (result := struct_result()):
            self.console.print(self.rich_table)
            raise ValueError(f"Mismatch in fields:\n{self.console.export_text()}")
        return result

    def struct(self, struct: StructType, field_results: List[Callable[[], bool]]) -> bool:
        results = [result() for result in field_results]
        return all(results)

    def field(self, field: NestedField, field_result: Callable[[], bool]) -> bool:
        return self._is_field_compatible(field) and field_result()

    def list(self, list_type: ListType, element_result: Callable[[], bool]) -> bool:
        return self._is_field_compatible(list_type.element_field) and element_result()

    def map(self, map_type: MapType, key_result: Callable[[], bool], value_result: Callable[[], bool]) -> bool:
        return all(
            [
                self._is_field_compatible(map_type.key_field),
                self._is_field_compatible(map_type.value_field),
                key_result(),
                value_result(),
            ]
        )

    def primitive(self, primitive: PrimitiveType) -> bool:
        return True