deltacat 1.1.9__py3-none-any.whl → 1.1.11__py3-none-any.whl

deltacat/storage/model/delta.py

@@ -12,6 +12,7 @@ from deltacat.storage.model.stream import StreamLocator
 from deltacat.storage.model.table import TableLocator
 from deltacat.storage.model.table_version import TableVersionLocator
 from deltacat.storage.model.types import DeltaType
+from deltacat.storage.model.partition_spec import DeltaPartitionSpec, PartitionValues
 
 
 class Delta(dict):
@@ -24,6 +25,7 @@ class Delta(dict):
         manifest: Optional[Manifest],
         previous_stream_position: Optional[int] = None,
         delete_parameters: Optional[DeleteParameters] = None,
+        partition_spec: Optional[DeltaPartitionSpec] = None,
     ) -> Delta:
         """
         Creates a Delta metadata model with the given Delta Locator, Delta Type,
@@ -38,6 +40,7 @@ class Delta(dict):
         delta.manifest = manifest
         delta.previous_stream_position = previous_stream_position
         delta.delete_parameters = delete_parameters
+        delta.partition_spec = partition_spec
         return delta
 
     @staticmethod
@@ -90,6 +93,12 @@ class Delta(dict):
                 f"Deltas to merge must all share the same delta type "
                 f"(found {len(distinct_delta_types)} delta types)."
             )
+        distinct_partition_spec = set([d.partition_spec for d in deltas])
+        if len(distinct_partition_spec) > 1:
+            raise ValueError(
+                f"Deltas to merge must all share the same partition spec "
+                f"(found {len(distinct_partition_spec)} partition specs)."
+            )
         merged_manifest = Manifest.merge_manifests(
             manifests,
             manifest_author,
@@ -252,7 +261,7 @@ class Delta(dict):
         return None
 
     @property
-    def partition_values(self) -> Optional[List[Any]]:
+    def partition_values(self) -> Optional[PartitionValues]:
         delta_locator = self.locator
         if delta_locator:
             return delta_locator.partition_values
@@ -276,6 +285,17 @@ class Delta(dict):
     def delete_parameters(self, delete_parameters: Optional[DeleteParameters]) -> None:
         self["delete_parameters"] = delete_parameters
 
+    @property
+    def partition_spec(self) -> Optional[DeltaPartitionSpec]:
+        val: Dict[str, Any] = self.get("partitionSpec")
+        if val is not None and not isinstance(val, DeltaPartitionSpec):
+            self.partition_spec = val = DeltaPartitionSpec(val)
+        return val
+
+    @partition_spec.setter
+    def partition_spec(self, value: Optional[DeltaPartitionSpec]) -> None:
+        self["partitionSpec"] = value
+
 
 class DeltaLocator(Locator, dict):
     @staticmethod
@@ -299,7 +319,7 @@ class DeltaLocator(Locator, dict):
         table_version: Optional[str],
         stream_id: Optional[str],
         storage_type: Optional[str],
-        partition_values: Optional[List[Any]],
+        partition_values: Optional[PartitionValues],
         partition_id: Optional[str],
         stream_position: Optional[int],
     ) -> DeltaLocator:
@@ -365,7 +385,7 @@ class DeltaLocator(Locator, dict):
         return None
 
     @property
-    def partition_values(self) -> Optional[List[Any]]:
+    def partition_values(self) -> Optional[PartitionValues]:
         partition_locator = self.partition_locator
         if partition_locator:
             return partition_locator.partition_values

deltacat/storage/model/partition.py

@@ -1,10 +1,9 @@
 # Allow classes to use self-referencing Type hints in Python 3.7.
 from __future__ import annotations
-
 from typing import Any, Dict, List, Optional, Union
 
 import pyarrow as pa
-
+from deltacat.storage.model.partition_spec import PartitionValues
 from deltacat.storage.model.locator import Locator
 from deltacat.storage.model.namespace import NamespaceLocator
 from deltacat.storage.model.stream import StreamLocator
@@ -127,7 +126,7 @@ class Partition(dict):
         return None
 
     @property
-    def partition_values(self) -> Optional[List[Any]]:
+    def partition_values(self) -> Optional[PartitionValues]:
         partition_locator = self.locator
         if partition_locator:
             return partition_locator.partition_values
@@ -199,7 +198,7 @@ class PartitionLocator(Locator, dict):
     @staticmethod
     def of(
         stream_locator: Optional[StreamLocator],
-        partition_values: Optional[List[Any]],
+        partition_values: Optional[PartitionValues],
         partition_id: Optional[str],
     ) -> PartitionLocator:
         """
@@ -225,7 +224,7 @@ class PartitionLocator(Locator, dict):
         table_version: Optional[str],
         stream_id: Optional[str],
         storage_type: Optional[str],
-        partition_values: Optional[List[Any]],
+        partition_values: Optional[PartitionValues],
         partition_id: Optional[str],
     ) -> PartitionLocator:
         stream_locator = StreamLocator.at(
@@ -253,11 +252,11 @@ class PartitionLocator(Locator, dict):
         self["streamLocator"] = stream_locator
 
     @property
-    def partition_values(self) -> Optional[List[Any]]:
+    def partition_values(self) -> Optional[PartitionValues]:
         return self.get("partitionValues")
 
     @partition_values.setter
-    def partition_values(self, partition_values: Optional[List[Any]]) -> None:
+    def partition_values(self, partition_values: Optional[PartitionValues]) -> None:
         self["partitionValues"] = partition_values
 
     @property

deltacat/storage/model/partition_spec.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+from typing import List, Optional, Any
+from deltacat.storage.model.transform import Transform
+
+"""
+An ordered list of partition values determining the values of
+ordered transforms specified in the partition spec.
+"""
+PartitionValues = List[Any]
+
+
+class PartitionFilter(dict):
+    """
+    This class represents a filter for partitions.
+    It is used to filter partitions based on certain criteria.
+    """
+
+    @staticmethod
+    def of(
+        partition_values: Optional[PartitionValues] = None,
+    ) -> PartitionFilter:
+        """
+        Creates a new PartitionFilter instance with the specified partition key and value.
+        """
+        partition_filter = PartitionFilter()
+        partition_filter["partitionValues"] = partition_values
+        return partition_filter
+
+    @property
+    def partition_values(self) -> Optional[PartitionValues]:
+        return self.get("partitionValues")
+
+
+class PartitionSpec(dict):
+    """
+    This class determines how the underlying entities in the
+    hierarchy are partitioned. Stream partitions deltas and
+    delta partitions files.
+    """
+
+    @staticmethod
+    def of(ordered_transforms: List[Transform] = None) -> PartitionSpec:
+        partition_spec = PartitionSpec()
+        partition_spec.ordered_transforms = ordered_transforms
+        return partition_spec
+
+    @property
+    def ordered_transforms(self) -> List[Transform]:
+        return self.get("orderedTransforms")
+
+    @ordered_transforms.setter
+    def ordered_transforms(self, value: List[Transform]) -> None:
+        self["orderedTransforms"] = value
+
+
+class StreamPartitionSpec(PartitionSpec):
+    """
+    A class representing a stream partition specification.
+    A stream partitions deltas into multiple different Partition
+    """
+
+    pass
+
+
+class DeltaPartitionSpec(PartitionSpec):
+    """
+    A class representing delta partition specification.
+    The manifest entries in delta are partitioned based on this spec.
+    """
+
+    pass

deltacat/storage/model/stream.py

@@ -8,6 +8,7 @@ from deltacat.storage.model.namespace import NamespaceLocator
 from deltacat.storage.model.table import TableLocator
 from deltacat.storage.model.table_version import TableVersionLocator
 from deltacat.storage.model.types import CommitState
+from deltacat.storage.model.partition_spec import StreamPartitionSpec, PartitionValues
 
 
 class Stream(dict):
@@ -17,12 +18,14 @@ class Stream(dict):
         partition_keys: Optional[List[Dict[str, Any]]],
         state: Optional[CommitState] = None,
         previous_stream_digest: Optional[bytes] = None,
+        partition_spec: Optional[StreamPartitionSpec] = None,
     ) -> Stream:
         stream = Stream()
         stream.locator = locator
         stream.partition_keys = partition_keys
         stream.state = state
         stream.previous_stream_digest = previous_stream_digest
+        stream.partition_spec = partition_spec
         return stream
 
     @property
@@ -38,6 +41,14 @@ class Stream(dict):
 
     @property
     def partition_keys(self) -> Optional[List[Dict[str, Any]]]:
+        """
+        Ordered list of unique column names in the table schema on
+        which the underlying data is partitioned. Either partition_spec
+        or partition_keys must be specified but not both.
+
+        (Deprecated): Partition keys will be deprecated in the favor
+        of partition_spec in future releases.
+        """
         return self.get("partitionKeys")
 
     @partition_keys.setter
@@ -46,6 +57,9 @@ class Stream(dict):
 
     @property
     def previous_stream_digest(self) -> Optional[str]:
+        """
+        Previous stream digest
+        """
         return self.get("previousStreamDigest")
 
     @previous_stream_digest.setter
@@ -54,6 +68,9 @@ class Stream(dict):
 
     @property
     def state(self) -> Optional[CommitState]:
+        """
+        The commit state of a stream.
+        """
         state = self.get("state")
         return None if state is None else CommitState(state)
 
@@ -61,6 +78,26 @@ class Stream(dict):
     def state(self, state: Optional[CommitState]) -> None:
         self["state"] = state
 
+    @property
+    def partition_spec(self) -> Optional[StreamPartitionSpec]:
+        """
+        If a table uses complex partitioning instead of identity,
+        partition spec can be specified to define that strategy.
+        For example, a partition spec can define a bucketing strategy
+        on composite column values or can define iceberg compliant
+        bucketing.
+
+        Either partition_spec or partition_keys must be specified but not both.
+        """
+        val: Dict[str, Any] = self.get("partitionSpec")
+        if val is not None and not isinstance(val, StreamPartitionSpec):
+            self.partition_spec = val = StreamPartitionSpec(val)
+        return val
+
+    @partition_spec.setter
+    def partition_spec(self, spec: StreamPartitionSpec) -> None:
+        self["partitionSpec"] = spec
+
     @property
     def namespace_locator(self) -> Optional[NamespaceLocator]:
         stream_locator = self.locator
@@ -110,7 +147,7 @@ class Stream(dict):
             return stream_locator.table_version
         return None
 
-    def validate_partition_values(self, partition_values: Optional[List[Any]]):
+    def validate_partition_values(self, partition_values: Optional[PartitionValues]):
         # TODO (pdames): ensure value data types match key data types
         partition_keys = self.partition_keys
         num_keys = len(partition_keys) if partition_keys else 0

deltacat/storage/model/transform.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+from typing import List
+from enum import Enum
+
+
+class TransformName(str, Enum):
+    IDENTITY = "identity"
+    BUCKET = "bucket"
+
+
+class TransformParameters(dict):
+    """
+    This is a parent class that contains properties
+    to be passed to the corresponding transform
+    """
+
+    pass
+
+
+class IdentityTransformParameters(TransformParameters):
+    """
+    This class is used to pass parameters to the identity transform
+    """
+
+    @staticmethod
+    def of(column_name: str) -> IdentityTransformParameters:
+        identify_transform_parameters = IdentityTransformParameters()
+        identify_transform_parameters["columnName"] = column_name
+        return identify_transform_parameters
+
+    @property
+    def column_name(self) -> str:
+        """
+        The name of the column to use for identity transform
+        """
+        return self["columnName"]
+
+    @column_name.setter
+    def column_name(self, value: str) -> None:
+        self["columnName"] = value
+
+
+class BucketingStrategy(str, Enum):
+    """
+    A bucketing strategy for the transform
+    """
+
+    # Uses default deltacat bucketing strategy.
+    # This strategy supports hashing on composite keys
+    # and uses SHA1 hashing for determining the bucket.
+    # If no columns passed, it will use a random UUID
+    # for determining the bucket.
+    DEFAULT = "default"
+
+    # Uses iceberg compliant bucketing strategy.
+    # As indicated in the iceberg spec, it does not support
+    # composite keys and uses murmur3 hash for determining
+    # the bucket.
+    # See https://iceberg.apache.org/spec/#partitioning
+    ICEBERG = "iceberg"
+
+
+class BucketTransformParameters(TransformParameters):
+    """
+    Encapsulates parameters for the bucket transform.
+    """
+
+    def of(
+        self,
+        num_buckets: int,
+        column_names: List[str],
+        bucketing_strategy: BucketingStrategy,
+    ) -> BucketTransformParameters:
+        bucket_transform_parameters = BucketTransformParameters()
+        bucket_transform_parameters["numBuckets"] = num_buckets
+        bucket_transform_parameters["columnNames"] = column_names
+        bucket_transform_parameters["bucketingStrategy"] = bucketing_strategy
+
+        return bucket_transform_parameters
+
+    @property
+    def num_buckets(self) -> int:
+        """
+        The total number of buckets to create for values of the column
+        """
+        return self["numBuckets"]
+
+    @property
+    def column_names(self) -> List[str]:
+        """
+        An ordered list of unique column names from the table schema
+        to use for bucketings.
+        """
+        return self["columnNames"]
+
+    @property
+    def bucketing_strategy(self) -> BucketingStrategy:
+        """
+        The bucketing strategy to used.
+        """
+        return self["bucketingStrategy"]
+
+
+class Transform(dict):
+    """
+    A transform is represents how a particular column value can be
+    transformed into a new value. This is mostly used in the context
+    of partitioning the data files in a table.
+    """
+
+    @staticmethod
+    def of(
+        name: TransformName,
+        parameters: TransformParameters,
+    ) -> Transform:
+        partition_transform = Transform()
+        partition_transform["name"] = name
+        partition_transform["parameters"] = parameters
+        return partition_transform
+
+    @property
+    def name(self) -> TransformName:
+        return self["name"]
+
+    @property
+    def parameters(self) -> TransformParameters:
+        return self["parameters"]

deltacat/tests/aws/test_s3u.py

@@ -20,6 +20,7 @@ from botocore.exceptions import (
     ConnectTimeoutError,
     HTTPClientError,
 )
+from ray.data.datasource import FilenameProvider
 from deltacat.exceptions import NonRetryableError
 from moto import mock_s3
 from tenacity import RetryError
@@ -34,6 +35,7 @@ class TestUuidBlockWritePathProvider(unittest.TestCase):
 
         result = provider("base_path")
 
+        self.assertTrue(isinstance(provider, FilenameProvider))
         self.assertRegex(result, r"^base_path/[\w-]{36}$")
 
 

deltacat/tests/compute/compactor/utils/test_round_completion_file.py

@@ -0,0 +1,231 @@
+import pytest
+import os
+from moto import mock_s3
+import boto3
+from boto3.resources.base import ServiceResource
+from deltacat.compute.compactor.utils.round_completion_file import (
+    read_round_completion_file,
+    write_round_completion_file,
+)
+from deltacat.tests.compute.test_util_common import get_test_partition_locator
+from deltacat.compute.compactor import RoundCompletionInfo
+
+RCF_BUCKET_NAME = "rcf-bucket"
+
+
+@pytest.fixture(autouse=True, scope="module")
+def mock_aws_credential():
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_ID"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+    yield
+
+
+@pytest.fixture(autouse=True, scope="module")
+def s3_resource(mock_aws_credential):
+    with mock_s3():
+        yield boto3.resource("s3")
+
+
+@pytest.fixture(autouse=True, scope="function")
+def setup_compaction_artifacts_s3_bucket(s3_resource: ServiceResource):
+    s3_resource.create_bucket(
+        ACL="authenticated-read",
+        Bucket=RCF_BUCKET_NAME,
+    )
+    yield
+    s3_resource.Bucket(RCF_BUCKET_NAME).objects.all().delete()
+
+
+class TestReadWriteRoundCompletionFile:
+    def test_read_when_rcf_written_without_destination(self):
+        """
+        This test case tests the backward compatibility by successfully
+        reading the previously written rcf.
+        """
+
+        source_locator = get_test_partition_locator("source")
+        destination_locator = get_test_partition_locator("destination")
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        rcf_url = write_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, None, expected_rcf
+        )
+
+        rcf = read_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator
+        )
+
+        assert (
+            rcf_url == "s3://rcf-bucket/f9829af39770d904dbb811bd8f4e886dd307f507.json"
+        )
+        assert rcf == expected_rcf
+
+    def test_read_when_rcf_written_with_destination(self):
+        """
+        This test case tests the backward compatibility by successfully
+        reading the previously written rcf.
+        """
+
+        source_locator = get_test_partition_locator("source")
+        destination_locator = get_test_partition_locator("destination")
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        rcf_url = write_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf
+        )
+
+        rcf = read_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator
+        )
+
+        assert (
+            rcf_url
+            == "s3://rcf-bucket/f9829af39770d904dbb811bd8f4e886dd307f507/e9939deadc091b3289a2eb0ca56b1ba86b9892f4.json"
+        )
+        assert rcf == expected_rcf
+
+    def test_read_without_destination_when_rcf_written_with_destination(self):
+        """
+        This test case tests the backward compatibility by successfully
+        reading the previously written rcf.
+        """
+
+        source_locator = get_test_partition_locator("source")
+        destination_locator = get_test_partition_locator("destination")
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        write_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf
+        )
+
+        rcf = read_round_completion_file(RCF_BUCKET_NAME, source_locator, None)
+
+        assert rcf is None
+
+    def test_read_without_destination_when_rcf_written_without_destination(self):
+        """
+        This test case tests the backward compatibility by successfully
+        reading the previously written rcf.
+        """
+
+        source_locator = get_test_partition_locator("source")
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        write_round_completion_file(RCF_BUCKET_NAME, source_locator, None, expected_rcf)
+
+        rcf = read_round_completion_file(RCF_BUCKET_NAME, source_locator, None)
+
+        assert rcf == expected_rcf
+
+    def test_read_when_rcf_written_both_with_and_without_destination(self):
+        """
+        This test case tests the backward compatibility by successfully
+        reading the previously written rcf.
+        """
+
+        source_locator = get_test_partition_locator("source")
+        destination_locator = get_test_partition_locator("destination")
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        expected_rcf_2 = RoundCompletionInfo.of(
+            high_watermark=1223,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=1233,
+        )
+
+        write_round_completion_file(RCF_BUCKET_NAME, source_locator, None, expected_rcf)
+
+        write_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf_2
+        )
+
+        rcf = read_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator
+        )
+
+        assert rcf == expected_rcf_2
+
+    def test_read_when_none_destination_partition_id(self):
+
+        source_locator = get_test_partition_locator("source")
+        destination_locator = get_test_partition_locator(None)
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        write_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator, expected_rcf
+        )
+
+        rcf = read_round_completion_file(
+            RCF_BUCKET_NAME, source_locator, destination_locator
+        )
+
+        assert rcf == expected_rcf
+
+    def test_write_when_custom_url_is_passed(self):
+        """
+        This test case tests the backward compatibility by successfully
+        reading the previously written rcf.
+        """
+
+        source_locator = get_test_partition_locator("source")
+
+        expected_rcf = RoundCompletionInfo.of(
+            high_watermark=122,
+            compacted_delta_locator={},
+            compacted_pyarrow_write_result={},
+            sort_keys_bit_width=12,
+        )
+
+        completion_file_s3_url = f"s3://{RCF_BUCKET_NAME}/test.json"
+        rcf_url = write_round_completion_file(
+            RCF_BUCKET_NAME,
+            source_locator,
+            None,
+            expected_rcf,
+            completion_file_s3_url=completion_file_s3_url,
+        )
+
+        rcf = read_round_completion_file(RCF_BUCKET_NAME, source_locator, None)
+
+        assert rcf_url == completion_file_s3_url
+        assert rcf is None