diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index 786223d47e..6e3d35eb35 100644 --- a/mkdocs/docs/api.md +++ b/mkdocs/docs/api.md @@ -517,6 +517,21 @@ long: [[2.349014],[4.896029],[6.0989],[-122.431297]] PyIceberg supports upsert operations, meaning that it is able to merge an Arrow table into an Iceberg table. Rows are considered the same based on the [identifier field](https://iceberg.apache.org/spec/?column-projection#identifier-field-ids). If a row is already in the table, it will update that row. If a row cannot be found, it will insert that new row. + + + + +!!! note "Join columns" + Use `join_cols` to pick the columns to match on; when omitted, the table's identifier fields are used. Unsupported join columns are rejected with a descriptive error before anything is written. + + - **Supported**: boolean, integer, long, decimal, date, time, timestamp, string, and binary columns. + - **Not supported**: `float` and `double` columns, because floating-point equality is unreliable; nested columns (structs, lists, and maps), including identifier fields nested inside a struct; and UUID columns. + - **Input requirements**: each join column must be present, must not contain null values, and the input rows must be unique on the join columns. Dictionary-encoded, `string_view`, `binary_view`, extension-type, and `pa.null()` columns must be cast or decoded to a plain type first. + + + + + Consider the following table, with some data: ```python @@ -1081,11 +1096,16 @@ Expert Iceberg users may choose to commit existing parquet files to the Iceberg + + !!! note "Name Mapping and Field IDs" `add_files` can work with Parquet files both with and without field IDs in their metadata: + - **Files with field IDs**: When field IDs are present in the Parquet metadata, they must match the corresponding field IDs in the Iceberg table schema. This is common for files generated by tools like Spark or when using or other libraries with explicit field ID metadata. - **Files without field IDs**: When field IDs are absent, the table must have a [Name Mapping](https://iceberg.apache.org/spec/?h=name+mapping#name-mapping-serialization) to map field names to Iceberg field IDs. `add_files` will automatically create a Name Mapping based on the table's current schema if one doesn't already exist. + + !!! note "Partitions" `add_files` only requires the client to read the existing parquet files' metadata footer to infer the partition value of each file. This implementation also supports adding files to Iceberg tables with partition transforms like `MonthTransform`, and `TruncateTransform` which preserve the order of the values after the transformation (Any Transform that has the `preserves_order` property set to True is supported). Please note that if the column statistics of the `PartitionField`'s source column are not present in the parquet metadata, the partition value is inferred as `None`. diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 303b3db135..3046efe30e 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -914,10 +914,10 @@ def upsert( if not when_matched_update_all and not when_not_matched_insert_all: raise ValueError("no upsert options selected...exiting") - if upsert_util.has_duplicate_rows(df, join_cols): - raise ValueError("Duplicate rows found in source dataset based on the key columns. No upsert executed") + from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible, schema_to_pyarrow - from pyiceberg.io.pyarrow import _check_pyarrow_schema_compatible + table_arrow_schema = schema_to_pyarrow(self.table_metadata.schema(), include_field_ids=False) + upsert_util.validate_join_cols(df, join_cols, table_arrow_schema) downcast_ns_timestamp_to_us = Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False _check_pyarrow_schema_compatible( @@ -927,6 +927,10 @@ def upsert( format_version=self.table_metadata.format_version, ) + # Validate uniqueness after type checks to avoid comparing/hashing unsupported types. + if upsert_util.has_duplicate_rows(df, join_cols): + raise ValueError("Duplicate rows found in source dataset based on the key columns. No upsert executed") + # get list of rows that exist so we don't have to load the entire target table matched_predicate = upsert_util.create_match_filter(df, join_cols) diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..a9feb4691e 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -16,6 +16,7 @@ # under the License. import functools import operator +from collections import Counter import pyarrow as pa from pyarrow import Table as pyarrow_table @@ -30,6 +31,75 @@ ) +def validate_join_cols(df: pyarrow_table, join_cols: list[str], table_schema: pa.Schema) -> None: + """Validate join-key presence and types before Arrow comparison or hashing.""" + if not isinstance(join_cols, (list, tuple)): + raise ValueError(f"join_cols must be a list of column names, got {type(join_cols).__name__}.") + + duplicates = sorted(col for col, count in Counter(join_cols).items() if count > 1) + if duplicates: + raise ValueError(f"join_cols contains duplicates: {', '.join(duplicates)}.") + + for col in join_cols: + _validate_table_join_col(col, table_schema) + _validate_input_join_col(col, df) + + +def _validate_table_join_col(col: str, table_schema: pa.Schema) -> None: + """Reject types that are unreliable or unsupported as join keys regardless of the input.""" + if col not in table_schema.names: + parent = col.split(".", 1)[0] + if "." in col and parent in table_schema.names and pa.types.is_nested(table_schema.field(parent).type): + raise ValueError(f"Join column '{col}' is a field inside struct '{parent}' and cannot be used as a join key.") + raise ValueError(f"Join column '{col}' does not exist in the table. Available columns: {', '.join(table_schema.names)}.") + + field_type = table_schema.field(col).type + + if pa.types.is_floating(field_type): + raise ValueError( + f"Join column '{col}' is floating point and cannot be used as a join key " + "because floating point equality is unreliable." + ) + + if pa.types.is_nested(field_type): + raise ValueError(f"Join column '{col}' has nested type '{field_type}'; only primitive columns can be join keys.") + + if isinstance(field_type, pa.BaseExtensionType): + raise NotImplementedError(f"Join column '{col}' has type '{field_type}', which is not yet supported as a join key.") + + +def _validate_input_join_col(col: str, df: pyarrow_table) -> None: + """Reject input representations that are unsupported even when the table type is valid.""" + # Schema compatibility permits missing optional fields, but upsert needs every join key. + if col not in df.schema.names: + raise ValueError(f"Join column '{col}' is missing from the input.") + + arr = df.column(col) + + if pa.types.is_dictionary(arr.type): + raise NotImplementedError( + f"Input column '{col}' is dictionary-encoded, which is not yet supported for join keys. Decode it first." + ) + + if pa.types.is_null(arr.type): + raise ValueError(f"Input column '{col}' has the null type and cannot be used as a join key.") + + if pa.types.is_string_view(arr.type) or pa.types.is_binary_view(arr.type): + plain = "string" if pa.types.is_string_view(arr.type) else "binary" + raise NotImplementedError( + f"Input column '{col}' has type '{arr.type}', which is not yet supported for join keys. Cast it to '{plain}' first." + ) + + if isinstance(arr.type, pa.BaseExtensionType): + raise NotImplementedError( + f"Input column '{col}' has extension type '{arr.type}', which is not yet supported for join keys." + ) + + # Null keys cannot be expressed as Iceberg literals in the match filter. + if arr.null_count > 0: + raise ValueError(f"Input column '{col}' contains null values, which cannot be join keys.") + + def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression: unique_keys = df.select(join_cols).group_by(join_cols).aggregate([]) diff --git a/tests/table/test_upsert.py b/tests/table/test_upsert.py index 78ddbc7c5c..7e88189325 100644 --- a/tests/table/test_upsert.py +++ b/tests/table/test_upsert.py @@ -14,10 +14,14 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from datetime import datetime +import uuid +from datetime import date, datetime, time, timezone +from decimal import Decimal from pathlib import PosixPath +from typing import Any import pyarrow as pa +import pyarrow.compute as pc import pytest from datafusion import SessionContext from pyarrow import Table as pa_table @@ -26,14 +30,29 @@ from pyiceberg.exceptions import NoSuchTableError from pyiceberg.expressions import AlwaysTrue, And, EqualTo, Reference from pyiceberg.expressions.literals import LongLiteral -from pyiceberg.io.pyarrow import schema_to_pyarrow +from pyiceberg.io.pyarrow import UnsupportedPyArrowTypeException, schema_to_pyarrow from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema from pyiceberg.table import Table, UpsertResult from pyiceberg.table.snapshots import Operation from pyiceberg.table.upsert_util import create_match_filter from pyiceberg.transforms import DayTransform -from pyiceberg.types import IntegerType, NestedField, StringType, StructType, TimestampType +from pyiceberg.types import ( + BinaryType, + BooleanType, + DateType, + DecimalType, + FixedType, + IntegerType, + LongType, + NestedField, + PrimitiveType, + StringType, + StructType, + TimestampType, + TimestamptzType, + TimeType, +) from tests.catalog.test_base import InMemoryCatalog @@ -326,7 +345,7 @@ def test_key_cols_misaligned(catalog: Catalog) -> None: df_src = ctx.sql("select 1 as item_id, date '2021-05-01' as order_date, 'B' as order_type").to_arrow_table() - with pytest.raises(Exception, match=r"""Field ".*" does not exist in schema"""): + with pytest.raises(ValueError, match="Join column 'order_id' is missing from the input"): table.upsert(df=df_src, join_cols=["order_id"]) @@ -669,7 +688,11 @@ def test_upsert_with_struct_field_as_join_key(catalog: Catalog) -> None: ) with pytest.raises( - pa.lib.ArrowNotImplementedError, match="Keys of type struct" + ValueError, + match=( + "Join column 'nested_type' has nested type 'struct'; " + "only primitive columns can be join keys" + ), ): _ = tbl.upsert(update_data, join_cols=["nested_type"]) @@ -927,3 +950,222 @@ def test_upsert_snapshot_properties(catalog: Catalog) -> None: for snapshot in snapshots[initial_snapshot_count:]: assert snapshot.summary is not None assert snapshot.summary.additional_properties.get("test_prop") == "test_value" + + +_UUID_BYTES = uuid.uuid4().bytes +_STRUCT = pa.struct([("a", pa.int32())]) +_MAP = pa.map_(pa.string(), pa.int32()) + + +@pytest.mark.parametrize( + "table_type, source_key, expected_error, match", + [ + pytest.param(pa.float32(), pa.array([1.0], pa.float32()), ValueError, "Join column 'k' is floating point", id="float32"), + pytest.param(pa.float64(), pa.array([1.0], pa.float64()), ValueError, "Join column 'k' is floating point", id="float64"), + pytest.param( + _STRUCT, pa.array([{"a": 1}], _STRUCT), ValueError, "Join column 'k' has nested type 'struct'", id="struct" + ), + pytest.param( + pa.list_(pa.int32()), + pa.array([[1]], pa.list_(pa.int32())), + ValueError, + "Join column 'k' has nested type 'large_list", + id="list", + ), + pytest.param( + _MAP, pa.array([[("a", 1)]], _MAP), ValueError, "Join column 'k' has nested type 'map'", id="map" + ), + pytest.param( + pa.uuid(), + pa.array([_UUID_BYTES], pa.uuid()), + NotImplementedError, + "Join column 'k' has type 'extension'", + id="uuid-table", + ), + pytest.param( + pa.uuid(), + pa.array([_UUID_BYTES], pa.binary(16)), + NotImplementedError, + "Join column 'k' has type 'extension'", + id="uuid-table-fixed-source", + ), + pytest.param( + pa.string(), + pa.array([_UUID_BYTES], pa.uuid()), + NotImplementedError, + "Input column 'k' has extension type 'extension'", + id="uuid-extension-source", + ), + pytest.param( + pa.string(), + pa.array(["a"]).dictionary_encode(), + NotImplementedError, + "Input column 'k' is dictionary-encoded", + id="dictionary-string", + ), + pytest.param( + pa.int64(), + pa.array([1]).dictionary_encode(), + NotImplementedError, + "Input column 'k' is dictionary-encoded", + id="dictionary-int", + ), + pytest.param(pa.int32(), pa.array([None], pa.null()), ValueError, "Input column 'k' has the null type", id="null-type"), + pytest.param( + pa.string(), + pa.array(["a"], pa.string_view()), + NotImplementedError, + "Input column 'k' has type 'string_view'", + id="string-view", + ), + pytest.param( + pa.binary(), + pa.array([b"a"], pa.binary_view()), + NotImplementedError, + "Input column 'k' has type 'binary_view'", + id="binary-view", + ), + pytest.param( + pa.int32(), + pc.run_end_encode(pa.array([1], pa.int32())), + UnsupportedPyArrowTypeException, + "unsupported type: run_end_encoded", + id="run-end-encoded", + ), + pytest.param( + pa.int32(), pa.array([1, None], pa.int32()), ValueError, "Input column 'k' contains null values", id="null-values" + ), + pytest.param( + pa.int32(), pa.array([None], pa.int32()), ValueError, "Input column 'k' contains null values", id="all-null-values" + ), + pytest.param(pa.int32(), pa.array(["1"]), ValueError, "Mismatch in fields", id="wrong-type-source"), + ], +) +def test_upsert_rejects_unsupported_join_key( + catalog: Catalog, table_type: pa.DataType, source_key: pa.Array, expected_error: type[Exception], match: str +) -> None: + """Unreliable or unsupported join keys fail with a descriptive error before anything is written.""" + identifier = "default.test_upsert_rejects_unsupported_join_key" + _drop_table(catalog, identifier) + table = catalog.create_table(identifier, pa.schema([("k", table_type), ("payload", pa.string())])) + source = pa.table({"k": source_key, "payload": ["val"] * len(source_key)}) + + with pytest.raises(expected_error, match=match): + table.upsert(source, join_cols=["k"]) + + assert table.current_snapshot() is None + + +@pytest.mark.parametrize( + "join_cols, identifier_field_ids, drop_source_columns, match", + [ + pytest.param(["missing"], [], [], "Join column 'missing' does not exist in the table", id="not-in-table"), + pytest.param(["K"], [], [], "Join column 'K' does not exist in the table", id="case-mismatch"), + pytest.param(["k"], [], ["k"], "Join column 'k' is missing from the input", id="required-not-in-source"), + pytest.param(["opt"], [], ["opt"], "Join column 'opt' is missing from the input", id="optional-not-in-source"), + pytest.param(["k", "opt"], [], ["opt"], "Join column 'opt' is missing from the input", id="composite-second-missing"), + pytest.param(["k", "k"], [], [], "join_cols contains duplicates: k", id="duplicate-join-cols"), + pytest.param(["s.x"], [], [], "Join column 's.x' is a field inside struct 's'", id="nested-path"), + pytest.param(None, [4], [], "Join column 's.x' is a field inside struct 's'", id="nested-identifier-field"), + pytest.param([], [], [], "Join columns could not be found", id="empty-join-cols"), + pytest.param(None, [], [], "Join columns could not be found", id="no-identifier-fields"), + pytest.param("k", [], [], "join_cols must be a list of column names, got str", id="string-not-list"), + pytest.param({"k"}, [], [], "join_cols must be a list of column names, got set", id="set-not-list"), + ], +) +def test_upsert_rejects_invalid_join_cols( + catalog: Catalog, join_cols: Any, identifier_field_ids: list[int], drop_source_columns: list[str], match: str +) -> None: + """Join column resolution fails clearly for missing, duplicate, nested, mistyped, or unresolvable columns.""" + identifier = "default.test_upsert_rejects_invalid_join_cols" + _drop_table(catalog, identifier) + schema = Schema( + NestedField(1, "k", IntegerType(), required=True), + NestedField(2, "opt", IntegerType(), required=False), + NestedField(3, "s", StructType(NestedField(4, "x", IntegerType(), required=True)), required=True), + NestedField(5, "payload", StringType(), required=False), + identifier_field_ids=identifier_field_ids, + ) + table = catalog.create_table(identifier, schema) + source = pa.Table.from_pylist( + [{"k": 1, "opt": 1, "s": {"x": 1}, "payload": "val"}], + schema=pa.schema( + [ + pa.field("k", pa.int32(), nullable=False), + pa.field("opt", pa.int32(), nullable=True), + pa.field("s", pa.struct([pa.field("x", pa.int32(), nullable=False)]), nullable=False), + pa.field("payload", pa.string(), nullable=True), + ] + ), + ).drop_columns(drop_source_columns) + + with pytest.raises(ValueError, match=match): + table.upsert(source, join_cols=join_cols) + + assert table.current_snapshot() is None + + +_UTC = timezone.utc + + +@pytest.mark.parametrize( + "iceberg_type, arrow_type, existing_key, new_key", + [ + pytest.param(BooleanType(), pa.bool_(), False, True, id="boolean"), + pytest.param(IntegerType(), pa.int32(), 1, 2, id="int"), + pytest.param(LongType(), pa.int64(), 1, 2, id="long"), + pytest.param(LongType(), pa.int32(), 1, 2, id="long-from-int32-source"), + pytest.param(DecimalType(10, 2), pa.decimal128(10, 2), Decimal("1.50"), Decimal("2.50"), id="decimal"), + pytest.param(DateType(), pa.date32(), date(2024, 1, 1), date(2024, 1, 2), id="date"), + pytest.param(TimeType(), pa.time64("us"), time(1, 2, 3), time(4, 5, 6), id="time"), + pytest.param(TimestampType(), pa.timestamp("us"), datetime(2024, 1, 1), datetime(2024, 1, 2), id="timestamp"), + pytest.param( + TimestamptzType(), + pa.timestamp("us", "UTC"), + datetime(2024, 1, 1, tzinfo=_UTC), + datetime(2024, 1, 2, tzinfo=_UTC), + id="timestamptz", + ), + pytest.param(StringType(), pa.string(), "a", "b", id="string"), + pytest.param(StringType(), pa.large_string(), "a", "b", id="string-from-large-string-source"), + pytest.param(BinaryType(), pa.binary(), b"a", b"b", id="binary"), + pytest.param(BinaryType(), pa.large_binary(), b"a", b"b", id="binary-from-large-binary-source"), + pytest.param(FixedType(2), pa.binary(2), b"aa", b"bb", id="fixed"), + ], +) +def test_upsert_accepts_supported_join_key( + catalog: Catalog, iceberg_type: PrimitiveType, arrow_type: pa.DataType, existing_key: Any, new_key: Any +) -> None: + """Every supported primitive key updates matched rows and inserts new ones, including from a multi-chunk source.""" + identifier = "default.test_upsert_accepts_supported_join_key" + _drop_table(catalog, identifier) + schema = Schema(NestedField(1, "k", iceberg_type, required=True), NestedField(2, "payload", StringType(), required=True)) + arrow_schema = pa.schema([pa.field("k", arrow_type, nullable=False), pa.field("payload", pa.string(), nullable=False)]) + table = catalog.create_table(identifier, schema) + table.append(pa.Table.from_pylist([{"k": existing_key, "payload": "old"}], schema=arrow_schema)) + + source = pa.concat_tables( + [ + pa.Table.from_pylist([{"k": existing_key, "payload": "updated"}], schema=arrow_schema), + pa.Table.from_pylist([{"k": new_key, "payload": "inserted"}], schema=arrow_schema), + ] + ) + result = table.upsert(source, join_cols=["k"]) + + assert (result.rows_updated, result.rows_inserted) == (1, 1) + rows = table.scan().to_arrow().to_pylist() + assert {(row["k"], row["payload"]) for row in rows} == {(existing_key, "updated"), (new_key, "inserted")} + + +def test_upsert_allows_null_join_key_values_in_target(catalog: Catalog) -> None: + """Null key values are only rejected in the source; existing null-key rows in the table are left untouched.""" + identifier = "default.test_upsert_allows_null_join_key_values_in_target" + _drop_table(catalog, identifier) + table = catalog.create_table(identifier, pa.schema([("k", pa.int32()), ("payload", pa.string())])) + table.append(pa.table({"k": pa.array([None, 1], pa.int32()), "payload": ["orphan", "old"]})) + + result = table.upsert(pa.table({"k": pa.array([1], pa.int32()), "payload": ["new"]}), join_cols=["k"]) + + assert (result.rows_updated, result.rows_inserted) == (1, 0) + rows = table.scan().to_arrow().to_pylist() + assert {(row["k"], row["payload"]) for row in rows} == {(None, "orphan"), (1, "new")}