deltacat 2.0.0b11__py3-none-any.whl → 2.0.0b12__py3-none-any.whl

deltacat/compute/converter/utils/converter_session_utils.py

@@ -1,97 +1,102 @@
 from collections import defaultdict
 import logging
 from deltacat import logs
-from deltacat.compute.converter.model.convert_input_files import ConvertInputFiles
+from deltacat.compute.converter.model.convert_input_files import (
+    ConvertInputFiles,
+    DataFileList,
+    DataFileListGroup,
+)
+from typing import List, Dict, Tuple, Any
+from enum import Enum
+from pyiceberg.manifest import DataFile
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
-def check_data_files_sequence_number(data_files_list, equality_delete_files_list):
+def check_data_files_sequence_number(
+    data_files_list: DataFileList,
+    equality_delete_files_list: DataFileList,
+) -> Tuple[DataFileListGroup, DataFileListGroup]:
     # Sort by file sequence number
     data_files_list.sort(key=lambda file_tuple: file_tuple[0])
     equality_delete_files_list.sort(key=lambda file_tuple: file_tuple[0])
-
-    equality_delete_files = []
-    result_data_file = []
-
-    # Pointer for list data_file
-    data_file_pointer = 0
+    data_file_delete_applicable = []
+    result_eq_files_list = []
 
     # Loop through each value in equality_delete_file
-    for equality_file_tuple in equality_delete_files_list:
-        # Find all values in data_file that are smaller than val_equality
-        valid_values = []
+    for data_file_tuple in data_files_list:
+
+        # Find all values in equality delete file that having a larger sequence number than current data file
+        valid_values_eq = []
 
+        # Pointer for equality delete file
+        eq_file_pointer = 0
         # Move data_file_pointer to the first value in data_file that is smaller than val_equality
         while (
-            data_file_pointer < len(data_files_list)
-            and data_files_list[data_file_pointer][0] < equality_file_tuple[0]
+            eq_file_pointer < len(equality_delete_files_list)
+            and equality_delete_files_list[eq_file_pointer][0] > data_file_tuple[0]
         ):
-            valid_values.append(data_files_list[data_file_pointer])
-            data_file_pointer += 1
-            equality_delete_files.append(equality_file_tuple)
-
-        # Append the value from equality_delete_file and the corresponding valid values from data_file
-        if valid_values:
-            result_data_file.append(valid_values)
-
-    result_equality_delete_file = append_larger_sequence_number_data_files(
-        equality_delete_files
+            valid_values_eq.append(equality_delete_files_list[eq_file_pointer])
+            eq_file_pointer += 1
+
+        if valid_values_eq:
+            # Append the value for both applicable eq files list and applicable data files list
+            data_file_delete_applicable.append(data_file_tuple)
+            result_eq_files_list.append(valid_values_eq)
+
+    res_data_file_list = []
+    res_equality_delete_file_list = []
+    merged_file_dict = defaultdict(list)
+    for data_file_sublist, eq_delete_sublist in zip(
+        data_file_delete_applicable, result_eq_files_list
+    ):
+        merged_file_dict[tuple(eq_delete_sublist)].append(data_file_sublist)
+    for eq_file_list, data_file_list in merged_file_dict.items():
+        res_data_file_list.append(list(set(data_file_list)))
+        res_equality_delete_file_list.append(list(set(eq_file_list)))
+
+    assert len(res_data_file_list) == len(res_equality_delete_file_list), (
+        f"length of applicable data files list: {len(res_data_file_list)} "
+        f"should equal to length of equality delete files list:{len(res_equality_delete_file_list)}"
     )
 
-    return result_equality_delete_file, result_data_file
-
-
-def append_larger_sequence_number_data_files(data_files_list):
-    result = []
-    # Iterate over the input list
-    for i in range(len(data_files_list)):
-        sublist = data_files_list[i:]
-        sublist_file_list = []
-        for file in sublist:
-            sublist_file_list.append(file)
-        result.append(sublist_file_list)
-    return result
+    return res_equality_delete_file_list, res_data_file_list
 
 
 def construct_iceberg_table_prefix(
-    iceberg_warehouse_bucket_name, table_name, iceberg_namespace
-):
+    iceberg_warehouse_bucket_name: str, table_name: str, iceberg_namespace: str
+) -> str:
     return f"{iceberg_warehouse_bucket_name}/{iceberg_namespace}/{table_name}/data"
 
 
-def partition_value_record_to_partition_value_string(partition):
+def partition_value_record_to_partition_value_string(partition: Any) -> str:
     # Get string representation of partition value out of Record[partition_value]
     partition_value_str = partition.__repr__().split("[", 1)[1].split("]")[0]
     return partition_value_str
 
 
 def group_all_files_to_each_bucket(
-    data_file_dict, equality_delete_dict, pos_delete_dict
-):
+    data_file_dict: Dict[Any, DataFileList],
+    equality_delete_dict: Dict[Any, DataFileList],
+    pos_delete_dict: Dict[Any, DataFileList],
+) -> List[ConvertInputFiles]:
     convert_input_files_for_all_buckets = []
     files_for_each_bucket_for_deletes = defaultdict(tuple)
     if equality_delete_dict:
         for partition_value, equality_delete_file_list in equality_delete_dict.items():
-            (
-                result_equality_delete_file,
-                result_data_file,
-            ) = check_data_files_sequence_number(
-                data_files_list=data_file_dict[partition_value],
-                equality_delete_files_list=equality_delete_dict[partition_value],
-            )
-            files_for_each_bucket_for_deletes[partition_value] = (
-                result_data_file,
-                result_equality_delete_file,
-                [],
-            )
-            if partition_value not in data_file_dict:
-                convert_input_file = ConvertInputFiles.of(
-                    partition_value=partition_value,
-                    applicable_data_files=result_data_file,
-                    applicable_equalitu_delete_files=result_equality_delete_file,
+            if partition_value in data_file_dict:
+                (
+                    result_equality_delete_file,
+                    result_data_file,
+                ) = check_data_files_sequence_number(
+                    data_files_list=data_file_dict[partition_value],
+                    equality_delete_files_list=equality_delete_dict[partition_value],
+                )
+                files_for_each_bucket_for_deletes[partition_value] = (
+                    result_data_file,
+                    result_equality_delete_file,
+                    [],
                 )
-                convert_input_files_for_all_buckets.append(convert_input_file)
 
     for partition_value, all_data_files_for_each_bucket in data_file_dict.items():
         convert_input_file = ConvertInputFiles.of(
@@ -102,8 +107,69 @@ def group_all_files_to_each_bucket(
             convert_input_file.applicable_data_files = (
                 files_for_each_bucket_for_deletes[partition_value][0]
             )
-            convert_input_file.applicable_delete_files = (
+            convert_input_file.applicable_equality_delete_files = (
                 files_for_each_bucket_for_deletes[partition_value][1]
             )
         convert_input_files_for_all_buckets.append(convert_input_file)
     return convert_input_files_for_all_buckets
+
+
+def sort_data_files_maintaining_order(data_files: DataFileList) -> DataFileList:
+    """
+    Sort data files deterministically based on two criterias:
+    1. Sequence number: Newly added files will have a higher sequence number
+    2. File path: If file sequence is the same, files are guaranteed to be returned in a deterministic order since file path is unique.
+    """
+    if data_files:
+        data_files = sorted(data_files, key=lambda f: (f[0], f[1].file_path))
+    return data_files
+
+
+class SnapshotType(Enum):
+    """Enumeration of possible snapshot types."""
+
+    NONE = "none"
+    APPEND = "append"
+    REPLACE = "replace"
+    DELETE = "delete"
+
+
+def _get_snapshot_action_description(
+    snapshot_type: SnapshotType,
+    files_to_delete: List[List[DataFile]],
+    files_to_add: List[DataFile],
+) -> str:
+    """Get a human-readable description of the snapshot action."""
+    descriptions = {
+        SnapshotType.NONE: "No changes needed",
+        SnapshotType.APPEND: f"Adding {len(files_to_add)} new files",
+        SnapshotType.REPLACE: f"Replacing {sum(len(files) for files in files_to_delete)} files with {len(files_to_add)} new files",
+        SnapshotType.DELETE: f"Deleting {sum(len(files) for files in files_to_delete)} files",
+    }
+    return descriptions[snapshot_type]
+
+
+def _determine_snapshot_type(
+    to_be_deleted_files: List[List[DataFile]], to_be_added_files: List[DataFile]
+) -> SnapshotType:
+    """
+    Determine the snapshot type based on file changes.
+
+    Args:
+        to_be_deleted_files: List of files to be deleted
+        to_be_added_files: List of files to be added
+
+    Returns:
+        SnapshotType indicating what kind of snapshot to commit
+    """
+    has_files_to_delete = bool(to_be_deleted_files)
+    has_files_to_add = bool(to_be_added_files)
+
+    if not has_files_to_delete and not has_files_to_add:
+        return SnapshotType.NONE
+    elif not has_files_to_delete and has_files_to_add:
+        return SnapshotType.APPEND
+    elif has_files_to_delete and has_files_to_add:
+        return SnapshotType.REPLACE
+    else:  # has_files_to_delete and not has_files_to_add
+        return SnapshotType.DELETE

deltacat/compute/converter/utils/iceberg_columns.py

@@ -1,5 +1,5 @@
 import pyarrow as pa
-from typing import Union
+from typing import Union, Iterator, Any
 import numpy as np
 
 # Refer to: https://iceberg.apache.org/spec/#reserved-field-ids for reserved field ids
@@ -9,7 +9,7 @@ ICEBERG_RESERVED_FIELD_ID_FOR_FILE_PATH_COLUMN = 2147483546
 ICEBERG_RESERVED_FIELD_ID_FOR_POS_COLUMN = 2147483545
 
 
-def _get_iceberg_col_name(suffix):
+def _get_iceberg_col_name(suffix: str) -> str:
     return suffix
 
 
@@ -26,15 +26,16 @@ _ORDERED_RECORD_IDX_COLUMN_FIELD = pa.field(
 )
 
 
-def get_record_index_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
+def get_record_index_column_array(obj: Any) -> Union[pa.Array, pa.ChunkedArray]:
     return pa.array(
         obj,
         _ORDERED_RECORD_IDX_COLUMN_TYPE,
     )
 
 
-def append_record_idx_col(table: pa.Table, ordered_record_indices) -> pa.Table:
-
+def append_record_idx_col(
+    table: pa.Table, ordered_record_indices: Iterator[int]
+) -> pa.Table:
     table = table.append_column(
         _ORDERED_RECORD_IDX_COLUMN_FIELD,
         get_record_index_column_array(ordered_record_indices),
@@ -55,7 +56,7 @@ _FILE_PATH_COLUMN_FIELD = pa.field(
 )
 
 
-def append_file_path_column(table: pa.Table, file_path: str):
+def append_file_path_column(table: pa.Table, file_path: str) -> pa.Table:
     table = table.append_column(
         _FILE_PATH_COLUMN_FIELD,
         pa.array(np.repeat(file_path, len(table)), _FILE_PATH_COLUMN_TYPE),
@@ -72,9 +73,8 @@ _GLOBAL_RECORD_IDX_COLUMN_FIELD = pa.field(
 
 
 def append_global_record_idx_column(
-    table: pa.Table, ordered_record_indices
+    table: pa.Table, ordered_record_indices: Iterator[int]
 ) -> pa.Table:
-
     table = table.append_column(
         _GLOBAL_RECORD_IDX_COLUMN_NAME,
         pa.array(ordered_record_indices, _GLOBAL_RECORD_IDX_COLUMN_TYPE),

deltacat/compute/converter/utils/io.py

@@ -1,25 +1,36 @@
 import logging
+
+from fsspec import AbstractFileSystem
 from deltacat import logs
 import deltacat.compute.converter.utils.iceberg_columns as sc
 import daft
 from deltacat.utils.daft import _get_s3_io_config
-from daft import TimeUnit
+from daft import TimeUnit, DataFrame
 import pyarrow as pa
+from typing import Callable, Optional, List, Dict, Any
 from deltacat.utils.pyarrow import sliced_string_cast
 from deltacat.compute.converter.constants import IDENTIFIER_FIELD_DELIMITER
-
+from deltacat.compute.converter.utils.s3u import upload_table_with_retry
+from pyiceberg.manifest import DataFile
 import pyarrow.compute as pc
+from deltacat.types.media import ContentType
+from deltacat.types.tables import (
+    get_table_writer,
+    get_table_slicer,
+    write_sliced_table as types_write_sliced_table,
+)
+from deltacat.storage import LocalTable, DistributedDataset
+from typing import Union
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
 def download_data_table_and_append_iceberg_columns(
-    file,
-    columns_to_download,
-    additional_columns_to_append,
-    sequence_number,
-    s3_client_kwargs,
-):
+    file: DataFile,
+    columns_to_download: List[str],
+    additional_columns_to_append: Optional[List[str]] = [],
+    s3_client_kwargs: Optional[Dict[str, Any]] = None,
+) -> pa.Table:
     table = download_parquet_with_daft_hash_applied(
         identifier_columns=columns_to_download,
         file=file,
@@ -36,8 +47,11 @@ def download_data_table_and_append_iceberg_columns(
 
 
 def download_parquet_with_daft_hash_applied(
-    identifier_columns, file, s3_client_kwargs, **kwargs
-):
+    identifier_columns: List[str],
+    file: DataFile,
+    s3_client_kwargs: Optional[Dict[str, Any]],
+    **kwargs: Any,
+) -> pa.Table:
 
     # TODO: Add correct read kwargs as in:
     #  https://github.com/ray-project/deltacat/blob/383855a4044e4dfe03cf36d7738359d512a517b4/deltacat/utils/daft.py#L97
@@ -65,7 +79,9 @@ def download_parquet_with_daft_hash_applied(
     return table
 
 
-def daft_read_parquet(path, io_config, coerce_int96_timestamp_unit):
+def daft_read_parquet(
+    path: str, io_config: Dict[str, Any], coerce_int96_timestamp_unit: TimeUnit
+) -> DataFrame:
     df = daft.read_parquet(
         path=path,
         io_config=io_config,
@@ -74,7 +90,9 @@ def daft_read_parquet(path, io_config, coerce_int96_timestamp_unit):
     return df
 
 
-def concatenate_hashed_identifier_columns(df, identifier_columns):
+def concatenate_hashed_identifier_columns(
+    df: DataFrame, identifier_columns: List[str]
+) -> pa.Array:
     pk_hash_columns = []
     previous_hash_column_length = None
     for i in range(len(identifier_columns)):
@@ -112,3 +130,74 @@ def concatenate_hashed_identifier_columns(df, identifier_columns):
     )
 
     return pk_hash_columns_concatenated
+
+
+def write_sliced_table(
+    table: Union[LocalTable, DistributedDataset],
+    base_path: str,
+    table_writer_kwargs: Optional[Dict[str, Any]],
+    content_type: ContentType = ContentType.PARQUET,
+    max_records_per_file: Optional[int] = 4000000,
+    filesystem: Optional[Union[AbstractFileSystem, pa.fs.FileSystem]] = None,
+    **kwargs,
+) -> List[str]:
+    """
+    Writes the given table to 1 or more files and return the paths
+    of the files written.
+    """
+    if isinstance(filesystem, pa.fs.FileSystem):
+        table_writer_fn = get_table_writer(table)
+        table_slicer_fn = get_table_slicer(table)
+
+        # Create a wrapper for the table writer that ensures directory creation
+        def table_writer_with_dir_creation(
+            dataframe: Any,
+            base_path: str,
+            filesystem: Optional[Union[AbstractFileSystem, pa.fs.FileSystem]],
+            block_path_provider: Callable,
+            content_type: str = ContentType.PARQUET.value,
+            **kwargs,
+        ):
+            try:
+                # Ensure base path directory exists
+                if isinstance(base_path, str):
+                    # Normalize the base path and ensure it's treated as a directory path
+                    base_dir = base_path.rstrip("/")
+                    filesystem.create_dir(base_dir, recursive=True)
+            except Exception:
+                # Directory might already exist or there might be permission issues
+                # Let the original write attempt proceed
+                pass
+            return table_writer_fn(
+                dataframe,
+                base_path,
+                filesystem,
+                block_path_provider,
+                content_type,
+                **kwargs,
+            )
+
+        # TODO(pdames): Disable redundant file info fetch currently
+        #   used to construct unused manifest entry metadata.
+        manifest_entry_list = types_write_sliced_table(
+            table=table,
+            base_path=base_path,
+            filesystem=filesystem,
+            max_records_per_entry=max_records_per_file,
+            table_writer_fn=table_writer_with_dir_creation,
+            table_slicer_fn=table_slicer_fn,
+            table_writer_kwargs=table_writer_kwargs,
+            content_type=content_type,
+        )
+        paths = [entry.uri for entry in manifest_entry_list]
+        return paths
+    else:
+        return upload_table_with_retry(
+            table=table,
+            s3_url_prefix=base_path,
+            s3_table_writer_kwargs=table_writer_kwargs,
+            content_type=content_type,
+            max_records_per_file=max_records_per_file,
+            s3_file_system=filesystem,
+            **kwargs,
+        )

deltacat/compute/converter/utils/s3u.py

@@ -4,14 +4,16 @@ from tenacity import (
     stop_after_delay,
     wait_random_exponential,
 )
-from typing import Union
-from deltacat.aws.s3u import CapturedBlockWritePaths, UuidBlockWritePathProvider
+from typing import Union, Optional, Dict, Any, List, Callable
+from deltacat.types.tables import (
+    CapturedBlockWritePaths,
+    UuidBlockWritePathProvider,
+)
 from deltacat.types.tables import (
     get_table_writer,
     get_table_length,
     TABLE_CLASS_TO_SLICER_FUNC,
 )
-from typing import Optional, Dict, Any
 from deltacat.exceptions import RetryableError
 from deltacat.storage import (
     DistributedDataset,
@@ -21,19 +23,22 @@ from deltacat.types.media import (
     ContentEncoding,
     ContentType,
 )
-from deltacat.aws.s3u import UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY
+from deltacat.constants import UPLOAD_SLICED_TABLE_RETRY_STOP_AFTER_DELAY
 import s3fs
+import boto3
+from boto3.session import Session
+from botocore.credentials import Credentials
 
 
-def get_credential():
-    import boto3
-
-    boto3_session = boto3.Session()
-    credentials = boto3_session.get_credentials()
+def get_credential() -> Credentials:
+    boto3_session: Session = boto3.Session()
+    credentials: Credentials = boto3_session.get_credentials()
     return credentials
 
 
 def get_s3_file_system(content_type):
+    import s3fs  # noqa: F401
+
     token_holder = get_credential()
     content_encoding = ContentEncoding.IDENTITY
 
@@ -57,12 +62,12 @@ def upload_table_with_retry(
     s3_table_writer_kwargs: Optional[Dict[str, Any]],
     content_type: ContentType = ContentType.PARQUET,
     max_records_per_file: Optional[int] = 4000000,
-    s3_file_system=None,
-    **s3_client_kwargs,
-):
+    filesystem: Optional[s3fs.S3FileSystem] = None,
+    **s3_client_kwargs: Any,
+) -> List[str]:
     """
-    Writes the given table to 1 or more S3 files and return Redshift
-    manifest entries describing the uploaded files.
+    Writes the given table to 1 or more S3 files and return the paths
+    of the S3 files written.
     """
     retrying = Retrying(
         wait=wait_random_exponential(multiplier=1, max=60),
@@ -73,8 +78,8 @@ def upload_table_with_retry(
     if s3_table_writer_kwargs is None:
         s3_table_writer_kwargs = {}
 
-    if not s3_file_system:
-        s3_file_system = get_s3_file_system(content_type=content_type)
+    if not filesystem:
+        filesystem = get_s3_file_system(content_type=content_type)
     capture_object = CapturedBlockWritePaths()
     block_write_path_provider = UuidBlockWritePathProvider(
         capture_object=capture_object, base_path=s3_url_prefix
@@ -86,7 +91,7 @@ def upload_table_with_retry(
             fn=upload_table,
             table_slices=table,
             s3_base_url=f"{s3_url_prefix}",
-            s3_file_system=s3_file_system,
+            s3_file_system=filesystem,
             s3_table_writer_func=s3_table_writer_func,
             s3_table_writer_kwargs=s3_table_writer_kwargs,
             block_write_path_provider=block_write_path_provider,
@@ -101,7 +106,7 @@ def upload_table_with_retry(
                 fn=upload_table,
                 table_slices=table_slice,
                 s3_base_url=f"{s3_url_prefix}",
-                s3_file_system=s3_file_system,
+                s3_file_system=filesystem,
                 s3_table_writer_func=s3_table_writer_func,
                 s3_table_writer_kwargs=s3_table_writer_kwargs,
                 block_write_path_provider=block_write_path_provider,
@@ -117,20 +122,21 @@ def upload_table_with_retry(
     return s3_write_paths
 
 
-def construct_s3_url(path):
+def construct_s3_url(path: Optional[str]) -> Optional[str]:
     if path:
         return f"s3://{path}"
+    return None
 
 
 def upload_table(
-    table_slices,
-    s3_base_url,
-    s3_file_system,
-    s3_table_writer_func,
-    block_write_path_provider,
-    content_type,
-    s3_table_writer_kwargs,
-):
+    table_slices: Union[LocalTable, DistributedDataset],
+    s3_base_url: str,
+    s3_file_system: s3fs.S3FileSystem,
+    s3_table_writer_func: Callable,
+    block_write_path_provider: UuidBlockWritePathProvider,
+    content_type: ContentType,
+    s3_table_writer_kwargs: Dict[str, Any],
+) -> None:
     s3_table_writer_func(
         table_slices,
         s3_base_url,