vastdb 0.1.11__py3-none-any.whl → 1.1.0__py3-none-any.whl

vastdb/_internal.py

@@ -35,6 +35,7 @@ from ibis.expr.operations.logical import (
 )
 from ibis.expr.operations.relations import Field
 from ibis.expr.operations.strings import StringContains
+from ibis.expr.operations.structs import StructField
 
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.BinaryLiteral as fb_binary_lit
 import vast_flatbuf.org.apache.arrow.computeir.flatbuf.BooleanLiteral as fb_bool_lit
@@ -182,7 +183,7 @@ class Predicate:
                 _logger.debug('OR args: %s op %s', or_args, op)
                 inner_offsets = []
 
-                prev_field_name = None
+                prev_field_path = None
                 for inner_op in or_args:
                     _logger.debug('inner_op %s', inner_op)
                     op_type = type(inner_op)
@@ -216,28 +217,38 @@ class Predicate:
                             if not isinstance(literal, Literal):
                                 raise NotImplementedError(self.expr)
 
+                    field_path = []
+                    while isinstance(column, StructField):
+                        column, subfield_name = column.args
+                        field_path.append(subfield_name)
+
                     if not isinstance(column, Field):
                         raise NotImplementedError(self.expr)
 
-                    field_name = column.name
-                    if prev_field_name is None:
-                        prev_field_name = field_name
-                    elif prev_field_name != field_name:
+                    field_path.append(column.name)
+                    field_path.reverse()  # first entry should be the top-level column name
+
+                    if prev_field_path is None:
+                        prev_field_path = field_path
+                    elif prev_field_path != field_path:
                         raise NotImplementedError(self.expr)
 
-                    node = self.nodes_map[field_name]
+                    nodes_map = self.nodes_map
+                    for name in field_path:
+                        node = nodes_map[name]
+                        nodes_map = node.children_map
+
                     # TODO: support predicate pushdown for leaf nodes (ORION-160338)
                     if node.children:
                         raise NotImplementedError(node.field)  # no predicate pushdown for nested columns
                     column_offset = self.build_column(position=node.index)
-                    field = self.schema.field(field_name)
                     for literal in literals:
                         args_offsets = [column_offset]
                         if literal is not None:
-                            args_offsets.append(self.build_literal(field=field, value=literal.value))
+                            args_offsets.append(self.build_literal(field=node.field, value=literal.value))
                         if builder_func == self.build_between:
-                            args_offsets.append(self.build_literal(field=field, value=lower.value))
-                            args_offsets.append(self.build_literal(field=field, value=upper.value))
+                            args_offsets.append(self.build_literal(field=node.field, value=lower.value))
+                            args_offsets.append(self.build_literal(field=node.field, value=upper.value))
 
                         inner_offsets.append(builder_func(*args_offsets))
 
@@ -572,6 +583,8 @@ class FieldNode:
         else:
             self.children = []  # for non-nested types
 
+        self.children_map = {c.field.name: c for c in self.children}
+
     def _iter_to_root(self) -> Iterator['FieldNode']:
         yield self
         if self.parent is not None:
@@ -748,7 +761,6 @@ class VastdbApi:
 
     def __init__(self, endpoint, access_key, secret_key,
             *,
-            auth_type=AuthType.SIGV4,
             ssl_verify=True,
             timeout=None,
             backoff_config: Optional[BackoffConfig] = None):
@@ -766,15 +778,15 @@ class VastdbApi:
         self._session.verify = ssl_verify
         self._session.headers['user-agent'] = self.client_sdk_version
 
-        backoff_config = backoff_config or BackoffConfig()
+        self.backoff_config = backoff_config or BackoffConfig()
         self._backoff_decorator = backoff.on_exception(
-            wait_gen=backoff_config.wait_gen,
+            wait_gen=self.backoff_config.wait_gen,
             exception=_RETRIABLE_EXCEPTIONS,
             giveup=_backoff_giveup,
-            max_tries=backoff_config.max_tries,
-            max_time=backoff_config.max_time,
-            max_value=backoff_config.max_value,  # passed to `backoff_config.wait_gen`
-            backoff_log_level=backoff_config.backoff_log_level)
+            max_tries=self.backoff_config.max_tries,
+            max_time=self.backoff_config.max_time,
+            max_value=self.backoff_config.max_value,  # passed to `self.backoff_config.wait_gen`
+            backoff_log_level=self.backoff_config.backoff_log_level)
         self._request = self._backoff_decorator(self._single_request)
 
         if url.port in {80, 443, None}:
@@ -812,6 +824,23 @@ class VastdbApi:
         _logger.critical(msg)
         raise NotImplementedError(msg)
 
+    def __enter__(self):
+        """Allow using this session as a context manager."""
+        return self
+
+    def __exit__(self, *args):
+        """Make sure that the connections closed."""
+        self._session.close()
+
+    def with_endpoint(self, endpoint):
+        """Open a new session for targeting a specific endpoint."""
+        return VastdbApi(endpoint=endpoint,
+            access_key=self.access_key,
+            secret_key=self.secret_key,
+            ssl_verify=self._session.verify,
+            timeout=self.timeout,
+            backoff_config=self.backoff_config)
+
     def _single_request(self, *, method, url, skip_status_check=False, **kwargs):
         _logger.debug("Sending request: %s %s %s timeout=%s", method, url, kwargs, self.timeout)
         try:

vastdb/bench/perf_bench/bench_repo/mega_combo.py

@@ -0,0 +1,87 @@
+from typing import List, Optional
+
+from vastdb.bench.perf_bench.common.constants import (
+    LOCAL_FS_DS_PATH,
+    NFS_DS_PATH,  # noqa: F401
+    S3_DS_PATH,  # noqa: F401
+    ParquetCompression,
+    VastConnDetails,
+)
+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.orchestrate.scenario_generator import (
+    generate_perf_bench_scenarios,
+)
+from vastdb.bench.perf_bench.query.query import QueryBackend
+
+
+def build_scenarios(
+    base_key: str,
+    conn_details: Optional[VastConnDetails] = None,
+) -> List[BenchScenario]:
+    return generate_perf_bench_scenarios(
+        base_key=base_key,
+        conn_details=conn_details or VastConnDetails(),
+        query_backends=[
+            QueryBackend.pyarrow,
+            # QueryBackend.vastdb,
+        ],
+        columns_choices=(DEFAULT_BARS_COLUMNS,),
+        universe_choices=(
+            "Single",
+            "Tiny",
+            "SmallSeq",
+            "Medium",
+            "Medium2",
+            "Large",
+        ),
+        num_bdays=[
+            1,  # 1d
+            5,  # 1w
+            # 22,  # 1m
+            65,  # 3m
+            # 130,  # 6m
+            252,  # 1y
+        ],
+
+        # Arrow-specific options
+        fs_path_choices=[
+            # NFS_DS_PATH,
+            LOCAL_FS_DS_PATH,
+            # S3_DS_PATH,
+        ],
+        rowgroup_size_choices=[  # make sure you have previously generated the respective datasets
+            # 64 * 1024,
+            # 128 * 1024,
+            256 * 1024,
+            # 512 * 1024,
+            # DEFAULT_ROW_GROUP_SIZE,
+            # int(1.5 * 1024 * 1024),
+        ],
+        compression_choices=[
+            ParquetCompression.LZ4,
+        ],
+        arrow_batching_spec_choices=[
+            # {"batch_size": 2*2**16, "batch_readahead": 16, "fragment_readahead": 4},
+            # {"batch_size": 6*2**16, "batch_readahead": 12, "fragment_readahead": 4},
+            # DEFAULT_ARROW_KWARGS,
+            {"batch_size": 16 * 2 ** 16, "batch_readahead": 16, "fragment_readahead": 4},
+            # {"batch_size": 24 * 2 ** 16, "batch_readahead": 12, "fragment_readahead": 4},
+            # {"batch_size": 32*2**16, "batch_readahead": 12, "fragment_readahead": 4},
+            # {"batch_size": 64*2**16, "batch_readahead": 12, "fragment_readahead": 4},
+            # {"batch_size": 128*2**16, "batch_readahead": 12, "fragment_readahead": 4},
+        ],
+
+        # VastDB-specific options
+        vdb_num_sub_splits_choices=(
+            # 1,  # Default
+            # 4,
+            8,
+            # 16,
+        ),
+        vdb_num_row_groups_per_sub_split_choices=(
+            # 1,
+            # 4,
+            8,  # Default
+        ),
+    )

vastdb/bench/perf_bench/cli.py

@@ -0,0 +1,225 @@
+from pathlib import Path
+from typing import Annotated, List, Optional
+
+import pandas as pd
+import typer
+
+from vastdb.bench.perf_bench.common.constants import (
+    DEFAULT_END_T,
+    DEFAULT_RESULTS_DIR,
+    DEFAULT_START_T,
+    DFAULT_PARQUET_COMPRESSION,
+    LOCAL_FS_DS_PATH,
+    LogLevel,
+    ParquetCompression,
+)
+from vastdb.bench.perf_bench.common.log_utils import (
+    get_logger,
+    set_log_file,
+    set_log_level,
+)
+from vastdb.bench.perf_bench.common.utils import getenv_flag, load_module_from_path
+from vastdb.bench.perf_bench.dataset.generate_secmaster import (
+    SM_PATH,
+    generate_secmaster,
+)
+from vastdb.bench.perf_bench.dataset.generate_stocks_dataset import (
+    generate_concurrent_synthetic_stock_1m_bars,
+)
+from vastdb.bench.perf_bench.orchestrate.scenario import BenchScenario
+from vastdb.bench.perf_bench.run import run_scenarios
+
+app = typer.Typer(pretty_exceptions_enable=getenv_flag("TYPER_PRETTY_EXCEPTIONS"))
+
+_MY_DIR = Path(__file__).parent
+
+
+# noinspection PyUnusedLocal
+@app.callback()
+def cli_common(
+    ctx: typer.Context,
+    verbose: Annotated[
+        bool,
+        typer.Option(
+            "--verbose",
+            is_flag=True,
+        ),
+    ] = False,
+    log_level: Annotated[
+        Optional[LogLevel],
+        typer.Option(
+            "--log-level",
+            case_sensitive=False,
+        ),
+    ] = None,
+    log_file: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--log-file",
+            writable=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+        ),
+    ] = None,
+):
+    if verbose:
+        log_level = LogLevel.DEBUG
+    if log_level:
+        set_log_level(log_level)
+    if log_file:
+        set_log_file(log_file)
+    get_logger(__name__).info("CLI common setup done.")
+
+
+def _positive_int(value: str) -> int:
+    i_value = int(value)
+    if i_value <= 0:
+        raise typer.BadParameter(f"Must be a positive integer: {value}.")
+    return i_value
+
+
+# noinspection PyUnusedLocal
+@app.command()
+def run_bench(
+    ctx: typer.Context,
+    bench_name: Annotated[
+        str,
+        typer.Option(
+            "--bench-name",
+        ),
+    ],
+    parallelism: Annotated[
+        List[int],
+        typer.Option(
+            "--parallelism",
+            callback=lambda par: [_positive_int(p) for p in par],
+        ),
+    ],
+    runs_per_bench: Annotated[
+        int,
+        typer.Option(
+            "--runs-per-bench",
+            callback=_positive_int,
+        ),
+    ] = 3,
+    bench_generator_path: Annotated[
+        Path,
+        typer.Option(
+            "--bench-generator-path",
+            readable=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+        ),
+    ] = _MY_DIR / "bench_repo" / "mega_combo.py",
+    results_base_dir: Annotated[
+        Path,
+        typer.Option(
+            "--log-file",
+            writable=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+    ] = DEFAULT_RESULTS_DIR,
+):
+    if not (bench_name := bench_name.strip()):
+        raise typer.BadParameter("Bench name must be non-empty.")
+
+    mod = load_module_from_path(bench_generator_path)
+    scenarios: List[BenchScenario] = mod.build_scenarios(base_key=bench_name)
+    for para in parallelism:
+        run_scenarios(
+            scenarios=scenarios,
+            runs_per_bench=runs_per_bench,
+            parallelism=para,
+            results_base_dir=str(results_base_dir),
+        )
+
+
+# noinspection PyUnusedLocal
+@app.command()
+def build_secmaster(
+    ctx: typer.Context,
+):
+    generate_secmaster()
+
+
+# noinspection PyUnusedLocal
+@app.command()
+def build_dataset(
+    ctx: typer.Context,
+    start_date: Annotated[
+        str,
+        typer.Option(
+            "--start-date",
+            help="Start date for the dataset.",
+            callback=lambda d: pd.Timestamp(d).normalize(),
+        ),
+    ] = DEFAULT_START_T.strftime("%Y%m%d"),
+    end_date: Annotated[
+        str,
+        typer.Option(
+            "--end-date",
+            help="Start date for the dataset.",
+            callback=lambda d: pd.Timestamp(d).normalize(),
+        ),
+    ] = DEFAULT_END_T.strftime("%Y%m%d"),
+    output_dir: Annotated[
+        Path,
+        typer.Option(
+            "--output-dir",
+            writable=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+    ] = LOCAL_FS_DS_PATH,
+    parallelism: Annotated[
+        int,
+        typer.Option(
+            "--parallelism",
+            callback=_positive_int,
+        ),
+    ] = 6,
+    row_group_size: Annotated[
+        int,
+        typer.Option(
+            "--row-group-size",
+            callback=_positive_int,
+            help=(
+                "Row group size for the dataset, some common values are: 64 * 1024, 128 * 1024, 256"
+                " * 1024, 512 * 1024,1024 * 1024, 1.5 * 1024 * 1024."
+            ),
+        ),
+    ] = 256 * 1024,
+    compression: Annotated[
+        ParquetCompression,
+        typer.Option(
+            "--compression",
+            help="Parquet compression algorithm.",
+        ),
+    ] = DFAULT_PARQUET_COMPRESSION,
+):
+    if row_group_size < 1024:
+        raise typer.BadParameter("Row group size must be at least 1024.")
+    if parallelism < 1:
+        raise typer.BadParameter("Parallelism must be at least 1.")
+    if start_date > end_date:
+        raise typer.BadParameter("Start date must be before the end date.")
+    if not SM_PATH.is_file():
+        generate_secmaster()
+    generate_concurrent_synthetic_stock_1m_bars(
+        from_t=start_date,
+        to_t=end_date,
+        output_dir=output_dir,
+        num_workers=parallelism,
+        row_group_size=row_group_size,
+        compression=compression,
+    )
+
+
+if __name__ == "__main__":
+    # Set the metadata only if we execute the main (not on just importing this module)
+    app()

vastdb/bench/perf_bench/common/constants.py

@@ -0,0 +1,96 @@
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+from urllib3.util import parse_url
+
+from vastdb.bench.perf_bench.common.types import StrEnum
+
+_MY_DIR = Path(__file__).parent
+
+# VastDB details
+VASTDB_ENDPOINT = os.getenv("VASTDB_ENDPOINT", "")
+VASTDB_BUCKET_NAME = os.getenv("VASTDB_BUCKET_NAME", "")
+VASTDB_TEST_SCHEMA_NAME = os.getenv("VASTDB_TEST_SCHEMA_NAME", "")
+VASTDB_TEST_TABLE_NAME = os.getenv("VASTDB_TEST_TABLE_NAME", "")
+
+# Regular S3 details
+S3_BUCKET_NAME = "my-s3-bucket"
+DEFAULT_S3_SSL_PORT = int(os.getenv("DEFAULT_S3_SSL_PORT", 443))
+if DEFAULT_S3_ENDPOINT_URL := os.getenv("AWS_S3_ENDPOINT_URL", ""):
+    _parsed = parse_url(DEFAULT_S3_ENDPOINT_URL)
+    DEFAULT_S3_HOST = str(_parsed.host)
+    DEFAULT_S3_PORT = int(_parsed.port or 80)
+else:
+    DEFAULT_S3_HOST = os.getenv("DEFAULT_S3_HOST", "1.1.1.1")
+    DEFAULT_S3_PORT = int(os.getenv("DEFAULT_S3_PORT", 80))
+    DEFAULT_S3_ENDPOINT_URL = f"http://{DEFAULT_S3_HOST}:{DEFAULT_S3_PORT}"
+
+# Paths
+DEFAULT_RESULTS_DIR = Path(__file__).parent.parent / "benchmark_results"
+NFS_DS_PATH = f"/mnt/data/{VASTDB_BUCKET_NAME}/data"
+LOCAL_FS_DS_PATH = _MY_DIR.parent / "dataset" / "test_dataset"
+S3_DS_PATH = f"s3://{VASTDB_BUCKET_NAME}/data"
+
+# Access keys
+DEFAULT_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY_ID", "some_access_key")
+DEFAULT_SECRET_KEY = os.getenv("AWS_SECRET_ACCESS_KEY", "some_secret_key")
+
+# Dataset start/end dates
+DEFAULT_START_T = pd.Timestamp("20180101")
+DEFAULT_END_T = pd.Timestamp("20200201")
+
+# Arrow related constants
+DEFAULT_ROW_GROUP_SIZE = 1024 * 1024
+DEFAULT_ARROW_KWARGS = {
+    "batch_size": (DEFAULT_ARROW_BATCH_SIZE := 786432),
+    "batch_readahead": (DEFAULT_ARROW_BATCH_READAHEAD := 12),
+    "fragment_readahead": (DEFAULT_ARROW_BATCH_FRAGMENT_READAHEAD := 4),
+}
+
+
+class ParquetCompression(StrEnum):
+    NONE = "NONE"
+    SNAPPY = "SNAPPY"
+    GZIP = "GZIP"
+    # LZO = "LZO"
+    BROTLI = "BROTLI"
+    LZ4 = "LZ4"
+    ZSTD = "ZSTD"
+
+
+DFAULT_PARQUET_COMPRESSION = ParquetCompression.LZ4
+
+
+class LogLevel(StrEnum):
+    CRITICAL = "CRITICAL"
+    FATAL = "FATAL"
+    ERROR = "ERROR"
+    WARN = "WARNING"
+    WARNING = "WARNING"
+    INFO = "INFO"
+    DEBUG = "DEBUG"
+    NOTSET = "NOTSET"
+
+    def to_int(self) -> int:
+        # noinspection PyUnresolvedReferences,PyProtectedMember
+        return logging._nameToLevel[self]  # noqa: SLF001
+
+
+@dataclass(frozen=True)
+class VastConnDetails:
+    """VAST Config."""
+
+    access: str = DEFAULT_ACCESS_KEY
+    secret: str = DEFAULT_SECRET_KEY
+    vastdb_bucket: str = VASTDB_BUCKET_NAME
+    vastdb_endpoint: str = VASTDB_ENDPOINT
+    vastdb_ssl_verify: bool = True
+    vastdb_schema: str = VASTDB_TEST_SCHEMA_NAME
+    vastdb_table: str = VASTDB_TEST_TABLE_NAME
+    s3_host: str = DEFAULT_S3_HOST
+    s3_bucket: str = VASTDB_BUCKET_NAME
+    s3_port: int = DEFAULT_S3_PORT
+    s3_ssl_port: int = DEFAULT_S3_SSL_PORT

vastdb/bench/perf_bench/common/log_utils.py

@@ -0,0 +1,67 @@
+import logging
+from pathlib import Path
+from threading import RLock
+from typing import Optional, Union
+
+from vastdb.bench.perf_bench.common.constants import LogLevel
+from vastdb.bench.perf_bench.common.types import PathLikeT
+
+
+class LogConfigError(Exception):
+    pass
+
+
+_logging_configured: bool = False
+_log_level: int = logging.INFO
+_log_file: Optional[Path] = None
+_lock: RLock = RLock()
+
+
+def set_log_file(log_file: PathLikeT) -> None:
+    global _log_file  # noqa: PLW0603
+    with _lock:
+        if _logging_configured:
+            raise LogConfigError(
+                "Cannot change log file after logging has been configured."
+            )
+
+        log_file = Path(str(log_file))
+        log_file.parent.mkdir(parents=True, exist_ok=True)
+        log_file.touch(exist_ok=True)
+        _log_file = log_file
+
+
+def get_log_level() -> int:
+    return _log_level
+
+
+def set_log_level(level: Union[str, LogLevel, int]) -> None:
+    global _log_level  # noqa: PLW0603
+    if isinstance(level, str):
+        level = LogLevel[level].to_int()
+    with _lock:
+        if level != _log_level and _logging_configured:
+            raise LogConfigError(
+                "Cannot change log level after logging has been configured."
+            )
+        _log_level = level
+
+
+def get_logger(name: Optional[str]) -> logging.Logger:
+    global _logging_configured  # noqa: PLW0603
+    with _lock:
+        if not _logging_configured:
+            logging.basicConfig(
+                level=_log_level,
+                format=(
+                    "%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s"
+                ),
+                datefmt="%Y-%m-%d %H:%M:%S",
+            )
+            if _log_file:
+                root_logger = logging.getLogger()
+                fh = logging.FileHandler(_log_file, mode="a")
+                fh.setFormatter(root_logger.handlers[0].formatter)
+                root_logger.addHandler(fh)
+            _logging_configured = True
+    return logging.getLogger(name)

vastdb/bench/perf_bench/common/types.py

@@ -0,0 +1,34 @@
+import datetime as dt
+import os
+import sys
+from typing import Union
+
+PathLikeT = Union[str, bytes, os.PathLike]
+DateLikeT = Union[str, dt.date]
+
+if sys.version_info >= (3, 11):
+    # noinspection PyUnresolvedReferences
+    from enum import StrEnum
+else:
+    from enum import Enum, auto
+
+    class StrEnum(str, Enum):
+        # noinspection PyTypeChecker
+        def __new__(cls, value: Union[auto, str], *args, **kwargs):
+            if not isinstance(value, (str, auto)):
+                raise TypeError(
+                    f"Not a string/auto type: {value=!r} [type={type(value)}]"
+                )
+            return super().__new__(cls, value, *args, **kwargs)
+
+        def __repr__(self):
+            """Return a string representation of the enumeration member."""
+            return f"<{self.__class__.__name__}.{self.name}: '{self.value}'>"
+
+        def __str__(self):
+            """Return the value of the enumeration member."""
+            return str(self.value)
+
+        @staticmethod
+        def _generate_next_value_(name: str, *args, **kwargs):
+            return name