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

deltacat/storage/rivulet/logical_plan.py

@@ -0,0 +1,105 @@
+from dataclasses import dataclass
+from typing import List, Callable, Any, Protocol
+
+from deltacat.storage.rivulet.dataset_executor import DatasetExecutor
+from deltacat.storage.rivulet.mvp.Table import MvpTable
+from deltacat.storage.rivulet import Schema
+
+
+class DatasetOperation(Protocol):
+    def visit(self, executor: DatasetExecutor):
+        ...
+
+
+@dataclass
+class SelectOperation(DatasetOperation):
+    """
+    Select a subset of fields within the schema
+
+    TODO need better interface for defining selection
+            (e.g. "all fields except X")
+
+    TODO in the future this should support basic filters (e.g. on primary key)
+    """
+
+    fields: List[str]
+
+    def visit(self, executor: DatasetExecutor):
+        executor.select(self.fields)
+
+
+@dataclass
+class MapOperation(DatasetOperation):
+    """
+    Map a function over each record in the dataset
+
+    TODO need more sophistication in the interface of the callable function
+    For now we will be super simple and just call the transform on each record
+    """
+
+    transform: Callable[[Any], Any]
+
+    def visit(self, executor: DatasetExecutor):
+        executor.map(self.transform)
+
+
+class CollectOperation(DatasetOperation):
+    """
+    Materialize dataset
+    """
+
+    def visit(self, executor: DatasetExecutor):
+        executor.collect()
+
+
+class LogicalPlan:
+    """
+    A fluent builder for constructing a sequence of dataset operations.
+
+    This class allows chaining of different dataset operations such as select and map.
+    The actual implementation of these operations is delegated to the Dataset class
+    using the visitor pattern.
+
+    Example:
+        plan = LogicalPlan().select(lambda x: x['age'] > 30).map(lambda x: x['name'])
+
+    The plan can then be executed on a Dataset object, which will apply the
+    operations in the order they were added.
+    """
+
+    def __init__(self, schema: Schema):
+        self.operations: List[DatasetOperation] = []
+        self.schema = schema
+        # Tracks effective schema to perform each operation on
+        self.effective_schema: Schema = schema.__deepcopy__()
+
+    def select(self, filter: List[str]) -> "LogicalPlan":
+        # Validate that select statement is allowed and mutate effective schema for future validations
+        invalid_fields = [
+            field for field in filter if field not in self.effective_schema.fields
+        ]
+        if invalid_fields:
+            raise ValueError(f"Invalid fields: {', '.join(invalid_fields)}")
+
+        # remove fields from effective schema if they are not in chosen fields
+        remove_fields = [
+            field for field in self.effective_schema.keys() if field not in filter
+        ]
+        for field in remove_fields:
+            self.effective_schema.__delitem__(field)
+
+        self.operations.append(SelectOperation(filter))
+        return self
+
+    def map(self, transform: Callable[[dict], dict]) -> "LogicalPlan":
+        self.operations.append(MapOperation(transform))
+        return self
+
+    def collect(self) -> "LogicalPlan":
+        self.operations.append(CollectOperation())
+        return self
+
+    def execute(self, executor: DatasetExecutor) -> "MvpTable":
+        for operation in self.operations:
+            operation.visit(executor)
+        return executor.output

deltacat/storage/rivulet/metastore/delta.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+from typing import Protocol, NamedTuple, List
+import time
+
+from deltacat.storage import (
+    ManifestMeta,
+    EntryType,
+    DeltaLocator,
+    Delta,
+    DeltaType,
+    Transaction,
+    TransactionType,
+    TransactionOperation,
+    TransactionOperationType,
+)
+from deltacat.storage.model.manifest import Manifest, ManifestEntryList, ManifestEntry
+from deltacat.storage.model.partition import PartitionLocator
+from deltacat.storage.model.transaction import TransactionOperationList
+
+from deltacat.storage.model.types import StreamFormat
+from deltacat.storage.rivulet import Schema
+
+StreamPosition = int
+"""The stream position for creating a consistent ordering of manifests."""
+TreeLevel = int
+"""The level of the manifest in the LSM-tree."""
+
+
+class DeltaContext(NamedTuple):
+    """Minimal amount of manifest context that may need to be circulated independently or alongside individual files"""
+
+    # Schema needed to understand which field group was added when writing manifest
+    # TODO in the future we should use something like a field group id and keep schema in dataset-level metadata
+    schema: Schema
+    stream_position: StreamPosition
+    level: TreeLevel
+
+
+class RivuletDelta(dict):
+    """
+    Temporary class during merging of deltacat/rivulet metadata formats
+
+    This class currently serves two purposes:
+    1. Avoid big bang refactor in which consumers of RivuletDelta have to update their code to consume deltacat Delta/Manifest
+    2. Provide more time to figure out how to represent SST files / schema / etc within deltacat constructs
+
+    """
+
+    context: DeltaContext
+
+    @staticmethod
+    def of(delta: Delta) -> RivuletDelta:
+        riv_delta = RivuletDelta()
+        riv_delta["dcDelta"] = delta
+        schema = Schema.from_dict(delta.get("schema"))
+        riv_delta["DeltaContext"] = DeltaContext(
+            schema, delta.stream_position, delta.get("level")
+        )
+
+        return riv_delta
+
+    @property
+    def dcDelta(self) -> Delta:
+        return self.get("dcDelta")
+
+    @property
+    def sst_files(self) -> List[str]:
+        if "sst_files" not in self.keys():
+            self["sst_files"] = [m.uri for m in self.dcDelta.manifest.entries]
+        return self["sst_files"]
+
+    @sst_files.setter
+    def sst_files(self, files: List[str]):
+        self["sst_files"] = files
+
+    @property
+    def context(self) -> DeltaContext:
+        return self["DeltaContext"]
+
+    @context.setter
+    def context(self, mc: DeltaContext):
+        self["DeltaContext"] = mc
+
+
+class ManifestIO(Protocol):
+    """
+    Minimal interface for reading and writing manifest files
+    """
+
+    def write(
+        self,
+        sst_files: List[str],
+        schema: Schema,
+        level: TreeLevel,
+    ) -> str:
+        ...
+
+    def read(self, file: str) -> RivuletDelta:
+        ...
+
+
+class DeltacatManifestIO(ManifestIO):
+    """
+    Writes manifest data, but by writing to a Deltacat metastore using Deltacat delta/manifest classes.
+    """
+
+    def __init__(self, root: str, locator: PartitionLocator):
+        self.root = root
+        self.locator = locator
+
+    def write(
+        self,
+        sst_files: List[str],
+        schema: Schema,
+        level: TreeLevel,
+    ) -> str:
+        entry_list = ManifestEntryList()
+        """
+        Currently, we use the "data files" manifest entry field for SST files
+        This is a bit of a hack - we should consider how to better model SST files
+        (e.g.: add Manifest entry of type "SST") and decide whether we also need to record data files separately
+         even though they're referenced by SST
+        Ticket: https://github.com/ray-project/deltacat/issues/469
+        """
+        for sst_uri in sst_files:
+            entry_list.append(
+                ManifestEntry.of(
+                    url=sst_uri,
+                    # TODO have rivulet writer populate these values
+                    # see: https://github.com/ray-project/deltacat/issues/476
+                    meta=ManifestMeta.of(
+                        record_count=None,  # or known
+                        content_length=None,
+                        content_type=None,
+                        content_encoding=None,
+                        entry_type=EntryType.DATA,
+                    ),
+                )
+            )
+        dc_manifest = Manifest.of(entries=entry_list)
+
+        # Create delta and transaction which writes manifest to root
+        # TODO replace this with higher level storage interface for deltacat
+        delta_locator = DeltaLocator.at(
+            namespace=self.locator.namespace,
+            table_name=self.locator.table_name,
+            table_version=self.locator.table_version,
+            partition_id=self.locator.partition_id,
+            partition_values=self.locator.partition_values,
+            stream_id=self.locator.stream_id,
+            stream_format=StreamFormat.DELTACAT,
+            # Using microsecond precision timestamp as stream position
+            # TODO consider having storage interface auto assign stream position
+            stream_position=time.time_ns(),
+        )
+
+        delta = Delta.of(
+            locator=delta_locator,
+            delta_type=DeltaType.APPEND,
+            meta=None,
+            properties={},
+            manifest=dc_manifest,
+        )
+        # TODO later support multiple schemas (https://github.com/ray-project/deltacat/issues/468)
+        delta["schema"] = schema.to_dict()
+        # TODO consider if level should be added as first class key to delta or
+        # kept as specific to storage interface
+        delta["level"] = level
+
+        tx_results = Transaction.of(
+            txn_type=TransactionType.APPEND,
+            txn_operations=TransactionOperationList.of(
+                [
+                    TransactionOperation.of(
+                        operation_type=TransactionOperationType.CREATE,
+                        dest_metafile=delta,
+                    )
+                ]
+            ),
+        ).commit(self.root)
+        paths = tx_results[0]
+        assert (
+            len(paths) == 1
+        ), "expected delta commit transaction to write exactly 1 metafile"
+        return paths[0]
+
+    def read(self, file: str):
+        delta = Delta.read(file)
+        return RivuletDelta.of(delta)

deltacat/storage/rivulet/metastore/json_sst.py

@@ -0,0 +1,105 @@
+import logging
+import json
+
+from itertools import zip_longest
+from typing import List
+
+from deltacat.storage.rivulet.fs.input_file import InputFile
+from deltacat.storage.rivulet.fs.output_file import OutputFile
+from deltacat.storage.rivulet.metastore.sst import (
+    SSTWriter,
+    SSTableRow,
+    SSTReader,
+    SSTable,
+)
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class JsonSstWriter(SSTWriter):
+    """
+    Class for reading and writing Json SST files
+
+    TODO use a more efficient format or compression. Also can factor out URI prefix across rows
+    We can also optimize by omitting offset_end if sequential rows use the same uri
+    """
+
+    def write(self, file: OutputFile, rows: List[SSTableRow]) -> None:
+        """
+        Writes SST file
+        """
+        if len(rows) == 0:
+            return
+
+        # File-level metadata for key min/max and offset end
+        min_key = rows[0].key_min
+        max_key = rows[-1].key_max
+        offset_end = rows[-1].offset_end
+
+        # Convert to dict for json serialization
+        file_rows = [
+            {
+                "key_min": row.key_min,
+                "key_max": row.key_max,
+                "offset": row.offset_start,
+                "uri": row.uri,
+            }
+            for row in rows
+        ]
+
+        file_as_dict = {
+            "key_min": min_key,
+            "key_max": max_key,
+            "offset_end": offset_end,
+            "metadata": file_rows,
+        }
+
+        try:
+            with file.create() as f:
+                f.write(json.dumps(file_as_dict).encode())
+            logger.debug(f"SSTable data successfully written to {file.location}")
+        except Exception as e:
+            # TODO better error handling for IO
+            logger.debug(f"Unexpected error occurred while writing SSTable data: {e}")
+            raise e
+
+
+class JsonSstReader(SSTReader):
+    """
+    interface for reading SST files
+    """
+
+    def read(self, file: InputFile) -> SSTable:
+        with file.open() as f:
+            data = json.loads(f.read())
+        sst_rows: List[SSTableRow] = []
+        file_offset_end = data["offset_end"]
+
+        # each row only has "key", "offset", "uri"
+        # need to get "key_end", "offset_end" from either next row
+        # OR from top level metadata
+        for row1, row2 in zip_longest(data["metadata"], data["metadata"][1:]):
+            # if not row2, we are on the very last row and need to use file metadata for key and offset end
+            if not row2:
+                sst_rows.append(
+                    SSTableRow(
+                        row1["key_min"],
+                        row1["key_max"],
+                        row1["uri"],
+                        row1["offset"],
+                        file_offset_end,
+                    )
+                )
+            else:
+                sst_rows.append(
+                    SSTableRow(
+                        row1["key_min"],
+                        row1["key_max"],
+                        row1["uri"],
+                        row1["offset"],
+                        file_offset_end,
+                    )
+                )
+
+        return SSTable(sst_rows, data["key_min"], data["key_max"])

deltacat/storage/rivulet/metastore/sst.py

@@ -0,0 +1,82 @@
+from dataclasses import dataclass
+from typing import Protocol, Any, List
+
+from deltacat.storage.rivulet.fs.input_file import InputFile
+from deltacat.storage.rivulet.fs.output_file import OutputFile
+
+
+@dataclass(frozen=True)
+class SSTableRow:
+    """
+    Row of Sorted String Table.
+
+    The metadata for a SSTable row referencing a offset-range of a data file containing data sorted by primary key.
+
+    Note that the actual file format for SSTables can omit some of these fields (e.g. key_end and offset_end) taking
+    advantage of sorted nature of file.
+    """
+
+    key_min: Any
+    """The first primary key found in referenced data range, inclusive."""
+    key_max: Any
+    """the last primary key found in referenced data range, inclusive"""
+    uri: str
+    """The URI of the data file containing this row's data.
+    Will be format dependent, e.g. file://<absolute_path> or s3://<bucket>/<key>
+    Note that if uri_prefix is specified in SSTable, this will just be a postfix
+    """
+    offset_start: int
+    """
+    offset start for data range within uri.
+    Note this offset is format dependent - e.g. for Parquet it will be zero-indexed row group
+    For other formats it will be byte offset into file
+    """
+    offset_end: int
+    """
+    offset end for data range within uri.
+    """
+
+    """The starting offset into the data file for data referenced by this row.
+    Note that offset is format dependent.
+        E.g. for parquet files it is row group, for other formats it is byte offset
+    """
+
+
+@dataclass(frozen=True)
+class SSTable:
+    """
+    In memory representation of Sorted String Table
+
+    List of references to data file ranges with statistics to enable pruning by primary key.
+    """
+
+    rows: List[SSTableRow]
+    """Sorted List of rows by key"""
+    min_key: Any
+    """Minimum observed primary key across all rows."""
+    max_key: Any
+    """Maximum observed primary key across all rows."""
+
+
+class SSTWriter(Protocol):
+    """
+    interface for writing SST files
+
+    Rows may be added iteratively. Input rows (within add_rows batch or across batches) MUST be sorted
+        by key_min
+    """
+
+    def write(self, file: OutputFile, rows: List[SSTableRow]) -> None:
+        """
+        Writes SST file
+        """
+        ...
+
+
+class SSTReader(Protocol):
+    """
+    interface for reading SST files
+    """
+
+    def read(self, file: InputFile) -> SSTable:
+        ...