diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..2ccdfc7b01 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -1626,6 +1626,136 @@ def _get_column_projection_values( return projected_missing_fields +def _fill_null_value(data_type: pa.DataType) -> Any: + if pa.types.is_integer(data_type) or pa.types.is_floating(data_type) or pa.types.is_decimal(data_type): + return 0 + if pa.types.is_boolean(data_type): + return False + if pa.types.is_string(data_type) or pa.types.is_large_string(data_type): + return "" + if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type): + return b"" + if pa.types.is_fixed_size_binary(data_type): + return b"\x00" * data_type.byte_width + if isinstance(data_type, pa.UuidType): + return pa.scalar(b"\x00" * 16, type=data_type) + if pa.types.is_timestamp(data_type) or pa.types.is_date(data_type) or pa.types.is_time(data_type): + return pa.scalar(0, type=data_type) + storage = getattr(data_type, "storage_type", None) + if storage is not None: + return _fill_null_value(storage) + return "" + + +def _column_name_for_field_id(table: pa.Table, field_id: int, schema: Schema) -> str | None: + try: + name = schema.find_field(field_id).name + except ValueError: + name = None + if name is not None and name in table.column_names: + return name + for field in table.schema: + if _get_field_id(field) == field_id: + return field.name + return None + + +def _align_equality_delete_table( + data: pa.Table, deletes: pa.Table, equality_ids: Iterable[int], table_schema: Schema +) -> tuple[pa.Table, pa.Table, list[str]]: + join_keys: list[str] = [] + for field_id in equality_ids: + delete_name: str | None = None + data_name = _column_name_for_field_id(data, field_id, table_schema) + if data_name is not None and data_name in deletes.column_names: + delete_name = data_name + else: + for field in deletes.schema: + if _get_field_id(field) == field_id: + delete_name = field.name + break + if delete_name is None: + try: + schema_name = table_schema.find_field(field_id).name + except ValueError: + schema_name = None + if schema_name is not None and schema_name in deletes.column_names: + delete_name = schema_name + if delete_name is None: + raise ValueError(f"Equality delete file is missing field id {field_id}") + if data_name is None: + delete_type = deletes.schema.field(delete_name).type + data = data.append_column(delete_name, pa.nulls(data.num_rows, type=delete_type)) + data_name = delete_name + if delete_name != data_name: + names = list(deletes.column_names) + names[names.index(delete_name)] = data_name + deletes = deletes.rename_columns(names) + data_type = data.schema.field(data_name).type + delete_type = deletes.schema.field(data_name).type + if delete_type != data_type: + deletes = deletes.set_column( + deletes.schema.get_field_index(data_name), data_name, pc.cast(deletes[data_name], data_type) + ) + join_keys.append(data_name) + return data, deletes, join_keys + + +def _null_safe_anti_join(data: pa.Table, deletes: pa.Table, join_keys: list[str]) -> pa.Table: + data_join = data + delete_join = deletes + join_cols: list[str] = [] + extra_cols: list[str] = [] + for key in join_keys: + null_col = f"__iceberg_eq_null_{key}" + val_col = f"__iceberg_eq_val_{key}" + key_type = data_join.schema.field(key).type + data_join = data_join.append_column(null_col, pc.is_null(data_join[key])) + delete_join = delete_join.append_column(null_col, pc.is_null(delete_join[key])) + data_vals = data_join[key] + delete_vals = delete_join[key] + fill_type = key_type + storage = getattr(key_type, "storage_type", None) + if storage is not None: + data_vals = data_vals.cast(storage) + delete_vals = delete_vals.cast(storage) + fill_type = storage + fill = _fill_null_value(fill_type) + data_join = data_join.append_column(val_col, pc.fill_null(data_vals, fill)) + delete_join = delete_join.append_column(val_col, pc.fill_null(delete_vals, fill)) + join_cols.extend([null_col, val_col]) + extra_cols.extend([null_col, val_col]) + if pa.types.is_floating(key_type): + nan_col = f"__iceberg_eq_nan_{key}" + data_join = data_join.append_column(nan_col, pc.fill_null(pc.is_nan(data_join[key]), False)) + delete_join = delete_join.append_column(nan_col, pc.fill_null(pc.is_nan(delete_join[key]), False)) + join_cols.append(nan_col) + extra_cols.append(nan_col) + delete_keys = delete_join.select(join_cols) + joined = data_join.join(delete_keys, keys=join_cols, join_type="left anti") + return joined.drop(extra_cols) + + +def _apply_equality_deletes( + data: pa.Table, equality_groups: dict[frozenset[int], list[pa.Table]], table_schema: Schema +) -> pa.Table: + if data.num_rows == 0: + return data + for equality_ids, delete_tables in equality_groups.items(): + if not equality_ids: + continue + aligned: list[pa.Table] = [] + join_keys: list[str] | None = None + for delete_table in delete_tables: + data, one, keys = _align_equality_delete_table(data, delete_table, equality_ids, table_schema) + if join_keys is None: + join_keys = keys + aligned.append(one.select(keys)) + deletes = pa.concat_tables(aligned, promote_options="permissive") + data = _null_safe_anti_join(data, deletes, join_keys or []) + return data + + def _task_to_record_batches( io: FileIO, task: FileScanTask, @@ -1640,6 +1770,7 @@ def _task_to_record_batches( format_version: TableVersion = TableProperties.DEFAULT_FORMAT_VERSION, downcast_ns_timestamp_to_us: bool | None = None, dictionary_columns: tuple[str, ...] = (), + equality_delete_tables: dict[str, pa.Table] | None = None, ) -> Iterator[pa.RecordBatch]: format_kwargs: dict[str, Any] = {"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8} if dictionary_columns and task.file.file_format == FileFormat.PARQUET: @@ -1672,7 +1803,20 @@ def _task_to_record_batches( bound_file_filter = bind(file_schema, translated_row_filter, case_sensitive=case_sensitive) pyarrow_filter = expression_to_pyarrow(bound_file_filter, file_schema) - file_project_schema = prune_columns(file_schema, projected_field_ids, select_full_types=False) + read_field_ids = set(projected_field_ids) + equality_groups: dict[frozenset[int], list[pa.Table]] = {} + if equality_delete_tables: + for delete_file in task.delete_files: + if delete_file.content != DataFileContent.EQUALITY_DELETES: + continue + eq_table = equality_delete_tables.get(delete_file.file_path) + if eq_table is None: + continue + eq_ids = frozenset(delete_file.equality_ids or []) + equality_groups.setdefault(eq_ids, []).append(eq_table) + read_field_ids.update(eq_ids) + + file_project_schema = prune_columns(file_schema, read_field_ids, select_full_types=False) fragment_scanner = ds.Scanner.from_fragment( fragment=fragment, @@ -1684,6 +1828,7 @@ def _task_to_record_batches( ) next_index = 0 + file_batches: list[pa.RecordBatch] = [] batches = fragment_scanner.to_batches() for batch in batches: next_index = next_index + len(batch) @@ -1709,24 +1854,61 @@ def _task_to_record_batches( if current_batch.num_rows == 0: continue - yield _to_requested_schema( - projected_schema, - file_project_schema, - current_batch, - downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, - projected_missing_fields=projected_missing_fields, - allow_timestamp_tz_mismatch=True, - ) + if not equality_groups: + yield _to_requested_schema( + projected_schema, + file_project_schema, + current_batch, + downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, + projected_missing_fields=projected_missing_fields, + allow_timestamp_tz_mismatch=True, + ) + else: + file_batches.append(current_batch) + + if equality_groups and file_batches: + data = pa.Table.from_batches(file_batches) + data = _apply_equality_deletes(data, equality_groups, table_schema) + for out_batch in data.to_batches(): + if out_batch.num_rows == 0: + continue + yield _to_requested_schema( + projected_schema, + file_project_schema, + out_batch, + downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us, + projected_missing_fields=projected_missing_fields, + allow_timestamp_tz_mismatch=True, + ) -def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[str, list[ChunkedArray]]: +def _read_equality_delete_table(io: FileIO, data_file: DataFile) -> pa.Table: + if data_file.file_format not in {FileFormat.PARQUET, FileFormat.ORC}: + raise ValueError(f"Equality delete file format not supported: {data_file.file_format}") + with io.new_input(data_file.file_path).open() as fi: + fragment = _get_file_format(data_file.file_format, pre_buffer=True, buffer_size=ONE_MEGABYTE).make_fragment(fi) + return ds.Scanner.from_fragment(fragment=fragment).to_table() + + +def _read_all_delete_files( + io: FileIO, tasks: Iterable[FileScanTask] +) -> tuple[dict[str, list[ChunkedArray]], dict[str, pa.Table]]: deletes_per_file: dict[str, list[ChunkedArray]] = {} - unique_deletes = set(itertools.chain.from_iterable([task.delete_files for task in tasks])) - if len(unique_deletes) > 0: - executor = ExecutorFactory.get_or_create() + equality_tables: dict[str, pa.Table] = {} + unique_pos: set[DataFile] = set() + unique_eq: set[DataFile] = set() + for task in tasks: + for delete_file in task.delete_files: + if delete_file.content == DataFileContent.EQUALITY_DELETES: + unique_eq.add(delete_file) + else: + unique_pos.add(delete_file) + + executor = ExecutorFactory.get_or_create() + if unique_pos: deletes_per_files: Iterator[dict[str, ChunkedArray]] = executor.map( lambda args: _read_deletes(*args), - [(io, delete_file) for delete_file in unique_deletes], + [(io, delete_file) for delete_file in unique_pos], ) for delete in deletes_per_files: for file, arr in delete.items(): @@ -1735,7 +1917,14 @@ def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[st else: deletes_per_file[file] = [arr] - return deletes_per_file + if unique_eq: + for path, table in executor.map( + lambda args: (args[1].file_path, _read_equality_delete_table(*args)), + [(io, delete_file) for delete_file in unique_eq], + ): + equality_tables[path] = table + + return deletes_per_file, equality_tables class ArrowScan: @@ -1841,7 +2030,7 @@ def to_record_batches(self, tasks: Iterable[FileScanTask]) -> Iterator[pa.Record ResolveError: When a required field cannot be found in the file ValueError: When a field type in the file cannot be projected to the schema type """ - deletes_per_file = _read_all_delete_files(self._io, tasks) + deletes_per_file, equality_delete_tables = _read_all_delete_files(self._io, tasks) total_row_count = 0 executor = ExecutorFactory.get_or_create() @@ -1850,7 +2039,7 @@ def batches_for_task(task: FileScanTask) -> list[pa.RecordBatch]: # Materialize the iterator here to ensure execution happens within the executor. # Otherwise, the iterator would be lazily consumed later (in the main thread), # defeating the purpose of using executor.map. - return list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file)) + return list(self._record_batches_from_scan_tasks_and_deletes([task], deletes_per_file, equality_delete_tables)) limit_reached = False for batches in executor.map(batches_for_task, tasks): @@ -1870,7 +2059,10 @@ def batches_for_task(task: FileScanTask) -> list[pa.RecordBatch]: break def _record_batches_from_scan_tasks_and_deletes( - self, tasks: Iterable[FileScanTask], deletes_per_file: dict[str, list[ChunkedArray]] + self, + tasks: Iterable[FileScanTask], + deletes_per_file: dict[str, list[ChunkedArray]], + equality_delete_tables: dict[str, pa.Table] | None = None, ) -> Iterator[pa.RecordBatch]: total_row_count = 0 for task in tasks: @@ -1890,6 +2082,7 @@ def _record_batches_from_scan_tasks_and_deletes( self._table_metadata.format_version, self._downcast_ns_timestamp_to_us, self._dictionary_columns, + equality_delete_tables, ) for batch in batches: if self._limit is not None: diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 3dffc2270c..35eca72a78 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2254,21 +2254,13 @@ def from_rest_response( Returns: A FileScanTask with the converted data and delete files. - - Raises: - NotImplementedError: If equality delete files are encountered. """ - from pyiceberg.catalog.rest.scan_planning import RESTEqualityDeleteFile - data_file = _rest_file_to_data_file(rest_task.data_file) resolved_deletes: set[DataFile] = set() if rest_task.delete_file_references: for idx in rest_task.delete_file_references: - delete_file = delete_files[idx] - if isinstance(delete_file, RESTEqualityDeleteFile): - raise NotImplementedError(f"PyIceberg does not yet support equality deletes: {delete_file.file_path}") - resolved_deletes.add(_rest_file_to_data_file(delete_file)) + resolved_deletes.add(_rest_file_to_data_file(delete_files[idx])) return FileScanTask( data_file=data_file, @@ -2279,7 +2271,7 @@ def from_rest_response( def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile: """Convert a REST content file to a manifest DataFile.""" - from pyiceberg.catalog.rest.scan_planning import RESTDataFile + from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTEqualityDeleteFile if isinstance(rest_file, RESTDataFile): column_sizes = rest_file.column_sizes.to_dict() if rest_file.column_sizes else None @@ -2292,6 +2284,8 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile: null_value_counts = None nan_value_counts = None + equality_ids = rest_file.equality_ids if isinstance(rest_file, RESTEqualityDeleteFile) else None + data_file = DataFile.from_args( content=DataFileContent.from_rest_type(rest_file.content), file_path=rest_file.file_path, @@ -2305,6 +2299,7 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile: nan_value_counts=nan_value_counts, split_offsets=rest_file.split_offsets, sort_order_id=rest_file.sort_order_id, + equality_ids=equality_ids, ) data_file.spec_id = rest_file.spec_id return data_file @@ -2804,7 +2799,7 @@ def plan_files( List of FileScanTasks that contain both data and delete files. """ data_entries: list[ManifestEntry] = [] - delete_index = DeleteFileIndex() + delete_index = DeleteFileIndex(self.table_metadata.schema()) residual_evaluators: dict[int, Callable[[DataFile], ResidualEvaluator]] = KeyDefaultDict(self._build_residual_evaluator) @@ -2818,7 +2813,7 @@ def plan_files( elif data_file.content == DataFileContent.POSITION_DELETES: delete_index.add_delete_file(manifest_entry, partition_key=data_file.partition) elif data_file.content == DataFileContent.EQUALITY_DELETES: - raise ValueError("PyIceberg does not yet support equality deletes: https://github.com/apache/iceberg/issues/6568") + delete_index.add_delete_file(manifest_entry, partition_key=data_file.partition) else: raise ValueError(f"Unknown DataFileContent ({data_file.content}): {manifest_entry}") diff --git a/pyiceberg/table/delete_file_index.py b/pyiceberg/table/delete_file_index.py index 3f513aabe5..4e5235f52e 100644 --- a/pyiceberg/table/delete_file_index.py +++ b/pyiceberg/table/delete_file_index.py @@ -17,11 +17,17 @@ from __future__ import annotations from bisect import bisect_left +from typing import TYPE_CHECKING +from pyiceberg.conversions import from_bytes from pyiceberg.expressions import EqualTo from pyiceberg.expressions.visitors import _InclusiveMetricsEvaluator -from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, ManifestEntry +from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, DataFileContent, ManifestEntry from pyiceberg.typedef import Record +from pyiceberg.types import NestedField + +if TYPE_CHECKING: + from pyiceberg.schema import Schema PATH_FIELD_ID = 2147483546 @@ -59,6 +65,15 @@ def referenced_delete_files(self) -> list[DataFile]: return [data_file for data_file, _ in self._files] +class EqualityDeletes(PositionDeletes): + """Collects equality delete files and indexes them by sequence number.""" + + def add(self, delete_file: DataFile, seq_num: int) -> None: + # Equality deletes are indexed by sequence number - 1 to ensure they only + # apply to data files added in strictly earlier snapshots. + super().add(delete_file, seq_num - 1) + + def _has_path_bounds(delete_file: DataFile) -> bool: lower = delete_file.lower_bounds upper = delete_file.upper_bounds @@ -76,6 +91,83 @@ def _applies_to_data_file(delete_file: DataFile, data_file: DataFile) -> bool: return evaluator.eval(delete_file) +def _is_all_null(data_file: DataFile, field_id: int) -> bool: + null_counts = data_file.null_value_counts + value_counts = data_file.value_counts + if not null_counts or not value_counts: + return False + null_count = null_counts.get(field_id) + value_count = value_counts.get(field_id) + return null_count is not None and value_count is not None and null_count == value_count + + +def _has_no_nulls(data_file: DataFile, field_id: int) -> bool: + null_counts = data_file.null_value_counts + if not null_counts: + return False + return null_counts.get(field_id) == 0 + + +def _contains_null(data_file: DataFile, field: NestedField) -> bool: + if field.required: + return False + null_counts = data_file.null_value_counts + if not null_counts: + return True + null_count = null_counts.get(field.field_id) + if null_count is None: + return True + return null_count > 0 + + +def _eq_applies_to_data_file(eq_delete_file: DataFile, data_file: DataFile, schema: Schema) -> bool: + if not eq_delete_file.equality_ids: + return True + + for field_id in eq_delete_file.equality_ids: + try: + field = schema.find_field(field_id) + except ValueError: + # Equality field is missing from the current schema (dropped column). + # Cannot prune; the delete may still apply. + continue + if not field.field_type.is_primitive: + continue + + # Iceberg bounds exclude nulls. If both files can contain nulls, the delete + # may match via null equality even when the value ranges are disjoint. + if _contains_null(data_file, field) and _contains_null(eq_delete_file, field): + continue + + # Data is only nulls for this field, but the delete has no null rows. + if _is_all_null(data_file, field_id) and _has_no_nulls(eq_delete_file, field_id): + return False + # Delete only removes null rows, but the data has none. + if _is_all_null(eq_delete_file, field_id) and _has_no_nulls(data_file, field_id): + return False + + if ( + eq_delete_file.lower_bounds + and eq_delete_file.upper_bounds + and data_file.lower_bounds + and data_file.upper_bounds + and field_id in eq_delete_file.lower_bounds + and field_id in eq_delete_file.upper_bounds + and field_id in data_file.lower_bounds + and field_id in data_file.upper_bounds + ): + field_type = field.field_type + eq_lower = from_bytes(field_type, eq_delete_file.lower_bounds[field_id]) + eq_upper = from_bytes(field_type, eq_delete_file.upper_bounds[field_id]) + data_lower = from_bytes(field_type, data_file.lower_bounds[field_id]) + data_upper = from_bytes(field_type, data_file.upper_bounds[field_id]) + + if eq_upper < data_lower or eq_lower > data_upper: + return False + + return True + + def _referenced_data_file_path(delete_file: DataFile) -> str | None: """Return the path, if the path bounds evaluate to the same location.""" lower_bounds = delete_file.lower_bounds @@ -103,26 +195,33 @@ def _partition_key(spec_id: int, partition: Record | None) -> tuple[int, Record] class DeleteFileIndex: - """Indexes position delete files by partition and by exact data file path.""" + """Indexes position and equality delete files by partition and by exact data file path.""" - def __init__(self) -> None: + def __init__(self, schema: Schema | None = None) -> None: + self._schema = schema self._by_partition: dict[tuple[int, Record], PositionDeletes] = {} self._by_path: dict[str, PositionDeletes] = {} + self._eq_deletes: dict[tuple[int, Record] | None, EqualityDeletes] = {} def is_empty(self) -> bool: - return not self._by_partition and not self._by_path + return not self._by_partition and not self._by_path and not self._eq_deletes def add_delete_file(self, manifest_entry: ManifestEntry, partition_key: Record | None = None) -> None: delete_file = manifest_entry.data_file seq = manifest_entry.sequence_number or INITIAL_SEQUENCE_NUMBER - target_path = _referenced_data_file_path(delete_file) - if target_path: - deletes = self._by_path.setdefault(target_path, PositionDeletes()) - deletes.add(delete_file, seq) - else: - key = _partition_key(delete_file.spec_id or 0, partition_key) - deletes = self._by_partition.setdefault(key, PositionDeletes()) + if delete_file.content == DataFileContent.POSITION_DELETES: + target_path = _referenced_data_file_path(delete_file) + if target_path: + deletes = self._by_path.setdefault(target_path, PositionDeletes()) + deletes.add(delete_file, seq) + else: + key = _partition_key(delete_file.spec_id or 0, partition_key) + deletes = self._by_partition.setdefault(key, PositionDeletes()) + deletes.add(delete_file, seq) + elif delete_file.content == DataFileContent.EQUALITY_DELETES: + eq_key = _partition_key(delete_file.spec_id or 0, partition_key) if partition_key else None + deletes = self._eq_deletes.setdefault(eq_key, EqualityDeletes()) deletes.add(delete_file, seq) def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> set[DataFile]: @@ -131,17 +230,33 @@ def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record deletes: set[DataFile] = set() spec_id = data_file.spec_id or 0 - key = _partition_key(spec_id, partition_key) - partition_deletes = self._by_partition.get(key) - if partition_deletes: - for delete_file in partition_deletes.filter_by_seq(seq_num): + + # Add position deletes + partition_pos_deletes = self._by_partition.get(key) + if partition_pos_deletes: + for delete_file in partition_pos_deletes.filter_by_seq(seq_num): if _applies_to_data_file(delete_file, data_file): deletes.add(delete_file) - path_deletes = self._by_path.get(data_file.file_path) - if path_deletes: - deletes.update(path_deletes.filter_by_seq(seq_num)) + path_pos_deletes = self._by_path.get(data_file.file_path) + if path_pos_deletes: + deletes.update(path_pos_deletes.filter_by_seq(seq_num)) + + # Add equality deletes + candidate_eq_deletes: list[DataFile] = [] + partition_eq_deletes = self._eq_deletes.get(key) + if partition_eq_deletes: + candidate_eq_deletes.extend(partition_eq_deletes.filter_by_seq(seq_num)) + + global_eq_deletes = self._eq_deletes.get(None) + if global_eq_deletes: + candidate_eq_deletes.extend(global_eq_deletes.filter_by_seq(seq_num)) + + for eq_delete_file in candidate_eq_deletes: + if self._schema and not _eq_applies_to_data_file(eq_delete_file, data_file, self._schema): + continue + deletes.add(eq_delete_file) return deletes @@ -154,4 +269,7 @@ def referenced_delete_files(self) -> list[DataFile]: for deletes in self._by_path.values(): data_files.extend(deletes.referenced_delete_files()) + for deletes in self._eq_deletes.values(): + data_files.extend(deletes.referenced_delete_files()) + return data_files diff --git a/pyiceberg/table/update/validate.py b/pyiceberg/table/update/validate.py index df8506aab4..f80887b362 100644 --- a/pyiceberg/table/update/validate.py +++ b/pyiceberg/table/update/validate.py @@ -248,13 +248,13 @@ def _added_delete_files( DeleteFileIndex """ if table.format_version < 2: - return DeleteFileIndex() + return DeleteFileIndex(table.schema()) manifests, snapshot_ids = _validation_history( table, parent_snapshot, starting_snapshot, VALIDATE_ADDED_DELETE_FILES_OPERATIONS, ManifestContent.DELETES ) - dfi = DeleteFileIndex() + dfi = DeleteFileIndex(table.schema()) for manifest in manifests: for entry in manifest.fetch_manifest_entry(table.io, discard_deleted=True): diff --git a/tests/catalog/test_scan_planning_models.py b/tests/catalog/test_scan_planning_models.py index f2c80cfb9b..c8aa4da110 100644 --- a/tests/catalog/test_scan_planning_models.py +++ b/tests/catalog/test_scan_planning_models.py @@ -38,7 +38,7 @@ ValueMap, ) from pyiceberg.expressions import AlwaysTrue, EqualTo, Reference -from pyiceberg.manifest import FileFormat +from pyiceberg.manifest import DataFileContent, FileFormat TEST_URI = "https://iceberg-test-catalog/" @@ -545,7 +545,7 @@ def test_plan_scan_cancelled(rest_scan_catalog: RestCatalog, requests_mock: Mock list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) -def test_plan_scan_equality_deletes_not_supported(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: +def test_plan_scan_with_equality_deletes(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: file_one = _rest_data_file(file_path="s3://bucket/tbl/data/file1.parquet") equality_delete = _rest_equality_delete_file(equality_ids=[1, 2]) requests_mock.post( @@ -566,5 +566,11 @@ def test_plan_scan_equality_deletes_not_supported(rest_scan_catalog: RestCatalog ) request = PlanTableScanRequest() - with pytest.raises(NotImplementedError, match="PyIceberg does not yet support equality deletes"): - list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + tasks = list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + + assert len(tasks) == 1 + assert tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" + assert len(tasks[0].delete_files) == 1 + delete_file = next(iter(tasks[0].delete_files)) + assert delete_file.content == DataFileContent.EQUALITY_DELETES + assert delete_file.equality_ids == [1, 2] diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..a4fc2babcd 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -117,6 +117,7 @@ TimestampType, TimestamptzType, TimeType, + UUIDType, ) from tests.catalog.test_base import InMemoryCatalog from tests.conftest import UNIFIED_AWS_SESSION_PROPERTIES @@ -1914,6 +1915,322 @@ def test_delete_duplicates(deletes_file: str, request: pytest.FixtureRequest, ta assert str(with_deletes) == expected_str +def test_scan_table_with_equality_deletes(tmp_path: str, catalog: InMemoryCatalog, table_schema_simple: Schema) -> None: + catalog.create_namespace("default") + table = catalog.create_table("default.eq_deletes", schema=table_schema_simple) + table.append( + pa.table( + { + "foo": ["a", "b", "c"], + "bar": pa.array([1, 2, 3], type=pa.int32()), + "baz": [True, False, True], + }, + schema=schema_to_pyarrow(table_schema_simple), + ) + ) + + deletes_path = f"{tmp_path}/eq-deletes.parquet" + pq.write_table(pa.table({"bar": pa.array([2], type=pa.int32())}), deletes_path) + + eq_file = DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_path, + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=1, + file_size_in_bytes=os.path.getsize(deletes_path), + equality_ids=[2], + spec_id=table.metadata.default_spec_id, + ) + + with table.transaction() as transaction: + with transaction.update_snapshot().fast_append() as update: + update.append_data_file(eq_file) + + result = table.scan().to_arrow() + assert result.column("bar").to_pylist() == [1, 3] + + +def test_equality_delete(example_task: FileScanTask, tmp_path: str, table_schema_simple: Schema) -> None: + deletes_path = f"{tmp_path}/eq-deletes.parquet" + pq.write_table(pa.table({"bar": pa.array([2], type=pa.int32())}), deletes_path) + + task = FileScanTask( + data_file=example_task.file, + delete_files={ + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_path, + file_format=FileFormat.PARQUET, + equality_ids=[2], + ) + }, + ) + result = ArrowScan( + table_metadata=TableMetadataV2( + location="file://a/b/c.json", + last_column_id=1, + format_version=2, + current_schema_id=1, + schemas=[table_schema_simple], + partition_specs=[PartitionSpec()], + ), + io=load_file_io(), + projected_schema=table_schema_simple, + row_filter=AlwaysTrue(), + ).to_table(tasks=[task]) + + assert result.column("bar").to_pylist() == [1, 3] + assert result.column("foo").to_pylist() == ["a", "c"] + + +def test_equality_delete_projected_without_equality_column( + example_task: FileScanTask, tmp_path: str, table_schema_simple: Schema +) -> None: + deletes_path = f"{tmp_path}/eq-deletes.parquet" + pq.write_table(pa.table({"bar": pa.array([2], type=pa.int32())}), deletes_path) + projected = Schema( + NestedField(field_id=1, name="foo", field_type=StringType(), required=False), + schema_id=1, + ) + task = FileScanTask( + data_file=example_task.file, + delete_files={ + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_path, + file_format=FileFormat.PARQUET, + equality_ids=[2], + ) + }, + ) + result = ArrowScan( + table_metadata=TableMetadataV2( + location="file://a/b/c.json", + last_column_id=1, + format_version=2, + current_schema_id=1, + schemas=[table_schema_simple], + partition_specs=[PartitionSpec()], + ), + io=load_file_io(), + projected_schema=projected, + row_filter=AlwaysTrue(), + ).to_table(tasks=[task]) + + assert result.column_names == ["foo"] + assert result.column("foo").to_pylist() == ["a", "c"] + + +def test_equality_delete_treats_nulls_as_equal(tmp_path: str) -> None: + schema = Schema( + NestedField(field_id=1, name="id", field_type=IntegerType(), required=False), + schema_id=1, + ) + data_path = f"{tmp_path}/data.parquet" + deletes_path = f"{tmp_path}/eq-deletes.parquet" + pq.write_table( + pa.table({"id": pa.array([1, None, 3], type=pa.int32())}, schema=schema_to_pyarrow(schema)), + data_path, + ) + pq.write_table(pa.table({"id": pa.array([None], type=pa.int32())}), deletes_path) + + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=data_path, + file_format=FileFormat.PARQUET, + record_count=3, + file_size_in_bytes=os.path.getsize(data_path), + ) + data_file.spec_id = 0 + task = FileScanTask( + data_file=data_file, + delete_files={ + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_path, + file_format=FileFormat.PARQUET, + equality_ids=[1], + ) + }, + ) + result = ArrowScan( + table_metadata=TableMetadataV2( + location="file://a/b/c.json", + last_column_id=1, + format_version=2, + current_schema_id=1, + schemas=[schema], + partition_specs=[PartitionSpec()], + ), + io=load_file_io(), + projected_schema=schema, + row_filter=AlwaysTrue(), + ).to_table(tasks=[task]) + + assert result.column("id").to_pylist() == [1, 3] + + +def _scan_equality_task(schema: Schema, data_path: str, deletes_path: str, equality_ids: list[int]) -> pa.Table: + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=data_path, + file_format=FileFormat.PARQUET, + record_count=1, + file_size_in_bytes=os.path.getsize(data_path), + ) + data_file.spec_id = 0 + task = FileScanTask( + data_file=data_file, + delete_files={ + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_path, + file_format=FileFormat.PARQUET, + equality_ids=equality_ids, + ) + }, + ) + return ArrowScan( + table_metadata=TableMetadataV2( + location="file://a/b/c.json", + last_column_id=1, + format_version=2, + current_schema_id=1, + schemas=[schema], + partition_specs=[PartitionSpec()], + ), + io=load_file_io(), + projected_schema=schema, + row_filter=AlwaysTrue(), + ).to_table(tasks=[task]) + + +def test_equality_delete_dropped_field_does_not_join_on_subset(tmp_path: str) -> None: + current_schema = Schema( + NestedField(field_id=1, name="id", field_type=IntegerType(), required=True), + schema_id=1, + ) + delete_schema = Schema( + NestedField(field_id=1, name="id", field_type=IntegerType(), required=True), + NestedField(field_id=2, name="extra", field_type=StringType(), required=False), + schema_id=1, + ) + data_path = f"{tmp_path}/data.parquet" + deletes_path = f"{tmp_path}/eq-deletes.parquet" + pq.write_table( + pa.table({"id": pa.array([1], type=pa.int32())}, schema=schema_to_pyarrow(current_schema)), + data_path, + ) + pq.write_table( + pa.table( + { + "id": pa.array([1], type=pa.int32()), + "extra": pa.array(["x"], type=pa.string()), + }, + schema=schema_to_pyarrow(delete_schema), + ), + deletes_path, + ) + result = _scan_equality_task(current_schema, data_path, deletes_path, [1, 2]) + assert result.column("id").to_pylist() == [1] + + +def test_equality_delete_treats_uuid_nulls_as_equal(tmp_path: str) -> None: + schema = Schema( + NestedField(field_id=1, name="id", field_type=UUIDType(), required=False), + schema_id=1, + ) + data_path = f"{tmp_path}/data.parquet" + deletes_path = f"{tmp_path}/eq-deletes.parquet" + kept = uuid.UUID("00000000-0000-0000-0000-000000000001") + pq.write_table( + pa.table({"id": pa.array([kept, None], type=pa.uuid())}, schema=schema_to_pyarrow(schema)), + data_path, + ) + pq.write_table(pa.table({"id": pa.array([None], type=pa.uuid())}), deletes_path) + result = _scan_equality_task(schema, data_path, deletes_path, [1]) + assert result.column("id").to_pylist() == [kept] + + +def test_equality_delete_treats_nan_as_equal(tmp_path: str) -> None: + schema = Schema( + NestedField(field_id=1, name="v", field_type=FloatType(), required=False), + schema_id=1, + ) + data_path = f"{tmp_path}/data.parquet" + deletes_path = f"{tmp_path}/eq-deletes.parquet" + pq.write_table( + pa.table({"v": pa.array([1.0, float("nan")], type=pa.float32())}, schema=schema_to_pyarrow(schema)), + data_path, + ) + pq.write_table(pa.table({"v": pa.array([float("nan")], type=pa.float32())}), deletes_path) + result = _scan_equality_task(schema, data_path, deletes_path, [1]) + values = result.column("v").to_pylist() + assert len(values) == 1 + assert values[0] == 1.0 + + +def test_equality_delete_renamed_field_across_delete_files(tmp_path: str) -> None: + schema = Schema( + NestedField(field_id=1, name="id", field_type=IntegerType(), required=True), + schema_id=1, + ) + renamed = Schema( + NestedField(field_id=1, name="pk", field_type=IntegerType(), required=True), + schema_id=1, + ) + data_path = f"{tmp_path}/data.parquet" + deletes_a = f"{tmp_path}/eq-a.parquet" + deletes_b = f"{tmp_path}/eq-b.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2, 3], type=pa.int32())}, schema=schema_to_pyarrow(schema)), + data_path, + ) + pq.write_table(pa.table({"id": pa.array([2], type=pa.int32())}, schema=schema_to_pyarrow(schema)), deletes_a) + pq.write_table(pa.table({"pk": pa.array([3], type=pa.int32())}, schema=schema_to_pyarrow(renamed)), deletes_b) + + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=data_path, + file_format=FileFormat.PARQUET, + record_count=3, + file_size_in_bytes=os.path.getsize(data_path), + ) + data_file.spec_id = 0 + task = FileScanTask( + data_file=data_file, + delete_files={ + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_a, + file_format=FileFormat.PARQUET, + equality_ids=[1], + ), + DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=deletes_b, + file_format=FileFormat.PARQUET, + equality_ids=[1], + ), + }, + ) + result = ArrowScan( + table_metadata=TableMetadataV2( + location="file://a/b/c.json", + last_column_id=1, + format_version=2, + current_schema_id=1, + schemas=[schema], + partition_specs=[PartitionSpec()], + ), + io=load_file_io(), + projected_schema=schema, + row_filter=AlwaysTrue(), + ).to_table(tasks=[task]) + assert result.column("id").to_pylist() == [1] + + def test_pyarrow_wrap_fsspec(example_task: FileScanTask, table_schema_simple: Schema) -> None: metadata_location = "file://a/b/c.json" diff --git a/tests/table/test_delete_file_index.py b/tests/table/test_delete_file_index.py index 09dd9ac81b..3adf7373eb 100644 --- a/tests/table/test_delete_file_index.py +++ b/tests/table/test_delete_file_index.py @@ -16,12 +16,22 @@ # under the License. import pytest +from pyiceberg.conversions import to_bytes from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus +from pyiceberg.schema import Schema from pyiceberg.table.delete_file_index import PATH_FIELD_ID, DeleteFileIndex, PositionDeletes from pyiceberg.typedef import Record +from pyiceberg.types import IntegerType, LongType, NestedField, StringType -def _create_data_file(file_path: str = "s3://bucket/data.parquet", spec_id: int = 0) -> DataFile: +def _create_data_file( + file_path: str = "s3://bucket/data.parquet", + spec_id: int = 0, + lower_bounds: dict[int, bytes] | None = None, + upper_bounds: dict[int, bytes] | None = None, + null_value_counts: dict[int, int] | None = None, + value_counts: dict[int, int] | None = None, +) -> DataFile: data_file = DataFile.from_args( content=DataFileContent.DATA, file_path=file_path, @@ -29,6 +39,10 @@ def _create_data_file(file_path: str = "s3://bucket/data.parquet", spec_id: int partition=Record(), record_count=100, file_size_in_bytes=1000, + lower_bounds=lower_bounds, + upper_bounds=upper_bounds, + null_value_counts=null_value_counts, + value_counts=value_counts, ) data_file._spec_id = spec_id return data_file @@ -81,6 +95,31 @@ def _create_deletion_vector( return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file) +def _create_equality_delete( + sequence_number: int = 1, + spec_id: int = 0, + lower_bounds: dict[int, bytes] | None = None, + upper_bounds: dict[int, bytes] | None = None, + null_value_counts: dict[int, int] | None = None, + value_counts: dict[int, int] | None = None, +) -> ManifestEntry: + delete_file = DataFile.from_args( + content=DataFileContent.EQUALITY_DELETES, + file_path=f"s3://bucket/eq-delete-{sequence_number}.parquet", + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=10, + file_size_in_bytes=100, + equality_ids=[1], + lower_bounds=lower_bounds, + upper_bounds=upper_bounds, + null_value_counts=null_value_counts, + value_counts=value_counts, + ) + delete_file._spec_id = spec_id + return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file) + + def test_empty_index() -> None: index = DeleteFileIndex() data_file = _create_data_file() @@ -187,3 +226,221 @@ def test_record_equality_for_partition_lookup() -> None: assert len(index.for_data_file(1, data_file, partition_b)) == 1 assert len(index.for_data_file(1, data_file, partition_c)) == 0 + + +def test_equality_delete_sequence_number_filtering() -> None: + index = DeleteFileIndex() + + # Equality delete with sequence number 2 + index.add_delete_file(_create_equality_delete(sequence_number=2)) + + data_file = _create_data_file() + + # Data file with sequence number 1 should be affected by equality delete with sequence number 2 + assert len(index.for_data_file(1, data_file)) == 1 + + # Data file with sequence number 2 should NOT be affected by equality delete with sequence number 2 + # Equality deletes apply only to data files added in strictly earlier snapshots (seq - 1) + assert len(index.for_data_file(2, data_file)) == 0 + + # Data file with sequence number 3 should NOT be affected + assert len(index.for_data_file(3, data_file)) == 0 + + +def test_equality_delete_sequence_number_unpartitioned() -> None: + data_file = _create_data_file() + + # Create both types of deletes at sequence number 10 + pos_delete = _create_positional_delete(sequence_number=10) + eq_delete = _create_equality_delete(sequence_number=10) + + index = DeleteFileIndex() + index.add_delete_file(pos_delete) + index.add_delete_file(eq_delete) + + # Sequence 10 + deletes = index.for_data_file(10, data_file) + assert len(deletes) == 1 + # Position deletes will apply (applies to seq <= 10) + assert pos_delete.data_file in deletes + # Equality deletes will not (applies to seq < 10) + assert eq_delete.data_file not in deletes + + # Sequence 9 + # Both deletes will apply. + deletes = index.for_data_file(9, data_file) + assert len(deletes) == 2 + assert pos_delete.data_file in deletes + assert eq_delete.data_file in deletes + + # At sequence 11 + # Neither should apply. + deletes = index.for_data_file(11, data_file) + assert len(deletes) == 0 + + +def test_global_equality_deletes() -> None: + index = DeleteFileIndex() + + # Global equality delete (unpartitioned) + index.add_delete_file(_create_equality_delete(sequence_number=10)) + + partition_1 = Record(1) + partition_2 = Record(2) + + # Partitioned equality delete for partition 1 + index.add_delete_file(_create_equality_delete(sequence_number=20), partition_1) + + file_1 = _create_data_file(file_path="s3://bucket/file_1.parquet") + file_2 = _create_data_file(file_path="s3://bucket/file_2.parquet") + + # Partition 1 should have 2 equality deletes (1 global, 1 partitioned) + assert len(index.for_data_file(1, file_1, partition_1)) == 2 + # Partition 2 should have 1 equality delete (1 global) + assert len(index.for_data_file(1, file_2, partition_2)) == 1 + + +def test_equality_delete_metrics_filtering() -> None: + schema = Schema(NestedField(1, "id", IntegerType(), required=True)) + index = DeleteFileIndex(schema=schema) + + # Equality delete for rows where id is between 10 and 20 + index.add_delete_file( + _create_equality_delete( + sequence_number=100, + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + ) + ) + + # Data file with id between 0 and 5 (no overlap) + file_no_overlap = _create_data_file( + "s3://bucket/no_overlap.parquet", + lower_bounds={1: to_bytes(IntegerType(), 0)}, + upper_bounds={1: to_bytes(IntegerType(), 5)}, + ) + assert len(index.for_data_file(1, file_no_overlap)) == 0 + + # Data file with id between 15 and 25 (overlap) + file_overlap = _create_data_file( + "s3://bucket/overlap.parquet", + lower_bounds={1: to_bytes(IntegerType(), 15)}, + upper_bounds={1: to_bytes(IntegerType(), 25)}, + ) + assert len(index.for_data_file(1, file_overlap)) == 1 + + # Data file with id between 25 and 30 (no overlap) + file_no_overlap_2 = _create_data_file( + "s3://bucket/no_overlap_2.parquet", + lower_bounds={1: to_bytes(IntegerType(), 25)}, + upper_bounds={1: to_bytes(IntegerType(), 30)}, + ) + assert len(index.for_data_file(1, file_no_overlap_2)) == 0 + + +def test_equality_delete_pruned_when_delete_all_null_data_no_nulls() -> None: + schema = Schema(NestedField(1, "id", IntegerType(), required=False)) + index = DeleteFileIndex(schema=schema) + + # Equality delete targets only null rows for field 1. + eq = _create_equality_delete( + sequence_number=10, + null_value_counts={1: 10}, + value_counts={1: 10}, + ) + index.add_delete_file(eq) + + # Data file has zero nulls for field 1, so the delete cannot match. + data = _create_data_file( + null_value_counts={1: 0}, + value_counts={1: 100}, + ) + + assert len(index.for_data_file(1, data)) == 0 + + +def test_equality_delete_pruned_data_all_null_delete_no_nulls() -> None: + schema = Schema(NestedField(1, "id", IntegerType(), required=False)) + index = DeleteFileIndex(schema=schema) + + # Equality delete has no null rows for field 1. + eq = _create_equality_delete( + sequence_number=10, + null_value_counts={1: 0}, + value_counts={1: 10}, + ) + index.add_delete_file(eq) + + # Data file is all nulls for field 1, so the delete cannot match. + data = _create_data_file( + null_value_counts={1: 100}, + value_counts={1: 100}, + ) + + assert len(index.for_data_file(1, data)) == 0 + + +def test_equality_delete_metrics_after_int_to_long_promotion() -> None: + schema = Schema(NestedField(1, "id", LongType(), required=True)) + index = DeleteFileIndex(schema=schema) + + index.add_delete_file( + _create_equality_delete( + sequence_number=100, + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + ) + ) + + file_no_overlap = _create_data_file( + "s3://bucket/no_overlap.parquet", + lower_bounds={1: to_bytes(IntegerType(), 0)}, + upper_bounds={1: to_bytes(IntegerType(), 5)}, + ) + assert len(index.for_data_file(1, file_no_overlap)) == 0 + + file_overlap = _create_data_file( + "s3://bucket/overlap.parquet", + lower_bounds={1: to_bytes(IntegerType(), 15)}, + upper_bounds={1: to_bytes(IntegerType(), 25)}, + ) + assert len(index.for_data_file(1, file_overlap)) == 1 + + +def test_equality_delete_dropped_equality_id_does_not_crash() -> None: + schema = Schema(NestedField(2, "other", StringType(), required=True)) + index = DeleteFileIndex(schema=schema) + index.add_delete_file( + _create_equality_delete( + sequence_number=10, + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + ) + ) + data = _create_data_file( + lower_bounds={1: to_bytes(IntegerType(), 15)}, + upper_bounds={1: to_bytes(IntegerType(), 25)}, + ) + assert len(index.for_data_file(1, data)) == 1 + + +def test_equality_delete_not_pruned_when_both_contain_nulls() -> None: + schema = Schema(NestedField(1, "id", IntegerType(), required=False)) + index = DeleteFileIndex(schema=schema) + index.add_delete_file( + _create_equality_delete( + sequence_number=100, + lower_bounds={1: to_bytes(IntegerType(), 10)}, + upper_bounds={1: to_bytes(IntegerType(), 20)}, + null_value_counts={1: 1}, + value_counts={1: 10}, + ) + ) + data = _create_data_file( + "s3://bucket/mixed.parquet", + lower_bounds={1: to_bytes(IntegerType(), 0)}, + upper_bounds={1: to_bytes(IntegerType(), 5)}, + null_value_counts={1: 1}, + value_counts={1: 100}, + ) + assert len(index.for_data_file(1, data)) == 1