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

deltacat/storage/model/shard.py

@@ -0,0 +1,51 @@
+from abc import abstractmethod
+from typing import Iterable, Optional, Protocol, TypeVar, Union
+
+from deltacat.experimental.storage.rivulet.reader.dataset_metastore import (
+    DatasetMetastore,
+)
+
+# TODO: Add type validation in dataset/schema classes
+T = TypeVar("T", bound=Union[int, str])
+
+
+class Shard(Protocol[T]):
+    """
+    Abstract base class representing a shard with defined inclusive boundaries.
+
+    A shard represents a logical partition of data, defined by its
+    minimum and maximum keys. These keys determine the range of data
+    within a dataset that the shard encompasses.
+    """
+
+    min_key: Optional[T]
+    max_key: Optional[T]
+
+
+class ShardingStrategy(Protocol):
+    """
+    A sharding strategy determines how the dataset is divided into shards.
+    """
+
+    @staticmethod
+    def from_string(strategy: str) -> "ShardingStrategy":
+        """
+        Factory method to create the appropriate ShardingStrategy from a string.
+
+        param: strategy: The string representation of the sharding strategy.
+        return: ShardingStrategy class.
+        """
+        if strategy == "range":
+            from deltacat.experimental.storage.rivulet.shard.range_shard import (
+                RangeShardingStrategy,
+            )
+
+            return RangeShardingStrategy()
+        else:
+            raise ValueError(f"Unsupported sharding strategy type: {strategy}")
+
+    @abstractmethod
+    def shards(self, num_shards: int, metastore: DatasetMetastore) -> Iterable[Shard]:
+        """
+        Generate the shards based on the chosen strategy.
+        """

deltacat/storage/model/sort_key.py

@@ -1,33 +1,230 @@
 # Allow classes to use self-referencing Type hints in Python 3.7.
 from __future__ import annotations
 
-from enum import Enum
+from typing import Optional, Any, List, Tuple, Dict
 
+from pyarrow.compute import SortOptions
 
-class SortOrder(str, Enum):
-    ASCENDING = "ascending"
-    DESCENDING = "descending"
+from deltacat.storage.model.types import (
+    SortOrder,
+    NullOrder,
+)
+from deltacat.storage.model.schema import FieldLocator
+from deltacat.storage.model.transform import Transform
+
+UNSORTED_SCHEME_NAME = "unsorted_scheme"
+UNSORTED_SCHEME_ID = "deadbeef-7277-49a4-a195-fdc8ed235d42"
 
 
 class SortKey(tuple):
     @staticmethod
-    def of(key_name: str, sort_order: SortOrder = SortOrder.ASCENDING) -> SortKey:
+    def of(
+        key: Optional[List[FieldLocator]],
+        sort_order: SortOrder = SortOrder.ASCENDING,
+        null_order: NullOrder = NullOrder.AT_END,
+        transform: Optional[Transform] = None,
+        native_object: Optional[Any] = None,
+    ) -> SortKey:
         """
         Create a sort key from a field name to use as the sort key, and
         the sort order for this key. If no sort order is specified, then the
-        data will be sorted in ascending order by default. Note that compaction
-        always keeps the LAST occurrence of this key post-sort. For example, if
-        you used an integer column as your sort key which contained the values
-        [2, 1, 3] specifying SortOrder.ASCENDING would ensure that the
-        value [3] is kept over [2, 1], and specifying SortOrder.DESCENDING
-        would ensure that [1] is kept over [2, 3].
+        data will be sorted in ascending order by default.
         """
-        return SortKey((key_name, sort_order.value))
+        return SortKey(
+            (
+                key,
+                sort_order.value if isinstance(sort_order, SortOrder) else sort_order,
+                null_order.value if isinstance(null_order, NullOrder) else null_order,
+                transform,
+                native_object,
+            )
+        )
+
+    def equivalent_to(
+        self,
+        other: SortKey,
+    ):
+        if other is None:
+            return False
+        if not isinstance(other, tuple):
+            return False
+        if not isinstance(other, SortKey):
+            other = SortKey(other)
+        return (
+            self.key == other.key
+            and self.transform == other.transform
+            and self.sort_order == other.sort_order
+            and self.null_order == other.null_order
+        )
 
     @property
-    def key_name(self) -> str:
+    def key(self) -> Optional[List[FieldLocator]]:
         return self[0]
 
     @property
     def sort_order(self) -> SortOrder:
         return SortOrder(self[1])
+
+    @property
+    def null_order(self) -> NullOrder:
+        return NullOrder(self[2])
+
+    @property
+    def transform(self) -> Optional[Transform]:
+        val: Dict[str, Any] = (
+            Transform(self[3]) if len(self) >= 4 and self[3] is not None else None
+        )
+        return val
+
+    @property
+    def arrow(self) -> List[Tuple[str, str]]:
+        # TODO(pdames): Convert unsupported field locators to arrow field names,
+        #   and transforms/multi-key-sorts to pyarrow compute expressions. Add
+        #   null order via SortOptions when supported per field by Arrow.
+        return (
+            [(field_locator, self[1]) for field_locator in self[0]] if self[0] else []
+        )
+
+    @property
+    def native_object(self) -> Optional[Any]:
+        return self[4] if len(self) >= 5 else None
+
+
+class SortKeyList(List[SortKey]):
+    @staticmethod
+    def of(items: List[SortKey]) -> SortKeyList:
+        typed_items = SortKeyList()
+        for item in items:
+            if item is not None and not isinstance(item, SortKey):
+                item = SortKey(item)
+            typed_items.append(item)
+        return typed_items
+
+    def __getitem__(self, item):
+        val = super().__getitem__(item)
+        if val is not None and not isinstance(val, SortKey):
+            self[item] = val = SortKey(val)
+        return val
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]  # This triggers __getitem__ conversion
+
+
+class SortScheme(dict):
+    @staticmethod
+    def of(
+        keys: Optional[SortKeyList],
+        name: Optional[str] = None,
+        scheme_id: Optional[str] = None,
+        native_object: Optional[Any] = None,
+    ) -> SortScheme:
+        # Validate keys if provided
+        if keys is not None:
+            # Check for empty keys list
+            if len(keys) == 0:
+                raise ValueError("Sort scheme cannot have empty keys list")
+
+            # Check for duplicate keys
+            key_names = []
+            for key in keys:
+                if key.key[0] in key_names:
+                    raise ValueError(f"Duplicate sort key found: {key.key[0]}")
+                key_names.append(key.key[0])
+
+        return SortScheme(
+            {
+                "keys": keys,
+                "name": name,
+                "id": scheme_id,
+                "nativeObject": native_object,
+            }
+        )
+
+    def equivalent_to(
+        self,
+        other: SortScheme,
+        check_identifiers: bool = False,
+    ) -> bool:
+        if other is None:
+            return False
+        if not isinstance(other, dict):
+            return False
+        if not isinstance(other, SortScheme):
+            other = SortScheme(other)
+        # If both have None keys, they are equivalent (for unsorted schemes)
+        if self.keys is None and other.keys is None:
+            return not check_identifiers or (
+                self.name == other.name and self.id == other.id
+            )
+        # If only one has None keys, they are not equivalent
+        if self.keys is None or other.keys is None:
+            return False
+        # Compare keys if both have them
+        for i in range(len(self.keys)):
+            if not self.keys[i].equivalent_to(other.keys[i]):
+                return False
+        return not check_identifiers or (
+            self.name == other.name and self.id == other.id
+        )
+
+    @property
+    def keys(self) -> Optional[SortKeyList]:
+        val: List[SortKey] = self.get("keys")
+        if val is not None and not isinstance(val, SortKeyList):
+            self["keys"] = val = SortKeyList.of(val)
+        return val
+
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def id(self) -> Optional[str]:
+        return self.get("id")
+
+    @property
+    def arrow(self) -> SortOptions:
+        # TODO(pdames): Remove homogenous null ordering when supported by Arrow.
+        if self.keys:
+            if len(set([key.null_order for key in self.keys])) == 1:
+                return SortOptions(
+                    sort_keys=[pa_key for k in self.keys for pa_key in k.arrow],
+                    null_placement=self.keys[0].null_order.value,
+                )
+            else:
+                err_msg = "All arrow sort keys must use the same null order."
+                raise ValueError(err_msg)
+        return SortOptions()
+
+    @property
+    def native_object(self) -> Optional[Any]:
+        return self.get("nativeObject")
+
+
+UNSORTED_SCHEME = SortScheme.of(
+    keys=None,
+    name=UNSORTED_SCHEME_NAME,
+    scheme_id=UNSORTED_SCHEME_ID,
+)
+
+
+class SortSchemeList(List[SortScheme]):
+    @staticmethod
+    def of(items: List[SortScheme]) -> SortSchemeList:
+        typed_items = SortSchemeList()
+        for item in items:
+            if item is not None and not isinstance(item, SortScheme):
+                item = SortScheme(item)
+            typed_items.append(item)
+        return typed_items
+
+    def __getitem__(self, item):
+        val = super().__getitem__(item)
+        if val is not None and not isinstance(val, SortScheme):
+            self[item] = val = SortScheme(val)
+        return val
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield self[i]  # This triggers __getitem__ conversion

deltacat/storage/model/stream.py

@@ -1,31 +1,54 @@
 # Allow classes to use self-referencing Type hints in Python 3.7.
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional
+import posixpath
 
-from deltacat.storage.model.locator import Locator
+import pyarrow
+
+import deltacat.storage.model.partition as partition
+
+from typing import Any, Dict, Optional, List
+
+from deltacat.storage.model.metafile import Metafile, MetafileRevisionInfo
+from deltacat.constants import TXN_DIR_NAME
+from deltacat.storage.model.locator import (
+    Locator,
+    LocatorName,
+)
 from deltacat.storage.model.namespace import NamespaceLocator
-from deltacat.storage.model.table import TableLocator
+from deltacat.storage.model.table import (
+    TableLocator,
+    Table,
+)
 from deltacat.storage.model.table_version import TableVersionLocator
-from deltacat.storage.model.types import CommitState
-from deltacat.storage.model.partition_spec import StreamPartitionSpec, PartitionValues
+from deltacat.storage.model.types import (
+    CommitState,
+    StreamFormat,
+)
+
 
+class Stream(Metafile):
+    """
+    An unbounded stream of Deltas, where each delta's records are optionally
+    partitioned according to the given partition scheme.
+    """
 
-class Stream(dict):
     @staticmethod
     def of(
         locator: Optional[StreamLocator],
-        partition_keys: Optional[List[Dict[str, Any]]],
+        partition_scheme: Optional[partition.PartitionScheme],
         state: Optional[CommitState] = None,
-        previous_stream_digest: Optional[bytes] = None,
-        partition_spec: Optional[StreamPartitionSpec] = None,
+        previous_stream_id: Optional[str] = None,
+        watermark: Optional[int] = None,
+        native_object: Optional[Any] = None,
     ) -> Stream:
         stream = Stream()
         stream.locator = locator
-        stream.partition_keys = partition_keys
+        stream.partition_scheme = partition_scheme
         stream.state = state
-        stream.previous_stream_digest = previous_stream_digest
-        stream.partition_spec = partition_spec
+        stream.previous_stream_id = previous_stream_id
+        stream.watermark = watermark
+        stream.native_object = native_object
         return stream
 
     @property
@@ -40,31 +63,44 @@ class Stream(dict):
         self["streamLocator"] = stream_locator
 
     @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.
+    def locator_alias(self) -> Optional[StreamLocatorAlias]:
+        return StreamLocatorAlias.of(self)
 
-        (Deprecated): Partition keys will be deprecated in the favor
-        of partition_spec in future releases.
+    @property
+    def partition_scheme(self) -> Optional[partition.PartitionScheme]:
         """
-        return self.get("partitionKeys")
+        A table's partition keys are defined within the context of a
+        Partition Scheme, which supports defining both fields to partition
+        a table by and optional transforms to apply to those fields to
+        derive the Partition Values that a given field, and its corresponding
+        record, belong to.
+        """
+        val: Dict[str, Any] = self.get("partitionScheme")
+        if val is not None and not isinstance(val, partition.PartitionScheme):
+            self.partition_scheme = val = partition.PartitionScheme(val)
+        return val
 
-    @partition_keys.setter
-    def partition_keys(self, partition_keys: Optional[List[Dict[str, Any]]]) -> None:
-        self["partitionKeys"] = partition_keys
+    @partition_scheme.setter
+    def partition_scheme(
+        self, partition_scheme: Optional[partition.PartitionScheme]
+    ) -> None:
+        self["partitionScheme"] = partition_scheme
 
     @property
-    def previous_stream_digest(self) -> Optional[str]:
-        """
-        Previous stream digest
-        """
-        return self.get("previousStreamDigest")
+    def previous_stream_id(self) -> Optional[str]:
+        return self.get("previousStreamId")
+
+    @previous_stream_id.setter
+    def previous_stream_id(self, previous_stream_id: Optional[str]) -> None:
+        self["previousStreamId"] = previous_stream_id
 
-    @previous_stream_digest.setter
-    def previous_stream_digest(self, previous_stream_digest: Optional[str]) -> None:
-        self["previousStreamDigest"] = previous_stream_digest
+    @property
+    def watermark(self) -> Optional[int]:
+        return self.get("watermark")
+
+    @watermark.setter
+    def watermark(self, watermark: Optional[int]) -> None:
+        self["watermark"] = watermark
 
     @property
     def state(self) -> Optional[CommitState]:
@@ -79,24 +115,12 @@ class Stream(dict):
         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.
+    def native_object(self) -> Optional[Any]:
+        return self.get("nativeObject")
 
-        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
+    @native_object.setter
+    def native_object(self, native_object: Optional[Any]) -> None:
+        self["nativeObject"] = native_object
 
     @property
     def namespace_locator(self) -> Optional[NamespaceLocator]:
@@ -126,6 +150,13 @@ class Stream(dict):
             return stream_locator.stream_id
         return None
 
+    @property
+    def stream_format(self) -> Optional[str]:
+        stream_locator = self.locator
+        if stream_locator:
+            return stream_locator.format
+        return None
+
     @property
     def namespace(self) -> Optional[str]:
         stream_locator = self.locator
@@ -147,16 +178,72 @@ class Stream(dict):
             return stream_locator.table_version
         return None
 
-    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
-        num_values = len(partition_values) if partition_values else 0
-        if num_values != num_keys:
-            raise ValueError(
-                f"Found {num_values} partition values but "
-                f"{num_keys} partition keys: {self}"
+    def url(self, catalog_name: Optional[str] = None) -> str:
+        return (
+            f"dc://{catalog_name}/{self.namespace}/{self.table_name}/{self.table_version}/{self.stream_format}/"
+            if catalog_name
+            else f"table://{self.namespace}/{self.table_name}/{self.table_version}/{self.stream_format}/"
+        )
+
+    def to_serializable(self) -> Stream:
+        serializable = self
+        if serializable.table_locator:
+            serializable: Stream = Stream.update_for(self)
+            # remove the mutable table locator
+            serializable.table_version_locator.table_locator = TableLocator.at(
+                namespace=self.id,
+                table_name=self.id,
+            )
+        return serializable
+
+    def from_serializable(
+        self,
+        path: str,
+        filesystem: Optional[pyarrow.fs.FileSystem] = None,
+    ) -> Stream:
+        # restore the table locator from its mapped immutable metafile ID
+        if self.table_locator and self.table_locator.table_name == self.id:
+            parent_rev_dir_path = Metafile._parent_metafile_rev_dir_path(
+                base_metafile_path=path,
+                parent_number=2,
             )
+            txn_log_dir = posixpath.join(
+                posixpath.dirname(
+                    posixpath.dirname(
+                        posixpath.dirname(parent_rev_dir_path),
+                    )
+                ),
+                TXN_DIR_NAME,
+            )
+            table = Table.read(
+                MetafileRevisionInfo.latest_revision(
+                    revision_dir_path=parent_rev_dir_path,
+                    filesystem=filesystem,
+                    success_txn_log_dir=txn_log_dir,
+                ).path,
+                filesystem,
+            )
+            self.table_version_locator.table_locator = table.locator
+        return self
+
+
+class StreamLocatorName(LocatorName):
+    def __init__(self, locator: StreamLocator):
+        self.locator = locator
+
+    @property
+    def immutable_id(self) -> Optional[str]:
+        return self.locator.stream_id
+
+    @immutable_id.setter
+    def immutable_id(self, immutable_id: Optional[str]):
+        self.locator.stream_id = immutable_id
+
+    def parts(self) -> List[str]:
+        return [
+            self.locator.stream_id,
+            self.locator.format,
+        ]
 
 
 class StreamLocator(Locator, dict):
@@ -164,7 +251,7 @@ class StreamLocator(Locator, dict):
     def of(
         table_version_locator: Optional[TableVersionLocator],
         stream_id: Optional[str],
-        storage_type: Optional[str],
+        stream_format: Optional[StreamFormat],
     ) -> StreamLocator:
         """
         Creates a table version Stream Locator. All input parameters are
@@ -173,7 +260,11 @@ class StreamLocator(Locator, dict):
         stream_locator = StreamLocator()
         stream_locator.table_version_locator = table_version_locator
         stream_locator.stream_id = stream_id
-        stream_locator.storage_type = storage_type
+        stream_locator.format = (
+            stream_format.value
+            if isinstance(stream_format, StreamFormat)
+            else stream_format
+        )
         return stream_locator
 
     @staticmethod
@@ -182,19 +273,31 @@ class StreamLocator(Locator, dict):
         table_name: Optional[str],
         table_version: Optional[str],
         stream_id: Optional[str],
-        storage_type: Optional[str],
+        stream_format: Optional[StreamFormat],
     ) -> StreamLocator:
-        table_version_locator = TableVersionLocator.at(
-            namespace,
-            table_name,
-            table_version,
+        table_version_locator = (
+            TableVersionLocator.at(
+                namespace,
+                table_name,
+                table_version,
+            )
+            if table_version
+            else None
         )
         return StreamLocator.of(
             table_version_locator,
             stream_id,
-            storage_type,
+            stream_format,
         )
 
+    @property
+    def name(self) -> StreamLocatorName:
+        return StreamLocatorName(self)
+
+    @property
+    def parent(self) -> Optional[TableVersionLocator]:
+        return self.table_version_locator
+
     @property
     def table_version_locator(self) -> Optional[TableVersionLocator]:
         val: Dict[str, Any] = self.get("tableVersionLocator")
@@ -217,12 +320,12 @@ class StreamLocator(Locator, dict):
         self["streamId"] = stream_id
 
     @property
-    def storage_type(self) -> Optional[str]:
-        return self.get("storageType")
+    def format(self) -> Optional[str]:
+        return self.get("format")
 
-    @storage_type.setter
-    def storage_type(self, storage_type: Optional[str]) -> None:
-        self["storageType"] = storage_type
+    @format.setter
+    def format(self, stream_format: Optional[str]) -> None:
+        self["format"] = stream_format
 
     @property
     def namespace_locator(self) -> Optional[NamespaceLocator]:
@@ -259,13 +362,45 @@ class StreamLocator(Locator, dict):
             return table_version_locator.table_version
         return None
 
-    def canonical_string(self) -> str:
-        """
-        Returns a unique string for the given locator that can be used
-        for equality checks (i.e. two locators are equal if they have
-        the same canonical string).
-        """
-        tvl_hexdigest = self.table_version_locator.hexdigest()
-        stream_id = self.stream_id
-        storage_type = self.storage_type
-        return f"{tvl_hexdigest}|{stream_id}|{storage_type}"
+
+class StreamLocatorAliasName(LocatorName):
+    def __init__(self, locator: StreamLocatorAlias):
+        self.locator = locator
+
+    @property
+    def immutable_id(self) -> Optional[str]:
+        return None
+
+    def parts(self) -> List[str]:
+        return [self.locator.format]
+
+
+class StreamLocatorAlias(Locator, dict):
+    @staticmethod
+    def of(
+        parent_stream: Stream,
+    ) -> StreamLocatorAlias:
+        return (
+            StreamLocatorAlias(
+                {
+                    "format": parent_stream.stream_format,
+                    "parent": (
+                        parent_stream.locator.parent if parent_stream.locator else None
+                    ),
+                }
+            )
+            if parent_stream.state != CommitState.STAGED
+            else None  # staged streams cannot be resolved by alias
+        )
+
+    @property
+    def format(self) -> Optional[str]:
+        return self.get("format")
+
+    @property
+    def name(self) -> StreamLocatorAliasName:
+        return StreamLocatorAliasName(self)
+
+    @property
+    def parent(self) -> Optional[Locator]:
+        return self.get("parent")