deltacat 1.1.36__py3-none-any.whl → 2.0.0b2__py3-none-any.whl

deltacat/storage/rivulet/reader/dataset_metastore.py

@@ -0,0 +1,178 @@
+import logging
+import posixpath
+from typing import Generator, Optional
+
+import pyarrow
+import pyarrow.fs
+
+from deltacat.storage import Delta
+from deltacat.storage.model.partition import PartitionLocator
+from deltacat.storage.rivulet.fs.file_provider import FileProvider
+from deltacat.utils.filesystem import resolve_path_and_filesystem
+from deltacat.storage.rivulet.metastore.json_sst import JsonSstReader
+from deltacat.storage.rivulet.metastore.delta import (
+    ManifestIO,
+    DeltaContext,
+    RivuletDelta,
+    DeltacatManifestIO,
+)
+from deltacat.storage.rivulet.metastore.sst import SSTReader, SSTable
+from deltacat.utils.metafile_locator import _find_table_path
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class ManifestAccessor:
+    """Accessor for retrieving a manifest's SSTable entities."""
+
+    def __init__(
+        self, delta: RivuletDelta, file_provider: FileProvider, sst_reader: SSTReader
+    ):
+        self.manifest: RivuletDelta = delta
+        self.file_provider: FileProvider = file_provider
+        self._sst_reader = sst_reader
+
+    @property
+    def context(self) -> DeltaContext:
+        return self.manifest.context
+
+    def generate_sstables(self) -> Generator[SSTable, None, None]:
+        """
+        Generate the SortedString Tables from this Manifest
+
+        :return a generator of SSTables for this manifest
+        """
+        for sst_uri in self.manifest.sst_files:
+            sst_file = self.file_provider.provide_input_file(sst_uri)
+            yield self._sst_reader.read(sst_file)
+
+
+class DatasetMetastore:
+    """
+    Metastore implementation for manifests stored on a filesystem
+
+    TODO this will be replaced with Deltacat Storage interface - https://github.com/ray-project/deltacat/issues/477
+    """
+
+    def __init__(
+        self,
+        # URI at which we expect to find deltas
+        delta_root_uri: str,
+        file_provider: FileProvider,
+        locator: PartitionLocator,
+        *,
+        manifest_io: ManifestIO = None,
+        sst_reader: SSTReader = None,
+    ):
+        self._min_key = None
+        self._max_key = None
+        self.delta_root_uri = delta_root_uri
+        self.file_provider = file_provider
+        self.manifest_io = manifest_io or DeltacatManifestIO(delta_root_uri, locator)
+        self.sst_reader = sst_reader or JsonSstReader()
+        self.locator = locator
+
+    def _get_delta(
+        self, delta_dir: str, filesystem: pyarrow.fs.FileSystem
+    ) -> Optional[RivuletDelta]:
+        """
+        Find the latest revision in a delta directory.
+
+        param: delta_dir: The directory containing the revisions.
+        param: filesystem: The filesystem to search for the revisions.
+        returns: The latest revision as a RivuletDelta.
+        """
+        rev_directory = posixpath.join(delta_dir, "rev")
+        revisions = filesystem.get_file_info(
+            pyarrow.fs.FileSelector(rev_directory, allow_not_found=True)
+        )
+
+        if not revisions:
+            logger.warning(f"No revision files found in {rev_directory}")
+            return None
+
+        # Take lexicographical max to find the latest revision
+        latest_revision = max(revisions, key=lambda f: f.path)
+
+        return (
+            RivuletDelta.of(Delta.read(latest_revision.path))
+            if latest_revision
+            else None
+        )
+
+    def generate_manifests(self) -> Generator[ManifestAccessor, None, None]:
+        """
+        Generate all manifests within the Metastore
+        NOTE: this will be replaced by deltacat storage API.
+
+        TODO: Generate partition path using Deltacat Storage interface.
+
+        param: delta_root_uri: The URI at which we expect to find deltas.
+        returns: a generator of ManifestAccessors for all manifests in the dataset.
+        """
+
+        root_path, filesystem = resolve_path_and_filesystem(self.delta_root_uri)
+
+        partition_path = posixpath.join(
+            _find_table_path(root_path, filesystem),
+            self.locator.table_version,
+            self.locator.stream_id,
+            self.locator.partition_id,
+        )
+
+        partition_info = filesystem.get_file_info(partition_path)
+
+        if partition_info.type != pyarrow.fs.FileType.Directory:
+            logger.debug(f"Partition directory {partition_path} not found. Skipping.")
+            return
+
+        # Locate "rev" directory inside the partition
+        rev_directory = posixpath.join(partition_path, "rev")
+        rev_info = filesystem.get_file_info(rev_directory)
+
+        if rev_info.type != pyarrow.fs.FileType.Directory:
+            logger.debug(f"Revision directory {rev_directory} not found. Skipping.")
+            return
+
+        # Fetch all delta directories inside the partition
+        delta_dirs = filesystem.get_file_info(
+            pyarrow.fs.FileSelector(
+                partition_path, allow_not_found=True, recursive=False
+            )
+        )
+
+        delta_dirs = [
+            delta
+            for delta in delta_dirs
+            if delta.type == pyarrow.fs.FileType.Directory and delta.base_name.isdigit()
+        ]
+
+        for delta_dir in delta_dirs:
+            rivulet_delta = self._get_delta(delta_dir.path, filesystem)
+            if rivulet_delta:
+                yield ManifestAccessor(
+                    rivulet_delta, self.file_provider, self.sst_reader
+                )
+
+    def get_min_max_keys(self):
+        """
+        Compute and cache the minimum and maximum keys in the dataset.
+
+        returns: a tuple of the minimum and maximum keys in the dataset
+        """
+        if self._min_key is not None and self._max_key is not None:
+            return (self._min_key, self._max_key)
+
+        min_key = None
+        max_key = None
+        for manifest_accessor in self.generate_manifests():
+            for sstable in manifest_accessor.generate_sstables():
+                if min_key is None or sstable.min_key < min_key:
+                    min_key = sstable.min_key
+                if max_key is None or sstable.max_key > max_key:
+                    max_key = sstable.max_key
+
+        self._min_key = min_key
+        self._max_key = max_key
+        return (min_key, max_key)

deltacat/storage/rivulet/reader/dataset_reader.py

@@ -0,0 +1,156 @@
+import logging
+from typing import Generator, Optional, Set, Type, TypeVar, Any
+
+from deltacat.storage.model.shard import Shard
+from deltacat.storage.rivulet.metastore.sst import SSTableRow, SSTable
+from deltacat.storage.rivulet.metastore.sst_interval_tree import (
+    BlockIntervalTree,
+    OrderedBlockGroups,
+)
+from deltacat.storage.rivulet.reader.block_scanner import BlockScanner
+from deltacat.storage.rivulet.reader.dataset_metastore import (
+    DatasetMetastore,
+    ManifestAccessor,
+)
+from deltacat.storage.rivulet.reader.query_expression import QueryExpression
+from deltacat.storage.rivulet import Schema
+
+# The type of data returned to reader
+T = TypeVar("T")
+
+
+class DatasetReader:
+    """
+    DatasetReader is an internal class used to execute a scan
+
+    TODO - Currently, this reader is limited to reading a single field group
+       The next CR will fast follow to modify this to read and zipper multiple field groups
+
+    TODO currently this assumes all SST files are L0 files with overlapping key ranges
+        Future CR will support L1+ SSTs
+    """
+
+    BLOCK_READER_POOL_SIZE = 8
+
+    def __init__(self, metastore: DatasetMetastore):
+        self.metastore: DatasetMetastore = metastore
+        self.block_scanner = BlockScanner(self.metastore)
+
+    def scan(
+        self,
+        schema: Schema,
+        deserialize_to: Type[T],
+        query: QueryExpression[Any](),
+        shard: Optional[Shard] = None,
+    ) -> Generator[T, None, None]:
+        """
+        Scan records given query and deserialize to desired memory output format
+
+        # TODO handle "partial schema" use case, in which the query schema is a subset of full schema
+
+        # TODO this is where we will do the ziper merge when we support multiple field groups
+        # for each SST row which may overlap key range, read data chunk
+        # we will later improve and parallelize this when we do zipper merge work
+        """
+
+        # Read manifests and differentiate between "full schema" and "zipper merge" use case
+        manifests = set(self.metastore.generate_manifests())
+        schemas = set([manifest.context.schema for manifest in manifests])
+        levels = set([manifest.context.level for manifest in manifests])
+        # Must zipper if there are multiple schemas
+        cannot_avoid_zipper = len(schemas) > 1
+        # Must zipper if L0 is involved or if manifests span multiple levels
+        cannot_avoid_zipper |= 0 in levels or len(levels) > 0
+
+        if cannot_avoid_zipper:
+            logging.info(f"Done scanning manifests. Can avoid zipper-merge")
+            for scan_result in self.__scan_with_zipper(
+                schema, deserialize_to, manifests, query, shard=shard
+            ):
+                yield scan_result
+        else:
+            logging.info(f"Done scanning manifests. Must perform zipper-merge")
+            for scan_result in self.__scan_no_zipper(
+                schema, deserialize_to, manifests, query, shard=shard
+            ):
+                yield scan_result
+
+    def __scan_no_zipper(
+        self,
+        schema: Schema,
+        deserialize_to: Type[T],
+        manifests: Set[ManifestAccessor],
+        query: QueryExpression[Any](),
+        shard: Optional[Shard] = None,
+    ) -> Generator[T, None, None]:
+        # Build final query using user query and shard boundaries (ensures only blocks in shard and query range are read).
+        # TODO: improve query expression implementation to have a builder of some sort.
+        query = QueryExpression().with_shard(query, shard)
+        # Map manifests to all SST rows which match query
+        matching_sst_rows: Set[SSTableRow] = {
+            row
+            for manifest in manifests
+            for table in manifest.generate_sstables()
+            for row in self.__load_sst_rows(table, query)
+        }
+
+        for result_row in self.block_scanner.scan(
+            schema, deserialize_to, matching_sst_rows, query
+        ):
+            yield result_row
+
+    def __scan_with_zipper(
+        self,
+        schema: Schema,
+        deserialize_to: Type[T],
+        manifests: Set[ManifestAccessor],
+        query: QueryExpression[Any](),
+        shard: Optional[Shard] = None,
+    ) -> Generator[T, None, None]:
+        # Build final query using user query and shard boundaries (ensures only blocks in shard and query range are read).
+        # TODO: improve query expression implementation to have a builder of some sort.
+        query = QueryExpression().with_shard(query, shard)
+        # Build interval tree from manifests and plan scan
+        sst_interval_tree = BlockIntervalTree()
+        for manifest in manifests:
+            for table in manifest.generate_sstables():
+                rows = self.__load_sst_rows(table, query)
+                sst_interval_tree.add_sst_rows(rows, manifest.context)
+
+        scan_block_groups: OrderedBlockGroups = (
+            sst_interval_tree.get_sorted_block_groups(query.min_key, query.max_key)
+        )
+        for result_row in self.block_scanner.scan_with_zipper(
+            schema, deserialize_to, scan_block_groups, query
+        ):
+            yield result_row
+
+    def __load_sst_rows(
+        self, table: SSTable, query: QueryExpression
+    ) -> Set[SSTableRow]:
+        # Short circuit table if there isn't any overlap with min and max
+        if not self.__overlaps_primary_key_range(query, table.min_key, table.max_key):
+            return set()
+        return {
+            r
+            for r in table.rows
+            if self.__overlaps_primary_key_range(query, r.key_min, r.key_max)
+        }
+
+    def __overlaps_primary_key_range(
+        self, query: QueryExpression, min_key, max_key
+    ) -> bool:
+        """
+        Helper method to check whether a query expression has overlap with a primary key range
+        """
+        # If no PK range set, the query is across all primary keys, so return true
+        if not query.key_range:
+            return True
+
+        query_start, query_end = query.key_range
+        if query_end < min_key:
+            return False
+        elif query_start > max_key:
+            return False
+        else:
+            return True

deltacat/storage/rivulet/reader/pyarrow_data_reader.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from typing import Generator, Dict, Type, NamedTuple, List
+
+from pyarrow import RecordBatch
+
+from deltacat.storage.rivulet.reader.data_reader import DataReader, MEMORY_FORMAT
+import pyarrow as pa
+
+
+class RecordBatchRowIndex(NamedTuple):
+    """
+    Named tuple for a record batch with an index into a specific row
+    Note that record batches store data by column, so the row index should be
+    used to index into each column array
+    """
+
+    batch: RecordBatch
+    row_index: int
+
+
+class ArrowDataReader(DataReader[RecordBatchRowIndex]):
+    """
+    Parquet reader to iteratively load records from parquet files
+    """
+
+    def deserialize_records(
+        self, record: RecordBatchRowIndex, output_type: Type[MEMORY_FORMAT]
+    ) -> Generator[MEMORY_FORMAT, None, None]:
+        """
+        Deserialize records into the specified format.
+
+        Note that output_type gets set based on what a DataScan converts results to,
+            e.g. to_arrow, to_dict
+
+        :param record: Input data (generated by generate_records method)
+        :param output_type: Type to deserialize into
+        :returns: A generator yielding records of the specified type.
+        """
+        batch, row_idx = record[0].batch, record[0].row_index
+
+        if output_type == Dict:
+            yield {
+                column: batch.column(column_idx)[row_idx].as_py()
+                for column_idx, column in enumerate(batch.column_names)
+            }
+
+        elif output_type == RecordBatch:
+            # only yield full record batch if row_idx is 0.
+            # TODO this logic will need to change in zipper use case across data formats
+            if row_idx == 0:
+                yield batch
+
+    def join_deserialize_records(
+        self,
+        records: List[RecordBatchRowIndex],
+        output_type: Type[MEMORY_FORMAT],
+        join_key: str,
+    ) -> Generator[MEMORY_FORMAT, None, None]:
+        """
+        Deserialize records into the specified format.
+
+        Note that output_type gets set based on what a DataScan converts results to,
+            e.g. to_arrow, to_dict
+
+        :param records: Input data (generated by generate_records method)
+        :param output_type: Type to deserialize into
+        :returns: A generator yielding records of the specified type.
+        """
+
+        if output_type == Dict:
+            yield self.__join_records_as_dict(records)
+        elif output_type == RecordBatch:
+            yield self.__join_records_as_record_batch(records, join_key)
+
+    @staticmethod
+    def __join_records_as_dict(records: List[RecordBatchRowIndex]) -> Dict[str, any]:
+        """
+        Deserialize records into a PyDict
+
+        :param records: input record data
+        :returns: A PyDict that's joined the given records around the primary key.
+        """
+        batch: RecordBatch
+        row_idx: int
+        out = {}
+        for r in records:
+            batch, row_idx = r
+            # Note this stomps over join key but that's OK
+            for column_idx, column in enumerate(batch.schema.names):
+                col = batch.column(column_idx)
+                if len(col) <= row_idx:
+                    raise IndexError(
+                        f"row index {row_idx} out of bounds for column {column} with length {len(col)}"
+                    )
+
+                out.update({column: col[row_idx].as_py()})
+        return out
+
+    @staticmethod
+    def __join_records_as_record_batch(
+        records: List[RecordBatchRowIndex], join_key: str
+    ) -> RecordBatch:
+        """
+        Deserialize records into a RecordBatch
+
+        :param records: input record data
+        :returns: RecordBatch that's inner-joined the given records around the primary key.
+        """
+        batch: RecordBatch
+        row_idx: int
+        out: pa.Table | None = None
+        for record in records:
+            batch, row_idx = record
+            batch_slice: RecordBatch = batch.slice(row_idx, 1)
+            if not out:
+                out = pa.Table.from_batches([batch_slice])
+            else:
+                table2 = pa.Table.from_batches([batch_slice])
+                out = out.join(table2, keys=join_key, join_type="inner")
+        return out.to_batches()[0]

deltacat/storage/rivulet/reader/query_expression.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import typing
+from typing import Optional
+
+from deltacat.storage.model.shard import Shard
+
+T = typing.TypeVar("T")  # Type of primary key in query expression. Must be comparable
+
+
+class QueryExpression(typing.Generic[T]):
+    """
+    Top level class for creating representing queries on a riv dataset.
+
+    For now, this is a minimal implementation which just allows for different predicates.
+
+    FUTURE IMPROVEMENTS
+    1. Support builder using operator overloading or fluent builder pattern,e.g.
+       (operator overloading) query = Column("Foo") < 10 & Column("PK")==100
+       (fluent interface) query = builder.column("colA").less_than(10)
+          .and_()
+          .column("PK").equals(100)
+          .build()
+
+    2. Support better push down predicate integration end to end. Specifically,
+        scan operation will need to return which query predicates were honored
+    """
+
+    def __init__(self):
+        self.key_range: Optional[(T, T)] = None
+
+    def with_key(self, val: T) -> QueryExpression:
+        """
+        Syntactic sugar for setting key range to a single value
+        """
+        if self.key_range:
+            raise ValueError(
+                f"Query expression already has set key range to: {self.key_range}"
+            )
+        self.key_range = (val, val)
+        return self
+
+    def with_range(self, bound1: T, bound2: T) -> QueryExpression:
+        if self.key_range:
+            raise ValueError(f"Key range already set to {self.key_range}")
+        self.key_range = tuple(sorted([bound1, bound2]))
+        return self
+
+    @staticmethod
+    def with_shard(query: Optional[QueryExpression], shard: Shard):
+        """
+        Generate a query expression that accounts for the shard boundaries.
+        Shard boundaries are inclusive and mark the outer bounds of the query.
+        """
+        if shard is None:
+            return query
+
+        if query.key_range is None:
+            return QueryExpression().with_range(shard.min_key, shard.max_key)
+
+        min_key = shard.min_key
+        max_key = shard.max_key
+
+        if min_key > query.min_key:
+            min_key = query.min_key
+
+        if max_key < query.max_key:
+            max_key = query.max_key
+
+        return QueryExpression().with_range(min_key, max_key)
+
+    @property
+    def min_key(self) -> T | None:
+        if not self.key_range:
+            return None
+        return self.key_range[0]
+
+    @property
+    def max_key(self) -> T | None:
+        if not self.key_range:
+            return None
+        return self.key_range[1]
+
+    def matches_query(self, key: any) -> bool:
+        """
+        Returns true if the key is within the range of the query expression
+        """
+        if not self.key_range:
+            return True
+        return self.min_key <= key <= self.max_key
+
+    def below_query_range(self, key: any) -> bool:
+        """
+        Returns true if the key is below the range of the query expression
+        will return false if key range is not set
+        """
+        if not self.key_range:
+            return False
+        return self.min_key > key

deltacat/storage/rivulet/reader/reader_type_registrar.py

@@ -0,0 +1,84 @@
+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 FileReader
+from typing import Type, Dict
+
+from deltacat.storage.rivulet.schema.schema import Schema
+
+
+class FileReaderRegistrar:
+    """
+    Registrar for readers of rivulet data
+
+    Readers must adhere to the Protocol DataReader
+
+    Packages with extension classes should call into this registrar in __init__.py
+    """
+
+    _readers = {}
+
+    @classmethod
+    def register_reader(
+        cls,
+        extension: str,
+        reader_class: Type[FileReader],
+        allow_overwrite: bool = False,
+    ):
+        """
+        Register a file extension associated with a dataset reader
+
+        Parameters:
+        - extension: str, the file extension to register
+        - reader_class: Type[DataReader], the reader class to associate with the extension
+        - allow_overwrite: bool, if True, allows overwriting an existing reader for the extension
+        """
+        if extension in cls._readers and not allow_overwrite:
+            raise ValueError(
+                f"Reader for extension '{extension}' is already registered. "
+                f"Set allow_overwrite=True to replace the existing reader."
+            )
+        normalized_extension = extension.lower()
+        cls._readers[normalized_extension] = reader_class
+
+    @classmethod
+    def get_reader_class(cls, uri: str) -> Type[FileReader]:
+        """
+        Gets the reader class given a URI
+
+        :param uri: URI of file to be read. Note that we expect the URI to end in a file extension
+        :raises ValueError: if no registered data reader is found for the URI's extension type
+        """
+        # Find the file extension from the URI
+        extension = uri.split(".")[-1].lower()
+
+        # Return the reader class if the extension is registered, otherwise return None
+        return cls._readers.get(extension)
+
+    @classmethod
+    def construct_reader_instance(
+        cls,
+        sst_row: SSTableRow,
+        file_provider: FileProvider,
+        primary_key: str,
+        schema: Schema,
+        reader_cache: Dict[str, FileReader] = None,
+    ) -> FileReader:
+        """
+        Construct a data reader for an instance of a given uri
+
+        :param uri: URI of file to be read. Note that we expect the URI to end in a file extension
+        :param reader_cache: Optional cache of readers keyed by extension
+        :raises ValueError: if no registered data reader is found for the URI's extension type
+        """
+        extension = sst_row.uri.split(".")[-1].lower()
+
+        if reader_cache is not None and extension in reader_cache:
+            return reader_cache[extension]
+
+        reader_class = FileReaderRegistrar.get_reader_class(sst_row.uri)
+        reader_instance = reader_class(sst_row, file_provider, primary_key, schema)
+
+        if reader_cache:
+            reader_cache[extension] = reader_instance
+
+        return reader_instance