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

deltacat/storage/rivulet/metastore/sst_interval_tree.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+from bisect import bisect_left, bisect_right
+from collections import defaultdict
+from dataclasses import dataclass
+from itertools import tee
+from typing import Any, Dict, Set, List, FrozenSet, Iterable, TypeVar, NamedTuple
+
+from intervaltree import Interval, IntervalTree
+
+from deltacat.storage.rivulet.metastore.delta import DeltaContext
+from deltacat.storage.rivulet.metastore.sst import SSTable, SSTableRow
+from deltacat.storage.rivulet import Schema
+
+T = TypeVar("T")
+
+
+# Can be replaced with itertools.pairwise once we're on python 3.10+
+def pairwise(iterable):
+    a, b = tee(iterable)
+    next(b, None)
+    return zip(a, b)
+
+
+class Block(NamedTuple):
+    row: SSTableRow
+    context: DeltaContext
+    """Context from the manifest around the placement of this row in the LSM-Tree"""
+
+
+@dataclass(frozen=True)
+class BlockGroup:
+    """
+    TODO discuss "Block" terminology.
+
+    This data structure represents a set of data blocks which have to be traversed
+     together, because rows in the block is in an overlapping key range
+
+    The data and key ranges within a BlockGroup can be processed in parallel during a scan
+
+    Each block may show up in many BlockGroups. This is because as soon as the combination of blocks in the group
+    change, it becomes a different group.
+
+    Take the following example: A dataset has Field Group 1 (FG1) and Field Group 2 (FG2). The SSTs rows look like below:
+
+    Field group 1     Field group 2
+    ------------------|---------------
+    [0,100] -> Block1 | [0-10]  -> Block4
+    [3-90]  -> Block2 | [0-100] -> Block5
+    [10-95] -> Block3 |
+
+    If we want to scan this data, we can index of off the SST rows into an interval tree and observe the following
+    "boundaries" in the interval tree:
+
+    BlockGroup1 - Covers rows [0-3), includes Blocks 1,4,5
+    BlockGroup2 - Covers rows [3-10), includes Blocks 1, 2, 4, 5
+    BlockGroup3 - Covers rows [10-90), includes Blocks 1, 2, 3, 4*, 5
+    BlockGroup4 - Covers rows [90-95), includes Blocks 1, 2*, 3, 5
+    BlockGroup5 - Covers rows [95-100], includes Blocks 1, 3*, 5
+    *special case - interval end==block group start
+
+    Creating a sorted list of BlockGroups like this allows us to know exactly which blocks contain records
+    for any given point query, range query, or scan. For instance, if the user queries for key=3, we know
+    to read BlockGroup1, or key[0-10] to read BlockGroup1+BlockGroup2.
+    """
+
+    key_min: T
+    """
+    Key min is inclusive (the block group covers data where primary key>=key_min)
+    """
+    key_max: T
+    field_group_to_blocks: Dict[Schema, FrozenSet[Block]]
+    key_max_inclusive: bool = False
+    """
+    By default, key_min is inclusive and key_max is non-inclusive.
+    For the highest key in an SSTable, we need to set key_max_inclusive to True
+    """
+
+    def key_in_range(self, key: T) -> bool:
+        if self.key_max_inclusive:
+            return self.key_min <= key <= self.key_max
+        else:
+            return self.key_min <= key < self.key_max
+
+    def key_below_range(self, key: T) -> bool:
+        return key < self.key_min
+
+
+@dataclass(frozen=True)
+class OrderedBlockGroups:
+    """
+    Ordered block groups representing a sequential interval of a primary key range
+
+    Block groups have an inclusive
+
+    The block groups have a "boundary table" which represents primary key lower/upper ranges in table
+
+    For example, a boundary table of [1,3,5,10] has blocks: [1,3), [3,5), [5,10]
+    """
+
+    key_min: T
+    """
+    Key min is inclusive
+    """
+    key_max: T
+    """
+    Key max is inclusive
+    """
+    block_groups: List[BlockGroup]
+    boundary_table: List[T]
+
+
+class BlockIntervalTree:
+    """
+    This interval tree combines all L0 SSTables in order to effectively traverse key ranges
+    so that data can be zippered across field groups
+
+    This interval tree is SHARED across N different field groups. This is because
+        maintaining a different interval tree for N different field groups would require
+        more complex traversal to achieve the same result
+    """
+
+    def __init__(self):
+        self.tree: IntervalTree = IntervalTree()
+        self.max_key_map: Dict[Any, List[Interval]] = {}
+
+    def add_sst_table(self, sst: SSTable, context: DeltaContext):
+        """
+        Add intervals to SSTree which use primary key min and max as intervals
+        The data for each interval is a tuple of (schema, SSTableRow)
+        """
+        self.add_sst_rows(sst.rows, context)
+
+    def add_sst_rows(self, sst_rows: Iterable[SSTableRow], context: DeltaContext):
+        """
+        Add individual SSTable rows to tree
+        """
+        for row in sst_rows:
+            interval: Interval = Interval(row.key_min, row.key_max, Block(row, context))
+            self.tree.add(interval)
+            if row.key_max not in self.max_key_map:
+                self.max_key_map[row.key_max] = [interval]
+            else:
+                self.max_key_map[row.key_max].append(interval)
+
+    def get_sorted_block_groups(
+        self, min_key: Any | None = None, max_key: Any | None = None
+    ) -> OrderedBlockGroups:
+        """
+        Returns an ordered list of block group by primary key range
+        The IntervalTree boundary table contains each boundary where the set of intervals change
+        This function traverses the boundary table and builds a list
+
+        Edge case - incompatibility with Interval Tree Library
+
+        Note that the IntervalTree Library treats ALL ranges as min=inclusive, max=non-inclusive
+        This is a different from our SSTables, where min_key and max_key are both inclusive.
+
+        As a suboptimal workaround, we can fix this by adding an SSTRow to an interval when the
+        SST row's max key is equal to the lower bound of the interval. Take the following example:
+
+        Field group 1     Field group 2
+        ------------------|---------------
+        [0,100] -> Block1 | [0-10]  -> Block4
+        [3-90]  -> Block2 | [0-100] -> Block5
+        [10-95] -> Block3 |
+
+        Our workaround adds for example Block4 to the interval [10,90), because Block4 is inclusive
+        of key 10 but the IntervalTree library thinks it is non-inclusive.
+
+        BlockGroup1 - Covers rows [0-3), includes Blocks 1,4,5
+        BlockGroup2 - Covers rows [3-10), includes Blocks 1, 2, 4, 5
+        BlockGroup3 - Covers rows [10-90), includes Blocks 1, 2, 3, 4*, 5
+        BlockGroup4 - Covers rows [90-95), includes Blocks 1, 2*, 3, 5
+        BlockGroup5 - Covers rows [95-100], includes Blocks 1, 3*, 5
+        *special case - interval end==block group start
+
+        An ideal solution would produce block groups like below. To do this against the IntervalTree
+        library, we would need to know how to convert an inclusive range like ["bar", "foo"] into a
+        range like ["bar", "fooa") where the end range is non-inclusive. It is fine for our block groups
+        to be non-optimal, we just need code when zipper merging to detect if a block group's max key is
+        less than the current iterator and therfore not consider it.
+
+        Optimal block groups:
+        BlockGroup1 - Covers rows [0-3), includes Blocks 1,5,4
+        BlockGroup2 - Covers rows [3-10), includes Blocks 1,5,2,4
+        BlockGroup3 - Covers rows [10-11), includes Blocks 1, 2, 3, 4, 5
+        BlockGroup3 - Covers rows [11-91), includes Blocks 1,2,3,5
+        BlockGroup4 - Covers rows [91-96), includes Blocks 1,3,5
+        BlockGroup5 - Covers rows [96-100], includes Blocks 1,5
+
+        :param: min_key optionally restrict result so that they must overlap or be greater than min_key, INCLUSIVE
+            A range [0,200) will be included for min_key <=100 because it overlaps with the min key
+            A range [100, 200) will be included because it overlaps with the min key (100 key inclusive)
+        :param: max_key optionally restrict result to be less than or overlap with max_key, INCLUSIVE
+            A range like [100,200) will NOT included for max key=200 because range maxes are non-inclusive
+        """
+        key_boundaries = self.tree.boundary_table.keys()
+        block_groups: List[BlockGroup] = []
+        block_groups_min = None
+        block_groups_max = None
+
+        # Note that we need to expand min_key_idx and max_key_idx by 1 to cover cases where
+        # the pairwise traversal (x,y) has x>=min_key>=y and x<=max_key<=y
+        if min_key is not None and max_key is not None and min_key > max_key:
+            raise ValueError(
+                f"min_key {min_key} cannot be greater than max_key {max_key}"
+            )
+
+        min_key_idx = (
+            max(0, bisect_left(key_boundaries, min_key) - 1)
+            if min_key is not None
+            else None
+        )
+        max_key_idx = (
+            bisect_right(key_boundaries, max_key) + 1 if max_key is not None else None
+        )
+        boundary_table = key_boundaries[min_key_idx:max_key_idx]
+
+        for lower_bound, upper_bound in pairwise(boundary_table):
+            # Note that IntervalTree library treats lower bound of slice as inclusive and upper as exclusive
+            # We follow the same structure in our BlockGroup
+            intervals: Set[Interval] = self.tree.overlap(lower_bound, upper_bound)
+
+            # Special case for if max key is equal to lower_bound. See method pydoc for more details
+            for i in self.max_key_map.get(lower_bound, []):
+                intervals.add(i)
+
+            field_group_to_blocks = defaultdict(set)
+            for interval in intervals:
+                data: Block = interval.data
+                schema = data.context.schema
+                field_group_to_blocks[schema].add(data)
+
+            # freeze dict to make it hashable
+            field_group_to_blocks = {
+                k: frozenset(v) for k, v in field_group_to_blocks.items()
+            }
+
+            # Special case - if this is the very last iteration, set key_max_inclusive to True
+            max_key_inclusive = upper_bound == boundary_table[-1]
+
+            block_group = BlockGroup(
+                lower_bound, upper_bound, field_group_to_blocks, max_key_inclusive
+            )
+            block_groups_min = (
+                lower_bound
+                if block_groups_min is None
+                else min(block_groups_min, lower_bound)
+            )
+            block_groups_max = (
+                upper_bound
+                if block_groups_max is None
+                else max(block_groups_max, upper_bound)
+            )
+            block_groups.append(block_group)
+
+        return OrderedBlockGroups(
+            block_groups_min, block_groups_max, block_groups, boundary_table
+        )

deltacat/storage/rivulet/mvp/Table.py

@@ -0,0 +1,101 @@
+from dataclasses import dataclass
+from typing import List, Dict, Any
+from typing import Iterable
+
+from collections.abc import Mapping
+
+from pyarrow import RecordBatch
+
+
+@dataclass
+class MvpRow(Mapping):
+    data: Dict[str, Any]
+
+    def __getitem__(self, key):
+        return self.data[key]
+
+    def __iter__(self):
+        return iter(self.data)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __contains__(self, key):
+        return key in self.data
+
+    def keys(self):
+        return self.data.keys()
+
+    def values(self):
+        return self.data.values()
+
+    def items(self):
+        return self.data.items()
+
+    def get(self, key, default=None):
+        return self.data.get(key, default)
+
+    @staticmethod
+    def generate_from_arrow(batch: RecordBatch):
+        for row_idx in range(batch.num_rows):
+            out = {}
+            for column_idx, column in enumerate(batch.column_names):
+                col = batch.column(column_idx)
+                out.update({column: col[row_idx].as_py()})
+                yield MvpRow(out)
+
+
+@dataclass
+class MvpTable(Iterable[Dict[str, Any]]):
+    data: Dict[str, List[Any]]
+
+    def __iter__(self):
+        # Get the lengths of all columns (they should be the same)
+        row_count = len(next(iter(self.data.values())))
+
+        # Iterate over the rows
+        for i in range(row_count):
+            row_data = {
+                field_name: field_arr[i] for field_name, field_arr in self.data.items()
+            }
+            yield row_data
+
+    def to_rows_by_key(self, mk: str) -> Dict[str, "MvpRow"]:
+        # Find the provided key field in the schema
+        # build row data
+        pk_col = self.data[mk]
+        row_data: Dict[str, MvpRow] = {}
+        for i, value in enumerate(pk_col):
+            row_data[value] = MvpRow(
+                {
+                    field_name: field_arr[i]
+                    for field_name, field_arr in self.data.items()
+                }
+            )
+        return row_data
+
+    def to_rows_list(self) -> List[Dict[str, Any]]:
+        return [r for r in self]
+
+    @classmethod
+    def merge(cls, dataset1: "MvpTable", dataset2: "MvpTable", pk: str) -> "MvpTable":
+
+        merged_data: Dict[str, List[Any]] = {}
+        # Initialize merged_data with keys from both datasets
+        for k in set(dataset1.data.keys()) | set(dataset2.data.keys()):
+            merged_data[k] = []
+
+        # Create dictionaries for quick lookup
+        row_data_ds1: dict[str, MvpRow] = dataset1.to_rows_by_key(pk)
+        row_data_ds2: dict[str, MvpRow] = dataset2.to_rows_by_key(pk)
+
+        # Merge the datasets
+        all_keys = set(row_data_ds1.keys()) | set(row_data_ds2.keys())
+        for k in all_keys:
+            row1: MvpRow = row_data_ds1.get(k, MvpRow({}))
+            row2: MvpRow = row_data_ds2.get(k, MvpRow({}))
+            merged_row = {**row1, **row2}
+            for column, values in merged_data.items():
+                values.append(merged_row.get(column))
+
+        return cls(merged_data)

deltacat/storage/rivulet/mvp/__init__.py

@@ -0,0 +1,5 @@
+"""
+These classes are temporary for our super scrappy MVP. I chose to make this Mvp package to make it explicit what interfaces need to be replaced later
+
+In particular all the IO will use in memory python lists/dicts for the MVP
+"""

deltacat/storage/rivulet/parquet/__init__.py

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

deltacat/storage/rivulet/parquet/file_reader.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pyarrow import RecordBatch
+
+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
+import pyarrow.parquet as pq
+import pyarrow as pa
+
+
+class ParquetFileReader(FileReader[RecordBatchRowIndex]):
+    """
+    Parquet file reader
+
+    This class is not thread safe
+    """
+
+    def __init__(
+        self,
+        sst_row: SSTableRow,
+        file_provider: FileProvider,
+        key: str,
+        schema: Schema,
+        iter_batch_size=1000,
+    ):
+        self.sst_row = sst_row
+        self.input = file_provider.provide_input_file(sst_row.uri)
+
+        self.key = key
+        self.parquet_file: pa.parquet.ParquetFile | None = None
+        self.iter_batch_size = iter_batch_size
+
+        # Iterator from pyarrow iter_batches API call. Pyarrow manages state of traversal within parquet row groups
+        self._record_batch_iter = None
+
+        self.schema = schema
+
+        """
+        These variables keep state about where the iterator is current at. They are initialized in __enter__()
+        """
+        self._curr_batch: RecordBatch | None = None
+        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 RowAndKey
+        """
+        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.parquet_file = pq.ParquetFile(f)
+            # Initialize _curr_batch
+            row_groups = list(range(self.sst_row.offset_start, self.sst_row.offset_end))
+            self._record_batch_iter = self.parquet_file.iter_batches(
+                self.iter_batch_size, row_groups, columns=self.schema.keys()
+            )
+            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):
+        self.parquet_file.close()
+
+    def __is_initialized(self):
+        return self.parquet_file and 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
+        """
+        self._curr_batch = next(self._record_batch_iter)
+        self._curr_row_offset = 0
+        self._pk_col = self._curr_batch[self.key]

deltacat/storage/rivulet/parquet/serializer.py

@@ -0,0 +1,37 @@
+from typing import List, Any
+
+import pyarrow as pa
+from pyarrow.parquet import FileMetaData
+
+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 ParquetDataSerializer(ArrowSerializer):
+    """
+    Parquet data writer. Responsible for flushing rows to parquet and returning SSTable rows for any file(s) written
+    """
+
+    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("parquet")
+        with file.create() as outfile:
+            metadata_collector: list[Any] = []
+            pa.parquet.write_table(
+                table=table, where=outfile, metadata_collector=metadata_collector
+            )
+            # look for file metadata
+            file_metadata: FileMetaData = next(
+                item for item in metadata_collector if isinstance(item, FileMetaData)
+            )
+            row_group_count = file_metadata.num_row_groups
+
+        # Because ParquetWriter only writes 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)
+        return [SSTableRow(key_min, key_max, file.location, 0, 0 + row_group_count)]