From ddb3d3e7d157c6dd5af981be91bd1e66694bb254 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Wed, 19 Aug 2026 12:54:41 -0700 Subject: [PATCH] IO: Absorb upsert PyArrow logic into io/pyarrow.py Move PyArrow-specific table operations (joins, group_by, duplicate detection, row comparison) from table/upsert_util.py into pyiceberg/io/pyarrow.py. The upsert_util module now delegates to helpers in io/pyarrow rather than importing pyarrow directly. This consolidates PyArrow logic behind the io/pyarrow module boundary, which is a precondition for the decomposition proposed in #3737. Fixes #3812 (PR A) --- pyiceberg/io/pyarrow.py | 82 ++++++++++++++++++++++++++++++ pyiceberg/table/upsert_util.py | 92 +++++++--------------------------- 2 files changed, 100 insertions(+), 74 deletions(-) diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..14a889791c 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -3137,3 +3137,85 @@ def _get_field_from_arrow_table(arrow_table: pa.Table, field_path: str) -> pa.Ar field_array = arrow_table[path_parts[0]] # Navigate into the struct using the remaining path parts return pc.struct_field(field_array, path_parts[1:]) + + +def _upsert_unique_keys(df: pa.Table, join_cols: list[str]) -> pa.Table: + """Extract unique key combinations from a table. + + Returns a table containing one row per distinct combination of join_cols. + """ + return df.select(join_cols).group_by(join_cols).aggregate([]) + + +def _upsert_has_duplicate_rows(df: pa.Table, join_cols: list[str]) -> bool: + """Check for duplicate rows in a PyArrow table based on the join columns.""" + return len(df.select(join_cols).group_by(join_cols).aggregate([([], "count_all")]).filter(pc.field("count_all") > 1)) > 0 + + +def _upsert_get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table: + """Return rows from source_table whose non-key columns differ from target_table. + + Performs an inner join on join_cols, then compares non-key column values + row-by-row. Returns the subset of source rows that have at least one + changed non-key column. If target_table is empty, returns an empty table. + + Raises: + ValueError: If target_table has duplicate rows on join_cols. + ValueError: If join_cols use reserved index column names. + """ + all_columns = set(source_table.column_names) + join_cols_set = set(join_cols) + + non_key_cols = list(all_columns - join_cols_set) + + if _upsert_has_duplicate_rows(target_table, join_cols): + raise ValueError("Target table has duplicate rows, aborting upsert") + + if len(target_table) == 0: + return source_table.schema.empty_table() + + # We need to compare non_key_cols in Python as PyArrow + # 1. Cannot do a join when non-join columns have complex types + # 2. Cannot compare columns with complex types + # See: https://github.com/apache/arrow/issues/35785 + SOURCE_INDEX_COLUMN_NAME = "__source_index" + TARGET_INDEX_COLUMN_NAME = "__target_index" + + if SOURCE_INDEX_COLUMN_NAME in join_cols or TARGET_INDEX_COLUMN_NAME in join_cols: + raise ValueError( + f"{SOURCE_INDEX_COLUMN_NAME} and {TARGET_INDEX_COLUMN_NAME} are reserved for joining " + f"DataFrames, and cannot be used as column names" + ) from None + + # Cast to target table schema so types align for the join. + # See: https://github.com/apache/arrow/issues/37542 + source_index = ( + source_table.cast(target_table.schema) + .select(join_cols_set) + .append_column(SOURCE_INDEX_COLUMN_NAME, pa.array(range(len(source_table)))) + ) + + target_index = target_table.select(join_cols_set).append_column(TARGET_INDEX_COLUMN_NAME, pa.array(range(len(target_table)))) + + matching_indices = source_index.join(target_index, keys=list(join_cols_set), join_type="inner") + + to_update_indices = [] + for source_idx, target_idx in zip( + matching_indices[SOURCE_INDEX_COLUMN_NAME].to_pylist(), + matching_indices[TARGET_INDEX_COLUMN_NAME].to_pylist(), + strict=True, + ): + source_row = source_table.slice(source_idx, 1) + target_row = target_table.slice(target_idx, 1) + + for key in non_key_cols: + source_val = source_row.column(key)[0].as_py() + target_val = target_row.column(key)[0].as_py() + if source_val != target_val: + to_update_indices.append(source_idx) + break + + if to_update_indices: + return source_table.take(to_update_indices) + else: + return source_table.schema.empty_table() diff --git a/pyiceberg/table/upsert_util.py b/pyiceberg/table/upsert_util.py index 6f32826eb0..d104178a0f 100644 --- a/pyiceberg/table/upsert_util.py +++ b/pyiceberg/table/upsert_util.py @@ -16,10 +16,7 @@ # under the License. import functools import operator - -import pyarrow as pa -from pyarrow import Table as pyarrow_table -from pyarrow import compute as pc +from typing import TYPE_CHECKING from pyiceberg.expressions import ( AlwaysFalse, @@ -28,10 +25,19 @@ In, Or, ) +from pyiceberg.io.pyarrow import ( + _upsert_get_rows_to_update, + _upsert_has_duplicate_rows, + _upsert_unique_keys, +) +if TYPE_CHECKING: + import pyarrow as pa -def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression: - unique_keys = df.select(join_cols).group_by(join_cols).aggregate([]) + +def create_match_filter(df: "pa.Table", join_cols: list[str]) -> BooleanExpression: + """Build an Iceberg filter expression matching the unique keys in df.""" + unique_keys = _upsert_unique_keys(df, join_cols) if len(join_cols) == 1: return In(join_cols[0], unique_keys[0].to_pylist()) @@ -48,77 +54,15 @@ def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpre return Or(*filters) -def has_duplicate_rows(df: pyarrow_table, join_cols: list[str]) -> bool: - """Check for duplicate rows in a PyArrow table based on the join columns.""" - return len(df.select(join_cols).group_by(join_cols).aggregate([([], "count_all")]).filter(pc.field("count_all") > 1)) > 0 +def has_duplicate_rows(df: "pa.Table", join_cols: list[str]) -> bool: + """Check for duplicate rows in a table based on the join columns.""" + return _upsert_has_duplicate_rows(df, join_cols) -def get_rows_to_update(source_table: pa.Table, target_table: pa.Table, join_cols: list[str]) -> pa.Table: - """ - Return a table with rows that need to be updated in the target table based on the join columns. +def get_rows_to_update(source_table: "pa.Table", target_table: "pa.Table", join_cols: list[str]) -> "pa.Table": + """Return rows from source that need to be updated in the target table based on the join columns. The table is joined on the identifier columns, and then checked if there are any updated rows. Those are selected and everything is renamed correctly. """ - all_columns = set(source_table.column_names) - join_cols_set = set(join_cols) - - non_key_cols = list(all_columns - join_cols_set) - - if has_duplicate_rows(target_table, join_cols): - raise ValueError("Target table has duplicate rows, aborting upsert") - - if len(target_table) == 0: - # When the target table is empty, there is nothing to update :) - return source_table.schema.empty_table() - - # We need to compare non_key_cols in Python as PyArrow - # 1. Cannot do a join when non-join columns have complex types - # 2. Cannot compare columns with complex types - # See: https://github.com/apache/arrow/issues/35785 - SOURCE_INDEX_COLUMN_NAME = "__source_index" - TARGET_INDEX_COLUMN_NAME = "__target_index" - - if SOURCE_INDEX_COLUMN_NAME in join_cols or TARGET_INDEX_COLUMN_NAME in join_cols: - raise ValueError( - f"{SOURCE_INDEX_COLUMN_NAME} and {TARGET_INDEX_COLUMN_NAME} are reserved for joining " - f"DataFrames, and cannot be used as column names" - ) from None - - # Step 1: Prepare source index with join keys and a marker index - # Cast to target table schema, so we can do the join - # See: https://github.com/apache/arrow/issues/37542 - source_index = ( - source_table.cast(target_table.schema) - .select(join_cols_set) - .append_column(SOURCE_INDEX_COLUMN_NAME, pa.array(range(len(source_table)))) - ) - - # Step 2: Prepare target index with join keys and a marker - target_index = target_table.select(join_cols_set).append_column(TARGET_INDEX_COLUMN_NAME, pa.array(range(len(target_table)))) - - # Step 3: Perform an inner join to find which rows from source exist in target - matching_indices = source_index.join(target_index, keys=list(join_cols_set), join_type="inner") - - # Step 4: Compare all rows using Python - to_update_indices = [] - for source_idx, target_idx in zip( - matching_indices[SOURCE_INDEX_COLUMN_NAME].to_pylist(), - matching_indices[TARGET_INDEX_COLUMN_NAME].to_pylist(), - strict=True, - ): - source_row = source_table.slice(source_idx, 1) - target_row = target_table.slice(target_idx, 1) - - for key in non_key_cols: - source_val = source_row.column(key)[0].as_py() - target_val = target_row.column(key)[0].as_py() - if source_val != target_val: - to_update_indices.append(source_idx) - break - - # Step 5: Take rows from source table using the indices and cast to target schema - if to_update_indices: - return source_table.take(to_update_indices) - else: - return source_table.schema.empty_table() + return _upsert_get_rows_to_update(source_table, target_table, join_cols)