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

vastdb/bench/perf_bench/common/utils.py

@@ -0,0 +1,219 @@
+import importlib.util
+import logging
+import os
+import sys
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from types import MappingProxyType, ModuleType
+from typing import (
+    Any,
+    Dict,
+    Generator,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import exchange_calendars as xcal
+import fsspec
+import pandas as pd
+from fsspec import AbstractFileSystem
+
+# noinspection PyProtectedMember
+from pandas._typing import IntervalClosedType
+from pyarrow import dataset as ds
+
+from vastdb.bench.perf_bench.common.constants import (
+    DEFAULT_ACCESS_KEY,
+    DEFAULT_END_T,
+    DEFAULT_ROW_GROUP_SIZE,
+    DEFAULT_S3_HOST,
+    DEFAULT_SECRET_KEY,
+    DEFAULT_START_T,
+    DFAULT_PARQUET_COMPRESSION,
+    VASTDB_ENDPOINT,
+    ParquetCompression,
+)
+from vastdb.bench.perf_bench.common.log_utils import get_logger
+from vastdb.bench.perf_bench.common.types import DateLikeT, PathLikeT
+from vastdb.session import Session
+
+_FS_MAP = MappingProxyType({"/mnt": "nfs", "/": "fs", "s3://": "s3", "": "vastdb"})
+
+
+def config_ipython():
+    try:
+        # noinspection PyUnresolvedReferences
+        _ = get_ipython
+        pd.set_option("display.float_format", lambda x: "%.6f" % x)
+        pd.set_option("display.max_rows", 50)
+        pd.set_option("display.max_columns", 50)
+        pd.set_option("display.width", 1000)
+    except NameError:
+        ...
+
+
+# If in interactive mode, make it pretty
+config_ipython()
+
+
+def getenv_flag(name: str, default: bool = False) -> bool:
+    return os.getenv(name, str(default)).lower() in ("1", "true", "yes")
+
+
+def to_ts(d: Optional[DateLikeT], normalize: bool = True) -> Optional[pd.Timestamp]:
+    if not d:
+        return None
+    ret_d: pd.Timestamp = d if isinstance(d, pd.Timestamp) else pd.Timestamp(d)
+    return ret_d.normalize() if normalize else ret_d
+
+
+def _default_s3_kwargs(
+    aws_access_key_id: str = DEFAULT_ACCESS_KEY,
+    aws_secret_access_key: str = DEFAULT_SECRET_KEY,
+    host: str = DEFAULT_S3_HOST,
+    config: Optional[Dict[str, Any]] = None,
+    **kwargs,
+) -> Dict[str, Any]:
+    return {
+        "endpoint_url": f"http://{host}",
+        "aws_access_key_id": aws_access_key_id,
+        "aws_secret_access_key": aws_secret_access_key,
+        "use_ssl": False,
+        "verify": False,
+        "config": {
+            "signature_version": "s3v4",
+            "s3": {"addressing_style": "path"},
+            **(config or {}),
+        },
+        "region_name": "us-east-1",
+        **kwargs,
+    }
+
+
+def get_vastdb_session(
+    access: str = DEFAULT_ACCESS_KEY,
+    secret: str = DEFAULT_SECRET_KEY,
+    vastdb_endpoint: str = VASTDB_ENDPOINT,
+    ssl_verify: bool = True,
+) -> Session:
+    return Session(
+        endpoint=vastdb_endpoint, access=access, secret=secret, ssl_verify=ssl_verify
+    )
+
+
+def get_filesystem(
+    path: PathLikeT,
+    botocore_client_kwargs: Optional[Dict[str, Any]] = None,
+    **kwargs,
+) -> Tuple[str, AbstractFileSystem]:
+    path = str(path)
+    if path.startswith("s3://"):
+        botocore_client_kwargs = _default_s3_kwargs(
+            **(botocore_client_kwargs or {}),
+            config=kwargs.pop("config_kwargs", None),
+        )
+        fs = fsspec.filesystem(
+            protocol="s3",
+            client_kwargs=botocore_client_kwargs,
+            **kwargs,
+        )
+        path = path[5:]  # remove the s3:// prefix
+    else:
+        fs = fsspec.filesystem(protocol="file", **kwargs)
+    return path, fs
+
+
+# noinspection PyShadowingBuiltins
+def get_parquet_dataset(
+    path: PathLikeT,
+    filesystem: Optional[AbstractFileSystem] = None,
+    fs_kwargs: Optional[Dict[str, Any]] = None,
+    format: Optional[str] = "parquet",  # noqa: A002
+    partitioning: Optional[str] = "hive",
+    **kwargs,
+) -> ds.Dataset:
+    path, fs = filesystem or get_filesystem(path=path, **(fs_kwargs or {}))
+    return ds.dataset(
+        source=path,
+        format=format,
+        partitioning=partitioning,
+        filesystem=fs,
+        **kwargs,
+    )
+
+
+def get_parquet_dataset_root(
+    base_dir: PathLikeT,
+    row_group_size: int = DEFAULT_ROW_GROUP_SIZE,
+    compression: Union[str, ParquetCompression] = DFAULT_PARQUET_COMPRESSION,
+) -> Path:
+    return Path(str(base_dir)) / f"rg_{row_group_size!s}_c_{ParquetCompression[compression]}"
+
+
+@contextmanager
+def time_me(logger: Optional[logging.Logger] = None):
+    logger = logger or get_logger(__name__)
+    start_t = time.time_ns()
+    yield  # Yield control back to the context block
+    total_t_sec = (time.time_ns() - start_t) / 1e9
+    logger.debug(f"Execution time: {total_t_sec:.2f} s")
+
+
+S = TypeVar("S", bound=Sequence)
+
+
+def chunk_sequence(input_seq: S, chunk_size: int) -> Generator[Sequence, None, None]:
+    for i in range(0, len(input_seq), chunk_size):
+        yield input_seq[i: i + chunk_size]
+
+
+def get_dates_range(
+    from_t: Optional[DateLikeT] = DEFAULT_START_T,
+    to_t: Optional[DateLikeT] = DEFAULT_END_T,
+    only_bdays: bool = True,
+    inclusive: IntervalClosedType = "both",
+) -> pd.DatetimeIndex:
+    fun = pd.bdate_range if only_bdays else pd.date_range
+    return fun(  # type: ignore[operator]
+        start=to_ts(from_t),
+        end=to_ts(to_t),
+        inclusive=inclusive,
+    )
+
+
+def infer_fs_type(path: str) -> str:
+    return next(k for m, k in _FS_MAP.items() if str(path).lower().startswith(m))
+
+
+def get_session_minutes(
+    date: DateLikeT,
+    exchange: str = "XNYS",
+) -> pd.DatetimeIndex:
+    ecal = xcal.get_calendar(exchange)
+    date = pd.Timestamp(date).normalize()
+    sess = ecal.date_to_session(date=date, direction="next")
+    sess_minutes = ecal.session_minutes(sess)
+    return sess_minutes
+
+
+def load_module_from_path(module_file_path: PathLikeT) -> ModuleType:
+    file_path = Path(str(module_file_path)).resolve()
+    module_name = file_path.stem
+
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot find module {module_name} at {file_path}")
+
+    module = importlib.util.module_from_spec(spec)
+    loader = spec.loader
+    if loader is None or not hasattr(loader, "exec_module"):
+        raise ImportError(f"Cannot load module {module_name} from {file_path}")
+
+    loader.exec_module(module)
+    sys.modules[module_name] = module
+
+    return module

vastdb/bench/perf_bench/dataset/generate_secmaster.py

@@ -0,0 +1,105 @@
+from pathlib import Path
+from typing import NamedTuple, TypedDict
+
+import requests
+
+from vastdb.bench.perf_bench.common.log_utils import get_logger
+
+_MY_DIR = Path(__file__).parent
+SM_PATH = _MY_DIR / "secmaster.py"
+
+LOG = get_logger(__name__)
+
+
+class NasdaqRawRecord(TypedDict):
+    symbol: str
+    name: str
+    lastsale: str
+    netchange: str
+    pctchange: str
+    marketCap: str
+    url: str
+
+
+class NasdaqRecord(NamedTuple):
+    symbol: str
+    id: int
+    last_sale: float
+
+    @staticmethod
+    def id_from_ticker(ticker: str) -> int:
+        base = 100
+        max_width = 5
+        offset = ord(" ")
+        if len(ticker := ticker.strip()) > max_width:
+            raise ValueError(f"Ticker too long: {ticker}")
+        ticker = ticker.rjust(max_width, " ")
+        return sum(
+            (ord(c) - offset) * base ** (len(ticker) - i)
+            for i, c in enumerate(ticker.upper())
+        )
+
+    @classmethod
+    def from_raw_dict(cls, raw_dict: NasdaqRawRecord) -> "NasdaqRecord":
+        return cls(
+            symbol=(sym := raw_dict["symbol"].strip().upper()),
+            id=cls.id_from_ticker(sym),
+            last_sale=float(raw_dict["lastsale"].replace("$", "").replace(",", "")),
+        )
+
+
+def generate_secmaster():
+    resp = requests.get(
+        "http://api.nasdaq.com/api/screener/stocks?limit=3000",
+        headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"},
+    )
+    resp.raise_for_status()
+    resp_j = resp.json()
+    records = [
+        NasdaqRecord.from_raw_dict(r)
+        for r in resp_j["data"]["table"]["rows"]
+        if r and r["symbol"].strip()
+    ]
+    _ticker_to_sid = {r.symbol: r.id for r in records}
+    _sid_to_ticker = {r.id: r.symbol for r in records}
+    _indicative_px = {r.id: r.last_sale for r in records}
+
+    uni = sorted(_ticker_to_sid)
+    secm_file_contents = f"""
+_ticker_to_sid = {_ticker_to_sid}
+
+
+_sid_to_ticker = {_sid_to_ticker}
+
+
+_indicative_px = {_indicative_px}
+
+
+def to_sid(ticker: str) -> int:
+    return _ticker_to_sid[ticker]
+
+
+def to_ticker(sid: int) -> str:
+    return _sid_to_ticker[sid]
+
+
+def get_indicative_px(sid: int) -> float:
+    return _indicative_px[sid]
+
+
+UNI_SPEC = {{
+    "Large": (large_uni := {uni}),
+    "Single": large_uni[:1],
+    "Tiny": (tiny_uni := large_uni[1::50]),
+    "Small": (small_uni := large_uni[1::10]),
+    "SmallSeq": large_uni[: len(small_uni)],
+    "Medium": (med_uni := large_uni[1::6]),
+    "MediumSeq": large_uni[: (len(med_uni))],
+    "Medium2": (med2_uni := large_uni[1::8]),
+    "Medium2Seq": large_uni[: (len(med2_uni))],
+}}
+"""
+    with open(SM_PATH, "w") as f:
+        f.write(secm_file_contents)
+
+    LOG.info("Secmaster generated (total stocks: %d): %s", len(uni), SM_PATH)

vastdb/bench/perf_bench/dataset/generate_stocks_dataset.py

@@ -0,0 +1,242 @@
+import concurrent.futures
+import datetime as dt
+import os
+import time
+from pathlib import Path
+from typing import Optional, Sequence, Union
+
+import numpy as np
+import pandas as pd
+
+from vastdb.bench.perf_bench.common.constants import (
+    DEFAULT_END_T,
+    DEFAULT_ROW_GROUP_SIZE,
+    DEFAULT_START_T,
+    DFAULT_PARQUET_COMPRESSION,
+    LOCAL_FS_DS_PATH,
+    ParquetCompression,
+)
+from vastdb.bench.perf_bench.common.log_utils import get_logger
+from vastdb.bench.perf_bench.common.types import DateLikeT, PathLikeT
+from vastdb.bench.perf_bench.common.utils import (
+    get_parquet_dataset_root,
+    get_session_minutes,
+)
+from vastdb.bench.perf_bench.dataset import secmaster as sm
+from vastdb.bench.perf_bench.dataset.schemas import (
+    BF,
+    StockBarsPandasSchema,
+)
+
+LOG = get_logger(__name__)
+
+DEFAULT_NUM_WORKERS = max(1, int((os.cpu_count() or 1) * 0.4))
+DEFAULT_DS_ROOT_PATH = LOCAL_FS_DS_PATH
+
+ref_bars: Optional[pd.DataFrame] = None
+
+
+def generate_synthetic_stock_1m_bars_day(
+    date: Optional[DateLikeT] = None,
+    exchange: str = "XNYS",
+) -> pd.DataFrame:
+    sess_minutes = get_session_minutes(date=date or DEFAULT_START_T, exchange=exchange)
+    uni_tickers = sm.large_uni
+    sids = [sm.to_sid(t) for t in sm.large_uni]
+    ref_px = pd.Series(
+        [sm.get_indicative_px(sid) for sid in sids],
+        dtype=BF.trade_close.pd_type,
+    )
+    uni_sz = len(uni_tickers)
+
+    def _build_minute(m: pd.Timestamp) -> pd.DataFrame:
+        return pd.DataFrame(
+            {
+                BF.sid.name: sids,
+                BF.ts.name: [m.astimezone("UTC").tz_localize(None)] * uni_sz,  # type: ignore[arg-type]
+                BF.ticker.name: uni_tickers,
+                # Ask
+                BF.ask_open.name: (
+                    ao := ref_px * np.random.uniform(0.999, 0.001, uni_sz)
+                ),
+                BF.ask_high.name: ao * np.random.uniform(1.0, 1.001, uni_sz),
+                BF.ask_low.name: ao * np.random.uniform(0.999, 1.0, uni_sz),
+                BF.ask_close.name: ao * np.random.uniform(0.999, 1.001, uni_sz),
+                BF.ask_qty.name: np.random.randint(1, 10000, uni_sz),
+                # Bid
+                BF.bid_open.name: (
+                    bo := ref_px * np.random.uniform(0.999, 0.001, uni_sz)
+                ),
+                BF.bid_high.name: bo * np.random.uniform(1.0, 1.001, uni_sz),
+                BF.bid_low.name: bo * np.random.uniform(0.999, 1.0, uni_sz),
+                BF.bid_close.name: bo * np.random.uniform(0.999, 1.001, uni_sz),
+                BF.bid_qty.name: np.random.randint(1, 10000, uni_sz),
+                # Trades
+                BF.trade_open.name: (
+                    to := ref_px * np.random.uniform(0.999, 0.001, uni_sz)
+                ),
+                BF.trade_high.name: to * np.random.uniform(1.0, 1.001, uni_sz),
+                BF.trade_low.name: to * np.random.uniform(0.999, 1.0, uni_sz),
+                BF.trade_close.name: to * np.random.uniform(0.999, 1.001, uni_sz),
+                BF.trade_volume.name: np.random.randint(1, 10000, uni_sz),
+                # Vwap
+                BF.vwap.name: ref_px * np.random.uniform(0.999, 1.001, uni_sz),
+            }
+        )
+
+    # noinspection PyUnreachableCode
+    return (
+        pd.concat([_build_minute(m) for m in sess_minutes])
+        .astype(StockBarsPandasSchema)
+        .sort_values([BF.sid.name, BF.ts.name])
+        .reset_index(drop=True)
+    )
+
+
+def worker_init():
+    global ref_bars  # noqa: PLW0603
+    ref_bars = generate_synthetic_stock_1m_bars_day()
+
+
+# noinspection DuplicatedCode
+def build_bars(
+    dates: Union[Sequence[dt.date], pd.DatetimeIndex],
+    output_dir: PathLikeT = DEFAULT_DS_ROOT_PATH,
+    row_group_size: int = DEFAULT_ROW_GROUP_SIZE,
+    compression: Union[str, ParquetCompression] = DFAULT_PARQUET_COMPRESSION,
+):
+    global ref_bars  # noqa: PLW0602
+    if ref_bars is None:
+        raise ValueError("Reference bars not initialized")
+
+    output_dir = Path(str(output_dir))
+    compression = ParquetCompression[compression]
+    output_dir = get_parquet_dataset_root(output_dir, row_group_size, compression)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    LOG.info(
+        f"Building bars from {dates[0]:%Y-%m-%d} to {dates[-1]:%Y-%m-%d}, saving to {output_dir}"
+    )
+
+    # noinspection DuplicatedCode
+    def build_new_bar(new_date: dt.date) -> pd.DataFrame:
+        new_bars = ref_bars
+        schema = new_bars.dtypes
+        excluded_cols = {BF.ticker.value, BF.sid.value}
+        for c in new_bars:
+            if c in excluded_cols:
+                continue
+            kind: str = schema[c].kind  # type: ignore[call-overload]
+            if kind == "f":
+                new_bars[c] = new_bars[c] * np.random.uniform(0.8, 1.2)
+            elif kind == "i":
+                new_bars[c] = (new_bars[c] * np.random.uniform(0.2, 5)).astype(
+                    schema[c]  # type: ignore[call-overload]
+                )
+            elif kind == "M":
+                new_bars[c] = new_date + (new_bars[c] - new_bars[c].dt.normalize())  # type: ignore[operator]
+        return new_bars
+
+    for d in dates:
+        try:
+            start_t = time.time()
+
+            bars_df = build_new_bar(new_date=d)
+            generated_time = time.time() - start_t
+
+            fpth = (
+                output_dir
+                / f"year={d:%Y}"
+                / f"date={d:%Y%m%d}"
+                / f"{d:%Y%m%d}_equity_etf_1m_bars.pq"
+            )
+            fpth.parent.mkdir(parents=True, exist_ok=True)
+
+            # noinspection PyTypeChecker
+            bars_df.to_parquet(
+                path=fpth,
+                engine="pyarrow",
+                compression=compression,
+                index=False,
+                row_group_size=row_group_size,
+            )  # type: ignore[call-overload]
+            LOG.info(
+                f"Written [{row_group_size=} {compression=}]"
+                f"[total={(time.time() - start_t):.2f}, generate={generated_time:.2f}) seconds]: "
+                f"{fpth}"
+            )
+        except Exception as e:
+            LOG.exception(f"FAILURE with date {d}: {e}")
+
+
+def common(output_dir: PathLikeT = DEFAULT_DS_ROOT_PATH):
+    output_dir = Path(str(output_dir))
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    global ref_bars  # noqa: PLW0602
+    if ref_bars is None:
+        worker_init()
+
+
+def generate_synthetic_stock_1m_bars(
+    from_t: DateLikeT = DEFAULT_START_T,
+    to_t: DateLikeT = DEFAULT_END_T,
+    output_dir: PathLikeT = DEFAULT_DS_ROOT_PATH,
+    row_group_size: int = DEFAULT_ROW_GROUP_SIZE,
+    compression: Union[str, ParquetCompression] = DFAULT_PARQUET_COMPRESSION,
+):
+    common(output_dir=output_dir)
+    dates = pd.bdate_range(
+        pd.Timestamp(from_t).normalize(), pd.Timestamp(to_t).normalize()
+    )
+    build_bars(
+        dates,
+        output_dir=output_dir,
+        row_group_size=row_group_size,
+        compression=compression,
+    )
+
+
+def generate_concurrent_synthetic_stock_1m_bars(
+    from_t: DateLikeT = DEFAULT_START_T,
+    to_t: DateLikeT = DEFAULT_END_T,
+    output_dir: PathLikeT = DEFAULT_DS_ROOT_PATH,
+    num_workers: int = DEFAULT_NUM_WORKERS,
+    row_group_size: int = DEFAULT_ROW_GROUP_SIZE,
+    compression: Union[str, ParquetCompression] = DFAULT_PARQUET_COMPRESSION,
+):
+    if (num_workers := max(1, int(num_workers))) == 1:
+        return generate_synthetic_stock_1m_bars(
+            from_t=from_t,
+            to_t=to_t,
+            output_dir=output_dir,
+            row_group_size=row_group_size,
+            compression=compression,
+        )
+
+    common(output_dir=output_dir)
+    dates = pd.bdate_range(
+        pd.Timestamp(from_t).normalize(), pd.Timestamp(to_t).normalize()
+    )
+    batch_size = int(len(dates) / num_workers) + 1
+    batches = [dates[i: i + batch_size] for i in range(0, len(dates), batch_size)]
+
+    with concurrent.futures.ProcessPoolExecutor(
+        max_workers=num_workers,
+        initializer=worker_init,
+    ) as executor:
+        # Prepare a list of futures
+        futures = [
+            executor.submit(
+                build_bars,
+                dates=batch_dates,
+                output_dir=output_dir,
+                row_group_size=row_group_size,
+                compression=compression,
+            )
+            for batch_dates in batches
+        ]
+
+        # Iterate over futures to get their results
+        for future in concurrent.futures.as_completed(futures):
+            LOG.info(future.result())

vastdb/bench/perf_bench/dataset/schemas.py

@@ -0,0 +1,101 @@
+from functools import lru_cache
+from types import MappingProxyType
+from typing import Dict, Mapping, Union
+
+import numpy as np
+import pyarrow as pa
+from pandas.core.dtypes.base import ExtensionDtype
+
+from vastdb.bench.perf_bench.common.types import StrEnum
+
+PandasDTypeT = Union[ExtensionDtype, np.dtype]
+
+
+class StockBarField(StrEnum):
+    sid = "sid"
+    ts = "ts"
+    ticker = "ticker"
+
+    # Ask
+    ask_open = "ask_open"
+    ask_high = "ask_high"
+    ask_low = "ask_low"
+    ask_close = "ask_close"
+    ask_qty = "ask_qty"
+
+    # Bid
+    bid_open = "bid_open"
+    bid_high = "bid_high"
+    bid_low = "bid_low"
+    bid_close = "bid_close"
+    bid_qty = "bid_qty"
+
+    # Trades
+    trade_open = "trade_open"
+    trade_high = "trade_high"
+    trade_low = "trade_low"
+    trade_close = "trade_close"
+    trade_volume = "trade_volume"
+
+    # VWAP
+    vwap = "vwap"
+
+    @property
+    def pa_type(self) -> pa.DataType:
+        return _get_field_pyarrow_types()[self]
+
+    @property
+    def pd_type(self) -> PandasDTypeT:
+        return self.pa_type.to_pandas_dtype()
+
+
+BF = StockBarField
+
+
+@lru_cache
+def _get_field_pyarrow_types() -> Mapping[str, pa.DataType]:
+    return MappingProxyType(
+        {
+            BF.sid.value: pa.int64(),
+            BF.ts.value: pa.timestamp(unit="ns"),
+            BF.ticker.value: pa.string(),
+            BF.ask_open.value: pa.float64(),
+            BF.ask_high.value: pa.float64(),
+            BF.ask_low.value: pa.float64(),
+            BF.ask_close.value: pa.float64(),
+            BF.ask_qty.value: pa.int64(),
+            BF.bid_open.value: pa.float64(),
+            BF.bid_high.value: pa.float64(),
+            BF.bid_low.value: pa.float64(),
+            BF.bid_close.value: pa.float64(),
+            BF.bid_qty.value: pa.int64(),
+            BF.trade_open.value: pa.float64(),
+            BF.trade_high.value: pa.float64(),
+            BF.trade_low.value: pa.float64(),
+            BF.trade_close.value: pa.float64(),
+            BF.trade_volume.value: pa.int64(),
+            BF.vwap.value: pa.float64(),
+        }
+    )
+
+
+DEFAULT_BARS_COLUMNS = (
+    BF.sid.value,
+    BF.ts.value,
+    BF.ask_open.value,
+    BF.ask_close.value,
+    BF.bid_open.value,
+    BF.bid_close.value,
+    BF.bid_qty.value,
+)
+
+# noinspection PyUnresolvedReferences
+StockBarsArrowSchema: pa.Schema = pa.schema(
+    (fld.value, fld.pa_type) for fld in StockBarField
+)
+# noinspection PyUnresolvedReferences
+StockBarsPandasSchema: Dict[str, PandasDTypeT] = {
+    fld.value: fld.pd_type for fld in StockBarField
+}
+
+BarsSortFields = (BF.sid.value, BF.ts.value)