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

deltacat/{storage → experimental/storage}/rivulet/writer/memtable_dataset_writer.py

@@ -6,15 +6,26 @@ from typing import Any, List, Set, Protocol, TypeVar, Dict, Iterable
 
 from pyarrow import RecordBatch, Table
 from deltacat.storage.model.partition import PartitionLocator
-from deltacat.storage.rivulet.metastore.delta import ManifestIO, DeltacatManifestIO
-
-from deltacat.storage.rivulet import Schema
-from deltacat.storage.rivulet.metastore.json_sst import JsonSstWriter
-from deltacat.storage.rivulet.serializer import MEMTABLE_DATA, DataSerializer
-from deltacat.storage.rivulet.serializer_factory import DataSerializerFactory
-from deltacat.storage.rivulet.writer.dataset_writer import DatasetWriter, DATA
-from deltacat.storage.rivulet.metastore.sst import SSTWriter
-from deltacat.storage.rivulet.fs.file_provider import FileProvider
+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")
 

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}")
+"""