deltacat 1.1.36__py3-none-any.whl → 2.0.0b2__py3-none-any.whl

deltacat/compute/converter/pyiceberg/overrides.py

@@ -0,0 +1,135 @@
+from collections import defaultdict
+import logging
+from deltacat import logs
+import pyarrow.parquet as pq
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def parquet_files_dict_to_iceberg_data_files(io, table_metadata, files_dict_list):
+    from pyiceberg.io.pyarrow import (
+        _check_pyarrow_schema_compatible,
+        data_file_statistics_from_parquet_metadata,
+        compute_statistics_plan,
+        parquet_path_to_id_mapping,
+    )
+    from pyiceberg.manifest import (
+        DataFile,
+        DataFileContent,
+        FileFormat,
+    )
+
+    data_file_content_type = DataFileContent.POSITION_DELETES
+    iceberg_files = []
+    schema = table_metadata.schema()
+    for files_dict in files_dict_list:
+        for partition_value, file_paths in files_dict.items():
+            for file_path in file_paths:
+                input_file = io.new_input(file_path)
+                with input_file.open() as input_stream:
+                    parquet_metadata = pq.read_metadata(input_stream)
+                _check_pyarrow_schema_compatible(
+                    schema, parquet_metadata.schema.to_arrow_schema()
+                )
+
+                statistics = data_file_statistics_from_parquet_metadata(
+                    parquet_metadata=parquet_metadata,
+                    stats_columns=compute_statistics_plan(
+                        schema, table_metadata.properties
+                    ),
+                    parquet_column_mapping=parquet_path_to_id_mapping(schema),
+                )
+
+                data_file = DataFile(
+                    content=data_file_content_type,
+                    file_path=file_path,
+                    file_format=FileFormat.PARQUET,
+                    partition=partition_value,
+                    # partition=Record(**{"pk": "111", "bucket": 2}),
+                    file_size_in_bytes=len(input_file),
+                    sort_order_id=None,
+                    spec_id=table_metadata.default_spec_id,
+                    equality_ids=None,
+                    key_metadata=None,
+                    **statistics.to_serialized_dict(),
+                )
+                iceberg_files.append(data_file)
+    return iceberg_files
+
+
+def fetch_all_bucket_files(table):
+    # step 1: filter manifests using partition summaries
+    # the filter depends on the partition spec used to write the manifest file, so create a cache of filters for each spec id
+    from pyiceberg.typedef import (
+        KeyDefaultDict,
+    )
+
+    data_scan = table.scan()
+    snapshot = data_scan.snapshot()
+    if not snapshot:
+        return iter([])
+    manifest_evaluators = KeyDefaultDict(data_scan._build_manifest_evaluator)
+
+    manifests = [
+        manifest_file
+        for manifest_file in snapshot.manifests(data_scan.io)
+        if manifest_evaluators[manifest_file.partition_spec_id](manifest_file)
+    ]
+
+    # step 2: filter the data files in each manifest
+    # this filter depends on the partition spec used to write the manifest file
+    from pyiceberg.expressions.visitors import _InclusiveMetricsEvaluator
+    from pyiceberg.types import (
+        strtobool,
+    )
+    from pyiceberg.table import _min_sequence_number, _open_manifest
+    from pyiceberg.utils.concurrent import ExecutorFactory
+    from itertools import chain
+    from pyiceberg.manifest import DataFileContent
+
+    partition_evaluators = KeyDefaultDict(data_scan._build_partition_evaluator)
+    metrics_evaluator = _InclusiveMetricsEvaluator(
+        data_scan.table_metadata.schema(),
+        data_scan.row_filter,
+        data_scan.case_sensitive,
+        strtobool(data_scan.options.get("include_empty_files", "false")),
+    ).eval
+
+    min_sequence_number = _min_sequence_number(manifests)
+
+    # {"bucket_index": List[DataFile]}
+    data_entries = defaultdict(list)
+    equality_data_entries = defaultdict(list)
+    positional_delete_entries = defaultdict(list)
+
+    executor = ExecutorFactory.get_or_create()
+    for manifest_entry in chain(
+        *executor.map(
+            lambda args: _open_manifest(*args),
+            [
+                (
+                    data_scan.io,
+                    manifest,
+                    partition_evaluators[manifest.partition_spec_id],
+                    metrics_evaluator,
+                )
+                for manifest in manifests
+                if data_scan._check_sequence_number(min_sequence_number, manifest)
+            ],
+        )
+    ):
+        data_file = manifest_entry.data_file
+        file_sequence_number = manifest_entry.sequence_number
+        data_file_tuple = (file_sequence_number, data_file)
+        partition_value = data_file.partition
+        if data_file.content == DataFileContent.DATA:
+            data_entries[partition_value].append(data_file_tuple)
+        if data_file.content == DataFileContent.POSITION_DELETES:
+            positional_delete_entries[partition_value].append(data_file_tuple)
+        elif data_file.content == DataFileContent.EQUALITY_DELETES:
+            equality_data_entries[partition_value].append(data_file_tuple)
+        else:
+            logger.warning(
+                f"Unknown DataFileContent ({data_file.content}): {manifest_entry}"
+            )
+    return data_entries, equality_data_entries, positional_delete_entries

deltacat/compute/converter/pyiceberg/update_snapshot_overrides.py

@@ -0,0 +1,251 @@
+from typing import Optional, List
+import uuid
+from pyiceberg.table.snapshots import (
+    Operation,
+)
+from pyiceberg.manifest import (
+    DataFileContent,
+    ManifestContent,
+    ManifestEntry,
+    ManifestEntryStatus,
+    ManifestFile,
+    write_manifest,
+)
+import itertools
+from pyiceberg.utils.concurrent import ExecutorFactory
+from pyiceberg.table.update.snapshot import UpdateSnapshot, _SnapshotProducer
+
+
+class _ReplaceFiles(_SnapshotProducer["_ReplaceFiles"]):
+    """Overwrites data from the table. This will produce an OVERWRITE snapshot.
+
+    Data and delete files were added and removed in a logical overwrite operation.
+    """
+
+    def _existing_manifests(self) -> List[ManifestFile]:
+        """Determine if there are any existing manifest files."""
+        existing_files = []
+        snapshot = self._transaction.table_metadata.current_snapshot()
+        if snapshot:
+            for manifest_file in snapshot.manifests(io=self._io):
+                entries = manifest_file.fetch_manifest_entry(
+                    io=self._io, discard_deleted=True
+                )
+
+                found_deleted_data_files = [
+                    entry.data_file
+                    for entry in entries
+                    if entry.data_file in self._deleted_data_files
+                ]
+
+                if len(found_deleted_data_files) == 0:
+                    existing_files.append(manifest_file)
+                else:
+                    # We have to replace the manifest file without the deleted data files
+                    if any(
+                        entry.data_file not in found_deleted_data_files
+                        for entry in entries
+                    ):
+                        with write_manifest(
+                            format_version=self._transaction.table_metadata.format_version,
+                            spec=self._transaction.table_metadata.specs()[
+                                manifest_file.partition_spec_id
+                            ],
+                            schema=self._transaction.table_metadata.schema(),
+                            output_file=self.new_manifest_output(),
+                            snapshot_id=self._snapshot_id,
+                        ) as writer:
+                            [
+                                writer.add_entry(
+                                    ManifestEntry(
+                                        status=ManifestEntryStatus.EXISTING,
+                                        snapshot_id=entry.snapshot_id,
+                                        sequence_number=entry.sequence_number,
+                                        file_sequence_number=entry.file_sequence_number,
+                                        data_file=entry.data_file,
+                                    )
+                                )
+                                for entry in entries
+                                if entry.data_file not in found_deleted_data_files
+                            ]
+                        existing_files.append(writer.to_manifest_file())
+        return existing_files
+
+    def _deleted_entries(self) -> List[ManifestEntry]:
+        """To determine if we need to record any deleted entries.
+
+        With a full overwrite all the entries are considered deleted.
+        With partial overwrites we have to use the predicate to evaluate
+        which entries are affected.
+        """
+        if self._parent_snapshot_id is not None:
+            previous_snapshot = self._transaction.table_metadata.snapshot_by_id(
+                self._parent_snapshot_id
+            )
+            if previous_snapshot is None:
+                # This should never happen since you cannot overwrite an empty table
+                raise ValueError(
+                    f"Could not find the previous snapshot: {self._parent_snapshot_id}"
+                )
+
+            executor = ExecutorFactory.get_or_create()
+
+            def _get_entries(manifest: ManifestFile) -> List[ManifestEntry]:
+                return [
+                    ManifestEntry(
+                        status=ManifestEntryStatus.DELETED,
+                        snapshot_id=entry.snapshot_id,
+                        sequence_number=entry.sequence_number,
+                        file_sequence_number=entry.file_sequence_number,
+                        data_file=entry.data_file,
+                    )
+                    for entry in manifest.fetch_manifest_entry(
+                        self._io, discard_deleted=True
+                    )
+                    if entry.data_file.content == DataFileContent.DATA
+                    and entry.data_file in self._deleted_data_files
+                ]
+
+            list_of_entries = executor.map(
+                _get_entries, previous_snapshot.manifests(self._io)
+            )
+            return list(itertools.chain(*list_of_entries))
+        else:
+            return []
+
+
+def replace(
+    self,
+    commit_uuid: Optional[uuid.UUID] = None,
+    using_starting_sequence: Optional[bool] = False,
+) -> _ReplaceFiles:
+    return _ReplaceFiles(
+        commit_uuid=commit_uuid,
+        operation=Operation.REPLACE
+        if self._transaction.table_metadata.current_snapshot() is not None
+        else Operation.APPEND,
+        transaction=self._transaction,
+        io=self._io,
+        snapshot_properties=self._snapshot_properties,
+        using_starting_sequence=using_starting_sequence,
+    )
+
+
+UpdateSnapshot.replace = replace
+
+
+def commit_replace_snapshot(
+    iceberg_table, to_be_deleted_files_list, new_position_delete_files
+):
+    tx = iceberg_table.transaction()
+    snapshot_properties = {}
+    commit_uuid = uuid.uuid4()
+    update_snapshot = tx.update_snapshot(snapshot_properties=snapshot_properties)
+    replace_snapshot = replace(
+        self=update_snapshot, commit_uuid=commit_uuid, using_starting_sequence=False
+    )
+    for to_be_deleted_file in to_be_deleted_files_list:
+        replace_snapshot.append_data_file(to_be_deleted_file)
+    for to_be_added_file in new_position_delete_files:
+        replace_snapshot.delete_data_file(to_be_added_file)
+    replace_snapshot._commit()
+    tx.commit_transaction()
+
+
+def append_delete_files_override(update_snapshot):
+    commit_uuid = uuid.uuid4()
+    return _AppendDeleteFilesOverride(
+        commit_uuid=commit_uuid,
+        operation=Operation.APPEND,
+        transaction=update_snapshot._transaction,
+        io=update_snapshot._io,
+        snapshot_properties=update_snapshot._snapshot_properties,
+    )
+
+
+class _AppendDeleteFilesOverride(_SnapshotProducer):
+    def _manifests(self):
+        def _write_added_manifest():
+            if self._added_data_files:
+                with write_manifest(
+                    format_version=self._transaction.table_metadata.format_version,
+                    spec=self._transaction.table_metadata.spec(),
+                    schema=self._transaction.table_metadata.schema(),
+                    output_file=self.new_manifest_output(),
+                    snapshot_id=self._snapshot_id,
+                ) as writer:
+                    for data_file in self._added_data_files:
+                        writer.add(
+                            ManifestEntry(
+                                status=ManifestEntryStatus.ADDED,
+                                snapshot_id=self._snapshot_id,
+                                sequence_number=None,
+                                file_sequence_number=None,
+                                data_file=data_file,
+                            )
+                        )
+                        writer.content = self.writer_content
+                return [writer.to_manifest_file()]
+            else:
+                return []
+
+        executor = ExecutorFactory.get_or_create()
+
+        added_manifests = executor.submit(_write_added_manifest)
+        existing_manifests = executor.submit(self._existing_manifests)
+
+        return self._process_manifests(
+            added_manifests.result() + existing_manifests.result()
+        )
+
+    def writer_content(self):
+        return ManifestContent.DELETES
+
+    def _existing_manifests(self) -> List[ManifestFile]:
+        """To determine if there are any existing manifest files.
+
+        A fast append will add another ManifestFile to the ManifestList.
+        All the existing manifest files are considered existing.
+        """
+        existing_manifests = []
+
+        if self._parent_snapshot_id is not None:
+            previous_snapshot = self._transaction.table_metadata.snapshot_by_id(
+                self._parent_snapshot_id
+            )
+
+            if previous_snapshot is None:
+                raise ValueError(
+                    f"Snapshot could not be found: {self._parent_snapshot_id}"
+                )
+
+            for manifest in previous_snapshot.manifests(io=self._io):
+                if (
+                    manifest.has_added_files()
+                    or manifest.has_existing_files()
+                    or manifest.added_snapshot_id == self._snapshot_id
+                ):
+                    existing_manifests.append(manifest)
+
+        return existing_manifests
+
+    def _deleted_entries(self) -> List[ManifestEntry]:
+        """To determine if we need to record any deleted manifest entries.
+
+        In case of an append, nothing is deleted.
+        """
+        return []
+
+
+def commit_append_snapshot(iceberg_table, new_position_delete_files):
+    with iceberg_table.transaction() as tx:
+        if iceberg_table.metadata.name_mapping() is None:
+            tx.set_properties(
+                **{
+                    "schema.name-mapping.default": tx.table_metadata.schema().name_mapping.model_dump_json()
+                }
+            )
+        with append_delete_files_override(tx.update_snapshot()) as append_snapshot:
+            if new_position_delete_files:
+                for data_file in new_position_delete_files:
+                    append_snapshot.append_data_file(data_file)

deltacat/compute/converter/steps/convert.py

@@ -0,0 +1,211 @@
+import pyarrow.compute as pc
+
+import deltacat.compute.converter.utils.iceberg_columns as sc
+import pyarrow as pa
+
+from collections import defaultdict
+import ray
+import logging
+from deltacat.compute.converter.model.convert_input import ConvertInput
+from deltacat.compute.converter.steps.dedupe import dedupe_data_files
+from deltacat.compute.converter.utils.s3u import upload_table_with_retry
+from deltacat.compute.converter.utils.io import (
+    download_data_table_and_append_iceberg_columns,
+)
+from deltacat.compute.converter.utils.converter_session_utils import (
+    partition_value_record_to_partition_value_string,
+)
+
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+@ray.remote
+def convert(convert_input: ConvertInput):
+    convert_input_files = convert_input.convert_input_files
+    convert_task_index = convert_input.convert_task_index
+    iceberg_table_warehouse_prefix = convert_input.iceberg_table_warehouse_prefix
+    identifier_fields = convert_input.identifier_fields
+    compact_small_files = convert_input.compact_small_files
+    position_delete_for_multiple_data_files = (
+        convert_input.position_delete_for_multiple_data_files
+    )
+    max_parallel_data_file_download = convert_input.max_parallel_data_file_download
+    s3_file_system = convert_input.s3_file_system
+    if not position_delete_for_multiple_data_files:
+        raise NotImplementedError(
+            f"Distributed file level position delete compute is not supported yet"
+        )
+    if compact_small_files:
+        raise NotImplementedError(f"Compact previous position delete not supported yet")
+
+    logger.info(f"Starting convert task index: {convert_task_index}")
+
+    applicable_data_files = convert_input_files.applicable_data_files
+    applicable_equality_delete_files = (
+        convert_input_files.applicable_equality_delete_files
+    )
+    all_data_files_for_this_bucket = convert_input_files.all_data_files_for_dedupe
+
+    partition_value_str = partition_value_record_to_partition_value_string(
+        convert_input_files.partition_value
+    )
+    partition_value = convert_input_files.partition_value
+    iceberg_table_warehouse_prefix_with_partition = (
+        f"{iceberg_table_warehouse_prefix}/{partition_value_str}"
+    )
+    enforce_primary_key_uniqueness = convert_input.enforce_primary_key_uniqueness
+    total_pos_delete_table = []
+    if applicable_equality_delete_files:
+        (
+            pos_delete_after_converting_equality_delete
+        ) = compute_pos_delete_with_limited_parallelism(
+            data_files_list=applicable_data_files,
+            identifier_columns=identifier_fields,
+            equality_delete_files_list=applicable_equality_delete_files,
+            iceberg_table_warehouse_prefix_with_partition=iceberg_table_warehouse_prefix_with_partition,
+            max_parallel_data_file_download=max_parallel_data_file_download,
+            s3_file_system=s3_file_system,
+        )
+        if pos_delete_after_converting_equality_delete:
+            total_pos_delete_table.append(pos_delete_after_converting_equality_delete)
+
+    if enforce_primary_key_uniqueness:
+        data_files_to_dedupe = get_additional_applicable_data_files(
+            all_data_files=all_data_files_for_this_bucket,
+            data_files_downloaded=applicable_data_files,
+        )
+        pos_delete_after_dedupe = dedupe_data_files(
+            data_file_to_dedupe=data_files_to_dedupe,
+            identify_column_name_concatenated=identifier_fields[0],
+            identifier_columns=identifier_fields,
+            merge_sort_column=sc._ORDERED_RECORD_IDX_COLUMN_NAME,
+        )
+        total_pos_delete_table.append(pos_delete_after_dedupe)
+
+    total_pos_delete = pa.concat_tables(total_pos_delete_table)
+    to_be_added_files_list = upload_table_with_retry(
+        table=total_pos_delete,
+        s3_url_prefix=iceberg_table_warehouse_prefix_with_partition,
+        s3_table_writer_kwargs={},
+        s3_file_system=s3_file_system,
+    )
+
+    to_be_delete_files_dict = defaultdict()
+    if applicable_equality_delete_files:
+        to_be_delete_files_dict[partition_value] = [
+            equality_delete_file[1]
+            for equality_delete_file in applicable_equality_delete_files
+        ]
+    to_be_added_files_dict = defaultdict()
+    to_be_added_files_dict[partition_value] = to_be_added_files_list
+    return (to_be_delete_files_dict, to_be_added_files_dict)
+
+
+def get_additional_applicable_data_files(all_data_files, data_files_downloaded):
+    data_file_to_dedupe = all_data_files
+    if data_files_downloaded:
+        data_file_to_dedupe = list(set(all_data_files) - set(data_files_downloaded))
+    return data_file_to_dedupe
+
+
+def filter_rows_to_be_deleted(
+    equality_delete_table, data_file_table, identifier_columns
+):
+    identifier_column = identifier_columns[0]
+    if equality_delete_table and data_file_table:
+        equality_deletes = pc.is_in(
+            data_file_table[identifier_column],
+            equality_delete_table[identifier_column],
+        )
+        position_delete_table = data_file_table.filter(equality_deletes)
+        logger.info(f"positional_delete_table:{position_delete_table.to_pydict()}")
+        logger.info(f"data_file_table:{data_file_table.to_pydict()}")
+        logger.info(
+            f"length_pos_delete_table, {len(position_delete_table)}, length_data_table:{len(data_file_table)}"
+        )
+    return position_delete_table
+
+
+def compute_pos_delete_converting_equality_deletes(
+    equality_delete_table,
+    data_file_table,
+    identifier_columns,
+    iceberg_table_warehouse_prefix_with_partition,
+    s3_file_system,
+):
+    new_position_delete_table = filter_rows_to_be_deleted(
+        data_file_table=data_file_table,
+        equality_delete_table=equality_delete_table,
+        identifier_columns=identifier_columns,
+    )
+    if new_position_delete_table:
+        logger.info(
+            f"Length of position delete table after converting from equality deletes:{len(new_position_delete_table)}"
+        )
+    else:
+        return None
+    return new_position_delete_table
+
+
+def download_bucketed_table(data_files, equality_delete_files):
+    from deltacat.utils.pyarrow import s3_file_to_table
+
+    compacted_table = s3_file_to_table(
+        [data_file.file_path for data_file in data_files]
+    )
+    equality_delete_table = s3_file_to_table(
+        [eq_file.file_path for eq_file in equality_delete_files]
+    )
+    return compacted_table, equality_delete_table
+
+
+def compute_pos_delete_with_limited_parallelism(
+    data_files_list,
+    identifier_columns,
+    equality_delete_files_list,
+    iceberg_table_warehouse_prefix_with_partition,
+    max_parallel_data_file_download,
+    s3_file_system,
+):
+    for data_files, equality_delete_files in zip(
+        data_files_list, equality_delete_files_list
+    ):
+        data_table_total = []
+        for data_file in data_files:
+            data_table = download_data_table_and_append_iceberg_columns(
+                data_files=data_file[1],
+                columns_to_download=identifier_columns,
+                additional_columns_to_append=[
+                    sc._FILE_PATH_COLUMN_NAME,
+                    sc._ORDERED_RECORD_IDX_COLUMN_NAME,
+                ],
+                sequence_number=data_file[0],
+            )
+            data_table_total.append(data_table)
+        data_table_total = pa.concat_tables(data_table_total)
+
+        equality_delete_table_total = []
+        for equality_delete in equality_delete_files:
+            equality_delete_table = download_data_table_and_append_iceberg_columns(
+                data_files=equality_delete[1],
+                columns_to_download=identifier_columns,
+            )
+            equality_delete_table_total.append(equality_delete_table)
+        equality_delete_table_total = pa.concat_tables(equality_delete_table_total)
+
+    new_pos_delete_table = compute_pos_delete_converting_equality_deletes(
+        equality_delete_table=equality_delete_table_total,
+        data_file_table=data_table_total,
+        iceberg_table_warehouse_prefix_with_partition=iceberg_table_warehouse_prefix_with_partition,
+        identifier_columns=identifier_columns,
+        s3_file_system=s3_file_system,
+    )
+    if not new_pos_delete_table:
+        logger.info("No records deleted based on equality delete converstion")
+
+    logger.info(
+        f"Number of records to delete based on equality delete convertion:{len(new_pos_delete_table)}"
+    )
+    return new_pos_delete_table

deltacat/compute/converter/steps/dedupe.py

@@ -0,0 +1,60 @@
+import pyarrow as pa
+import pyarrow.compute as pc
+import deltacat.compute.converter.utils.iceberg_columns as sc
+from deltacat.compute.converter.utils.io import (
+    download_data_table_and_append_iceberg_columns,
+)
+
+
+def dedupe_data_files(
+    data_file_to_dedupe,
+    identify_column_name_concatenated,
+    identifier_columns,
+    merge_sort_column,
+):
+    data_file_table = []
+
+    # Sort data files by file sequence number first
+    data_file_to_dedupe = sorted(data_file_to_dedupe, key=lambda f: f[0])
+    for file_tuple in data_file_to_dedupe:
+        sequence_number = file_tuple[0]
+        data_file = file_tuple[1]
+        data_file_to_dedupe_table = download_data_table_and_append_iceberg_columns(
+            file=data_file,
+            columns_to_download=identifier_columns,
+            additional_columns_to_append=[
+                sc._FILE_PATH_COLUMN_NAME,
+                sc._ORDERED_RECORD_IDX_COLUMN_NAME,
+            ],
+            sequence_number=sequence_number,
+        )
+        data_file_table.append(data_file_to_dedupe_table)
+
+    final_data_to_dedupe = pa.concat_tables(data_file_table)
+
+    record_idx_iterator = iter(range(len(final_data_to_dedupe)))
+
+    # Append global record index to used as aggregate column
+    final_data_to_dedupe = sc.append_global_record_idx_column(
+        final_data_to_dedupe, record_idx_iterator
+    )
+
+    final_data_table_indices = final_data_to_dedupe.group_by(
+        identify_column_name_concatenated, use_threads=False
+    ).aggregate([(sc._GLOBAL_RECORD_IDX_COLUMN_NAME, "max")])
+
+    pos_delete_indices = pc.invert(
+        pc.is_in(
+            final_data_to_dedupe[sc._GLOBAL_RECORD_IDX_COLUMN_NAME],
+            value_set=final_data_table_indices[
+                f"{sc._GLOBAL_RECORD_IDX_COLUMN_NAME}_max"
+            ],
+        )
+    )
+
+    final_data_table_to_delete = final_data_to_dedupe.filter(pos_delete_indices)
+
+    final_data_table_to_delete = final_data_table_to_delete.drop(
+        [identify_column_name_concatenated, sc._GLOBAL_RECORD_IDX_COLUMN_NAME]
+    )
+    return final_data_table_to_delete