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

deltacat/tests/storage/rivulet/test_manifest.py

@@ -0,0 +1,67 @@
+import os
+
+import pytest
+
+from deltacat import Dataset
+from deltacat.storage.rivulet.fs.file_store import FileStore
+from deltacat.storage.rivulet.schema.datatype import Datatype
+from deltacat.storage.rivulet.metastore.delta import DeltacatManifestIO
+from deltacat.storage.rivulet import Schema, Field
+import pyarrow as pa
+import pyarrow.parquet
+
+
+@pytest.fixture
+def sample_schema():
+    return Schema(
+        fields=[
+            Field("id", Datatype.int32(), is_merge_key=True),
+            Field("name", Datatype.string()),
+            Field("age", Datatype.int32()),
+        ]
+    )
+
+
+@pytest.fixture
+def sample_pydict():
+    return {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"], "age": [25, 30, 35]}
+
+
+@pytest.fixture
+def path(tmp_path):
+    return tmp_path
+
+
+@pytest.fixture
+def sample_parquet_data(path, sample_pydict):
+    parquet_path = path / "test.parquet"
+    table = pa.Table.from_pydict(sample_pydict)
+    pyarrow.parquet.write_table(table, parquet_path)
+    return parquet_path
+
+
+def test_write_manifest_round_trip(sample_parquet_data, sample_schema):
+    dataset = Dataset.from_parquet(
+        file_uri=sample_parquet_data, name="dataset", merge_keys="id"
+    )
+
+    path, filesystem = FileStore.filesystem(dataset._metadata_path)
+    file_store = FileStore(path, filesystem=filesystem)
+    manifest_io = DeltacatManifestIO(path, dataset._locator)
+
+    sst_files = ["sst1.sst", "sst2.sst"]
+    schema = Schema(
+        {("id", Datatype.int32()), ("name", Datatype.string())},
+        "id",
+    )
+    level = 2
+
+    uri = os.path.join(path, "manifest.json")
+
+    file_store.create_output_file(uri)
+    written = manifest_io.write(sst_files, schema, level)
+    manifest = manifest_io.read(written)
+
+    assert manifest.context.schema == schema
+    assert manifest.context.level == level
+    assert manifest.sst_files == sst_files

deltacat/tests/storage/rivulet/test_sst_interval_tree.py

@@ -0,0 +1,232 @@
+from typing import List, FrozenSet, Dict
+
+import pytest
+
+from deltacat.storage.rivulet.metastore.delta import DeltaContext
+from deltacat.storage.rivulet.metastore.sst import SSTable, SSTableRow
+from deltacat.storage.rivulet.metastore.sst_interval_tree import (
+    BlockIntervalTree,
+    BlockGroup,
+    OrderedBlockGroups,
+    Block,
+)
+from deltacat.storage.rivulet.schema.datatype import Datatype
+from deltacat.storage.rivulet import Schema
+
+
+@pytest.fixture
+def schema1() -> Schema:
+    return Schema(
+        {
+            ("id", Datatype.int32()),
+            ("name", Datatype.string()),
+            ("age", Datatype.int32()),
+        },
+        "id",
+    )
+
+
+@pytest.fixture
+def schema2() -> Schema:
+    return Schema(
+        {
+            ("id", Datatype.int32()),
+            ("address", Datatype.string()),
+            ("zip", Datatype.string()),
+        },
+        "id",
+    )
+
+
+@pytest.fixture
+def sst_row_list() -> List[SSTableRow]:
+    return [
+        SSTableRow(0, 100, "block1", 0, 1),
+        SSTableRow(3, 90, "block2", 0, 1),
+        SSTableRow(10, 95, "block3", 0, 1),
+        SSTableRow(0, 10, "block4", 0, 1),
+        SSTableRow(0, 100, "block5", 0, 1),
+    ]
+
+
+@pytest.fixture
+def sst1(sst_row_list) -> SSTable:
+    return SSTable(sst_row_list[0:3], 0, 100)
+
+
+@pytest.fixture
+def sst2(sst_row_list) -> SSTable:
+    return SSTable(sst_row_list[3:5], 0, 100)
+
+
+@pytest.fixture
+def manifest_context1(schema1) -> DeltaContext:
+    return DeltaContext(schema1, "manifest-001", 0)
+
+
+@pytest.fixture
+def manifest_context2(schema2) -> DeltaContext:
+    return DeltaContext(schema2, "manifest-002", 1)
+
+
+def with_field_group(
+    context: DeltaContext, rows: List[SSTableRow], indexes: List[int]
+) -> Dict[Schema, FrozenSet[Block]]:
+    """Construct a BlockGroup dict for a singular field group"""
+    schema = context.schema
+    return {schema: frozenset([Block(rows[i], context) for i in indexes])}
+
+
+@pytest.fixture
+def expected_block_groups(
+    manifest_context1, manifest_context2, sst_row_list
+) -> List[BlockGroup]:
+    return [
+        BlockGroup(
+            0,
+            3,
+            with_field_group(manifest_context1, sst_row_list, [0])
+            | with_field_group(manifest_context2, sst_row_list, [3, 4]),
+        ),
+        BlockGroup(
+            3,
+            10,
+            with_field_group(manifest_context1, sst_row_list, [0, 1])
+            | with_field_group(manifest_context2, sst_row_list, [3, 4]),
+        ),
+        BlockGroup(
+            10,
+            90,
+            with_field_group(manifest_context1, sst_row_list, [0, 1, 2])
+            | with_field_group(manifest_context2, sst_row_list, [3, 4]),
+        ),
+        BlockGroup(
+            90,
+            95,
+            with_field_group(manifest_context1, sst_row_list, [0, 1, 2])
+            | with_field_group(manifest_context2, sst_row_list, [4]),
+        ),
+        BlockGroup(
+            95,
+            100,
+            with_field_group(manifest_context1, sst_row_list, [0, 2])
+            | with_field_group(manifest_context2, sst_row_list, [4]),
+        ),
+    ]
+
+
+def test_build_sst(
+    sst1,
+    sst2,
+    manifest_context1,
+    manifest_context2,
+    sst_row_list,
+    expected_block_groups,
+):
+    t = BlockIntervalTree()
+    t.add_sst_table(sst1, manifest_context1)
+    t.add_sst_table(sst2, manifest_context2)
+
+    block_groups = t.get_sorted_block_groups()
+    expected = _build_ordered_block_groups(expected_block_groups)
+    assert expected == block_groups
+
+
+def test_build_sst_with_bounds(
+    sst1,
+    sst2,
+    manifest_context1,
+    manifest_context2,
+    sst_row_list,
+    expected_block_groups,
+):
+    t = BlockIntervalTree()
+    t.add_sst_table(sst1, manifest_context1)
+    t.add_sst_table(sst2, manifest_context2)
+
+    block_groups_filtered = t.get_sorted_block_groups(20, 100)
+    expected = _build_ordered_block_groups(expected_block_groups[2:])
+    assert expected == block_groups_filtered
+
+    block_groups_filtered = t.get_sorted_block_groups(96, 100)
+    expected = _build_ordered_block_groups(expected_block_groups[4:])
+    assert expected == block_groups_filtered
+
+    block_groups_filtered = t.get_sorted_block_groups(0, 10)
+    expected = _build_ordered_block_groups(expected_block_groups[0:3])
+    assert expected == block_groups_filtered
+
+    # Max key of 95 is inclusive of last range so it is included
+    block_groups_filtered = t.get_sorted_block_groups(None, 95)
+    expected = _build_ordered_block_groups(expected_block_groups)
+    assert expected == block_groups_filtered
+
+    block_groups_filtered = t.get_sorted_block_groups(None, 94)
+    expected = _build_ordered_block_groups(expected_block_groups[0:4])
+    assert expected == block_groups_filtered
+
+    block_groups_filtered = t.get_sorted_block_groups(0, 10)
+    expected = _build_ordered_block_groups(expected_block_groups[0:3])
+    assert expected == block_groups_filtered
+
+    block_groups_filtered = t.get_sorted_block_groups(0, 0)
+    expected = _build_ordered_block_groups(expected_block_groups[0:1])
+    assert expected == block_groups_filtered
+
+
+def test_build_sst_with_non_zero_min_key_matching_global_min_key(manifest_context1):
+    # Using a non-0 value since 0 evaluates to False
+    min_key = 1
+    max_key = 95
+
+    sst_row = SSTableRow(min_key, max_key, "row-with-non-zero-min-key", 0, 1)
+    t = BlockIntervalTree()
+    t.add_sst_table(SSTable([sst_row], min_key, max_key), manifest_context1)
+
+    block_groups_filtered = t.get_sorted_block_groups(min_key, min_key + 1)
+    expected = _build_ordered_block_groups(
+        [
+            BlockGroup(
+                min_key,
+                max_key,
+                {
+                    manifest_context1.schema: frozenset(
+                        [Block(sst_row, manifest_context1)]
+                    )
+                },
+            )
+        ]
+    )
+    assert expected == block_groups_filtered
+
+
+def test_build_sst_invalid_bounds(
+    sst1, sst2, schema1, schema2, sst_row_list, expected_block_groups
+):
+    t = BlockIntervalTree()
+
+    with pytest.raises(ValueError):
+        t.get_sorted_block_groups(10, 0)
+
+
+def _build_ordered_block_groups(block_groups: List[BlockGroup]) -> OrderedBlockGroups:
+    """
+    Helper method to build OrderedBlockGroups from a sorted list of block groups
+
+    """
+    ordered_groups = []
+    boundary_table = []
+    for i, bg in enumerate(block_groups):
+        boundary_table.append(bg.key_min)
+        is_last = i == len(block_groups) - 1
+        if is_last:
+            bg = BlockGroup(bg.key_min, bg.key_max, bg.field_group_to_blocks, True)
+            boundary_table.append(bg.key_max)
+        ordered_groups.append(bg)
+
+    return OrderedBlockGroups(
+        ordered_groups[0].key_min,
+        ordered_groups[-1].key_max,
+        ordered_groups,
+        boundary_table,
+    )

deltacat/tests/storage/rivulet/test_utils.py

@@ -0,0 +1,122 @@
+import inspect
+import os
+
+from pyarrow import RecordBatch, Table
+
+from deltacat.storage.rivulet.dataset import Dataset
+from deltacat.storage.rivulet.reader.query_expression import QueryExpression
+from deltacat.storage.rivulet.writer.dataset_writer import DatasetWriter
+
+from deltacat.storage.rivulet.mvp.Table import MvpTable, MvpRow
+from deltacat.storage.rivulet import Schema
+from typing import Dict, List, Generator, Set
+
+FIXTURE_ROW_COUNT = 10000
+
+
+def write_mvp_table(writer: DatasetWriter, table: MvpTable):
+    writer.write(table.to_rows_list())
+
+
+def mvp_table_to_record_batches(table: MvpTable, schema: Schema) -> RecordBatch:
+    data = table.to_rows_list()
+    columns = {key: [d.get(key) for d in data] for key in schema.keys()}
+    record_batch = RecordBatch.from_pydict(columns, schema=schema.to_pyarrow())
+    return record_batch
+
+
+def compare_mvp_table_to_scan_results(
+    table: MvpTable, scan_results: List[dict], pk: str
+):
+    table_row_list = table.to_rows_list()
+    assert len(scan_results) == len(table_row_list)
+    rows_by_pk: Dict[str, MvpRow] = table.to_rows_by_key(pk)
+    assert len(rows_by_pk) == len(scan_results)
+    for record in scan_results:
+        pk_val = record[pk]
+        assert rows_by_pk[pk_val].data == record
+
+
+def validate_with_full_scan(dataset: Dataset, expected: MvpTable, schema: Schema):
+    # best way to validate is to use dataset reader and read records
+    read_records = list(dataset.scan(QueryExpression()).to_pydict())
+    compare_mvp_table_to_scan_results(
+        expected, read_records, list(dataset.get_merge_keys())[0]
+    )
+
+
+def generate_data_files(dataset: Dataset) -> Generator[str, None, None]:
+    for ma in dataset._metastore.generate_manifests():
+        for sstable in ma.generate_sstables():
+            for row in sstable.rows:
+                yield row.uri
+
+
+def assert_data_file_extension(dataset: Dataset, file_extension: str):
+    data_file_count = 0
+    for data_file in generate_data_files(dataset):
+        data_file_count += 1
+        assert data_file.endswith(file_extension)
+    assert data_file_count > 0, "No data files found in dataset"
+    print(f"Asserted that {data_file_count} data files end with {file_extension}")
+
+
+def assert_data_file_extension_set(dataset: Dataset, file_extension_set: Set[str]):
+    """
+    Asserts that each file extension in set appears at least once in dataset
+    """
+    data_file_count = 0
+    found_extensions = set()
+
+    for data_file in generate_data_files(dataset):
+        data_file_count += 1
+        for extension in file_extension_set:
+            if data_file.endswith(extension):
+                found_extensions.add(extension)
+                break
+
+    assert data_file_count > 0, "No data files found in dataset"
+    assert (
+        found_extensions == file_extension_set
+    ), f"Missing extensions: {file_extension_set - found_extensions}"
+    print(
+        f"Asserted that among {data_file_count} data files, all extensions {file_extension_set} were found"
+    )
+
+
+def create_dataset_for_method(temp_dir: str):
+    """
+    Given a temp directory, creates a directory within it based on the name of the function calling this.
+    Then returns a dataset based from that directory
+    """
+    caller_frame = inspect.getouterframes(inspect.currentframe())[1]
+    dataset_dir = os.path.join(temp_dir, caller_frame.function)
+    os.makedirs(dataset_dir)
+    return Dataset(
+        dataset_name=f"dataset-${caller_frame.function}", metadata_uri=dataset_dir
+    )
+
+
+def verify_pyarrow_scan(
+    scan_result: Generator[RecordBatch, None, None],
+    expected_schema: Schema,
+    expected_data: dict,
+):
+    record_batches = list(scan_result)
+    assert record_batches, "Scan should return at least one record batch."
+
+    combined_table = Table.from_batches(record_batches)
+
+    expected_fields = {field.name for field in expected_schema.values()}
+    scanned_fields = set(combined_table.schema.names)
+    assert (
+        scanned_fields == expected_fields
+    ), f"Scanned fields {scanned_fields} do not match expected fields {expected_fields}."
+
+    for field in expected_fields:
+        assert (
+            field in combined_table.column_names
+        ), f"Field '{field}' is missing in the scan result."
+        assert (
+            combined_table[field].to_pylist() == expected_data[field]
+        ), f"Field '{field}' data does not match expected values."