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

deltacat/benchmarking/benchmark_engine.py

@@ -0,0 +1,84 @@
+import sys
+import time
+from contextlib import contextmanager
+from typing import Generator, Tuple
+
+from deltacat.benchmarking.benchmark_report import BenchmarkMetric, BenchmarkStep
+from deltacat.experimental.storage.rivulet.dataset import Dataset
+from deltacat.experimental.storage.rivulet.reader.query_expression import (
+    QueryExpression,
+)
+
+
+@contextmanager
+def timed_step(description: str) -> Generator[BenchmarkStep, None, None]:
+    """Convenience for computing elapsed time of a block of code as a metric.
+
+    :param description: description of the step
+    :return: a benchmark operation populated with the elapsed time
+    """
+    metric = BenchmarkStep(description)
+    start_time = time.time()
+    yield metric
+    end_time = time.time()
+    metric.add(BenchmarkMetric("elapsed_time", 1000 * (end_time - start_time), "ms"))
+
+
+class BenchmarkEngine:
+    def __init__(self, dataset: Dataset):
+        self.dataset = dataset
+
+    def load_and_commit(
+        self, schema_name, generator, count
+    ) -> Tuple[str, BenchmarkStep]:
+        """Load count number of rows from the generator and commit.
+
+        :param generator: row generator
+        :param count: the number of rows to load into the dataset
+        :return: tuple of the manifest URI and a operation measurement
+        """
+        desc = f"load {count} from {generator}"
+        writer = self.dataset.writer(schema_name)
+        with timed_step(desc) as step:
+            rows = [generator.generate() for _ in range(count)]
+            writer.write(rows)
+            result = writer.flush()
+        step.add(BenchmarkMetric("loaded", count))
+        return result, step
+
+    def scan(self) -> Tuple[set[any], BenchmarkStep]:
+        """
+        Scans the rows in dataset and prints some basic statistics about the manifest
+
+        :return: Tuple[set[any], BenchmarkStep] - a tuple containing a set of merge keys and a benchmark step with metrics
+        """
+        keys = set()
+        object_count = 0
+        size_b = 0
+        # Note that we expect single col merge keys so we can return key set
+        # this will fail with validation error if dataset has multiple merge keys
+        merge_key_name = self.dataset.schemas["all"].get_merge_key()
+        with timed_step("full scan") as step:
+            for row in self.dataset.scan(QueryExpression()).to_pydict():
+                object_count += 1
+                size_b += sum([sys.getsizeof(x) for x in row.values()])
+                keys.add(row.get(merge_key_name))
+                # TODO replace with the actual metrics we want to measure
+        step.add(BenchmarkMetric("rows read", object_count))
+        step.add(BenchmarkMetric("size", size_b / (1024 * 1024), "MB"))
+        return keys, step
+
+    def run_queries(
+        self, description, manifest_uri, queries: list[QueryExpression]
+    ) -> BenchmarkStep:
+        object_count = 0
+        size_b = 0
+        with timed_step(description) as step:
+            for query in queries:
+                for row in self.dataset.scan(query).to_pydict():
+                    object_count += 1
+                    size_b += sum([sys.getsizeof(x) for x in row.values()])
+        # TODO replace with the actual metrics we want to measure
+        step.add(BenchmarkMetric("rows read", object_count))
+        step.add(BenchmarkMetric("size", size_b / (1024 * 1024), "MB"))
+        return step

deltacat/benchmarking/benchmark_report.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+from tabulate import tabulate
+from typing import Union, Optional
+
+
+@dataclass
+class BenchmarkMetric:
+    name: str
+    value: Union[float, int]
+    unit: Optional[str] = None
+
+
+class BenchmarkStep:
+    """Captures measurements from a given operation"""
+
+    def __init__(self, description):
+        self.description: str = description
+        """Description of the operation"""
+        self._metrics: dict[str, BenchmarkMetric] = {}
+        """Description of the operation"""
+
+    def add(self, metric: BenchmarkMetric):
+        self._metrics[metric.name] = metric
+
+    def list_metrics(self):
+        """List the metrics (sorted by name)"""
+        return sorted(self._metrics.values(), key=lambda x: x.name)
+
+
+class BenchmarkRun:
+    """Class for capturing measurements for a given test suite for comparison."""
+
+    def __init__(self, suite: str, description: str):
+        self.suite = suite
+        """The test suite associated with this report."""
+        self.description = description
+        """Description of the report"""
+        self.steps: list[BenchmarkStep] = []
+        """List of steps and their metrics"""
+
+    def add(self, operation):
+        self.steps.append(operation)
+
+
+class BenchmarkReport:
+    def __init__(self, name):
+        self.name = name
+        self.runs: list[BenchmarkRun] = []
+
+    def add(self, run):
+        self.runs.append(run)
+
+    def __str__(self):
+        """Pretty-print a table that compares the metrics across each report.
+
+        We want to transpose these such that each report gets their own column and each metric gets its own row
+        (ideally grouped by operation).
+        """
+        if not self.runs:
+            print("No runs to compare!")
+            return
+        suites = set(r.suite for r in self.runs)
+        if len(suites) > 1:
+            print("Found more than one type of suite")
+            return
+        suite = self.runs[0].suite
+
+        headers = [
+            f"{suite} Operation",
+            "Metric",
+            "Unit",
+            *[r.description for r in self.runs],
+        ]
+        rows = []
+        for step_tranche in zip(*[r.steps for r in self.runs]):
+            # TODO zip by metric name instead of assuming all metrics are being measured
+            step_name = step_tranche[0].description
+            for metric_tuple in zip(*[x.list_metrics() for x in step_tranche]):
+                row = [
+                    step_name,
+                    metric_tuple[0].name,
+                    metric_tuple[0].unit,
+                    *[p.value for p in metric_tuple],
+                ]
+                rows.append(row)
+        return tabulate(rows, headers=headers, tablefmt="fancy_outline")

deltacat/benchmarking/benchmark_suite.py

@@ -0,0 +1,11 @@
+from typing import Protocol
+
+from deltacat.benchmarking.benchmark_report import BenchmarkRun
+
+
+class BenchmarkSuite(Protocol):
+    def run(self) -> BenchmarkRun:
+        """Run the benchmark suite and produce a report.
+
+        Each report should be comparable against other reports by the same suite"""
+        ...

deltacat/benchmarking/conftest.py

@@ -4,12 +4,28 @@ import pyarrow as pa
 import pyarrow.fs as pafs
 import pyarrow.parquet as papq
 import pytest
+from _pytest.terminal import TerminalReporter
 
-from deltacat.utils.pyarrow import s3_file_to_table
-from deltacat.types.media import (
-    ContentEncoding,
-    ContentType,
-)
+from deltacat.benchmarking.benchmark_report import BenchmarkReport
+
+
+@pytest.fixture(autouse=True, scope="function")
+def report(request):
+    report = BenchmarkReport(request.node.name)
+
+    def final_callback():
+        terminal_reporter: TerminalReporter = request.config.pluginmanager.get_plugin(
+            "terminalreporter"
+        )
+        capture_manager = request.config.pluginmanager.get_plugin("capturemanager")
+        with capture_manager.global_and_fixture_disabled():
+            terminal_reporter.ensure_newline()
+            terminal_reporter.section(request.node.name, sep="-", blue=True, bold=True)
+            terminal_reporter.write(str(report))
+            terminal_reporter.ensure_newline()
+
+    request.addfinalizer(final_callback)
+    return report
 
 
 def pyarrow_read(path: str, columns: list[str] | None = None) -> pa.Table:
@@ -21,17 +37,6 @@ def pyarrow_read(path: str, columns: list[str] | None = None) -> pa.Table:
     return papq.read_table(path, columns=columns, filesystem=fs)
 
 
-def deltacat_read(path: str, columns: list[str] | None = None) -> pa.Table:
-    assert path.startswith("s3://")
-    return s3_file_to_table(
-        path,
-        content_type=ContentType.PARQUET,
-        content_encoding=ContentEncoding.IDENTITY,
-        column_names=None,  # Parquet files are schemaful
-        include_columns=columns,
-    )
-
-
 def daft_table_read(path: str, columns: list[str] | None = None) -> pa.Table:
     try:
         import daft
@@ -40,7 +45,7 @@ def daft_table_read(path: str, columns: list[str] | None = None) -> pa.Table:
             "Daft not installed. Install Daft using pip to run these benchmarks: `pip install getdaft`"
         )
 
-    tbl = daft.table.Table.read_parquet(path, columns=columns)
+    tbl = daft.read_parquet(path, columns=columns)
     return tbl.to_arrow()
 
 
@@ -48,12 +53,10 @@ def daft_table_read(path: str, columns: list[str] | None = None) -> pa.Table:
     params=[
         daft_table_read,
         pyarrow_read,
-        deltacat_read,
     ],
     ids=[
         "daft_table",
         "pyarrow",
-        "deltacat",
     ],
 )
 def read_fn(request):

deltacat/benchmarking/data/random_row_generator.py

@@ -0,0 +1,94 @@
+import math
+import os
+from enum import Enum
+from io import BytesIO
+from typing import Any, Dict
+
+import faker
+from faker_file.providers.png_file import GraphicPngFileProvider
+from faker_file.storages.filesystem import FileSystemStorage
+from PIL import Image
+
+from deltacat.benchmarking.data.row_generator import RowGenerator
+
+
+class ImageStyle(Enum):
+    RANDOM_BYTES = 1
+    """Generate random bytes to simulate an image.
+
+    This is the fastest option (if you want to test correctness).
+    """
+    PILLOW = 2
+    """Generate actual PNG files in-memory directly using Pillow"""
+    FAKER_FILE = 3
+    """Generate PNG files on-disk using with some random elements"""
+
+
+class RandomRowGenerator(RowGenerator):
+    """Generate rows with 'images' that are just randomly-generated bytes"""
+
+    def __init__(
+        self, seed=0, tmp_dir=None, style: ImageStyle = ImageStyle.RANDOM_BYTES
+    ):
+        self.seed = seed
+        self.fake = faker.Faker()
+        self.fake.seed_instance(seed)
+        self.fake.add_provider(GraphicPngFileProvider)
+        self.temp_dir = tmp_dir
+        self.style = style
+
+    def __str__(self):
+        return f"random source"
+
+    def _generate_image(self, width, height) -> bytes:
+        if self.style == ImageStyle.RANDOM_BYTES:
+            return self._generate_with_random_bytes(width, height)
+        elif self.style == ImageStyle.PILLOW:
+            return self._generate_with_pillow(width, height)
+        elif self.style == ImageStyle.FAKER_FILE:
+            return self._generate_with_faker(width, height)
+        else:
+            raise ValueError("Unknown ImageStyle")
+
+    @staticmethod
+    def _generate_with_random_bytes(width, height) -> bytes:
+        """Generate random bytes to simulate an image."""
+        target_size = math.floor(
+            width * height / 50
+        )  # this isn't actually how file size relates to image size
+        # Assumption: we don't actually need images. It suffices to generate arbitrary-length bytes of random characters.
+        return os.urandom(target_size)
+
+    @staticmethod
+    def _generate_with_pillow(width, height) -> bytes:
+        """Generate actual PNG files in-memory directly using Pillow"""
+        file = BytesIO()
+        image = Image.new("RGBA", size=(width, height), color=(155, 0, 0))
+        image.save(file, "png")
+        file.name = "test.png"
+        file.seek(0)
+        return file.read()
+
+    def _generate_with_faker(self, width, height) -> bytes:
+        """Generate PNG files on-disk using with some random elements"""
+        rel_name = self.fake.graphic_png_file(
+            storage=FileSystemStorage(
+                root_path=self.temp_dir,
+                rel_path="tmp",
+            ),
+            size=(width, height),
+        )
+        file_name = f"{self.temp_dir}/{rel_name}"
+        with open(file_name, "rb") as f:
+            return f.read()
+
+    def generate(self) -> Dict[str, Any]:
+        return {
+            "id": self.fake.random_int(0, 10_000_000),
+            "source": self.fake.image_url(),
+            "media": (
+                self._generate_image(
+                    self.fake.random_int(512, 2048), self.fake.random_int(512, 4096)
+                )
+            ),
+        }

deltacat/benchmarking/data/row_generator.py

@@ -0,0 +1,10 @@
+from typing import Protocol, Iterator, Dict, Any
+
+
+class RowGenerator(Protocol):
+    def generate(self) -> Dict[str, Any]:
+        ...
+
+    def generate_dataset(self, count) -> Iterator[Dict[str, Any]]:
+        """Generate a dataset with a given number of records"""
+        return map(lambda x: self.generate(), iter(range(count)))

deltacat/benchmarking/test_benchmark_pipeline.py

@@ -0,0 +1,108 @@
+import math
+from random import shuffle
+import pytest
+from deltacat.experimental.storage.rivulet.dataset import Dataset
+from deltacat.experimental.storage.rivulet.schema.datatype import Datatype
+from deltacat.experimental.storage.rivulet.reader.query_expression import (
+    QueryExpression,
+)
+from deltacat.experimental.storage.rivulet.schema.schema import Schema
+from deltacat.benchmarking.benchmark_engine import BenchmarkEngine
+from deltacat.benchmarking.benchmark_report import BenchmarkRun, BenchmarkReport
+from deltacat.benchmarking.benchmark_suite import BenchmarkSuite
+from deltacat.benchmarking.data.random_row_generator import RandomRowGenerator
+from deltacat.benchmarking.data.row_generator import RowGenerator
+from deltacat.tests.test_utils.filesystem import temp_dir_autocleanup
+
+pytestmark = pytest.mark.benchmark
+
+
+@pytest.fixture
+def schema():
+    return Schema(
+        [
+            ("id", Datatype.int32()),
+            ("source", Datatype.string()),
+            ("media", Datatype.image("png")),
+        ],
+        "id",
+    )
+
+
+class LoadAndScanSuite(BenchmarkSuite):
+    """Load some number of rows and scan"""
+
+    schema_name = "LoadAndScanSuite"
+
+    def __init__(self, dataset: Dataset, schema: Schema, generator, description=None):
+        self.suite = "ReadSuite"
+        self.dataset: Dataset = dataset
+        self.schema = schema
+        self.dataset.add_schema(schema, LoadAndScanSuite.schema_name)
+        self.generator: RowGenerator = generator
+        self.description: str = description or f"{self.dataset} x {self.generator}"
+
+    def run(self) -> BenchmarkRun:
+        container = BenchmarkEngine(self.dataset)
+        run = BenchmarkRun(self.suite, self.description)
+        # load a large number of rows
+        manifest_uri, step = container.load_and_commit(
+            LoadAndScanSuite.schema_name, self.generator, 1000
+        )
+        run.add(step)
+        # do a full scan of all rows (and eagerly load them)
+        keys, step = container.scan()
+        run.add(step)
+        # randomly retrieve all keys one-by-one from the dataset
+        random_keys = list(keys)
+        shuffle(random_keys)
+        step = container.run_queries(
+            "load all keys individually",
+            manifest_uri,
+            [QueryExpression().with_key(k) for k in random_keys],
+        )
+        run.add(step)
+        # split into 4 key ranges and get them individually
+        quartiles = self._generate_quartiles(keys)
+        expressions = [
+            QueryExpression().with_range(start, end) for (start, end) in quartiles
+        ]
+        step = container.run_queries(
+            "load key ranges by quartile", manifest_uri, expressions
+        )
+        run.add(step)
+        return run
+
+    @staticmethod
+    def _generate_quartiles(keys):
+        sorted_keys = sorted(keys)
+        size = len(keys)
+        starts = list(range(0, size, math.ceil(size / 4)))
+        ends = list([x - 1 for x in starts[1:]])
+        ends.append(size - 1)
+        quartiles = list(zip(starts, ends))
+        return [(sorted_keys[start], sorted_keys[end]) for (start, end) in quartiles]
+
+
+def test_suite1(schema: Schema, report: BenchmarkReport):
+    with temp_dir_autocleanup() as temp_dir:
+        generator = RandomRowGenerator(123, temp_dir)
+        report.add(
+            LoadAndScanSuite(
+                Dataset(dataset_name="test_suite1_ds1", metadata_uri=temp_dir),
+                schema,
+                generator,
+                "SST (rand)",
+            ).run()
+        )
+
+    with temp_dir_autocleanup() as temp_dir:
+        generator = RandomRowGenerator(123, temp_dir)
+        report.add(
+            LoadAndScanSuite(
+                Dataset(dataset_name="test_suite1_ds2", metadata_uri=temp_dir),
+                schema,
+                generator,
+                "dupe",
+            ).run()
+        )

deltacat/catalog/__init__.py

@@ -0,0 +1,73 @@
+from deltacat.catalog.delegate import (
+    alter_namespace,
+    alter_table,
+    create_namespace,
+    create_table,
+    default_namespace,
+    drop_namespace,
+    drop_table,
+    get_namespace,
+    get_table,
+    list_namespaces,
+    list_tables,
+    namespace_exists,
+    read_table,
+    refresh_table,
+    rename_table,
+    table_exists,
+    truncate_table,
+    write_to_table,
+)
+from deltacat.catalog.model.catalog import (  # noqa: F401
+    all_catalogs,
+    init,
+    init_local,
+    is_initialized,
+    clear_catalogs,
+    get_catalog,
+    pop_catalog,
+    put_catalog,
+    raise_if_not_initialized,
+    Catalog,
+)
+from deltacat.catalog.model.properties import (  # noqa: F401
+    CatalogProperties,
+    get_catalog_properties,
+)
+from deltacat.catalog.model.table_definition import TableDefinition
+from deltacat.catalog.main import impl as dcat
+
+__all__ = [
+    "alter_namespace",
+    "alter_table",
+    "create_namespace",
+    "create_table",
+    "default_namespace",
+    "drop_namespace",
+    "drop_table",
+    "get_namespace",
+    "get_table",
+    "list_namespaces",
+    "list_tables",
+    "namespace_exists",
+    "read_table",
+    "refresh_table",
+    "rename_table",
+    "table_exists",
+    "truncate_table",
+    "write_to_table",
+    "all_catalogs",
+    "init",
+    "init_local",
+    "is_initialized",
+    "clear_catalogs",
+    "get_catalog",
+    "get_catalog_properties",
+    "pop_catalog",
+    "put_catalog",
+    "raise_if_not_initialized",
+    "dcat",
+    "Catalog",
+    "CatalogProperties",
+    "TableDefinition",
+]