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

deltacat/storage/rivulet/dataset_executor.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from typing import List, Callable, Any
+
+from deltacat.storage.rivulet.field_group import FieldGroup
+from deltacat.storage.rivulet.mvp.Table import MvpTable
+from deltacat.storage.rivulet import Schema
+from deltacat.storage.rivulet.reader.data_scan import DataScan
+from deltacat.storage.rivulet.reader.dataset_metastore import DatasetMetastore
+from deltacat.storage.rivulet.reader.dataset_reader import DatasetReader
+from deltacat.storage.rivulet.reader.query_expression import QueryExpression
+
+
+class DatasetExecutor:
+    """
+    Executor class which runs operations such as select, map, take, save
+
+    This class may store intermediate state while it is executing.
+
+    LogicalPlan is responsible for constructor an executor and ordering operations appropriately
+    """
+
+    def __init__(
+        self,
+        field_groups: List[FieldGroup],
+        schema: Schema,
+        metastore: DatasetMetastore,
+    ):
+        self.effective_schema: Schema = schema.__deepcopy__()
+        self.field_groups = field_groups
+        self.output: MvpTable | None = None
+        self._metastore = metastore
+
+    def collect(self) -> MvpTable:
+        if not self.output:
+            self.output = self._read(self.effective_schema)
+        return self.output
+
+    def select(self, fields: List[str]) -> "DatasetExecutor":
+        """
+        Reads data and selects a subset of fields
+        Note that this implementation is super inefficient (does not push down filters to read, copies data to new MvpTable). That is OK since this will all be replaced
+        """
+        # Read data from original input sources if not already read
+        if not self.output:
+            self.output = self._read(self.effective_schema)
+        # Calculate effective schema and apply it to data
+        self.effective_schema.filter(fields)
+        self.output = MvpTable(
+            {
+                key: value
+                for key, value in self.output.data.items()
+                if key in self.effective_schema
+            },
+        )
+        return self
+
+    def map(self, transform: Callable[[Any], Any]) -> "DatasetExecutor":
+        raise NotImplementedError
+
+    def _read(self, schema: Schema) -> MvpTable:
+        """
+        Internal helper method to read data
+
+        TODO for now this is doing dumb in-memory implementation and later this is going to be replaced by rust library
+        """
+        if len(self.field_groups) == 1:
+            return self._read_as_mvp_table(schema, self.field_groups[0])
+        else:
+            ds1 = self._read_as_mvp_table(schema, self.field_groups[0])
+            ds2 = self._read_as_mvp_table(schema, self.field_groups[1])
+            merged = MvpTable.merge(ds1, ds2, schema.primary_key.name)
+            for i in range(2, len(self.field_groups)):
+                ds_i = self._read_as_mvp_table(schema, self.field_groups[i])
+                merged = MvpTable.merge(merged, ds_i, schema.primary_key.name)
+            return merged
+
+    def _read_as_mvp_table(self, schema: Schema, field_group: FieldGroup):
+        data = list(
+            DataScan(
+                schema, QueryExpression(), DatasetReader(self._metastore)
+            ).to_pydict()
+        )
+        output = {}
+        for key in schema.fields.keys():
+            output[key] = [d.get(key) for d in data]
+        return MvpTable(output)

deltacat/storage/rivulet/feather/__init__.py

@@ -0,0 +1,5 @@
+# TODO later on this will be moved to a dedicated package
+from deltacat.storage.rivulet.feather.file_reader import FeatherFileReader
+from deltacat.storage.rivulet.reader.reader_type_registrar import FileReaderRegistrar
+
+FileReaderRegistrar.register_reader("feather", FeatherFileReader)

deltacat/storage/rivulet/feather/file_reader.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import pyarrow.ipc
+from pyarrow import RecordBatch, RecordBatchFileReader
+
+from deltacat.storage.rivulet.fs.file_provider import FileProvider
+from deltacat.storage.rivulet.metastore.sst import SSTableRow
+from deltacat.storage.rivulet.reader.data_reader import (
+    RowAndKey,
+    FileReader,
+    FILE_FORMAT,
+)
+from deltacat.storage.rivulet.reader.pyarrow_data_reader import RecordBatchRowIndex
+from deltacat.storage.rivulet.schema.schema import Schema
+
+
+class FeatherFileReader(FileReader[RecordBatchRowIndex]):
+    """
+    Feather file reader. This class is not thread safe
+
+    This is mostly a copy-pasta from ParquetFileReader
+    TODO can consider abstracting code between this and ParquetFileReader
+    """
+
+    def __init__(
+        self,
+        sst_row: SSTableRow,
+        file_provider: FileProvider,
+        primary_key: str,
+        schema: Schema,
+    ):
+        self.sst_row = sst_row
+        self.input = file_provider.provide_input_file(sst_row.uri)
+
+        self.key = primary_key
+        self.feather_file = sst_row.uri
+
+        self.schema = schema
+
+        # Iterator from pyarrow iter_batches API call. Pyarrow manages state of traversal within parquet row groups
+
+        """
+        These variables keep state about where the iterator is current at. They are initialized in __enter__()
+        """
+        self._curr_batch: RecordBatch | None = None
+        self._feather_reader: RecordBatchFileReader | None = None
+        # Arrow only lets you read feather files chunk by chunk using
+        # RecordBatchFileReader.get_batch(index)
+        self._curr_batch_index = 0
+        self._curr_row_offset = 0
+        self._pk_col = None
+
+    def peek(self) -> Optional[RowAndKey[FILE_FORMAT]]:
+        """
+        Peek next record
+
+        Note that there is an edge case where peek() is called on the bounary between record batches
+        This only happens curr_row_offset == curr_batch.num_rows, meaning next() or peek() would need to advance
+        to the next record batch. When this happens, peek() increments _curr_batch and sets _curr_row_offset to 0
+
+        :return: Optional of RowAndPrimaryKey
+        """
+        if not self.__is_initialized():
+            raise RuntimeError(
+                "ParquetFileReader must be initialized with __enter__ before reading"
+            )
+
+        if self.__need_to_advance_record_batch():
+            try:
+                self.__advance_record_batch()
+            except StopIteration:
+                return None
+
+        pk = self._pk_col[self._curr_row_offset].as_py()
+        return RowAndKey(
+            RecordBatchRowIndex(self._curr_batch, self._curr_row_offset), pk
+        )
+
+    def __next__(self) -> RowAndKey[FILE_FORMAT]:
+        if not self.__is_initialized():
+            raise RuntimeError(
+                "ParquetFileReader must be initialized with __enter__ before reading"
+            )
+
+        if self.__need_to_advance_record_batch():
+            self.__advance_record_batch()
+            pk = self._pk_col[0].as_py()
+            return RowAndKey(RecordBatchRowIndex(self._curr_batch, 0), pk)
+        else:
+            pk = self._pk_col[self._curr_row_offset].as_py()
+            offset = self._curr_row_offset
+            self._curr_row_offset += 1
+            return RowAndKey(RecordBatchRowIndex(self._curr_batch, offset), pk)
+
+    def __enter__(self):
+        with self.input.open() as f:
+            self._feather_reader = pyarrow.ipc.RecordBatchFileReader(f)
+            self.__advance_record_batch()
+
+    def __exit__(self, __exc_type, __exc_value, __traceback):
+        self.close()
+        # return False to propagate up error messages
+        return False
+
+    def close(self):
+        # no op
+        return
+
+    def __is_initialized(self):
+        return self._curr_batch and self._pk_col
+
+    def __need_to_advance_record_batch(self):
+        return not self._curr_row_offset < self._curr_batch.num_rows
+
+    def __advance_record_batch(self):
+        """
+        Advance to next record batch
+        :raise StopIteration: If there are no more record batches
+        """
+        try:
+            self._curr_batch = self._feather_reader.get_batch(self._curr_batch_index)
+            self._curr_batch_index += 1
+            self._curr_row_offset = 0
+            self._pk_col = self._curr_batch[self.key]
+            # Filter the batch to only include fields in the schema
+            # Pyarrow select will throw a ValueError if the field is not in the schema
+            fields = [
+                field
+                for field in self.schema.keys()
+                if field in self._curr_batch.schema.names
+            ]
+            self._curr_batch = self._curr_batch.select(fields)
+        except ValueError:
+            raise StopIteration(f"Ended iteration at batch {self._curr_batch_index}")

deltacat/storage/rivulet/feather/serializer.py

@@ -0,0 +1,35 @@
+from typing import List
+
+import pyarrow as pa
+from pyarrow import feather
+
+from deltacat.storage.rivulet.metastore.sst import SSTableRow
+from deltacat.storage.rivulet import Schema
+from deltacat.storage.rivulet.arrow.serializer import ArrowSerializer
+from deltacat.storage.rivulet.fs.file_provider import FileProvider
+
+
+class FeatherDataSerializer(ArrowSerializer):
+    """
+    Feather data writer. Responsible for flushing rows to feather files and returning SSTable rows for any file(s) written
+
+    TODO Support recording byte range offsets. Deferring this for now
+        We may need to provide a wrapper class over fsspec which introspects how many bytes written
+        when .write is called on output stream
+    """
+
+    def __init__(self, file_provider: FileProvider, schema: Schema):
+        super().__init__(file_provider, schema)
+
+    def serialize(self, table: pa.Table) -> List[SSTableRow]:
+        file = self.file_provider.provide_data_file("feather")
+
+        with file.create() as outfile:
+            # Note that write_feather says that dest is a string, but it is really any object implementing write()
+            feather.write_feather(table, dest=outfile)
+
+        # Because we only write one row group, it only creates one SSTableRow
+        #  we may have more granular SST indexes for other formats
+        key_min, key_max = self._get_min_max_key(table)
+        # TODO need to populate byte offsets. For now, we are writing single files per SSTableRow
+        return [SSTableRow(key_min, key_max, file.location, 0, 0)]

deltacat/storage/rivulet/fs/file_provider.py

@@ -0,0 +1,105 @@
+import posixpath
+import time
+from typing import List, Generator
+
+from deltacat.storage.model.partition import PartitionLocator
+from deltacat.storage.rivulet.fs.file_store import FileStore
+from deltacat.storage.rivulet.fs.input_file import InputFile
+from deltacat.storage.rivulet.fs.output_file import OutputFile
+from deltacat.utils.metafile_locator import _find_partition_path
+
+
+class FileProvider:
+    """
+    Manages the generation of URIs for data and metadata files and facilitates the creation of files at those URIs.
+    All files are generated relative to the root of the storage location.
+
+    This class is inspired by the Iceberg `LocationProvider` and provides methods
+    to generate paths for various types of files (e.g., data files, SSTs, and manifests)
+    while maintaining a clear structure within the dataset.
+
+    TODO (deltacat): FileProvider will be replaced/refactored once we are able to integrate with Deltacat.
+    TODO: Incorporate additional file naming conventions, such as including
+          partitionId, taskId, and operationId, to improve traceability and
+          idempotency.
+    """
+
+    uri: str
+
+    def __init__(self, uri: str, locator: PartitionLocator, file_store: FileStore):
+        """
+        Initializes the file provider.
+
+        param: uri: Base URI of the dataset.
+        param: file_store: FileStore instance for creating and reading files.
+        """
+        self.uri = uri
+        self._locator = locator
+        self._file_store = file_store
+
+    def provide_data_file(self, extension: str) -> OutputFile:
+        """
+        Creates a new data file.
+
+        TODO: Ensure storage interface can provide data files.
+
+        param: extension: File extension (e.g., "parquet").
+        returns: OutputFile instance pointing to the created data file.
+        """
+        partition_path = _find_partition_path(self.uri, self._locator)
+        uri = posixpath.join(
+            partition_path, "data", f"{int(time.time_ns())}.{extension}"
+        )
+        return self._file_store.create_output_file(uri)
+
+    def provide_l0_sst_file(self) -> OutputFile:
+        """
+        Creates a new L0 SST file.
+
+        TODO: Ensure storage interface can provide sst files.
+
+        returns: OutputFile instance pointing to the created SST file.
+        """
+        partition_path = _find_partition_path(self.uri, self._locator)
+        uri = posixpath.join(
+            partition_path, "metadata", "ssts", "0", f"{int(time.time_ns())}.json"
+        )
+        return self._file_store.create_output_file(uri)
+
+    def provide_input_file(self, uri: str) -> InputFile:
+        """
+        Reads an existing file.
+
+        param sst_uri: URI of the file to read.
+        returns: InputFile instance for the specified URI.
+        """
+        return self._file_store.create_input_file(uri)
+
+    def provide_manifest_file(self) -> OutputFile:
+        """
+        Creates a new manifest file.
+
+        returns: OutputFile instance pointing to the created manifest file.
+        """
+        uri = f"{self.uri}/metadata/manifests/{int(time.time_ns())}.json"
+        return self._file_store.create_output_file(uri)
+
+    def get_sst_scan_directories(self) -> List[str]:
+        """
+        Retrieves SST scan directories.
+
+        returns: List of directories containing SSTs.
+        """
+        partition_path = _find_partition_path(self.uri, self._locator)
+        return [f"{partition_path}/metadata/ssts/0/"]
+
+    def generate_sst_uris(self) -> Generator[InputFile, None, None]:
+        """
+        Generates all SST URIs.
+
+        returns: Generator of InputFile instances for SSTs.
+        """
+        sst_directories = self.get_sst_scan_directories()
+        for directory in sst_directories:
+            for file in self._file_store.list_files(directory):
+                yield file

deltacat/storage/rivulet/fs/file_store.py

@@ -0,0 +1,130 @@
+from typing import Tuple, Iterator, Optional
+from pyarrow.fs import FileSystem, FileType, FileSelector
+
+# TODO(deltacat): Rely on deltacat implementation to resolve path and filesystem.
+from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+
+from deltacat.storage.rivulet.fs.input_file import FSInputFile
+from deltacat.storage.rivulet.fs.output_file import FSOutputFile
+
+
+class FileStore:
+    """
+    Manages the filesystem and low-level file operations.
+    This class is designed to work with any filesystem supported by PyArrow; local, S3, HDFS, GCP,
+    and other fsspec-compatible filesystems.
+
+    TODO: Add better error consolidation between filesystems. Will be handled by deltacat implementation?
+
+    method: `filesystem`: Resolves and normalizes a given path and filesystem.
+    method: `file_exists`: Checks if a file exists at a given URI.
+    method: `create_file`: Creates a new file for writing at a specified URI.
+    method: `read_file`: Reads an existing file from a specified URI.
+    method: `list_files`: Lists all files within a specified directory URI.
+    """
+
+    def __init__(self, path: str, filesystem: Optional[FileSystem] = None):
+        """
+        Serves as the source of truth for all file operations, ensuring that
+        all paths and operations are relative to the specified filesystem,
+        providing consistency and compatibility across fsspec supported backends.
+
+        TODO (deltacat): maybe rely on deltacat catalog as a source of truth for rivulet filesystem.
+
+        param: path (str): The base URI or path for the filesystem.
+        param: filesystem (FileSystem): A PyArrow filesystem instance.
+        """
+        _, fs = FileStore.filesystem(path, filesystem)
+        self.filesystem = filesystem or fs
+
+    @staticmethod
+    def filesystem(
+        path: str, filesystem: Optional[FileSystem] = None
+    ) -> Tuple[str, FileSystem]:
+        """
+        Resolves and normalizes the given path and filesystem.
+
+        param: path (str): The URI or path to resolve.
+        param: filesystem (Optional[FileSystem]): An optional filesystem instance.
+        returns: Tuple[str, FileSystem]: The normalized path and filesystem.
+        raises: AssertionError: If multiple paths are resolved.
+        """
+        paths, filesystem = _resolve_paths_and_filesystem(
+            paths=path, filesystem=filesystem
+        )
+        assert len(paths) == 1, "Multiple paths not supported"
+        return paths[0], filesystem
+
+    def file_exists(
+        self, data_uri: str, filesystem: Optional[FileSystem] = None
+    ) -> bool:
+        """
+        Checks if a file exists at the specified URI.
+
+        param: data_uri (str): The URI of the file to check.
+        param: filesystem (Optional[FileSystem]): Filesystem to use. Defaults to the instance filesystem.
+        returns: bool: True if the file exists, False otherwise.
+        """
+        path, filesystem = FileStore.filesystem(data_uri, filesystem or self.filesystem)
+        return filesystem.get_file_info(path).type != FileType.NotFound
+
+    def create_output_file(
+        self, data_uri: str, filesystem: Optional[FileSystem] = None
+    ) -> FSOutputFile:
+        """
+        Creates a new output file for writing at the specified URI.
+
+        param: data_uri (str): The URI where the file will be created.
+        param: filesystem (Optional[FileSystem]): Filesystem to use. Defaults to the instance filesystem.
+        returns: FSOutputFile: An object for writing to the file.
+        raises: IOError: If file creation fails.
+        """
+        try:
+            path, filesystem = FileStore.filesystem(
+                data_uri, filesystem or self.filesystem
+            )
+            return FSOutputFile(path, filesystem)
+        except Exception as e:
+            raise IOError(f"Failed to create file '{data_uri}': {e}")
+
+    def create_input_file(
+        self, data_uri: str, filesystem: Optional[FileSystem] = None
+    ) -> FSInputFile:
+        """
+        Create a new input file for reading at the specified URI.
+
+        param: data_uri (str): The URI of the file to read.
+        param: filesystem (Optional[FileSystem]): Filesystem to use. Defaults to the instance filesystem.
+        returns: FSInputFile: An object for reading the file.
+        raises: IOError: If file reading fails.
+        """
+        try:
+            path, filesystem = FileStore.filesystem(
+                data_uri, filesystem or self.filesystem
+            )
+            return FSInputFile(path, filesystem)
+        except Exception as e:
+            raise IOError(f"Failed to read file '{data_uri}': {e}")
+
+    def list_files(
+        self, data_uri: str, filesystem: Optional[FileSystem] = None
+    ) -> Iterator[FSInputFile]:
+        """
+        Lists all files in the specified directory URI.
+
+        param: data_uri (str): The URI of the directory to list files from.
+        param: filesystem (Optional[FileSystem]): Filesystem to use. Defaults to the instance filesystem.
+        returns: Iterator[FSInputFile]: An iterator of FSInputFile objects representing the files.
+        raises: IOError: If listing files fails.
+        """
+        try:
+            path, filesystem = FileStore.filesystem(
+                data_uri, filesystem or self.filesystem
+            )
+            file_info = filesystem.get_file_info(FileSelector(path, recursive=False))
+
+            for file in file_info:
+                if file.type == FileType.File:
+                    yield FSInputFile(file.path, filesystem)
+        except Exception as e:
+            raise IOError(f"Failed to list files in '{data_uri}': {e}")

deltacat/storage/rivulet/fs/input_file.py

@@ -0,0 +1,76 @@
+from contextlib import contextmanager
+import io
+from abc import ABC, abstractmethod
+from typing import Protocol
+
+from pyarrow.fs import FileSystem, FileType
+
+
+class InputStream(Protocol):
+    """A protocol with a subset of IOBase for file-like input objects"""
+
+    @abstractmethod
+    def read(self, size: int = -1) -> bytes:
+        ...
+
+    @abstractmethod
+    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
+        ...
+
+    @abstractmethod
+    def tell(self) -> int:
+        ...
+
+    @abstractmethod
+    def close(self) -> None:
+        ...
+
+    def __enter__(self) -> "InputStream":
+        ...
+
+    @abstractmethod
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        ...
+
+
+class InputFile(ABC):
+    """Abstraction for interacting with input files"""
+
+    def __init__(self, location: str):
+        self._location = location
+
+    @property
+    def location(self) -> str:
+        return self._location
+
+    @abstractmethod
+    def exists(self) -> bool:
+        """Return whether the location exists.
+
+        Raises:
+            PermissionError: If this has insufficient permissions to access the file at location.
+        """
+
+    @abstractmethod
+    def open(self) -> InputStream:
+        """Return a file-like object for input
+
+        Raises:
+            FileNotFoundError: If the file does not exist at self.location.
+            PermissionError: If this has insufficient permissions to access the file at location.
+        """
+
+
+class FSInputFile(InputFile):
+    def __init__(self, location: str, fs: FileSystem):
+        self._location = location
+        self.fs = fs
+
+    def exists(self) -> bool:
+        file_info = self.fs.get_file_info(self._location)
+        return file_info.type != FileType.NotFound
+
+    @contextmanager
+    def open(self, seekable: bool = False):
+        with self.fs.open_input_file(self._location) as input_stream:
+            yield input_stream

deltacat/storage/rivulet/fs/output_file.py

@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+import posixpath
+from typing import Protocol
+
+from pyarrow.fs import FileSystem, FileType
+
+from deltacat.storage.rivulet.fs.input_file import FSInputFile, InputFile
+
+
+class OutputStream(Protocol):  # pragma: no cover
+    """A protocol with a subset of IOBase for file-like output objects"""
+
+    @abstractmethod
+    def write(self, b: bytes) -> int:
+        ...
+
+    @abstractmethod
+    def close(self) -> None:
+        ...
+
+    @abstractmethod
+    def __enter__(self) -> "OutputStream":
+        ...
+
+    @abstractmethod
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        ...
+
+
+class OutputFile(ABC):
+    """Abstraction for interacting with output files"""
+
+    def __init__(self, location: str):
+        self._location = location
+
+    @property
+    def location(self) -> str:
+        return self._location
+
+    @abstractmethod
+    def exists(self) -> bool:
+        """Return whether the location exists.
+
+        Raises:
+            PermissionError: If this has insufficient permissions to access the file at location.
+        """
+
+    @abstractmethod
+    def to_input_file(self) -> InputFile:
+        """Return an InputFile for this output file's location"""
+
+    @abstractmethod
+    def create(self) -> OutputStream:
+        """Return a file-like object for output
+
+        TODO: overwrite protection (FileExistsError?)
+        Raises:
+            PermissionError: If this has insufficient permissions to access the file at location.
+        """
+
+
+class FSOutputFile(OutputFile):
+    def __init__(self, location: str, fs: FileSystem):
+        self._location = location
+        self.fs = fs
+
+    def exists(self) -> bool:
+        file_info = self.fs.get_file_info(self._location)
+        return file_info.type != FileType.NotFound
+
+    def to_input_file(self) -> "FSInputFile":
+        return FSInputFile(self._location, self.fs)
+
+    @contextmanager
+    def create(self):
+        """Create and open the file for writing."""
+        try:
+            parent_dir = posixpath.dirname(self._location)
+            if parent_dir:  # Check if there's a parent directory to create
+                self.fs.create_dir(parent_dir, recursive=True)
+
+            with self.fs.open_output_stream(self._location) as output_stream:
+                yield output_stream
+        except Exception as e:
+            raise IOError(f"Failed to create or write to file '{self._location}': {e}")