deltacat 1.1.36__py3-none-any.whl → 2.0__py3-none-any.whl

deltacat/tests/compute/converter/test_convert_session.py

@@ -0,0 +1,478 @@
+import pytest
+import ray
+from typing import List
+from pyiceberg.catalog.rest import RestCatalog
+from pyiceberg.expressions import EqualTo
+from pyiceberg.schema import Schema
+from pyiceberg.types import (
+    NestedField,
+    StringType,
+    LongType,
+)
+from pyiceberg.partitioning import PartitionSpec, PartitionField
+from pyiceberg.transforms import IdentityTransform
+import pyarrow as pa
+
+from deltacat.compute.converter.steps.convert import convert
+from deltacat.compute.converter.model.convert_input import ConvertInput
+from deltacat.compute.converter.pyiceberg.overrides import (
+    fetch_all_bucket_files,
+    parquet_files_dict_to_iceberg_data_files,
+)
+from collections import defaultdict
+from deltacat.compute.converter.utils.converter_session_utils import (
+    group_all_files_to_each_bucket,
+)
+from deltacat.tests.compute.converter.utils import (
+    get_s3_file_system,
+    drop_table_if_exists,
+)
+from deltacat.compute.converter.pyiceberg.update_snapshot_overrides import (
+    commit_append_snapshot,
+)
+
+
+def run_spark_commands(spark, sqls: List[str]) -> None:
+    for sql in sqls:
+        spark.sql(sql)
+
+
+@pytest.mark.integration
+def test_pyiceberg_spark_setup_sanity(spark, session_catalog: RestCatalog) -> None:
+    """
+    This Test was copied over from Pyiceberg integ test: https://github.com/apache/iceberg-python/blob/main/tests/integration/test_deletes.py#L62
+    First sanity check to ensure all integration with Pyiceberg and Spark are working as expected.
+    """
+    identifier = "default.table_partitioned_delete"
+
+    run_spark_commands(
+        spark,
+        [
+            f"DROP TABLE IF EXISTS {identifier}",
+            f"""
+            CREATE TABLE {identifier} (
+                number_partitioned  int,
+                number              int
+            )
+            USING iceberg
+            PARTITIONED BY (number_partitioned)
+            TBLPROPERTIES('format-version' = 2)
+        """,
+            f"""
+            INSERT INTO {identifier} VALUES (10, 20), (10, 30)
+        """,
+            f"""
+            INSERT INTO {identifier} VALUES (11, 20), (11, 30)
+        """,
+        ],
+    )
+
+    tbl = session_catalog.load_table(identifier)
+    tbl.delete(EqualTo("number_partitioned", 10))
+
+    # No overwrite operation
+    assert [snapshot.summary.operation.value for snapshot in tbl.snapshots()] == [
+        "append",
+        "append",
+        "delete",
+    ]
+    assert tbl.scan().to_arrow().to_pydict() == {
+        "number_partitioned": [11, 11],
+        "number": [20, 30],
+    }
+
+
+@pytest.mark.integration
+def test_spark_position_delete_production_sanity(
+    spark, session_catalog: RestCatalog
+) -> None:
+    """
+    Sanity test to ensure Spark position delete production is successful with `merge-on-read` spec V2.
+    Table has two partition levels. 1. BucketTransform on primary key
+    """
+    identifier = "default.table_spark_position_delete_production_sanity"
+
+    run_spark_commands(
+        spark,
+        [
+            f"DROP TABLE IF EXISTS {identifier}",
+            f"""
+            CREATE TABLE {identifier} (
+                number_partitioned INT,
+                primary_key STRING
+            )
+            USING iceberg
+            PARTITIONED BY (bucket(3, primary_key), number_partitioned)
+            TBLPROPERTIES(
+                'format-version' = 2,
+                'write.delete.mode'='merge-on-read',
+                'write.update.mode'='merge-on-read',
+                'write.merge.mode'='merge-on-read'
+            )
+            """,
+            f"""
+            INSERT INTO {identifier} VALUES (0, 'pk1'), (0, 'pk2'), (0, 'pk3')
+            """,
+            f"""
+            INSERT INTO {identifier} VALUES (1, 'pk1'), (1, 'pk2'), (1, 'pk3')
+            """,
+        ],
+    )
+
+    run_spark_commands(
+        spark,
+        [
+            f"""
+                DELETE FROM {identifier} WHERE primary_key in ("pk1")
+            """,
+        ],
+    )
+
+    tbl = session_catalog.load_table(identifier)
+    tbl.refresh()
+
+    assert [snapshot.summary.operation.value for snapshot in tbl.snapshots()] == [
+        "append",
+        "append",
+        "delete",
+    ]
+
+    assert tbl.scan().to_arrow().to_pydict() == {
+        "number_partitioned": [1, 1, 0, 0],
+        "primary_key": ["pk2", "pk3", "pk2", "pk3"],
+    }
+
+
+@pytest.mark.integration
+def test_converter_drop_duplicates_success(
+    spark, session_catalog: RestCatalog, setup_ray_cluster, mocker
+) -> None:
+    """
+    Test for convert compute remote function happy case. Download file results are mocked.
+    """
+
+    # 1. Create Iceberg table
+    namespace = "default"
+    table_name = "table_converter_ray_pos_delete_drop_duplicates_compute"
+    identifier = f"{namespace}.{table_name}"
+
+    schema = Schema(
+        NestedField(
+            field_id=1, name="number_partitioned", field_type=LongType(), required=False
+        ),
+        NestedField(
+            field_id=2, name="primary_key", field_type=StringType(), required=False
+        ),
+        # Explicitly define "file_path" and "pos" for assertion of deterministic record after dedupe
+        NestedField(
+            field_id=2147483546,
+            name="file_path",
+            field_type=StringType(),
+            required=False,
+        ),
+        NestedField(
+            field_id=2147483545, name="pos", field_type=LongType(), required=False
+        ),
+        schema_id=0,
+    )
+
+    partition_field_identity = PartitionField(
+        source_id=1,
+        field_id=101,
+        transform=IdentityTransform(),
+        name="number_partitioned",
+    )
+    partition_spec = PartitionSpec(partition_field_identity)
+
+    properties = dict()
+    properties["write.format.default"] = "parquet"
+    properties["write.delete.mode"] = "merge-on-read"
+    properties["write.update.mode"] = "merge-on-read"
+    properties["write.merge.mode"] = "merge-on-read"
+    properties["format-version"] = "2"
+
+    drop_table_if_exists(identifier, session_catalog)
+    session_catalog.create_table(
+        identifier,
+        schema=schema,
+        partition_spec=partition_spec,
+        properties=properties,
+    )
+
+    # 2. Use Spark to generate initial data files
+    tbl = session_catalog.load_table(identifier)
+    tbl.refresh()
+    run_spark_commands(
+        spark,
+        [
+            f"""
+            INSERT INTO {identifier} VALUES (0, "pk1", "path1", 1), (0, "pk2", "path2", 2), (0, "pk3", "path3", 3)
+            """
+        ],
+    )
+    run_spark_commands(
+        spark,
+        [
+            f"""
+                INSERT INTO {identifier} VALUES (0, "pk1", "path1", 4), (0, "pk2", "path2", 5), (0, "pk3", "path3", 6)
+                """
+        ],
+    )
+    run_spark_commands(
+        spark,
+        [
+            f"""
+                INSERT INTO {identifier} VALUES (0, "pk4", "path4", 7), (0, "pk2", "path2", 8), (0, "pk3", "path3", 9)
+                """
+        ],
+    )
+
+    tbl = session_catalog.load_table(identifier)
+    # 3. Use convert.remote() function to compute position deletes
+    data_file_dict, equality_delete_dict, pos_delete_dict = fetch_all_bucket_files(tbl)
+
+    convert_input_files_for_all_buckets = group_all_files_to_each_bucket(
+        data_file_dict=data_file_dict,
+        equality_delete_dict=equality_delete_dict,
+        pos_delete_dict=pos_delete_dict,
+    )
+
+    s3_file_system = get_s3_file_system()
+    for i, one_bucket_files in enumerate(convert_input_files_for_all_buckets):
+        convert_input = ConvertInput.of(
+            convert_input_files=one_bucket_files,
+            convert_task_index=i,
+            iceberg_table_warehouse_prefix="warehouse/default",
+            identifier_fields=["primary_key"],
+            compact_small_files=False,
+            enforce_primary_key_uniqueness=True,
+            position_delete_for_multiple_data_files=True,
+            max_parallel_data_file_download=10,
+            s3_file_system=s3_file_system,
+        )
+
+    number_partitioned_array_1 = pa.array([0, 0, 0], type=pa.int32())
+    primary_key_array_1 = pa.array(["pk1", "pk2", "pk3"])
+    names = ["number_partitioned", "primary_key"]
+    data_table_1 = pa.Table.from_arrays(
+        [number_partitioned_array_1, primary_key_array_1], names=names
+    )
+
+    number_partitioned_array_2 = pa.array([0, 0, 0], type=pa.int32())
+    primary_key_array_2 = pa.array(["pk1", "pk2", "pk3"])
+    names = ["number_partitioned", "primary_key"]
+    data_table_2 = pa.Table.from_arrays(
+        [number_partitioned_array_2, primary_key_array_2], names=names
+    )
+
+    number_partitioned_array_3 = pa.array([0, 0, 0], type=pa.int32())
+    primary_key_array_3 = pa.array(["pk4", "pk2", "pk3"])
+    names = ["number_partitioned", "primary_key"]
+    data_table_3 = pa.Table.from_arrays(
+        [number_partitioned_array_3, primary_key_array_3], names=names
+    )
+
+    download_data_mock = mocker.patch(
+        "deltacat.compute.converter.utils.io.download_parquet_with_daft_hash_applied"
+    )
+    download_data_mock.side_effect = (data_table_1, data_table_2, data_table_3)
+
+    convert_ref = convert.remote(convert_input)
+
+    to_be_deleted_files_list = []
+    to_be_added_files_dict_list = []
+    convert_result = ray.get(convert_ref)
+
+    partition_value = convert_input.convert_input_files.partition_value
+
+    if convert_result[0]:
+        to_be_deleted_files_list.extend(convert_result[0].values())
+
+    file_location = convert_result[1][partition_value][0]
+    to_be_added_files = f"s3://{file_location}"
+
+    to_be_added_files_dict = defaultdict()
+    to_be_added_files_dict[partition_value] = [to_be_added_files]
+    to_be_added_files_dict_list.append(to_be_added_files_dict)
+
+    # 4. Commit position delete, delete equality deletes from table
+    new_position_delete_files = parquet_files_dict_to_iceberg_data_files(
+        io=tbl.io,
+        table_metadata=tbl.metadata,
+        files_dict_list=to_be_added_files_dict_list,
+    )
+    commit_append_snapshot(
+        iceberg_table=tbl,
+        new_position_delete_files=new_position_delete_files,
+    )
+    tbl.refresh()
+
+    # 5. Only primary key 2 and 3 should exist in table, as primary key 1 is deleted.
+    pyiceberg_scan_table_rows = tbl.scan().to_arrow().to_pydict()
+
+    # Only one unique record for each pk exists
+    all_pk = sorted(pyiceberg_scan_table_rows["primary_key"])
+    assert all_pk == ["pk1", "pk2", "pk3", "pk4"]
+
+    # Expected unique record to keep for each pk
+    expected_pk_to_pos_mapping = {"pk1": 4, "pk2": 8, "pk3": 9, "pk4": 7}
+    for pk, pos in zip(
+        pyiceberg_scan_table_rows["primary_key"], pyiceberg_scan_table_rows["pos"]
+    ):
+        assert pos == expected_pk_to_pos_mapping[pk]
+
+
+@pytest.mark.integration
+def test_converter_pos_delete_read_by_spark_success(
+    spark, session_catalog: RestCatalog, setup_ray_cluster, mocker
+) -> None:
+    """
+    Test for convert compute remote function happy case. Download file results are mocked.
+    """
+
+    # 1. Create Iceberg table
+    namespace = "default"
+    table_name = "table_converter_ray_pos_delete_read_by_spark_success"
+    identifier = f"{namespace}.{table_name}"
+
+    schema = Schema(
+        NestedField(
+            field_id=1, name="number_partitioned", field_type=LongType(), required=False
+        ),
+        NestedField(
+            field_id=2, name="primary_key", field_type=StringType(), required=False
+        ),
+        schema_id=0,
+    )
+
+    partition_field_identity = PartitionField(
+        source_id=1,
+        field_id=101,
+        transform=IdentityTransform(),
+        name="number_partitioned",
+    )
+    partition_spec = PartitionSpec(partition_field_identity)
+
+    properties = dict()
+    properties["write.format.default"] = "parquet"
+    properties["write.delete.mode"] = "merge-on-read"
+    properties["write.update.mode"] = "merge-on-read"
+    properties["write.merge.mode"] = "merge-on-read"
+    properties["format-version"] = "2"
+
+    drop_table_if_exists(identifier, session_catalog)
+    session_catalog.create_table(
+        identifier,
+        schema=schema,
+        partition_spec=partition_spec,
+        properties=properties,
+    )
+
+    # 2. Use Spark to generate initial data files
+    tbl = session_catalog.load_table(identifier)
+
+    run_spark_commands(
+        spark,
+        [
+            f"""
+               INSERT INTO {identifier} VALUES (0, "pk1"), (0, "pk2"), (0, "pk3")
+               """
+        ],
+    )
+    run_spark_commands(
+        spark,
+        [
+            f"""
+                   INSERT INTO {identifier} VALUES (0, "pk1"), (0, "pk2"), (0, "pk3")
+                   """
+        ],
+    )
+    run_spark_commands(
+        spark,
+        [
+            f"""
+                   INSERT INTO {identifier} VALUES (0, "pk4"), (0, "pk2"), (0, "pk3")
+                   """
+        ],
+    )
+    tbl.refresh()
+
+    # 3. Use convert.remote() function to compute position deletes
+    data_file_dict, equality_delete_dict, pos_delete_dict = fetch_all_bucket_files(tbl)
+
+    convert_input_files_for_all_buckets = group_all_files_to_each_bucket(
+        data_file_dict=data_file_dict,
+        equality_delete_dict=equality_delete_dict,
+        pos_delete_dict=pos_delete_dict,
+    )
+
+    s3_file_system = get_s3_file_system()
+    for i, one_bucket_files in enumerate(convert_input_files_for_all_buckets):
+        convert_input = ConvertInput.of(
+            convert_input_files=one_bucket_files,
+            convert_task_index=i,
+            iceberg_table_warehouse_prefix="warehouse/default",
+            identifier_fields=["primary_key"],
+            compact_small_files=False,
+            enforce_primary_key_uniqueness=True,
+            position_delete_for_multiple_data_files=True,
+            max_parallel_data_file_download=10,
+            s3_file_system=s3_file_system,
+        )
+
+    primary_key_array_1 = pa.array(["pk1", "pk2", "pk3"])
+    names = ["primary_key"]
+    data_table_1 = pa.Table.from_arrays([primary_key_array_1], names=names)
+
+    primary_key_array_2 = pa.array(["pk1", "pk2", "pk3"])
+    names = ["primary_key"]
+    data_table_2 = pa.Table.from_arrays([primary_key_array_2], names=names)
+
+    primary_key_array_3 = pa.array(["pk4", "pk2", "pk3"])
+    names = ["primary_key"]
+    data_table_3 = pa.Table.from_arrays([primary_key_array_3], names=names)
+
+    download_data_mock = mocker.patch(
+        "deltacat.compute.converter.utils.io.download_parquet_with_daft_hash_applied"
+    )
+    download_data_mock.side_effect = (data_table_1, data_table_2, data_table_3)
+
+    convert_ref = convert.remote(convert_input)
+
+    to_be_deleted_files_list = []
+    to_be_added_files_dict_list = []
+    convert_result = ray.get(convert_ref)
+
+    partition_value = convert_input.convert_input_files.partition_value
+
+    if convert_result[0]:
+        to_be_deleted_files_list.extend(convert_result[0].values())
+
+    file_location = convert_result[1][partition_value][0]
+    to_be_added_files = f"s3://{file_location}"
+
+    to_be_added_files_dict = defaultdict()
+    to_be_added_files_dict[partition_value] = [to_be_added_files]
+    to_be_added_files_dict_list.append(to_be_added_files_dict)
+
+    # 4. Commit position delete, delete equality deletes from table
+    new_position_delete_files = parquet_files_dict_to_iceberg_data_files(
+        io=tbl.io,
+        table_metadata=tbl.metadata,
+        files_dict_list=to_be_added_files_dict_list,
+    )
+
+    commit_append_snapshot(
+        iceberg_table=tbl,
+        new_position_delete_files=new_position_delete_files,
+    )
+    tbl.refresh()
+
+    # 5. Result assertion: Spark read table contains unique primary key
+    spark_read_pos_delete = spark.sql(f"select * from {identifier}").collect()
+    all_pk = [
+        spark_read_pos_delete[row_idx][1]
+        for row_idx in range(len(spark_read_pos_delete))
+    ]
+    all_pk_sorted = sorted(all_pk)
+    assert all_pk_sorted == ["pk1", "pk2", "pk3", "pk4"]

deltacat/tests/compute/converter/utils.py

@@ -0,0 +1,123 @@
+import uuid
+import logging
+from pyiceberg.exceptions import NoSuchTableError
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def get_s3_file_system():
+    import pyarrow
+
+    return pyarrow.fs.S3FileSystem(
+        access_key="admin",
+        secret_key="password",
+        endpoint_override="http://localhost:9000",
+    )
+    #        'region="us-east-1", proxy_options={'scheme': 'http', 'host': 'localhost',
+    # 'port': 9000, 'username': 'admin',
+    # 'password': 'password'})
+
+
+def write_equality_data_table(
+    file_link_prefix, table, partition_value, equality_delete_table
+):
+    import pyarrow.parquet as pq
+
+    uuid_path = uuid.uuid4()
+    deletes_file_path = f"{file_link_prefix}/{uuid_path}_deletes.parquet"
+    file_system = get_s3_file_system()
+    pq.write_table(equality_delete_table, deletes_file_path, filesystem=file_system)
+    return f"s3://{deletes_file_path}"
+
+
+def add_equality_data_files(file_paths, table, partition_value):
+    with table.transaction() as tx:
+        if table.metadata.name_mapping() is None:
+            tx.set_properties(
+                **{
+                    "schema.name-mapping.default": table.metadata.schema().name_mapping.model_dump_json()
+                }
+            )
+        with tx.update_snapshot().fast_append() as update_snapshot:
+            data_files = parquet_files_to_equality_data_files(
+                table_metadata=table.metadata,
+                file_paths=file_paths,
+                io=table.io,
+                partition_value=partition_value,
+            )
+            for data_file in data_files:
+                update_snapshot.append_data_file(data_file)
+
+
+def parquet_files_to_equality_data_files(
+    io, table_metadata, file_paths, partition_value
+):
+    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,
+    )
+    import pyarrow.parquet as pq
+
+    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)
+
+        schema = table_metadata.schema()
+        _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=DataFileContent.EQUALITY_DELETES,
+            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(),
+        )
+
+        yield data_file
+
+
+def commit_equality_delete_to_table(
+    table, file_link_prefix, partition_value, equality_delete_table
+):
+
+    data_files = [
+        write_equality_data_table(
+            table=table,
+            file_link_prefix=file_link_prefix,
+            partition_value=partition_value,
+            equality_delete_table=equality_delete_table,
+        )
+    ]
+
+    add_equality_data_files(
+        file_paths=data_files, partition_value=partition_value, table=table
+    )
+    return data_files
+
+
+def drop_table_if_exists(table, catalog):
+    try:
+        catalog.drop_table(table)
+    except NoSuchTableError:
+        logger.warning(f"table:{table} doesn't exist, not dropping table.")

deltacat/tests/compute/resource_estimation/test_delta.py

@@ -1,6 +1,5 @@
 import deltacat.tests.local_deltacat_storage as ds
 from deltacat.types.media import ContentType
-import os
 import pytest
 from deltacat.storage import Delta
 from deltacat.compute.resource_estimation.delta import (
@@ -21,21 +20,6 @@ Function scoped fixtures
 """
 
 
-@pytest.fixture(scope="function")
-def local_deltacat_storage_kwargs():
-    DATABASE_FILE_PATH_KEY, DATABASE_FILE_PATH_VALUE = (
-        "db_file_path",
-        "deltacat/tests/local_deltacat_storage/db_test.sqlite",
-    )
-    # see deltacat/tests/local_deltacat_storage/README.md for documentation
-    kwargs_for_local_deltacat_storage = {
-        DATABASE_FILE_PATH_KEY: DATABASE_FILE_PATH_VALUE,
-    }
-    yield kwargs_for_local_deltacat_storage
-    if os.path.exists(DATABASE_FILE_PATH_VALUE):
-        os.remove(DATABASE_FILE_PATH_VALUE)
-
-
 @pytest.fixture(scope="function")
 def parquet_delta_with_manifest(local_deltacat_storage_kwargs):
     """