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

deltacat/storage/rivulet/schema/datatype.py

@@ -0,0 +1,128 @@
+# Similar to daft's datatype, this is a big ole enum of all possible types
+#   In the long term, this will have to be interoperable with pandas/daft/spark/parquet/iceberg/etc type systems
+#   Our Spec will need to publish data type mappings, such as Iceberg's data type mappings: https://iceberg.apache.org/spec/#file-system-operations
+#   It also has the unique responsibility of representing multi-modal (e.g. image) types
+from dataclasses import dataclass
+from typing import Optional
+
+import pyarrow as pa
+
+
+# OPEN QUESTIONS:
+# Do we want to support the notion of logical vs physical type like parquet?
+
+# TODO turn into an interface or otherwise allow pluggable datatypes
+@dataclass(frozen=True)
+class Datatype:
+    type_name: str
+
+    @property
+    def subtype(self) -> Optional[str]:
+        """
+        Higher level formats like binary or image will have "subtype", such as image(jpg) or binary(np_array)
+        TODO - Note that we are replacing this schema system with DeltaCat schema model, which supports extended/decorated pyarrow types
+        For now going to do a super minimal/hacky implementation of types like binary and image, where
+        :return: Subtype if it exists, or None
+        """
+        if not self.type_name.endswith(")"):
+            return None
+        if self.type_name.startswith("binary(") or self.type_name.startswith("image("):
+            return self.type_name[self.type_name.find("(") + 1 : -1]
+        return None
+
+    @classmethod
+    def binary(cls, binary_format):
+        """
+        :param binary_format:
+        :return:
+        """
+        return cls(type_name=f"binary({binary_format})")
+
+    @classmethod
+    def image(cls, image_format):
+        return cls(type_name=f"image({image_format})")
+
+    @classmethod
+    def string(cls):
+        return cls(type_name="string")
+
+    @classmethod
+    def float(cls):
+        return cls(type_name="float")
+
+    @classmethod
+    def int16(cls):
+        return cls(type_name="int16")
+
+    @classmethod
+    def int32(cls):
+        return cls(type_name="int32")
+
+    @classmethod
+    def int64(cls):
+        return cls(type_name="int64")
+
+    @classmethod
+    def bool(cls):
+        return cls(type_name="bool")
+
+    @classmethod
+    def from_pyarrow(cls, pa_type: pa.DataType) -> "Datatype":
+        """
+        Convert a pa type to a Rivulet Datatype.
+
+        Args:
+            pa_type: pa DataType to convert
+
+        Returns:
+            Datatype: Corresponding Rivulet Datatype
+
+        Raises:
+            ValueError: If the pa type is not supported
+        """
+        if pa.types.is_string(pa_type):
+            return cls.string()
+        elif pa.types.is_float64(pa_type):
+            return cls.float()
+        elif pa.types.is_int16(pa_type):
+            return cls.int16()
+        elif pa.types.is_int32(pa_type):
+            return cls.int32()
+        elif pa.types.is_int64(pa_type):
+            return cls.int64()
+        elif pa.types.is_boolean(pa_type):
+            return cls.bool()
+        elif pa.types.is_binary(pa_type):
+            # TODO: Use pyarrow metadata on schema field to map correctly into image and other binary types
+            return cls.binary("binary")  # Default binary format
+        else:
+            raise ValueError(f"Unsupported pa type: {pa_type}")
+
+    def to_pyarrow(self) -> pa.field:
+        """
+        In the future we want to be more thoughtful about how we do type conversions
+
+        For now, just build a simple mapping of every time to pyarrow
+        For what it's worth, Daft schema types have a giant if/else like this
+
+        :return: pyarrow type
+        """
+        if self.type_name == "string":
+            return pa.string()
+        elif self.type_name == "float":
+            return pa.float64()
+        elif self.type_name == "int16":
+            return pa.int16()
+        elif self.type_name == "int32":
+            return pa.int32()
+        elif self.type_name == "int64":
+            return pa.int64()
+        elif self.type_name == "bool":
+            return pa.bool_()
+        elif self.type_name.startswith("image(") or self.type_name.startswith(
+            "binary("
+        ):
+            # TODO we will need to think about how custom types work with tabular libraries
+            return pa.binary()
+        else:
+            raise ValueError(f"Unsupported type conversion to pa: {self.type_name}")

deltacat/storage/rivulet/schema/schema.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, asdict
+from typing import MutableMapping, Dict, Iterable, Tuple, Optional
+
+import pyarrow as pa
+
+from deltacat.storage.rivulet.schema.datatype import Datatype
+
+
+@dataclass(frozen=True)
+class Field:
+    name: str
+    datatype: Datatype
+    is_merge_key: bool = False
+
+
+class Schema(MutableMapping[str, Field]):
+    """
+    A mutable mapping representing a schema for structured data, requiring at least one merge key field.
+
+    TODO FUTURE ITERATIONS
+    1. We may use Deltacat for schema
+    2. We almost certainly want our schema system based on arrow types,
+       since many libraries we are integrating with (e.g. daft) are
+       interoperable with arrow schemas
+
+    Attributes:
+        name: The name of the schema (for storing in dict/map)
+       _fields (dict): Maps field names to Field objects.
+
+    Methods:
+       from_pyarrow(pyarrow_schema: pa.Schema, key: str) -> Schema:
+           Creates a Schema instance from a PyArrow schema.
+
+       __len__() -> int: Returns number of fields.
+       __getitem__(key: str) -> Field: Gets field by name.
+       __setitem__(key: str, value: Field | Datatype): Adds/updates field.
+       __delitem__(key: str): Deletes field if not a merge key.
+       __iter__(): Iterates over fields.
+
+       add_field(field: Field): Adds a Field using its name as the key.
+       to_pyarrow() -> pa.Schema:
+           Converts schema to PyArrow format.
+
+       keys(): Returns field names.
+       values(): Returns Field objects.
+       items(): Returns (name, Field) pairs.
+    """
+
+    def __init__(
+        self,
+        fields: Iterable[Tuple[str, Datatype] | Field] = None,
+        merge_keys: Optional[Iterable[str]] = None,
+    ):
+        self._fields: Dict[str, Field] = {}
+        merge_keys = merge_keys or {}
+        if len(fields or []) == 0:
+            if len(merge_keys) > 0:
+                raise TypeError(
+                    "It is invalid to specify merge keys when no fields are specified. Add fields or remove the merge keys."
+                )
+            return
+        # Convert all input tuples to Field objects and add to fields
+        for field in fields:
+            if isinstance(field, tuple):
+                name, datatype = field
+                processed_field = Field(
+                    name=name, datatype=datatype, is_merge_key=(name in merge_keys)
+                )
+            elif isinstance(field, Field):
+                processed_field = field
+                name = field.name
+                # Check if merge key status conflicts
+                if len(merge_keys) > 0:
+                    expected_merge_key_status = name in merge_keys
+                    if processed_field.is_merge_key != expected_merge_key_status:
+                        raise TypeError(
+                            f"Merge key status conflict for field '{name}': "
+                            f"Provided as merge key: {expected_merge_key_status}, "
+                            f"Field's current status: {processed_field.is_merge_key}. "
+                            f"Merge keys should only be defined if raw (name, Datatype) tuples are used."
+                        )
+            else:
+                raise TypeError(f"Unexpected field type: {type(field)}")
+            self.add_field(processed_field)
+
+    @classmethod
+    def from_dict(cls, data) -> Schema:
+        fields = [
+            Field(
+                name=field_data["name"],
+                datatype=Datatype(**field_data["datatype"])
+                if isinstance(field_data["datatype"], dict)
+                else field_data["datatype"],
+                is_merge_key=field_data["is_merge_key"],
+            )
+            for field_data in data["fields"]
+        ]
+        return cls(fields)
+
+    @classmethod
+    def from_pyarrow(
+        cls, pyarrow_schema: pa.Schema, merge_keys: str | Iterable[str] = None
+    ) -> Schema:
+        """
+        Create a Schema instance from a PyArrow schema.
+
+        Args:
+            pyarrow_schema: PyArrow Schema to convert
+            merge_keys: The optional set of merge keys to add to the schema as it's being translated.
+                        These keys must be present in the schema.
+
+        Returns:
+            Schema: New Schema instance
+
+        Raises:
+            ValueError: If key is not found in schema
+        """
+        merge_keys = [merge_keys] if isinstance(merge_keys, str) else merge_keys
+        fields = {}
+
+        for field in pyarrow_schema:
+            dtype = Datatype.from_pyarrow(field.type)
+            fields[field.name] = Field(
+                field.name, dtype, is_merge_key=(field.name in merge_keys)
+            )
+
+        # Validate that the defined merge_keys are present in the fields being added
+        missing_keys = merge_keys - fields.keys()
+        if missing_keys:
+            raise ValueError(
+                f"The following merge keys not found in the provided schema: {', '.join(missing_keys)}"
+            )
+
+        return cls(fields.values())
+
+    @classmethod
+    def merge_all(cls, schemas: Iterable[Schema]) -> Schema:
+        """Merges a list of schemas into a new schema"""
+        merged = cls({})
+        for schema in schemas:
+            merged.merge(schema)
+        return merged
+
+    def __getitem__(self, key: str) -> Field:
+        return self._fields[key]
+
+    def __setitem__(
+        self, key: str, value: Field | Datatype | Tuple[Datatype, bool]
+    ) -> None:
+        # Create field from [str, Datatype, bool] where bool is merge_key
+        if isinstance(value, Field):
+            processed_field = value
+        elif isinstance(value, Datatype):
+            processed_field = Field(
+                key, value
+            )  # is_merge_key is always false in this case
+        elif isinstance(value, tuple):
+            (datatype, merge_key) = value
+            processed_field = Field(key, datatype, merge_key)
+        else:
+            raise TypeError(
+                "The field must be an instance of the Field class, Datatype, or Tuple[Datatype, bool], where bool is whether the field is a merge key."
+            )
+        processed_field: Field = processed_field
+        # if len(self._fields) == 0 and not processed_field.is_merge_key:
+        #    raise TypeError("The first field set on a Schema must be a merge key.")
+
+        self._fields[processed_field.name] = processed_field
+
+    def __delitem__(self, key: str) -> None:
+        field = self._fields[key]
+        if field.is_merge_key:
+            raise ValueError("Cannot delete a merge key field")
+        del self._fields[key]
+
+    def __len__(self) -> int:
+        return len(self._fields)
+
+    def __iter__(self) -> Iterable[str]:
+        return iter(self._fields.keys())
+
+    def __hash__(self) -> int:
+        return hash((frozenset(self._fields.items())))
+
+    def __eq__(self, other) -> bool:
+        if isinstance(other, Schema):
+            return self._fields == other._fields
+        return False
+
+    # Has a spurious type check problem in @dataclass + asdict(): https://youtrack.jetbrains.com/issue/PY-76059/Incorrect-Type-warning-with-asdict-and-Dataclass
+    def to_dict(self) -> dict[str, list[dict[str, Field]]]:
+        return {"fields": [asdict(field) for field in self._fields.values()]}
+
+    def add_field(self, field: Field) -> None:
+        """Adds a Field object using its name as the key, raises ValueError if it already exists"""
+        if field.name in self._fields:
+            raise ValueError(
+                f"Attempting to add a field with the same name as an existing field: {field.name}"
+            )
+        self[field.name] = field
+
+    def get_merge_keys(self) -> Iterable[str]:
+        """Return a list of all merge keys."""
+        return [field.name for field in self._fields.values() if field.is_merge_key]
+
+    def get_merge_key(self) -> str:
+        """Returns a single merge key if there is one, or raises if not. Used for simple schemas w/ a single key"""
+        # Get the merge key
+        merge_keys = list(self.get_merge_keys())
+        if len(merge_keys) != 1:
+            raise ValueError(
+                f"Schema must have exactly one merge key, but found {merge_keys}"
+            )
+        return merge_keys[0]
+
+    def merge(self, other: Schema) -> None:
+        """Merges another schema's fields into the current schema."""
+        if not other:
+            return
+        for name, field in other._fields.items():
+            if name in self._fields:
+                if self._fields[name] != field:
+                    raise ValueError(
+                        f"Field '{name}' already exists in the current schema with different definition"
+                    )
+            else:
+                self.add_field(field)
+
+    def to_pyarrow(self) -> pa.Schema:
+        """
+        Convert the Schema to a PyArrow schema.
+
+        Returns:
+            pyarrow.schema: A PyArrow schema representation of this Schema.
+        """
+        # TODO: Should we track merge_keys as it goes to/from pyarrow?
+        fields = []
+        for name, field in self._fields.items():
+            fields.append(pa.field(name, field.datatype.to_pyarrow()))
+        return pa.schema(fields)
+
+    def keys(self) -> Iterable[str]:
+        return self._fields.keys()
+
+    def values(self) -> Iterable[Field]:
+        return self._fields.values()
+
+    def items(self) -> Iterable[tuple[str, Field]]:
+        return self._fields.items()

deltacat/storage/rivulet/serializer.py

@@ -0,0 +1,40 @@
+from typing import Protocol, Iterable, List, Union, Any, Dict
+
+from deltacat.storage.rivulet.metastore.sst import SSTableRow
+import pyarrow as pa
+
+MEMTABLE_DATA = Union[Iterable[Dict[str, Any]], pa.Table]
+
+
+class DataSerializer(Protocol):
+    """
+    Interface for writing data only.
+
+    As data is written, it must emit sufficient metadata to build SSTable
+    Each format will have a specific data writer (e.g. ParquetDataWriter)
+
+    TODO future improvements:
+        1. How does data writer control how it chooses to write to existing files vs new files?
+            For now, we will not expose this configuration and always write each batch to
+            a new file
+        2. Related to 1, how should we expose URI(s) to write to? Probably DataWriter can
+            use FileProvider and needs to know relevant ids like task ID.
+    """
+
+    def flush_batch(self, sorted_records: MEMTABLE_DATA) -> List[SSTableRow]:
+        """
+        Flushes rows to file, and return appropriate metadata to build SSTable
+
+        TODO future improvements
+         1. Finalize type for input records (instead of MvpRow)
+
+            Options could be:
+            (a) Something like Iceberg "StructLike" which allows flexible integrations without memcopy for row-oriented formats, e.g. can make Spark InternalRow structlike
+            (b) use arrow. We will probably use arrow for writing parquet, although
+                probably it isn't ideal for row oriented formats
+         2. Keep in mind, most implementation of DataWriter will be written in rust.
+
+        :param sorted_records: Records sorted by key
+        :return: metadata used to build SS Table
+        """
+        ...

deltacat/storage/rivulet/serializer_factory.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from deltacat.storage.rivulet.parquet.serializer import ParquetDataSerializer
+from deltacat.storage.rivulet import Schema
+from deltacat.storage.rivulet.serializer import DataSerializer
+from deltacat.storage.rivulet.fs.file_provider import FileProvider
+
+from deltacat.storage.rivulet.feather.serializer import FeatherDataSerializer
+
+
+class DataSerializerFactory:
+    """
+    Simple factory class for getting the appropriate serializer given a schema
+    TODO make this more modular/pluggable like DatasetReaderRegistrar
+      This will be more challenging to make pluggable, because we should not rely on a simple 1:1 mapping of type to serializer
+      The actual logic for determining how to serialize a given schema may be complex
+      e.g.: if schema contains datatype X, you must use serializer Y. Otherwise, default to serializer Z
+    """
+
+    @classmethod
+    def get_serializer(
+        self,
+        schema: Schema,
+        file_provider: FileProvider,
+        user_provided_format: str | None = None,
+    ) -> DataSerializer:
+        if user_provided_format == "parquet":
+            return ParquetDataSerializer(file_provider, schema)
+        elif user_provided_format == "feather":
+            return FeatherDataSerializer(file_provider, schema)
+        elif user_provided_format is not None:
+            raise ValueError("Unsupported format. Must be 'parquet' or 'feather'.")
+
+        # Default engine logic. For now, if there is image or binary use feather
+        has_binary_or_image = any(
+            field.datatype.type_name.startswith(("binary", "image"))
+            for field in schema.values()
+        )
+        if has_binary_or_image:
+            return FeatherDataSerializer(file_provider, schema)
+        else:
+            return ParquetDataSerializer(file_provider, schema)

deltacat/storage/rivulet/writer/dataset_writer.py

@@ -0,0 +1,29 @@
+from typing import Protocol, Iterable, Union, Any, Dict
+import pyarrow as pa
+
+DATA = Union[Iterable[Dict[str, Any]], Iterable[pa.RecordBatch], pa.RecordBatch]
+
+
+class DatasetWriter(Protocol):
+    """
+    Top level interface for writing records to rivulet dataset. This is used by dataset.py
+
+    This writes both data AND metadata (SSTs, manifests).
+
+    The general paradigm is that records are written iteratively through write or write_batch. At configurable intervals (based on record count or size), data and metadata gets flushed.
+
+    When the user either closes the dataset writer or calls commit(), this triggers all buffered data and metadata to be flushed.
+    """
+
+    def write(self, record: DATA) -> None:
+        ...
+
+    def flush(self) -> str:
+        """
+        Explicitly flush any data and metadata and commit to dataset
+
+        This is a blocking operation
+
+        :return: URI of manifest written for commit
+        """
+        ...