deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/compute/converter/model/converter_session_params.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+from typing import Optional, Dict, Any, List
+from deltacat.compute.converter.constants import (
+    DEFAULT_CONVERTER_TASK_MAX_PARALLELISM,
+)
+from deltacat.constants import DEFAULT_NAMESPACE
+from fsspec import AbstractFileSystem
+from pyiceberg.catalog import Catalog
+
+
+class ConverterSessionParams(dict):
+    """
+    This class represents the parameters passed to convert_ (deltacat/compute/compactor/compaction_session.py)
+    """
+
+    @staticmethod
+    def of(params: Optional[Dict[str, Any]]) -> ConverterSessionParams:
+        params = {} if params is None else params
+        assert params.get("catalog") is not None, "catalog is a required arg"
+        assert (
+            params.get("iceberg_table_name") is not None
+        ), "iceberg_table_name is a required arg"
+        assert (
+            params.get("iceberg_warehouse_bucket_name") is not None
+        ), "iceberg_warehouse_bucket_name is a required arg"
+        result = ConverterSessionParams(params)
+
+        result.iceberg_namespace = params.get("iceberg_namespace", DEFAULT_NAMESPACE)
+        result.enforce_primary_key_uniqueness = params.get(
+            "enforce_primary_key_uniqueness", False
+        )
+        result.compact_previous_position_delete_files = params.get(
+            "compact_previous_position_delete_files", False
+        )
+
+        # For Iceberg v3 spec, option to produce delete vector that can establish 1:1 mapping with data files.
+        result.position_delete_for_multiple_data_files = params.get(
+            "position_delete_for_multiple_data_files", True
+        )
+        result.task_max_parallelism = params.get(
+            "task_max_parallelism", DEFAULT_CONVERTER_TASK_MAX_PARALLELISM
+        )
+        result.merge_keys = params.get("merge_keys", None)
+        result.s3_client_kwargs = params.get("s3_client_kwargs", {})
+        result.filesystem = params.get("filesystem", None)
+        result.s3_prefix_override = params.get("s3_prefix_override", None)
+
+        return result
+
+    @property
+    def catalog(self) -> Catalog:
+        return self["catalog"]
+
+    @property
+    def iceberg_table_name(self) -> str:
+        return self["iceberg_table_name"]
+
+    @property
+    def iceberg_warehouse_bucket_name(self) -> str:
+        return self["iceberg_warehouse_bucket_name"]
+
+    @property
+    def iceberg_namespace(self) -> str:
+        return self["iceberg_namespace"]
+
+    @iceberg_namespace.setter
+    def iceberg_namespace(self, iceberg_namespace: str) -> None:
+        self["iceberg_namespace"] = iceberg_namespace
+
+    @property
+    def enforce_primary_key_uniqueness(self) -> bool:
+        return self["enforce_primary_key_uniqueness"]
+
+    @enforce_primary_key_uniqueness.setter
+    def enforce_primary_key_uniqueness(
+        self, enforce_primary_key_uniqueness: bool
+    ) -> None:
+        self["enforce_primary_key_uniqueness"] = enforce_primary_key_uniqueness
+
+    @property
+    def compact_previous_position_delete_files(self) -> bool:
+        return self["compact_previous_position_delete_files"]
+
+    @compact_previous_position_delete_files.setter
+    def compact_previous_position_delete_files(
+        self, compact_previous_position_delete_files: bool
+    ) -> None:
+        self[
+            "compact_previous_position_delete_files"
+        ] = compact_previous_position_delete_files
+
+    @property
+    def position_delete_for_multiple_data_files(self) -> bool:
+        return self["position_delete_for_multiple_data_files"]
+
+    @position_delete_for_multiple_data_files.setter
+    def position_delete_for_multiple_data_files(
+        self, position_delete_for_multiple_data_files: bool
+    ) -> None:
+        self[
+            "position_delete_for_multiple_data_files"
+        ] = position_delete_for_multiple_data_files
+
+    @property
+    def task_max_parallelism(self) -> int:
+        return self["task_max_parallelism"]
+
+    @task_max_parallelism.setter
+    def task_max_parallelism(self, task_max_parallelism: int) -> None:
+        self["task_max_parallelism"] = task_max_parallelism
+
+    @property
+    def merge_keys(self) -> Optional[List[str]]:
+        return self["merge_keys"]
+
+    @merge_keys.setter
+    def merge_keys(self, merge_keys: Optional[List[str]]) -> None:
+        self["merge_keys"] = merge_keys
+
+    @property
+    def s3_client_kwargs(self) -> Dict[str, Any]:
+        return self["s3_client_kwargs"]
+
+    @s3_client_kwargs.setter
+    def s3_client_kwargs(self, s3_client_kwargs: Dict[str, Any]) -> None:
+        self["s3_client_kwargs"] = s3_client_kwargs
+
+    @property
+    def filesystem(self) -> Optional[AbstractFileSystem]:
+        return self["filesystem"]
+
+    @filesystem.setter
+    def filesystem(self, filesystem: Optional[AbstractFileSystem]) -> None:
+        self["filesystem"] = filesystem
+
+    @property
+    def location_provider_prefix_override(self) -> Optional[str]:
+        return self["location_provider_prefix_override"]
+
+    @location_provider_prefix_override.setter
+    def location_provider_prefix_override(
+        self, location_provider_prefix_override: Optional[str]
+    ) -> None:
+        self["location_provider_prefix_override"] = location_provider_prefix_override

deltacat/compute/converter/pyiceberg/catalog.py

@@ -0,0 +1,78 @@
+from typing import Optional, Dict, Any
+from pyiceberg.table import Table
+from pyiceberg.catalog import Catalog, load_catalog as pyiceberg_load_catalog
+from botocore.credentials import Credentials
+import boto3
+from boto3.session import Session
+
+
+def load_catalog(
+    iceberg_catalog_name: str, iceberg_catalog_properties: Dict[str, Any]
+) -> Catalog:
+    catalog = pyiceberg_load_catalog(
+        name=iceberg_catalog_name,
+        **iceberg_catalog_properties,
+    )
+    return catalog
+
+
+def get_s3_path(
+    bucket_name: str,
+    database_name: Optional[str] = None,
+    table_name: Optional[str] = None,
+) -> str:
+    result_path = f"s3://{bucket_name}"
+    if database_name is not None:
+        result_path += f"/{database_name}.db"
+
+    if table_name is not None:
+        result_path += f"/{table_name}"
+    return result_path
+
+
+def get_bucket_name() -> str:
+    return "test-bucket"
+
+
+def get_s3_prefix() -> str:
+    return get_s3_path(get_bucket_name())
+
+
+def get_credential() -> Credentials:
+    boto3_session: Session = boto3.Session()
+    credentials: Credentials = boto3_session.get_credentials()
+    return credentials
+
+
+def get_glue_catalog() -> Catalog:
+    credential = get_credential()
+    # Credentials are refreshable, so accessing your access key / secret key
+    # separately can lead to a race condition. Use this to get an actual matched
+    # set.
+    credential = credential.get_frozen_credentials()
+    access_key_id = credential.access_key
+    secret_access_key = credential.secret_key
+    session_token = credential.token
+    s3_path = get_s3_prefix()
+    glue_catalog = pyiceberg_load_catalog(
+        "glue",
+        **{
+            "warehouse": s3_path,
+            "type": "glue",
+            "aws_access_key_id": access_key_id,
+            "aws_secret_access_key": secret_access_key,
+            "aws_session_token": session_token,
+            "region_name": "us-east-1",
+            "s3.access-key-id": access_key_id,
+            "s3.secret-access-key": secret_access_key,
+            "s3.session-token": session_token,
+            "s3.region": "us-east-1",
+        },
+    )
+
+    return glue_catalog
+
+
+def load_table(catalog: Catalog, table_name: str) -> Table:
+    loaded_table = catalog.load_table(table_name)
+    return loaded_table

deltacat/compute/converter/pyiceberg/overrides.py

@@ -0,0 +1,263 @@
+from collections import defaultdict
+import logging
+from deltacat import logs
+import pyarrow
+import pyarrow.parquet as pq
+from pyiceberg.io.pyarrow import (
+    parquet_path_to_id_mapping,
+    StatisticsCollector,
+    MetricModeTypes,
+    DataFileStatistics,
+    MetricsMode,
+    StatsAggregator,
+)
+from typing import Dict, List, Set, Any, Tuple
+from deltacat.compute.converter.utils.iceberg_columns import (
+    ICEBERG_RESERVED_FIELD_ID_FOR_FILE_PATH_COLUMN,
+    ICEBERG_RESERVED_FIELD_ID_FOR_POS_COLUMN,
+)
+from pyiceberg.io.pyarrow import (
+    compute_statistics_plan,
+)
+from pyiceberg.manifest import (
+    DataFile,
+    DataFileContent,
+    FileFormat,
+)
+from pyiceberg.table import _min_sequence_number, _open_manifest, Table
+from pyiceberg.utils.concurrent import ExecutorFactory
+from itertools import chain
+from pyiceberg.typedef import (
+    KeyDefaultDict,
+)
+from pyiceberg.schema import Schema
+from pyiceberg.io import FileIO
+from deltacat.compute.converter.model.convert_input_files import (
+    DataFileList,
+)
+
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def parquet_path_to_id_mapping_override(schema: Schema) -> Dict[str, int]:
+    res = parquet_path_to_id_mapping(schema)
+    # Override here to insert position delete reserved column field IDs
+    res["file_path"] = ICEBERG_RESERVED_FIELD_ID_FOR_FILE_PATH_COLUMN
+    res["pos"] = ICEBERG_RESERVED_FIELD_ID_FOR_POS_COLUMN
+    return res
+
+
+def data_file_statistics_from_parquet_metadata(
+    parquet_metadata: pq.FileMetaData,
+    stats_columns: Dict[int, StatisticsCollector],
+    parquet_column_mapping: Dict[str, int],
+) -> DataFileStatistics:
+    """
+    Overrides original Pyiceberg function: Compute and return DataFileStatistics that includes the following.
+
+    - record_count
+    - column_sizes
+    - value_counts
+    - null_value_counts
+    - nan_value_counts
+    - column_aggregates
+    - split_offsets
+
+    Args:
+        parquet_metadata (pyarrow.parquet.FileMetaData): A pyarrow metadata object.
+        stats_columns (Dict[int, StatisticsCollector]): The statistics gathering plan. It is required to
+            set the mode for column metrics collection
+        parquet_column_mapping (Dict[str, int]): The mapping of the parquet file name to the field ID
+    """
+    column_sizes: Dict[int, int] = {}
+    value_counts: Dict[int, int] = {}
+    split_offsets: List[int] = []
+
+    null_value_counts: Dict[int, int] = {}
+    nan_value_counts: Dict[int, int] = {}
+
+    col_aggs = {}
+
+    invalidate_col: Set[int] = set()
+    for r in range(parquet_metadata.num_row_groups):
+        # References:
+        # https://github.com/apache/iceberg/blob/fc381a81a1fdb8f51a0637ca27cd30673bd7aad3/parquet/src/main/java/org/apache/iceberg/parquet/ParquetUtil.java#L232
+        # https://github.com/apache/parquet-mr/blob/ac29db4611f86a07cc6877b416aa4b183e09b353/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/metadata/ColumnChunkMetaData.java#L184
+
+        row_group = parquet_metadata.row_group(r)
+
+        data_offset = row_group.column(0).data_page_offset
+        dictionary_offset = row_group.column(0).dictionary_page_offset
+
+        if row_group.column(0).has_dictionary_page and dictionary_offset < data_offset:
+            split_offsets.append(dictionary_offset)
+        else:
+            split_offsets.append(data_offset)
+
+        for pos in range(parquet_metadata.num_columns):
+            column = row_group.column(pos)
+            field_id = parquet_column_mapping[column.path_in_schema]
+            if field_id in stats_columns:
+                stats_col = stats_columns[field_id]
+
+                column_sizes.setdefault(field_id, 0)
+                column_sizes[field_id] += column.total_compressed_size
+
+                if stats_col.mode == MetricsMode(MetricModeTypes.NONE):
+                    continue
+
+                value_counts[field_id] = (
+                    value_counts.get(field_id, 0) + column.num_values
+                )
+
+                if column.is_stats_set:
+                    try:
+                        statistics = column.statistics
+
+                        if statistics.has_null_count:
+                            null_value_counts[field_id] = (
+                                null_value_counts.get(field_id, 0)
+                                + statistics.null_count
+                            )
+
+                        if stats_col.mode == MetricsMode(MetricModeTypes.COUNTS):
+                            continue
+
+                        if field_id not in col_aggs:
+                            col_aggs[field_id] = StatsAggregator(
+                                stats_col.iceberg_type,
+                                statistics.physical_type,
+                                stats_col.mode.length,
+                            )
+
+                        col_aggs[field_id].update_min(statistics.min)
+                        col_aggs[field_id].update_max(statistics.max)
+
+                    except pyarrow.lib.ArrowNotImplementedError as e:
+                        invalidate_col.add(field_id)
+                        logger.warning(e)
+            else:
+                # Note: Removed original adding columns without stats to invalid column logic here
+                logger.warning(
+                    "PyArrow statistics missing for column %d when writing file", pos
+                )
+
+    split_offsets.sort()
+
+    for field_id in invalidate_col:
+        del col_aggs[field_id]
+        del null_value_counts[field_id]
+
+    return DataFileStatistics(
+        record_count=parquet_metadata.num_rows,
+        column_sizes=column_sizes,
+        value_counts=value_counts,
+        null_value_counts=null_value_counts,
+        nan_value_counts=nan_value_counts,
+        column_aggregates=col_aggs,
+        split_offsets=split_offsets,
+    )
+
+
+def parquet_files_dict_to_iceberg_data_files(
+    io: FileIO,
+    table_metadata: Any,
+    files_dict: Dict[Any, List[str]],
+    file_content_type: DataFileContent,
+) -> List[DataFile]:
+    iceberg_files = []
+    schema = table_metadata.schema()
+    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)
+
+            # Removed _check_pyarrow_schema_compatible() here since reserved columns does not comply to all rules.
+
+            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_override(schema),
+            )
+
+            data_file = DataFile(
+                content=file_content_type,
+                file_path=file_path,
+                file_format=FileFormat.PARQUET,
+                partition=partition_value,
+                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: Table,
+) -> Tuple[Dict[Any, DataFileList], Dict[Any, DataFileList], Dict[Any, DataFileList]]:
+    # 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
+    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
+    partition_evaluators = KeyDefaultDict(data_scan._build_partition_evaluator)
+    residual_evaluators = KeyDefaultDict(data_scan._build_residual_evaluator)
+    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],
+                    residual_evaluators[manifest.partition_spec_id],
+                    data_scan._build_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)
+        elif 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