vastdb 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

vastdb/bench/perf_bench/orchestrate/bench_spec.py

@@ -0,0 +1,91 @@
+import datetime as dt
+import os
+from typing import Dict, List, NamedTuple, Optional, Union
+
+import pandas as pd
+
+DEFAULT_CSV_SEPARATOR = ","
+
+
+class BenchResult(NamedTuple):
+    # The identifier for this result
+    key: str
+    round: int
+
+    # Host details and time
+    host: str
+    pid: int
+
+    # Results details
+    start_ts: pd.Timestamp
+    end_ts: pd.Timestamp
+    duration_sec: float
+    n_rows: int
+    n_cols: int
+    n_bytes: int
+
+    params: Optional[Dict[str, Union[str, int, float]]] = None
+
+    def to_dict(
+        self,
+        include_params: bool = True,
+        flatten_params: bool = True,
+    ) -> Dict[str, Union[str, int, float]]:
+        dc = {f: getattr(self, f) for f in self._fields if f != "params"}
+        if include_params:
+            params = self.params or {}
+            if flatten_params:
+                dc = {**dc, **{pk: pv for pk, pv in params.items() if pk not in dc}}
+            else:
+                dc["params"] = params
+        return dc
+
+    def fields(
+        self,
+        include_params: bool = True,
+        flatten_params: bool = True,
+    ) -> List[str]:
+        return list(
+            self.to_dict(include_params=include_params, flatten_params=flatten_params)
+        )
+
+    def csv_header(
+        self,
+        separator: str = DEFAULT_CSV_SEPARATOR,
+        include_params: bool = True,
+    ) -> str:
+        return separator.join(
+            self.to_dict(include_params=include_params, flatten_params=True)
+        )
+
+    def to_csv(
+        self,
+        separator: str = DEFAULT_CSV_SEPARATOR,
+        include_params: bool = True,
+        include_header: bool = False,
+    ) -> str:
+        def _to_str(v) -> str:
+            if isinstance(v, dt.date):
+                v = v.isoformat()
+            v_str = str(v) if not isinstance(v, str) else v
+            if separator in v_str:
+                if '"' in v_str:
+                    raise ValueError(f"Can't handle double quotes in value: {v_str}")
+                v_str = f'"{v_str}"'
+            return v_str
+
+        data = [
+            _to_str(d)
+            for d in self.to_dict(include_params=include_params, flatten_params=True)
+        ]
+        data_csv = separator.join(data)
+
+        if include_header:
+            data_csv = os.linesep.join(
+                [
+                    self.csv_header(separator=separator, include_params=include_params),
+                    data_csv,
+                ]
+            )
+
+        return data_csv

vastdb/bench/perf_bench/orchestrate/results_helpers.py

@@ -0,0 +1,126 @@
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+import pandas as pd
+
+from vastdb.bench.perf_bench.common.log_utils import get_logger
+from vastdb.bench.perf_bench.orchestrate.bench_spec import BenchResult
+
+LOG = get_logger(__name__)
+
+PARAM_NAMES_SEP = "|"
+
+
+def results_to_df(
+    results: List[BenchResult],
+) -> pd.DataFrame:
+    results_df = pd.DataFrame([r.to_dict() for r in results])
+    results_df["param_names"] = PARAM_NAMES_SEP.join(
+        [c for c in (results[0].params or ()) if c in results_df]
+    )
+    return results_df
+
+
+def save_results(
+    results: List[BenchResult],
+    results_path: Union[str, Path],
+    append: bool = False,
+):
+    if not results:
+        return
+
+    results_path = Path(results_path)
+    results_path.parent.mkdir(parents=True, exist_ok=True)
+    results_df = results_to_df(results)
+
+    if append and results_path.is_file():
+        prev_results = pd.read_csv(results_path).astype(results_df.dtypes)
+        results_df = pd.concat([prev_results, results_df], ignore_index=True)
+
+    results_df.to_csv(results_path, index=False)
+    LOG.info(f"Results written to {results_path}")
+
+
+def calc_total_time_coverage_seconds(results_df: pd.DataFrame) -> float:
+    # Step 1: Separate start and end times, tagging them differently
+    starts = results_df["start_ts"].rename("time").to_frame().assign(event=1)
+    ends = results_df["end_ts"].rename("time").to_frame().assign(event=-1)
+
+    # Combine and sort chronologically
+    timeline = pd.concat([starts, ends]).sort_values("time")
+
+    # Step 2 & 3: Calculate cumulative sum to identify active intervals
+    timeline["active_intervals"] = timeline["event"].cumsum()
+
+    # Identify transitions (changes in active_intervals)
+    timeline["change"] = timeline["active_intervals"].diff().fillna(1).astype(bool)
+
+    # Calculate durations between changes
+    timeline["duration"] = timeline["time"].diff().shift(-1).fillna(pd.Timedelta(0))
+
+    # Only consider periods where at least one interval is active
+    total_coverage: pd.Timedelta = timeline.query("active_intervals > 0")[
+        "duration"
+    ].sum()
+    return total_coverage.total_seconds()
+
+
+def calculate_aggregate_stats(
+    results: Optional[List[BenchResult]] = None,
+    results_df: Optional[pd.DataFrame] = None,
+    results_path: Union[None, str, Path] = None,
+) -> pd.DataFrame:
+    if sum(bool(x) for x in (results, results_df, results_path)) != 1:
+        raise ValueError(
+            "Exactly one of results, results_df, or results_path must be provided"
+        )
+
+    if results:
+        r_df = results_to_df(results)
+        group_flds = list(results[0].params or ())
+    else:
+        r_df: pd.DataFrame = (  # type: ignore[assignment,no-redef]
+            pd.read_csv(results_path) if results_path else results_df
+        )
+        p_names: pd.Series = results_df["param_names"]  # type: ignore[index]
+        group_flds = list(p_names.values[0].split(PARAM_NAMES_SEP))
+    group_flds = ["key", "n_cols", *group_flds]
+
+    if not group_flds:
+        raise ValueError("No group fields found")
+
+    def _get_agg(col: str) -> str:
+        mapper = {
+            "start_ts": "min",
+            "end_ts": "max",
+            "round": "max",
+            "duration_sec": "mean",
+        }
+        return mapper.get(
+            col, "sum" if np.issubdtype(r_df.dtypes[col], np.number) else "last"
+        )
+
+    non_group_flds = [c for c in r_df.columns if c not in group_flds]
+    agg_df = (
+        r_df.groupby(group_flds)
+        .aggregate(
+            {c: _get_agg(c) for c in non_group_flds},
+            total=("duration_sec", "count"),
+        )
+        .drop(columns=["pid", "host", "param_names", "round", "duration_sec"])
+        .sort_index()
+    )
+    agg_df["duration_sec"] = (
+        r_df.groupby(group_flds)
+        .apply(calc_total_time_coverage_seconds, include_groups=False)
+        .sort_index()
+    )
+    agg_df["M_rows_per_sec"] = (agg_df["n_rows"] / agg_df["duration_sec"] / 1e6).astype(
+        "float64"
+    )
+    agg_df["MB_per_sec"] = (
+        agg_df["n_bytes"] / 1024 ** 2 / agg_df["duration_sec"]
+    ).astype("float64")
+
+    return agg_df

vastdb/bench/perf_bench/orchestrate/scenario.py

@@ -0,0 +1,109 @@
+import concurrent
+import os
+import socket
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+
+from vastdb.bench.perf_bench.orchestrate.bench_spec import BenchResult
+from vastdb.bench.perf_bench.orchestrate.results_helpers import LOG
+
+
+@dataclass
+class BenchScenario:
+    key: str
+    func: Callable
+    kwargs: Dict[str, Any]
+    results: Optional[Tuple[BenchResult, ...]] = tuple()
+    default_result_params: Optional[Dict[str, Union[str, int, float]]] = None
+
+    # noinspection PyUnresolvedReferences
+    def run(
+        self,
+        n_runs: int = 1,
+        discard_first_run: bool = True,
+        parallelism: int = 1,
+        workers_init: Optional[Callable] = None,
+    ):
+        if parallelism == 1:
+            self.results = tuple(
+                self._do_run(n_runs=n_runs, discard_first_run=discard_first_run)
+            )
+        elif parallelism > 1:
+            with concurrent.futures.ProcessPoolExecutor(
+                max_workers=parallelism,
+                initializer=workers_init,
+            ) as executor:
+                futures = [
+                    executor.submit(
+                        self._do_run, n_runs=n_runs, discard_first_run=discard_first_run
+                    )
+                    for _ in range(parallelism)
+                ]
+
+                # Iterate over futures to get their results
+                results = []
+                for future in concurrent.futures.as_completed(futures):
+                    LOG.info("Future completed")
+                    results.extend(future.result())
+                self.results = tuple(results)
+        else:
+            raise ValueError("Can't use a negative number of parallel benches")
+
+    def _do_run(
+        self,
+        n_runs: int = 1,
+        discard_first_run: bool = True,
+    ) -> List[BenchResult]:
+        hostname = socket.gethostname()
+        pid = os.getpid()
+        results = []
+        kwargs = self.kwargs or {}
+
+        for i in range(n_runs + int(discard_first_run)):
+            key = self.key
+
+            start_ts = time.time_ns()
+            ret_v = self.func(**kwargs)
+            end_ts = time.time_ns()
+
+            if discard_first_run and i == 0:
+                LOG.info(f"Discarding first run ({key})")
+                continue
+
+            n_bytes, n_rows, n_cols = 0, 0, 0
+
+            if hasattr(ret_v, "shape"):
+                shape = ret_v.shape
+                if len(shape):
+                    n_rows = shape[0]
+                if len(shape) > 1:
+                    n_cols = shape[1]
+
+            if hasattr(ret_v, "nbytes"):
+                n_bytes = ret_v.nbytes
+            elif hasattr(ret_v, "memory_usage"):
+                n_bytes = ret_v.memory_usage().sum()
+            elif hasattr(ret_v, "estimated_size"):
+                n_bytes = ret_v.estimated_size()
+
+            result = BenchResult(
+                key=key,
+                round=i,
+                host=hostname,
+                pid=pid,
+                start_ts=pd.Timestamp(start_ts, unit="ns"),
+                end_ts=pd.Timestamp(end_ts, unit="ns"),
+                duration_sec=(end_ts - start_ts) / 1e9,
+                n_rows=n_rows,
+                n_cols=n_cols,
+                n_bytes=n_bytes,
+                params=self.default_result_params or {},
+            )
+            results.append(result)
+
+            LOG.info(f"{result.duration_sec} sec")
+
+        return results

vastdb/bench/perf_bench/orchestrate/scenario_generator.py

@@ -0,0 +1,144 @@
+import datetime as dt
+from itertools import product
+from typing import Dict, List, Optional, Sequence, Union
+
+import pandas as pd
+
+from vastdb.bench.perf_bench.common.constants import (
+    DEFAULT_ARROW_KWARGS,
+    DEFAULT_ROW_GROUP_SIZE,
+    DEFAULT_START_T,
+    DFAULT_PARQUET_COMPRESSION,
+    ParquetCompression,
+    VastConnDetails,
+)
+from vastdb.bench.perf_bench.common.types import PathLikeT
+from vastdb.bench.perf_bench.common.utils import get_parquet_dataset_root, infer_fs_type
+from vastdb.bench.perf_bench.dataset import secmaster as sm
+from vastdb.bench.perf_bench.dataset.schemas import DEFAULT_BARS_COLUMNS
+from vastdb.bench.perf_bench.orchestrate.scenario import BenchScenario
+from vastdb.bench.perf_bench.query.query import QueryBackend
+
+
+def generate_perf_bench_scenarios(
+    base_key: str,
+    conn_details: VastConnDetails,
+    # Common parameters
+    query_backends: Sequence[Union[str, QueryBackend]] = (QueryBackend.pyarrow,),
+    universe_choices: Sequence[Union[str, Sequence[str]]] = ("Single",),
+    columns_choices: Sequence[Sequence[str]] = (DEFAULT_BARS_COLUMNS,),
+    dt_start_t: Union[dt.date, str] = DEFAULT_START_T,
+    num_bdays: Sequence[int] = (1,),
+    # Arrow-specific options
+    fs_path_choices: Optional[Sequence[PathLikeT]] = None,
+    rowgroup_size_choices: Optional[Sequence[int]] = None,
+    compression_choices: Optional[Sequence[Union[str, ParquetCompression]]] = None,
+    arrow_batching_spec_choices: Optional[Sequence[Dict[str, int]]] = None,
+    # VastDB-specific options
+    vdb_num_sub_splits_choices: Optional[Sequence[int]] = None,
+    vdb_num_row_groups_per_sub_split_choices: Optional[Sequence[int]] = None,
+) -> List[BenchScenario]:
+    if not base_key:
+        raise ValueError("base_key must be provided")
+
+    uni_choices: List[List[str]] = []
+    for k in universe_choices:
+        if isinstance(k, str):
+            uni_choices.append(sm.UNI_SPEC.get(k, [k]))
+        else:
+            uni_choices.append(list(k))
+
+    columns_choices = [[c for c in cs if c] for cs in columns_choices if cs]
+    dt_start_t = pd.Timestamp(dt_start_t)
+    dt_range_choices = [
+        (dt_start_t, dt_start_t + pd.tseries.offsets.BDay(d))
+        for d in sorted(set(num_bdays))
+        if d > 0
+    ]
+
+    scenarios = []
+    for i, (  # type: ignore[misc]
+        uni,
+        columns,
+        (from_t, to_t),
+        path,
+        rowgroup_size,
+        compression,
+        arrow_batch_spec,
+        vdb_num_sub_splits,
+        vdb_num_row_groups_per_sub_split,
+        query_backend,
+    ) in enumerate(
+        product(
+            uni_choices,
+            columns_choices,
+            dt_range_choices,
+            # Arrow specific options
+            (*(fs_path_choices or []), None),
+            (*(rowgroup_size_choices or [DEFAULT_ROW_GROUP_SIZE]), None),
+            (*(compression_choices or [DFAULT_PARQUET_COMPRESSION]), None),
+            (*(arrow_batching_spec_choices or [DEFAULT_ARROW_KWARGS]), None),
+            # VastDB specific options
+            (*(vdb_num_sub_splits_choices or [1]), None),
+            (*(vdb_num_row_groups_per_sub_split_choices or [8]), None),
+            [QueryBackend[qb] for qb in query_backends],
+        )
+    ):
+        if query_backend is QueryBackend.vastdb:
+            if not all((vdb_num_sub_splits, vdb_num_row_groups_per_sub_split)) or any(
+                (path, rowgroup_size, compression, arrow_batch_spec)
+            ):
+                # ignore path and rg size for vastdb runs
+                continue
+            backend_kwargs = {
+                "num_sub_splits": vdb_num_sub_splits,
+                "num_row_groups_per_sub_split": vdb_num_row_groups_per_sub_split,
+                "conn_details": conn_details,
+            }
+            path = None
+        elif query_backend is QueryBackend.pyarrow:
+            if not all((path, rowgroup_size, compression, arrow_batch_spec)) or any(
+                (vdb_num_sub_splits, vdb_num_row_groups_per_sub_split)
+            ):
+                # ignore non-vastdb runs without a path or rg size
+                continue
+            backend_kwargs: Dict[str, Union[str, int, VastConnDetails]] = arrow_batch_spec  # type: ignore[no-redef]
+            backend_kwargs["conn_details"] = conn_details
+            path = get_parquet_dataset_root(path, rowgroup_size, compression)  # type: ignore[arg-type]
+        else:
+            raise NotImplementedError(f"Unsupported query backend: {query_backend=!r}")
+
+        scen = BenchScenario(
+            key=f"{base_key}_{query_backend.name}",
+            func=query_backend.run_query,
+            kwargs=dict(
+                universe=uni,
+                columns=columns,
+                from_t=from_t,
+                to_t=to_t,
+                path=path,
+                backend_kwargs=backend_kwargs,
+            ),
+            default_result_params={
+                "uni_sz": len(uni),  # type: ignore
+                "num_days": (to_t - from_t).days,  # type: ignore
+                "query_backend": query_backend.name,  # type: ignore
+                "fs_type": infer_fs_type(path or ""),  # type: ignore
+                "rowgroup_size": rowgroup_size or -1,  # type: ignore
+                "arrow_fragment_readahead": (  # type: ignore
+                    arrow_batch_spec["fragment_readahead"] if arrow_batch_spec else -1  # type: ignore
+                ),
+                "arrow_batch_readahead": (  # type: ignore
+                    arrow_batch_spec["batch_readahead"] if arrow_batch_spec else -1  # type: ignore
+                ),
+                "arrow_batch_size": (  # type: ignore
+                    arrow_batch_spec["batch_size"] if arrow_batch_spec else -1  # type: ignore
+                ),
+                "vdb_num_sub_splits": vdb_num_sub_splits or -1,  # type: ignore
+                "vdb_num_row_groups_per_sub_split": vdb_num_row_groups_per_sub_split or -1,  # type: ignore
+            },
+        )
+
+        scenarios.append(scen)
+
+    return scenarios

vastdb/bench/perf_bench/query/arrow_common.py

@@ -0,0 +1,59 @@
+# ruff: noqa: F841
+
+import datetime as dt
+from functools import reduce
+from typing import Optional, Sequence
+
+import pyarrow as pa
+from pyarrow import dataset as ds
+
+from vastdb.bench.perf_bench.common import utils as bu
+from vastdb.bench.perf_bench.dataset import secmaster as sm
+from vastdb.bench.perf_bench.dataset.schemas import BF
+
+
+def build_arrow_filter(
+    universe: Sequence[str],
+    from_t: Optional[dt.datetime] = None,
+    to_t: Optional[dt.datetime] = None,
+    use_sid: bool = True,
+) -> ds.Expression:
+    # Build the filters
+    y_m_filter = from_t_filter = to_t_filter = None
+    from_t, to_t = bu.to_ts(from_t, normalize=False), bu.to_ts(to_t, normalize=False)
+    if from_t or to_t:
+        all_dates = bu.get_dates_range(from_t, to_t, inclusive="left")
+        years = sorted({d.year for d in all_dates})
+
+        y_m_filter = ds.field("year").isin(years) & ds.field("date").isin(
+            tuple(d.strftime("%Y%m%d") for d in all_dates)
+        )
+
+        if from_t and from_t.time() != dt.time.min:
+            from_t_filter = ds.field("ts") >= pa.scalar(from_t)
+        if to_t and to_t.time() != dt.time.min:
+            to_t_filter = ds.field("ts") < pa.scalar(to_t)
+
+    if use_sid:
+        universe = [sm.to_sid(s) for s in universe]  # type: ignore[misc]
+        fld = BF.sid.value
+    else:
+        fld = BF.ticker.value
+    universe = sorted(universe)
+    if len(universe) <= 16:
+        tickers_filter = reduce(
+            lambda x, y: x | y, [ds.field(fld) == s for s in universe]
+        )
+    else:
+        tickers_filter = ds.field(fld).isin(universe)
+
+    ds_filter = reduce(
+        lambda x, y: x & y,
+        [
+            flt
+            for flt in (y_m_filter, tickers_filter, from_t_filter, to_t_filter)
+            if flt is not None
+        ],
+    )
+
+    return ds_filter

vastdb/bench/perf_bench/query/query.py

@@ -0,0 +1,42 @@
+import datetime as dt
+import logging
+from typing import Any, Dict, Optional, Protocol, Sequence
+
+import pyarrow as pa
+
+from vastdb.bench.perf_bench.common.types import StrEnum
+
+
+class QueryRunner(Protocol):
+    def __call__(
+        self,
+        universe: Sequence[str],
+        columns: Optional[Sequence[str]],
+        from_t: Optional[dt.datetime],
+        to_t: Optional[dt.datetime],
+        path: Optional[str],
+        use_sid: bool,
+        logger: Optional[logging.Logger],
+        backend_kwargs: Optional[Dict[str, Any]],
+    ) -> pa.Table:
+        ...
+
+
+class QueryBackend(StrEnum):
+    vastdb = "vastdb"
+    pyarrow = "pyarrow"
+
+    @property
+    def _query_fun(self) -> QueryRunner:
+        if self is QueryBackend.vastdb:
+            from vastdb.bench.perf_bench.query.query_vastdb import query_vastdb
+
+            return query_vastdb
+        elif self is QueryBackend.pyarrow:
+            from vastdb.bench.perf_bench.query.query_pyarrow import query_pyarrow
+
+            return query_pyarrow
+        raise NotImplementedError(f"Unsupported query backend: {self}")
+
+    def run_query(self, *args, **kwargs) -> pa.Table:
+        return self._query_fun(*args, **kwargs)

vastdb/bench/perf_bench/query/query_pyarrow.py

@@ -0,0 +1,70 @@
+# ruff: noqa: F841
+
+import datetime as dt
+import logging
+from typing import Any, Dict, Optional, Sequence
+
+import pyarrow as pa
+
+import vastdb.bench.perf_bench.common.log_utils
+from vastdb.bench.perf_bench.common import utils as bu
+from vastdb.bench.perf_bench.common.constants import (
+    DEFAULT_ARROW_KWARGS,
+    VastConnDetails,
+)
+from vastdb.bench.perf_bench.query.arrow_common import build_arrow_filter
+
+LOG = vastdb.bench.perf_bench.common.log_utils.get_logger(__name__)
+
+
+# noinspection PyUnusedLocal
+def query_pyarrow(
+    universe: Sequence[str],
+    columns: Optional[Sequence[str]] = None,
+    from_t: Optional[dt.datetime] = None,
+    to_t: Optional[dt.datetime] = None,
+    path: Optional[str] = None,
+    use_sid: bool = True,
+    logger: Optional[logging.Logger] = None,
+    backend_kwargs: Optional[Dict[str, Any]] = None,
+) -> pa.Table:
+    # ------------------------------------------------------------
+    #    Query the PyArrow's parquet dataset
+    # ------------------------------------------------------------
+    backend_kwargs = backend_kwargs or {}
+    conn_details: Optional[VastConnDetails] = backend_kwargs.pop("conn_details", None)
+    conn_details = conn_details or VastConnDetails()
+    fs_kwargs = {
+        "botocore_client_kwargs": {
+            "aws_access_key_id": conn_details.access,
+            "aws_secret_access_key": conn_details.secret,
+            "host": conn_details.s3_host,
+        }
+    }
+
+    # Note that one can optimize the below by passing the filesystem object and/or partitioning
+    #  scheme to avoid the slow penalty of discovering the partitioning upon init.
+    dset = bu.get_parquet_dataset(path=path, filesystem=None, fs_kwargs=fs_kwargs)  # type: ignore[arg-type]
+
+    ds_filter = build_arrow_filter(
+        universe=universe,
+        from_t=from_t,
+        to_t=to_t,
+        use_sid=use_sid,
+    )
+
+    # Cleanup the kwargs to be passed to the arrow scanner
+    kwargs = backend_kwargs or DEFAULT_ARROW_KWARGS
+    kwargs.pop("arrow_kwargs", None)
+    if flt := kwargs.pop("filter", None):
+        ds_filter = ds_filter & flt
+
+    # Read the data
+    # noinspection PyArgumentList
+    table = dset.to_table(
+        columns=columns,
+        filter=ds_filter,
+        **{**DEFAULT_ARROW_KWARGS, **(kwargs or {})},
+    )
+
+    return table