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

deltacat/experimental/storage/rivulet/writer/memtable_dataset_writer.py

@@ -0,0 +1,305 @@
+from __future__ import annotations
+
+import threading
+from threading import Thread
+from typing import Any, List, Set, Protocol, TypeVar, Dict, Iterable
+
+from pyarrow import RecordBatch, Table
+from deltacat.storage.model.partition import PartitionLocator
+from deltacat.experimental.storage.rivulet.metastore.delta import (
+    ManifestIO,
+    DeltacatManifestIO,
+)
+
+from deltacat.experimental.storage.rivulet import Schema
+from deltacat.experimental.storage.rivulet.metastore.json_sst import JsonSstWriter
+from deltacat.experimental.storage.rivulet.serializer import (
+    MEMTABLE_DATA,
+    DataSerializer,
+)
+from deltacat.experimental.storage.rivulet.serializer_factory import (
+    DataSerializerFactory,
+)
+from deltacat.experimental.storage.rivulet.writer.dataset_writer import (
+    DatasetWriter,
+    DATA,
+)
+from deltacat.experimental.storage.rivulet.metastore.sst import SSTWriter
+from deltacat.experimental.storage.rivulet.fs.file_provider import FileProvider
+
+INPUT_ROW = TypeVar("INPUT_ROW")
+
+
+class Memtable(Protocol[INPUT_ROW]):
+    """
+    Protocol defining the interface for a memtable that can store and sort records of type T.
+    """
+
+    def add_record(self, record: INPUT_ROW) -> bool:
+        """
+        Add a record to the memtable.
+
+        Args:
+            record: The record to add of type INPUT_ROW
+
+        Returns:
+            bool: True if the memtable is full after adding the record, False otherwise
+        """
+        ...
+
+    def get_sorted_records(self, schema: Schema) -> MEMTABLE_DATA:
+        """
+        Get all records in the memtable in sorted order.
+
+        Returns:
+            List[T]: A list of sorted records
+        """
+        ...
+
+
+class DictMemTable(Memtable[Dict[str, Any]]):
+    """
+    Unit of in memory buffering of sorted records before records are written to file
+
+    TODO future improvements:
+        1. build b+ tree of record indexes on insertion
+             OR If we end up using arrow as intermediate format, we can use
+          pyarrow compute sort
+        2. Probably we will re-write in rust
+    """
+
+    def __init__(self, merge_key: str):
+        self.row_size = 0
+        self.merge_key = merge_key
+
+        self._records: List[Dict[str, Any]] = []
+        self.lock = threading.Lock()
+
+    def add_record(self, record: Dict[str, Any]):
+        with self.lock:
+            self._records.append(record)
+            self.row_size += 1
+
+        if self.row_size >= MemtableDatasetWriter.MAX_ROW_SIZE:
+            return True
+        return False
+
+    def get_sorted_records(self, schema: Schema) -> List[Dict[str, Any]]:
+        """
+        Gets sorted records
+
+        :return: iterator over sorted record
+        """
+        with self.lock:
+            self._records.sort(key=lambda x: x.__getitem__(self.merge_key))
+            return self._records
+
+
+class RecordBatchMemTable(Memtable[RecordBatch]):
+    """
+    Note that this will not respect max row size.
+    """
+
+    def __init__(self, merge_key: str):
+        self.row_size = 0
+        self.merge_key = merge_key
+
+        # list of full record batches in memtable
+        self._records_batches: List[RecordBatch] = []
+        self.lock = threading.Lock()
+
+    def add_record(self, record: RecordBatch):
+        with self.lock:
+            self._records_batches.append(record)
+            self.row_size += record.num_rows
+
+        if self.row_size >= MemtableDatasetWriter.MAX_ROW_SIZE:
+            return True
+        return False
+
+    def get_sorted_records(self, schema: Schema) -> Table:
+        """
+        Gets sorted records
+
+        :return: iterator over sorted record
+        """
+        with self.lock:
+            # Note that we are providing schema so that pyarrow does not infer it
+            table = Table.from_batches(self._records_batches, schema.to_pyarrow())
+            return table.sort_by(self.merge_key)
+
+
+class MemtableDatasetWriter(DatasetWriter):
+    # Note that this max row size is not respected when PyArrow RecordBatches are used
+    # In that case, the entire record batch is written within one memtable even if the row count overflows
+    MAX_ROW_SIZE = 1000000
+    """
+    Buffers data into rotating memtables. When a memtable reaches a certain size, it is flushed to disk and a new memtable is allocated
+
+    Uses DataWriter which will be format specific for writing data
+    Uses MetadataWriter for writing metadata
+
+    TODO Future Improvements
+    1. Maybe we should re-write this class in Rust (pending testing)
+    """
+
+    def __init__(
+        self,
+        file_provider: FileProvider,
+        schema: Schema,
+        locator: PartitionLocator,
+        file_format: str | None = None,
+        sst_writer: SSTWriter = None,
+        manifest_io: ManifestIO = None,
+    ):
+
+        if not sst_writer:
+            sst_writer = JsonSstWriter()
+        if not manifest_io:
+            manifest_io = DeltacatManifestIO(file_provider.uri, locator)
+
+        self.schema = schema
+
+        self.file_provider = file_provider
+        self.data_serializer: DataSerializer = DataSerializerFactory.get_serializer(
+            self.schema, self.file_provider, file_format
+        )
+        self.sst_writer = sst_writer
+        self.manifest_io = manifest_io
+
+        self._sst_files: Set[str] = set()
+        self.__curr_memtable = None
+        self.__open_memtables = []
+        self.__rlock = threading.RLock()
+        self.__open_threads: List[Thread] = []
+        self._locator = locator
+
+    def write_dict(self, record: Dict[str, Any]) -> None:
+
+        # Construct memtable if doesn't exist. If previous memtable wrong type, rotate
+        memtable_ctor = lambda: DictMemTable(self.schema.get_merge_key())
+        if not self.__curr_memtable:
+            self.__curr_memtable = memtable_ctor()
+        try:
+            isinstance(self.__curr_memtable, DictMemTable)
+        except TypeError:
+            self.__rotate_memtable(memtable_ctor)
+
+        # Write record(s). If memtable is full, rotate
+        if self.__curr_memtable.add_record(record):
+            self.__rotate_memtable(memtable_ctor)
+
+    def write_record_batch(self, record: RecordBatch) -> None:
+        # Construct memtable if doesn't exist. If previous memtable wrong type, rotate
+        memtable_ctor = lambda: RecordBatchMemTable(self.schema.get_merge_key())
+        if not self.__curr_memtable:
+            self.__curr_memtable = memtable_ctor()
+
+        try:
+            isinstance(self.__curr_memtable, RecordBatchMemTable)
+        except TypeError:
+            self.__rotate_memtable(memtable_ctor)
+
+        # Write record(s). If memtable is full, rotate
+        if self.__curr_memtable.add_record(record):
+            self.__rotate_memtable(memtable_ctor)
+
+    def write(self, data: DATA) -> None:
+        if isinstance(data, RecordBatch):
+            self.write_record_batch(data)
+        elif isinstance(data, Iterable):
+            for x in data:
+                if isinstance(x, dict):
+                    self.write_dict(x)
+                elif isinstance(x, RecordBatch):
+                    self.write_record_batch(x)
+                else:
+                    raise ValueError(
+                        f"Iterable contained unsupported type {type(x).__name__}."
+                        f" Supported data types to write are: {DATA}"
+                    )
+        else:
+            raise ValueError(
+                f"Unsupported data type {type(data).__name__}. Supported data types to write are: {DATA}"
+            )
+
+    def flush(self) -> str:
+        """
+        Explicitly flush any data and metadata and commit to dataset
+        """
+        self.__flush_memtable(self.__curr_memtable)
+        for thread in [t for t in self.__open_threads if t.is_alive()]:
+            thread.join()
+
+        manifest_location = self.__write_manifest_file()
+        self._sst_files.clear()
+
+        return manifest_location
+
+    def __enter__(self) -> Any:
+        """
+        Enter and exit method allows python "with" statement
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Closes all open memtables and ensures all data is flushed.
+        """
+        self.flush()
+        # return False to propogate up error messages
+        return False
+
+    def __rotate_memtable(self, memtable_constructor_closure):
+        """
+        Replace the active memtable
+        :return:
+        """
+        with self.__rlock:
+            self.__flush_memtable(self.__curr_memtable)
+            self.__curr_memtable = memtable_constructor_closure()
+            self.__open_memtables.append(self.__curr_memtable)
+
+            # Reap dead threads
+            self.__open_threads = [t for t in self.__open_threads if t.is_alive()]
+
+    def __flush_memtable(self, memtable):
+        thread = threading.Thread(target=self.__flush_memtable_async, args=(memtable,))
+        thread.start()
+        with self.__rlock:
+            self.__open_threads.append(thread)
+
+    def __flush_memtable_async(self, memtable: Memtable):
+        """
+        Flushes data and metadata for a given memtable
+        Called asynchronously in background thread
+        """
+        if not memtable:
+            return
+
+        sst_metadata_list = self.data_serializer.flush_batch(
+            memtable.get_sorted_records(self.schema)
+        )
+
+        # short circuit if no data/metadata written
+        if not sst_metadata_list:
+            with self.__rlock:
+                self.__open_memtables.remove(memtable)
+            return
+
+        # Write SST. Each memtable is going to have a dedicated L0 SST file because that is the unit at which
+        # we have contiguously sorted data
+        sst_file = self.file_provider.provide_l0_sst_file()
+
+        with self.__rlock:
+            self.sst_writer.write(sst_file, sst_metadata_list)
+            self._sst_files.add(sst_file.location)
+
+            if memtable in self.__open_memtables:
+                self.__open_memtables.remove(memtable)
+
+    def __write_manifest_file(self) -> str:
+        """
+        Write the manifest file to the filesystem at the given URI.
+        """
+        return self.manifest_io.write(list(self._sst_files), self.schema, 0)

deltacat/io/__init__.py

@@ -0,0 +1,13 @@
+from deltacat.io.reader.deltacat_read_api import read_deltacat
+from deltacat.io.datasource.deltacat_datasource import DeltacatReadType
+from deltacat.io.datasource.deltacat_datasource import (
+    METAFILE_DATA_COLUMN_NAME,
+    METAFILE_TYPE_COLUMN_NAME,
+)
+
+__all__ = [
+    "read_deltacat",
+    "DeltacatReadType",
+    "METAFILE_DATA_COLUMN_NAME",
+    "METAFILE_TYPE_COLUMN_NAME",
+]

deltacat/io/dataset/deltacat_dataset.py

@@ -0,0 +1,91 @@
+# Allow classes to use self-referencing Type hints in Python 3.7.
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Optional, cast
+
+import pyarrow as pa
+from ray.data import Dataset
+
+from deltacat.utils.url import DeltaCatUrl
+from deltacat.io.datasink.deltacat_datasink import DeltaCatDatasink
+
+
+class DeltaCatDataset(Dataset):
+    @staticmethod
+    def from_dataset(dataset: Dataset) -> DeltaCatDataset:
+        # cast to DeltacatDataset in-place since it only adds new methods
+        dataset.__class__ = DeltaCatDataset
+        return cast(DeltaCatDataset, dataset)
+
+    def write_deltacat(
+        self,
+        url: DeltaCatUrl,
+        *,
+        # if the source dataset only contains DeltaCAT metadata, then only copy the metadata to the destination... if it contains external source file paths, then register them in a new Delta.
+        metadata_only: bool = False,
+        # merge all deltas as part of the write operation
+        copy_on_write: Optional[bool] = False,
+        filesystem: Optional[pa.fs.S3FileSystem] = None,
+        try_create_dir: bool = True,
+        arrow_open_stream_args: Optional[Dict[str, Any]] = None,
+        arrow_parquet_args_fn: Callable[[], Dict[str, Any]] = lambda: {},
+        min_rows_per_file: Optional[int] = None,
+        ray_remote_args: Dict[str, Any] = None,
+        concurrency: Optional[int] = None,
+        **arrow_parquet_args,
+    ) -> None:
+        """Writes the dataset to files and commits DeltaCAT metadata indexing
+        the files written.
+
+        This is only supported for datasets convertible to Arrow records.
+        To control the number of files, use ``.repartition()``.
+
+        Unless a custom block path provider is given, the format of the output
+        files will be {uuid}_{block_idx}.{extension}, where ``uuid`` is a
+        unique id for the dataset.
+
+        The DeltaCAT manifest will be written to ``f"{path}/manifest``
+
+        Examples:
+            >>> ds.write_deltacat("s3://catalog/root/path")
+
+        Time complexity: O(dataset size / parallelism)
+
+        Args:
+            url: The path to the root directory where materialized files and
+                DeltaCAT manifest will be written.
+            filesystem: The filesystem implementation to write to. This should
+                be either a PyArrow S3FileSystem.
+            try_create_dir: Try to create all directories in destination path
+                if True. Does nothing if all directories already exist.
+            arrow_open_stream_args: kwargs passed to
+                pyarrow.fs.S3FileSystem.open_output_stream
+            filename_provider: FilenameProvider implementation
+                to write each dataset block to a custom output path.
+            arrow_parquet_args_fn: Callable that returns a dictionary of write
+                arguments to use when writing each block to a file. Overrides
+                any duplicate keys from arrow_parquet_args. This should be used
+                instead of arrow_parquet_args if any of your write arguments
+                cannot be pickled, or if you'd like to lazily resolve the write
+                arguments for each dataset block.
+            arrow_parquet_args: Options to pass to
+                pyarrow.parquet.write_table(), which is used to write out each
+                block to a file.
+        """
+        datasink = DeltaCatDatasink(
+            url,
+            metadata_only=metadata_only,
+            copy_on_write=copy_on_write,
+            arrow_parquet_args_fn=arrow_parquet_args_fn,
+            arrow_parquet_args=arrow_parquet_args,
+            min_rows_per_file=min_rows_per_file,
+            filesystem=filesystem,
+            try_create_dir=try_create_dir,
+            open_stream_args=arrow_open_stream_args,
+            dataset_uuid=self._uuid,
+        )
+        self.write_datasink(
+            datasink,
+            ray_remote_args=ray_remote_args,
+            concurrency=concurrency,
+        )

deltacat/io/datasink/deltacat_datasink.py

@@ -0,0 +1,207 @@
+import logging
+
+from collections import OrderedDict
+from typing import Dict, Any, Optional, List, Iterable
+
+from ray.data import Datasink
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource import WriteResult
+
+from ray.data.datasource.filename_provider import (
+    FilenameProvider,
+)
+
+from deltacat import logs
+
+from deltacat.constants import METAFILE_FORMAT_MSGPACK
+from deltacat.storage import Metafile
+from deltacat.io.datasource.deltacat_datasource import (
+    METAFILE_DATA_COLUMN_NAME,
+    METAFILE_TYPE_COLUMN_NAME,
+)
+from deltacat.utils.url import DeltaCatUrl, DeltaCatUrlWriter
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class CapturingBlockWritePathProvider(FilenameProvider):
+    """Delegating block write path provider that saves an ordered dictionary of
+    input keyword arguments for every block write path returned."""
+
+    def __init__(
+        self,
+        block_write_path_provider: FilenameProvider,
+        base_path: Optional[str] = None,
+    ):
+        self.base_path = base_path
+        self.block_write_path_provider = block_write_path_provider
+        self.write_path_kwargs: Dict[str, Dict[str, Any]] = OrderedDict()
+
+    def get_filename_for_block(
+        self,
+        block: Any,
+        task_index: int,
+        block_index: int,
+    ) -> str:
+        if self.base_path is None:
+            raise ValueError(
+                "Base path must be provided to CapturingBlockWritePathProvider",
+            )
+        return self._get_write_path_for_block(
+            base_path=self.base_path,
+            block=block,
+            block_index=block_index,
+        )
+
+    def _get_write_path_for_block(
+        self,
+        base_path: str,
+        *args,
+        **kwargs,
+    ) -> str:
+        filename = self.block_write_path_provider.get_filename_for_block(
+            *args,
+            **kwargs,
+        )
+        write_path = f"{base_path}/{filename}"
+        kwargs["base_path"] = base_path
+        self.write_path_kwargs[write_path] = kwargs
+        return write_path
+
+
+class DeltaCatWriteResult:
+    def __init__(self):
+        self.metadata = None
+        self.path = None
+        self.dataset_uuid = None
+        self.block_write_path_provider = None
+        self.content_type = None
+        self.content_encoding = None
+        self.filesystem = None
+
+
+class DeltaCatDatasink(Datasink[List[Metafile]]):
+    def __init__(
+        self,
+        url: DeltaCatUrl,
+        *,
+        metadata_only: bool = False,
+        copy_on_write: Optional[bool] = False,
+    ):
+        self._url = url
+        self._metadata_only = metadata_only
+        self._copy_on_write = copy_on_write
+
+    def on_write_start(self) -> None:
+        pass
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> List[Metafile]:
+        for block in blocks:
+            pa_table = BlockAccessor.for_block(block).to_arrow()
+            if (
+                METAFILE_DATA_COLUMN_NAME in pa_table.column_names
+                and METAFILE_TYPE_COLUMN_NAME in pa_table.column_names
+            ):
+                for pa_scalar in pa_table[METAFILE_DATA_COLUMN_NAME]:
+                    metafile_msgpack_bytes = pa_scalar.as_py()
+                    metafile = Metafile.deserialize(
+                        serialized=metafile_msgpack_bytes,
+                        meta_format=METAFILE_FORMAT_MSGPACK,
+                    )
+                    # TODO(pdames): Add `metafile` to writer as a kwarg instead
+                    #  of constructing a new URL with the metafile as input.
+                    writer_url = DeltaCatUrlWriter(self._url, metafile=metafile)
+                    # TODO(pdames): Run writes in order from catalog -> delta
+                    #  by truncating the URL down to just dc://{catalog-name}
+                    #  and rebuilding all path elements from there.
+                    writer_url.write(metafile)
+            else:
+                raise NotImplementedError(
+                    f"Expected {METAFILE_DATA_COLUMN_NAME} and "
+                    f"{METAFILE_TYPE_COLUMN_NAME} columns in the input block, "
+                    f"but found {pa_table.column_names}."
+                )
+
+    def on_write_complete(
+        self,
+        write_result: WriteResult[List[Metafile]],
+    ):
+        pass
+
+
+"""
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> List[ObjectRef[DeltacatWriteResult]]:
+        paths, filesystem = resolve_paths_and_filesystem(
+            self.path,
+            self.filesystem,
+        )
+        assert len(paths) == 1, f"Expected 1 write path, found {len(paths)}."
+        path = paths[0]
+        write_results = super().write(blocks)
+        # append a summary of this write operation in the last write result
+        metadata = [BlockAccessor.for_block(_).get_metadata() for _ in blocks]
+        rwr = DeltacatWriteResult()
+        rwr.metadata = metadata
+        rwr.path = path
+        rwr.dataset_uuid = self.dataset_uuid
+        rwr.block_write_path_provider = self.filename_provider
+        rwr.content_type = ContentType.PARQUET.value
+        rwr.content_encoding = ContentEncoding.IDENTITY.value
+        rwr.filesystem = filesystem
+        rwr_obj_ref = ray.put(rwr)
+        write_results.append(rwr_obj_ref)
+        return write_results
+
+    def on_write_complete(self, write_results: List[Any], **kwargs) -> None:
+        # TODO (pdames): time latency of this operation - overall s3 write times
+        #  are 2-3x pure read_parquet_fast() times
+        # restore the write operation summary from the last write result
+        result: DeltacatWriteResult = write_results[len(write_results) - 1]
+        write_path_args = result.block_write_path_provider.write_path_kwargs
+        blocks_written = len(write_path_args)
+        expected_blocks_written = len(result.metadata)
+        # TODO(pdames): Corner cases where mismatch is expected? Emply blocks?
+        #  Blocks filtered/split/merged to more/less write paths?
+        assert blocks_written == expected_blocks_written, (
+            f"Dataset write result validation failed. Found "
+            f"{blocks_written}/{expected_blocks_written} Dataset blocks "
+            f"written. Refusing to commit DeltaCAT Manifest."
+        )
+        manifest_entries = ManifestEntryList()
+        for block_idx, path in enumerate(write_path_args.keys()):
+            file_info = result.filesystem.get_file_info(path)
+            if file_info.type == pyarrow.fs.FileType.File:
+                content_length = file_info.size
+            else:
+                raise FileNotFoundError(ENOENT, strerror(ENOENT), path)
+            num_rows = result.metadata[block_idx].num_rows
+            source_content_length = result.metadata[block_idx].size_bytes
+            manifest_entry_meta = ManifestMeta.of(
+                int(num_rows) if num_rows is not None else None,
+                int(content_length) if content_length is not None else None,
+                result.content_type,
+                result.content_encoding,
+                int(source_content_length) if source_content_length else None,
+            )
+            parsed_url = parse_s3_url(path)
+            manifest_entry = ManifestEntry.of(
+                parsed_url.url,
+                manifest_entry_meta,
+            )
+            manifest_entries.append(manifest_entry)
+        manifest = Manifest.of(manifest_entries)
+        manifest_path = f"{result.path}/manifest"
+        logger.debug(f"Write succeeded for Dataset ID: {result.dataset_uuid}")
+        with result.filesystem.open_output_stream(manifest_path) as f:
+            f.write(json.dumps(manifest).encode("utf-8"))
+        logger.debug(f"Manifest committed to: {manifest_path}")
+"""