You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by fo...@apache.org on 2024/02/25 22:33:58 UTC
(iceberg-python) branch main updated: Improve error message in case of a mismatch (#352)
This is an automated email from the ASF dual-hosted git repository.
fokko pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new 339ba53 Improve error message in case of a mismatch (#352)
339ba53 is described below
commit 339ba53fe2f1563c2ed2a99d09c0986b2e7e8caf
Author: Fokko Driesprong <fo...@tabular.io>
AuthorDate: Sun Feb 25 23:33:53 2024 +0100
Improve error message in case of a mismatch (#352)
* Nice error
* Simplify a bit
* Property
Co-authored-by: Honah J. <un...@gmail.com>
* Property
Co-authored-by: Honah J. <un...@gmail.com>
* Whitespace
---------
Co-authored-by: Honah J. <un...@gmail.com>
---
pyiceberg/io/pyarrow.py | 4 +++
pyiceberg/schema.py | 12 +++++++
pyiceberg/table/__init__.py | 45 +++++++++++++++++++++++-
pyiceberg/table/name_mapping.py | 6 +++-
tests/table/test_init.py | 78 +++++++++++++++++++++++++++++++++++++++++
5 files changed, 143 insertions(+), 2 deletions(-)
diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index 76bb041..d41f7a0 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -655,6 +655,10 @@ def pyarrow_to_schema(schema: pa.Schema, name_mapping: Optional[NameMapping] = N
return visit_pyarrow(schema, visitor)
+def _pyarrow_to_schema_without_ids(schema: pa.Schema) -> Schema:
+ return visit_pyarrow(schema, _ConvertToIcebergWithoutIDs())
+
+
@singledispatch
def visit_pyarrow(obj: Union[pa.DataType, pa.Schema], visitor: PyArrowSchemaVisitor[T]) -> T:
"""Apply a pyarrow schema visitor to any point within a schema.
diff --git a/pyiceberg/schema.py b/pyiceberg/schema.py
index 6dd174f..e805895 100644
--- a/pyiceberg/schema.py
+++ b/pyiceberg/schema.py
@@ -22,6 +22,7 @@ from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import cached_property, partial, singledispatch
from typing import (
+ TYPE_CHECKING,
Any,
Callable,
Dict,
@@ -62,6 +63,11 @@ from pyiceberg.types import (
UUIDType,
)
+if TYPE_CHECKING:
+ from pyiceberg.table.name_mapping import (
+ NameMapping,
+ )
+
T = TypeVar("T")
P = TypeVar("P")
@@ -221,6 +227,12 @@ class Schema(IcebergBaseModel):
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.
diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py
index b07da7e..060f137 100644
--- a/pyiceberg/table/__init__.py
+++ b/pyiceberg/table/__init__.py
@@ -84,7 +84,11 @@ from pyiceberg.table.metadata import (
TableMetadata,
TableMetadataUtil,
)
-from pyiceberg.table.name_mapping import NameMapping, parse_mapping_from_json, update_mapping
+from pyiceberg.table.name_mapping import (
+ NameMapping,
+ parse_mapping_from_json,
+ update_mapping,
+)
from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef
from pyiceberg.table.snapshots import (
Operation,
@@ -129,6 +133,41 @@ TABLE_ROOT_ID = -1
_JAVA_LONG_MAX = 9223372036854775807
+def _check_schema(table_schema: Schema, other_schema: "pa.Schema") -> None:
+ from pyiceberg.io.pyarrow import _pyarrow_to_schema_without_ids, pyarrow_to_schema
+
+ name_mapping = table_schema.name_mapping
+ try:
+ task_schema = pyarrow_to_schema(other_schema, name_mapping=name_mapping)
+ except ValueError as e:
+ other_schema = _pyarrow_to_schema_without_ids(other_schema)
+ additional_names = set(other_schema.column_names) - set(table_schema.column_names)
+ raise ValueError(
+ f"PyArrow table contains more columns: {', '.join(sorted(additional_names))}. Update the schema first (hint, use union_by_name)."
+ ) from e
+
+ if table_schema.as_struct() != task_schema.as_struct():
+ from rich.console import Console
+ from rich.table import Table as RichTable
+
+ console = Console(record=True)
+
+ rich_table = RichTable(show_header=True, header_style="bold")
+ rich_table.add_column("")
+ rich_table.add_column("Table field")
+ rich_table.add_column("Dataframe field")
+
+ for lhs in table_schema.fields:
+ try:
+ rhs = task_schema.find_field(lhs.field_id)
+ rich_table.add_row("✅" if lhs == rhs else "❌", str(lhs), str(rhs))
+ except ValueError:
+ rich_table.add_row("❌", str(lhs), "Missing")
+
+ console.print(rich_table)
+ raise ValueError(f"Mismatch in fields:\n{console.export_text()}")
+
+
class TableProperties:
PARQUET_ROW_GROUP_SIZE_BYTES = "write.parquet.row-group-size-bytes"
PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT = 128 * 1024 * 1024 # 128 MB
@@ -1015,6 +1054,8 @@ class Table:
if len(self.spec().fields) > 0:
raise ValueError("Cannot write to partitioned tables")
+ _check_schema(self.schema(), other_schema=df.schema)
+
with self.update_snapshot().fast_append() as update_snapshot:
# skip writing data files if the dataframe is empty
if df.shape[0] > 0:
@@ -1045,6 +1086,8 @@ class Table:
if len(self.spec().fields) > 0:
raise ValueError("Cannot write to partitioned tables")
+ _check_schema(self.schema(), other_schema=df.schema)
+
with self.update_snapshot().overwrite() as update_snapshot:
# skip writing data files if the dataframe is empty
if df.shape[0] > 0:
diff --git a/pyiceberg/table/name_mapping.py b/pyiceberg/table/name_mapping.py
index 9497483..baa15f1 100644
--- a/pyiceberg/table/name_mapping.py
+++ b/pyiceberg/table/name_mapping.py
@@ -26,7 +26,7 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from collections import ChainMap
from functools import cached_property, singledispatch
-from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
+from typing import Any, Dict, Generic, Iterator, List, Optional, TypeVar, Union
from pydantic import Field, conlist, field_validator, model_serializer
@@ -97,6 +97,10 @@ class NameMapping(IcebergRootModel[List[MappedField]]):
"""Return the number of mappings."""
return len(self.root)
+ def __iter__(self) -> Iterator[MappedField]:
+ """Iterate over the mapped fields."""
+ return iter(self.root)
+
def __str__(self) -> str:
"""Convert the name-mapping into a nicely formatted string."""
if len(self.root) == 0:
diff --git a/tests/table/test_init.py b/tests/table/test_init.py
index efee43b..7bf65a6 100644
--- a/tests/table/test_init.py
+++ b/tests/table/test_init.py
@@ -19,6 +19,7 @@ import uuid
from copy import copy
from typing import Dict
+import pyarrow as pa
import pytest
from sortedcontainers import SortedList
@@ -58,6 +59,7 @@ from pyiceberg.table import (
Table,
UpdateSchema,
_apply_table_update,
+ _check_schema,
_generate_snapshot_id,
_match_deletes_to_data_file,
_TableMetadataUpdateContext,
@@ -982,3 +984,79 @@ def test_correct_schema() -> None:
_ = t.scan(snapshot_id=-1).projection()
assert "Snapshot not found: -1" in str(exc_info.value)
+
+
+def test_schema_mismatch_type(table_schema_simple: Schema) -> None:
+ other_schema = pa.schema((
+ pa.field("foo", pa.string(), nullable=True),
+ pa.field("bar", pa.decimal128(18, 6), nullable=False),
+ pa.field("baz", pa.bool_(), nullable=True),
+ ))
+
+ expected = r"""Mismatch in fields:
+┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+┃ ┃ Table field ┃ Dataframe field ┃
+┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+│ ✅ │ 1: foo: optional string │ 1: foo: optional string │
+│ ❌ │ 2: bar: required int │ 2: bar: required decimal\(18, 6\) │
+│ ✅ │ 3: baz: optional boolean │ 3: baz: optional boolean │
+└────┴──────────────────────────┴─────────────────────────────────┘
+"""
+
+ with pytest.raises(ValueError, match=expected):
+ _check_schema(table_schema_simple, other_schema)
+
+
+def test_schema_mismatch_nullability(table_schema_simple: Schema) -> None:
+ other_schema = pa.schema((
+ pa.field("foo", pa.string(), nullable=True),
+ pa.field("bar", pa.int32(), nullable=True),
+ pa.field("baz", pa.bool_(), nullable=True),
+ ))
+
+ expected = """Mismatch in fields:
+┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+┃ ┃ Table field ┃ Dataframe field ┃
+┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+│ ✅ │ 1: foo: optional string │ 1: foo: optional string │
+│ ❌ │ 2: bar: required int │ 2: bar: optional int │
+│ ✅ │ 3: baz: optional boolean │ 3: baz: optional boolean │
+└────┴──────────────────────────┴──────────────────────────┘
+"""
+
+ with pytest.raises(ValueError, match=expected):
+ _check_schema(table_schema_simple, other_schema)
+
+
+def test_schema_mismatch_missing_field(table_schema_simple: Schema) -> None:
+ other_schema = pa.schema((
+ pa.field("foo", pa.string(), nullable=True),
+ pa.field("baz", pa.bool_(), nullable=True),
+ ))
+
+ expected = """Mismatch in fields:
+┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+┃ ┃ Table field ┃ Dataframe field ┃
+┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+│ ✅ │ 1: foo: optional string │ 1: foo: optional string │
+│ ❌ │ 2: bar: required int │ Missing │
+│ ✅ │ 3: baz: optional boolean │ 3: baz: optional boolean │
+└────┴──────────────────────────┴──────────────────────────┘
+"""
+
+ with pytest.raises(ValueError, match=expected):
+ _check_schema(table_schema_simple, other_schema)
+
+
+def test_schema_mismatch_additional_field(table_schema_simple: Schema) -> None:
+ other_schema = pa.schema((
+ pa.field("foo", pa.string(), nullable=True),
+ pa.field("bar", pa.int32(), nullable=True),
+ pa.field("baz", pa.bool_(), nullable=True),
+ pa.field("new_field", pa.date32(), nullable=True),
+ ))
+
+ expected = r"PyArrow table contains more columns: new_field. Update the schema first \(hint, use union_by_name\)."
+
+ with pytest.raises(ValueError, match=expected):
+ _check_schema(table_schema_simple, other_schema)