deltacat 1.1.36__py3-none-any.whl → 2.0.0b2__py3-none-any.whl

deltacat/storage/model/transform.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
-from typing import List
 from enum import Enum
+from typing import Dict, Any, Optional
 
 
 class TransformName(str, Enum):
     IDENTITY = "identity"
     BUCKET = "bucket"
+    YEAR = "year"
+    MONTH = "month"
+    DAY = "day"
+    HOUR = "hour"
+    TRUNCATE = "truncate"
+    VOID = "void"
+    UNKNOWN = "unknown"
 
 
 class TransformParameters(dict):
@@ -17,63 +24,30 @@ class TransformParameters(dict):
     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 DeltaCAT SHA-1 based hash bucketing strategy.
     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-compliant murmur3 based hash bucketing strategy.
     ICEBERG = "iceberg"
 
 
 class BucketTransformParameters(TransformParameters):
     """
-    Encapsulates parameters for the bucket transform.
+    Parameters for the bucket transform.
     """
 
+    @staticmethod
     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
@@ -81,47 +55,210 @@ class BucketTransformParameters(TransformParameters):
     @property
     def num_buckets(self) -> int:
         """
-        The total number of buckets to create for values of the column
+        The total number of buckets to create.
         """
         return self["numBuckets"]
 
     @property
-    def column_names(self) -> List[str]:
+    def bucketing_strategy(self) -> BucketingStrategy:
         """
-        An ordered list of unique column names from the table schema
-        to use for bucketings.
+        The bucketing strategy to use.
         """
-        return self["columnNames"]
+        return BucketingStrategy(self["bucketingStrategy"])
+
+
+class TruncateTransformParameters(TransformParameters):
+    """
+    Parameters for the truncate transform.
+    """
+
+    @staticmethod
+    def of(width: int) -> TruncateTransformParameters:
+        truncate_transform_parameters = TruncateTransformParameters()
+        truncate_transform_parameters["width"] = width
+        return truncate_transform_parameters
 
     @property
-    def bucketing_strategy(self) -> BucketingStrategy:
+    def width(self) -> int:
         """
-        The bucketing strategy to used.
+        The width to truncate the field to.
         """
-        return self["bucketingStrategy"]
+        return self["width"]
 
 
 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.
+    A transform represents how a particular column value can be
+    transformed into a new value. For example, transforms may be used
+    to determine partition or sort values for table records.
+    """
+
+    @property
+    def name(self) -> TransformName:
+        return TransformName(self["name"])
+
+    @name.setter
+    def name(self, name: TransformName) -> None:
+        self["name"] = name
+
+    @property
+    def parameters(self) -> Optional[TransformParameters]:
+        return NAME_TO_TRANSFORM[self.name].parameters
+
+    @parameters.setter
+    def parameters(
+        self,
+        parameters: Optional[TransformParameters] = None,
+    ) -> None:
+        NAME_TO_TRANSFORM[self.name].parameters = parameters
+
+
+class BucketTransform(Transform):
+    """
+    A transform that hashes field values into a fixed number of buckets.
     """
 
     @staticmethod
-    def of(
-        name: TransformName,
-        parameters: TransformParameters,
-    ) -> Transform:
-        partition_transform = Transform()
-        partition_transform["name"] = name
-        partition_transform["parameters"] = parameters
-        return partition_transform
+    def of(parameters: BucketTransformParameters) -> BucketTransform:
+        transform = BucketTransform()
+        transform.name = TransformName.BUCKET
+        transform.parameters = parameters
+        return transform
 
     @property
-    def name(self) -> TransformName:
-        return self["name"]
+    def parameters(self) -> BucketTransformParameters:
+        val: Dict[str, Any] = self.get("parameters")
+        if val is not None and not isinstance(val, BucketTransformParameters.__class__):
+            self["parameters"] = val = BucketTransformParameters(val)
+        return val
+
+    @parameters.setter
+    def parameters(
+        self,
+        parameters: Optional[BucketTransformParameters] = None,
+    ) -> None:
+        self["parameters"] = parameters
+
+
+class TruncateTransform(Transform):
+    """
+    A transform that truncates field values to a fixed width.
+    """
+
+    @staticmethod
+    def of(parameters: TruncateTransformParameters) -> TruncateTransform:
+        transform = TruncateTransform()
+        transform.name = TransformName.TRUNCATE
+        transform.parameters = parameters
+        return transform
 
     @property
-    def parameters(self) -> TransformParameters:
-        return self["parameters"]
+    def parameters(self) -> TruncateTransformParameters:
+        val: Dict[str, Any] = self.get("parameters")
+        if val is not None and not isinstance(val, TruncateTransformParameters):
+            self["parameters"] = val = TruncateTransformParameters(val)
+        return val
+
+    @parameters.setter
+    def parameters(
+        self,
+        parameters: Optional[TruncateTransformParameters] = None,
+    ) -> None:
+        self["parameters"] = parameters
+
+
+class IdentityTransform(Transform):
+    """
+    A no-op transform that returns unmodified field values.
+    """
+
+    @staticmethod
+    def of() -> IdentityTransform:
+        transform = IdentityTransform()
+        transform.name = TransformName.IDENTITY
+        return transform
+
+
+class HourTransform(Transform):
+    """
+    A transform that returns the hour of a datetime value.
+    """
+
+    @staticmethod
+    def of() -> HourTransform:
+        transform = HourTransform()
+        transform.name = TransformName.HOUR
+        return transform
+
+
+class DayTransform(Transform):
+    """
+    A transform that returns the day of a datetime value.
+    """
+
+    @staticmethod
+    def of() -> DayTransform:
+        transform = DayTransform()
+        transform.name = TransformName.DAY
+        return transform
+
+
+class MonthTransform(Transform):
+    """
+    A transform that returns the month of a datetime value.
+    """
+
+    @staticmethod
+    def of() -> MonthTransform:
+        transform = MonthTransform()
+        transform.name = TransformName.MONTH
+        return transform
+
+
+class YearTransform(Transform):
+    """
+    A transform that returns the year of a datetime value.
+    """
+
+    @staticmethod
+    def of() -> YearTransform:
+        transform = YearTransform()
+        transform.name = TransformName.YEAR
+        return transform
+
+
+class VoidTransform(Transform):
+    """
+    A transform that coerces all field values to None.
+    """
+
+    @staticmethod
+    def of() -> VoidTransform:
+        transform = VoidTransform()
+        transform.name = TransformName.VOID
+        return transform
+
+
+class UnknownTransform(Transform):
+    """
+    An unknown or invalid transform.
+    """
+
+    @staticmethod
+    def of() -> UnknownTransform:
+        transform = UnknownTransform()
+        transform.name = TransformName.UNKNOWN
+        return transform
+
+
+NAME_TO_TRANSFORM: Dict[TransformName, Transform] = {
+    TransformName.IDENTITY: IdentityTransform,
+    TransformName.BUCKET: BucketTransform,
+    TransformName.YEAR: YearTransform,
+    TransformName.MONTH: MonthTransform,
+    TransformName.DAY: DayTransform,
+    TransformName.HOUR: HourTransform,
+    TransformName.TRUNCATE: TruncateTransform,
+    TransformName.VOID: VoidTransform,
+    TransformName.UNKNOWN: UnknownTransform,
+}

deltacat/storage/model/types.py

@@ -1,25 +1,100 @@
+from __future__ import annotations
+
 from enum import Enum
 from typing import List, Union
 
-from pyarrow.parquet import ParquetFile
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 from ray.data.dataset import Dataset
 from daft import DataFrame as DaftDataFrame
 
-LocalTable = Union[pa.Table, pd.DataFrame, np.ndarray, ParquetFile]
+
+LocalTable = Union[
+    pa.Table,
+    pd.DataFrame,
+    np.ndarray,
+    pa.parquet.ParquetFile,
+]
 LocalDataset = List[LocalTable]
 DistributedDataset = Union[Dataset, DaftDataFrame]
 
 
+class StreamFormat(str, Enum):
+    DELTACAT = "deltacat"
+    ICEBERG = "iceberg"
+    HUDI = "hudi"
+    DELTA_LAKE = "delta_lake"
+    SQLITE3 = "SQLITE3"  # used by tests
+
+
 class DeltaType(str, Enum):
     APPEND = "append"
     UPSERT = "upsert"
     DELETE = "delete"
 
 
+class TransactionType(str, Enum):
+    # the transaction reads existing data
+    # does not conflict with any other transaction types
+    READ = "read"
+    # the transaction only appends new data
+    # conflicts with other transaction types can be auto-resolved
+    APPEND = "append"
+    # the transaction alters existing data
+    # (even if it also appends data)
+    # conflicts with other alters/overwrites/restates/deletes fail
+    ALTER = "alter"
+    # the transaction overwrites existing data
+    # (even if it also appends or alters data)
+    # conflicts with other alters/overwrites/restates/deletes fail
+    OVERWRITE = "overwrite"
+    # the transaction restates existing data with a new layout
+    # (even if it appends, alters, or overwrites data to do so)
+    # conflicts with other alters/overwrites/restates/deletes fail
+    RESTATE = "restate"
+    # the transaction deletes existing data
+    # (even if it also appends, alters, overwrites, or restates data)
+    # conflicts with other alters/overwrites/restates/deletes fail
+    DELETE = "delete"
+
+
+class TransactionOperationType(str, Enum):
+    CREATE = "create"
+    UPDATE = "update"
+    DELETE = "delete"
+
+    READ_SIBLINGS = "read_siblings"
+    READ_CHILDREN = "read_children"
+    READ_LATEST = "read_latest"
+    READ_EXISTS = "read_exists"
+
+    @staticmethod
+    def write_operations():
+        return {
+            TransactionOperationType.CREATE,
+            TransactionOperationType.UPDATE,
+            TransactionOperationType.DELETE,
+        }
+
+    @staticmethod
+    def read_operations():
+        return {
+            TransactionOperationType.READ_SIBLINGS,
+            TransactionOperationType.READ_CHILDREN,
+            TransactionOperationType.READ_LATEST,
+            TransactionOperationType.READ_EXISTS,
+        }
+
+    def is_write_operation(self) -> bool:
+        return self in TransactionOperationType.write_operations()
+
+    def is_read_operation(self) -> bool:
+        return self in TransactionOperationType.read_operatins()
+
+
 class LifecycleState(str, Enum):
+    CREATED = "created"
     UNRELEASED = "unreleased"
     ACTIVE = "active"
     DEPRECATED = "deprecated"
@@ -35,22 +110,45 @@ class CommitState(str, Enum):
 
 class SchemaConsistencyType(str, Enum):
     """
-    Schemas are optional for DeltaCAT tables and can be used to inform the data
-    consistency checks run for each field. If a schema is present, it can be
-    used to enforce the following column-level data consistency policies at
-    table load time:
+    DeltaCAT table schemas can be used to inform the data consistency checks
+    run for each field. When present, the schema can be used to enforce the
+    following field-level data consistency policies at table load time:
 
-    NONE: No consistency checks are run. May be mixed with the below two
-    policies by specifying column names to pass through together with
-    column names to coerce/validate.
+    NONE: No consistency checks are run.
 
-    COERCE: Coerce fields to fit the schema whenever possible. An explicit
-    subset of column names to coerce may optionally be specified.
+    COERCE: Coerce fields to fit the schema whenever possible.
 
-    VALIDATE: Raise an error for any fields that don't fit the schema. An
-    explicit subset of column names to validate may optionally be specified.
+    VALIDATE: Raise an error for any fields that don't fit the schema.
     """
 
     NONE = "none"
     COERCE = "coerce"
     VALIDATE = "validate"
+
+
+class SortOrder(str, Enum):
+    ASCENDING = "ascending"
+    DESCENDING = "descending"
+
+    @classmethod
+    def _missing_(cls, value: str):
+        # pyiceberg.table.sorting.SortDirection mappings
+        if value.lower() == "asc":
+            return SortOrder.ASCENDING
+        elif value.lower() == "desc":
+            return SortOrder.DESCENDING
+        return None
+
+
+class NullOrder(str, Enum):
+    AT_START = "at_start"
+    AT_END = "at_end"
+
+    @classmethod
+    def _missing_(cls, value: str):
+        # pyiceberg.table.sorting.NullOrder mappings
+        if value.lower() == "nulls-first":
+            return NullOrder.AT_START
+        elif value.lower() == "nulls-last":
+            return NullOrder.AT_END
+        return None

deltacat/storage/rivulet/__init__.py

@@ -0,0 +1,11 @@
+from .schema.schema import Schema
+from .schema.schema import Field
+from .dataset import Dataset
+from .schema.schema import Datatype
+
+__all__ = [
+    "Schema",
+    "Field",
+    "Dataset",
+    "Datatype",
+]

deltacat/storage/rivulet/arrow/serializer.py

@@ -0,0 +1,75 @@
+from abc import ABC, abstractmethod
+from typing import Iterator, List, Any
+import pyarrow as pa
+
+from deltacat.storage.rivulet.metastore.sst import SSTableRow
+from deltacat.storage.rivulet import Schema
+from deltacat.storage.rivulet.serializer import DataSerializer, MEMTABLE_DATA
+from deltacat.storage.rivulet.fs.file_provider import FileProvider
+
+
+class ArrowSerializer(DataSerializer, ABC):
+    """
+    Utility class which can serialize data by first converting to arrow as intermediate format
+    and then using the provided serialization function
+    """
+
+    def __init__(self, file_provider: FileProvider, schema: Schema):
+        self.schema = schema
+        self.file_provider = file_provider
+        self.arrow_schema = self.schema.to_pyarrow()
+
+    @abstractmethod
+    def serialize(self, table: pa.Table) -> List[SSTableRow]:
+        """
+        Write an Arrow table to the specified output file
+
+        :param table: PyArrow table to write
+        :return: Number of row groups in the written file
+        """
+        pass
+
+    def _to_arrow_table(self, sorted_records: MEMTABLE_DATA) -> pa.Table:
+        """
+        Convert input records to an Arrow table
+        """
+        if isinstance(sorted_records, pa.Table):
+            return sorted_records
+        elif isinstance(sorted_records, (Iterator, List)):
+            return pa.Table.from_pylist(sorted_records, schema=self.arrow_schema)
+        else:
+            raise ValueError(f"Unsupported record type: {type(sorted_records)}")
+
+    def _get_min_max_key(self, table: pa.Table) -> (Any, Any):
+        """
+        Get min and max values for the merge key from the table
+        """
+        key_col = table[self.schema.get_merge_key()]
+        return key_col[0].as_py(), key_col[len(key_col) - 1].as_py()
+
+    def flush_batch(self, sorted_records: MEMTABLE_DATA) -> List[SSTableRow]:
+        """
+        Write records to new parquet file as row group
+        For now, we will depend on pyarrow to write to parquet
+
+        :param sorted_records: record batch in SORTED ORDER
+        :return: metadata for constructing SSTable
+        """
+        if not sorted_records:
+            return []
+
+        table = self._to_arrow_table(sorted_records)
+        return self.serialize(table)
+
+    def write(self, sorted_records: MEMTABLE_DATA) -> List[SSTableRow]:
+        """
+        Write records using the provided serialization function
+
+        :param sorted_records: record batch in SORTED ORDER
+        :return: metadata for constructing SSTable
+        """
+        if not sorted_records:
+            return []
+
+        table = self._to_arrow_table(sorted_records)
+        return self.serialize(table)