vastdb 0.1.9__py3-none-any.whl → 0.1.10__py3-none-any.whl

vastdb/_internal.py

@@ -1807,7 +1807,7 @@ def _iter_query_data_response_columns(fileobj, stream_ids=None):
             batches.append(batch)
         except StopIteration:  # we got an end-of-stream IPC message for a given stream ID
             reader, batches = readers.pop(stream_id)  # end of column
-            table = pa.Table.from_batches(batches)  # concatenate all column chunks (as a single)
+            table = pa.Table.from_batches(batches=batches, schema=reader.schema)  # concatenate all column chunks (as a single)
             _logger.debug("stream_id=%d rows=%d column=%s", stream_id, len(table), table)
             yield (stream_id, next_row_id, table)
 

vastdb/bench/test_sample.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+
+import functools
+import itertools
+import logging
+import os
+import random
+import threading
+import time
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+import numpy as np
+import pyarrow as pa
+
+import vastdb.errors
+from vastdb.table import INTERNAL_ROW_ID
+from vastdb.tests import metrics
+
+logging.basicConfig(
+    level="INFO",
+    format="%(asctime)s %(levelname)-10s %(process)d/%(thread)d %(filename)s:%(lineno)d %(message)s")
+
+log = logging.getLogger()
+
+log.info("Python SDK version: %s", vastdb.util.version())
+
+NUM_COLUMNS = 10_000
+COLUMNS_BATCH = 10
+
+NUM_ROW_GROUPS = 100
+ROW_GROUP_SIZE = 100_000
+
+
+INTERNAL_ROWID_FIELD = pa.field(INTERNAL_ROW_ID, pa.uint64())  # used for UPDATE
+EXTERNAL_ROWID_FIELD = pa.field("vastdb_rowid", pa.int64())  # used for INSERT & SELECT
+
+SCHEMA = "perf"
+TABLE = "sample"
+
+SCHEMA_ARROW = pa.schema(
+    [pa.field(f'c{i}', pa.float32()) for i in range(NUM_COLUMNS)]
+)
+
+
+def load_batch(bucket, session_kwargs, offset, limit):
+    log.info('loading into [%d..%d)', offset, limit)
+
+    # Iterate over all row-groups in this file
+    rowids_range = range(offset, limit)
+    rowids = pa.array(rowids_range, INTERNAL_ROWID_FIELD.type)
+
+    session = vastdb.connect(**session_kwargs)
+    metrics_rows = []
+
+    with session.transaction() as tx:
+        table = tx.bucket(bucket).schema(SCHEMA).table(TABLE)
+
+        col = table[EXTERNAL_ROWID_FIELD.name]
+        pred = (col >= rowids_range[0]) & (col <= rowids_range[-1])
+        count = sum(len(rb) for rb in table.select(columns=[], predicate=pred))
+        log.info("%d rows exist at %s", count, rowids_range)
+        if count == len(rowids_range):
+            # skip already loaded rows
+            log.info('skipping [%d..%d)', offset, limit)
+
+        total_nbytes = 0
+        calls = 0
+        t0 = time.time()
+        # Insert/update every chunk of columns in this rowgroup
+        for j in range(0, len(SCHEMA_ARROW), COLUMNS_BATCH):
+            cols_batch = list(SCHEMA_ARROW)[j:j + COLUMNS_BATCH]
+            arrays = [
+                pa.array(np.float32(np.random.uniform(size=[ROW_GROUP_SIZE])))
+                for _ in cols_batch
+            ]
+            chunk = pa.table(data=arrays, schema=pa.schema(cols_batch))
+            nbytes = chunk.get_total_buffer_size()
+            start = time.perf_counter()
+            if j == 0:
+                chunk = chunk.add_column(0, EXTERNAL_ROWID_FIELD, rowids.cast(EXTERNAL_ROWID_FIELD.type))
+                op = 'insert'
+                table.insert(chunk)
+            else:
+                chunk = chunk.add_column(0, INTERNAL_ROWID_FIELD, rowids)
+                op = 'update'
+                table.update(chunk)
+            finish = time.perf_counter()
+
+            metrics_rows.append(metrics.Row(
+                start=start, finish=finish, table_path=table.path, op=op,
+                nbytes=nbytes, rows=len(chunk), cols=len(cols_batch),
+                pid=os.getpid(), tid=threading.get_native_id()))
+
+            total_nbytes += nbytes
+            calls += 1
+            log.debug("%s into %s: %d rows x %d cols, %.3f MB",
+                op, rowids_range, len(chunk), len(chunk.schema),
+                chunk.get_total_buffer_size() / 1e6)
+
+        dt = time.time() - t0
+
+    log.info('loaded into [%d..%d): %d rows x %d cols, %.3f MB, %d RPCs, %.3f seconds',
+             offset, limit, limit - offset, NUM_COLUMNS, total_nbytes / 1e6, calls, dt)
+    return metrics_rows
+
+
+def test_ingest(clean_bucket_name, session_kwargs, tabular_endpoint_urls, num_workers, perf_metrics_db):
+    session = vastdb.connect(**session_kwargs)
+    metrics_table = metrics.Table(perf_metrics_db, "ingest")
+
+    with session.transaction() as tx:
+        b = tx.bucket(clean_bucket_name)
+        try:
+            s = b.schema(SCHEMA)
+        except vastdb.errors.MissingSchema:
+            s = b.create_schema(SCHEMA)
+
+        try:
+            s.table(TABLE)
+        except vastdb.errors.MissingTable:
+            s.create_table(TABLE, pa.schema([EXTERNAL_ROWID_FIELD] + list(SCHEMA_ARROW)))
+
+    ranges = [
+        (i * ROW_GROUP_SIZE, (i + 1) * ROW_GROUP_SIZE)
+        for i in range(NUM_ROW_GROUPS)
+    ]
+
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = [
+            executor.submit(load_batch, clean_bucket_name, session_kwargs | {'endpoint': url}, offset, limit)
+            for (offset, limit), url in zip(ranges, itertools.cycle(tabular_endpoint_urls))
+        ]
+        log.info("spawned %d futures", len(futures))
+        for future in as_completed(futures):
+            metrics_table.insert(future.result())
+
+    with session.transaction() as tx:
+        t = tx.bucket(clean_bucket_name).schema(SCHEMA).table(TABLE)
+        count = sum(len(rb) for rb in t.select([]))
+        log.info("%s has %d rows: %s", t, count, t.stats)
+
+
+def run_query(session_kwargs, i, bucket_name, endpoint_url):
+    num_columns = 2000
+    row_groups_per_query = 10
+
+    config = vastdb.table.QueryConfig(
+        num_sub_splits=1,
+        num_splits=1,
+        limit_rows_per_sub_split=ROW_GROUP_SIZE,
+        num_row_groups_per_sub_split=1)
+
+    row_group_indices = list(range(NUM_ROW_GROUPS))
+    r = random.Random(i)
+    r.shuffle(row_group_indices)
+
+    session = vastdb.connect(**(session_kwargs | {"endpoint": endpoint_url}))
+    with session.transaction() as tx:
+        t = tx.bucket(bucket_name).schema(SCHEMA).table(TABLE)
+
+        fields = list(t.arrow_schema)[1:]
+        r.shuffle(fields)
+        cols = [f.name for f in fields[:num_columns]]
+
+        vastdb_rowid = t['vastdb_rowid']
+        preds = []
+        for offset in range(0, len(row_group_indices), row_groups_per_query):
+            rowid_ranges = (
+                vastdb_rowid.between(j * ROW_GROUP_SIZE, (j + 1) * ROW_GROUP_SIZE - 1)
+                for j in row_group_indices[offset:offset + row_groups_per_query]
+            )
+            pred = functools.reduce((lambda x, y: x | y), rowid_ranges)
+            preds.append(pred)
+
+        for j, pred in enumerate(preds):
+            log.info("%d) starting query #%d on %s", i, j, endpoint_url)
+            t0 = time.time()
+            res = t.select(columns=cols, predicate=pred, config=config)
+            rows = 0
+            data = 0
+            for rb in res:
+                rows += len(rb)
+                data += rb.nbytes
+                dt = time.time() - t0
+                log.info("%d) got query #%d batch %.3f[s], %.3f[GB] %.3f[MB/s], %.3f[Mrows]", i, j, dt, data / 1e9, data / 1e6 / dt, rows / 1e6)
+
+            dt = time.time() - t0
+            log.info("%d) finished query #%d %.3f[s], %.3f[GB], %.3f[MB/s], %.3f[Mrows]", i, j, dt, data / 1e9, data / 1e6 / dt, rows / 1e6)
+
+
+def test_scan(test_bucket_name, session, num_workers, session_kwargs, tabular_endpoint_urls):
+    log.info("starting %d workers, endpoints=%s", num_workers, tabular_endpoint_urls)
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        for i, url in zip(range(num_workers), itertools.cycle(tabular_endpoint_urls)):
+            executor.submit(run_query, session_kwargs, i, test_bucket_name, url)
+    log.info("finished %d workers", num_workers)

vastdb/conftest.py

@@ -1,4 +1,5 @@
 import os
+import sqlite3
 from pathlib import Path
 
 import boto3
@@ -13,27 +14,43 @@ def pytest_addoption(parser):
                      default=os.environ.get("AWS_ACCESS_KEY_ID", None))
     parser.addoption("--tabular-secret-key", help="Secret key with Tabular permissions (AWS_SECRET_ACCESS_KEY)",
                      default=os.environ.get("AWS_SECRET_ACCESS_KEY", None))
-    parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint", default="http://localhost:9090")
+    parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint", default=[], action="append")
     parser.addoption("--data-path", help="Data files location", default=None)
     parser.addoption("--crater-path", help="Save benchmark results in a dedicated location", default=None)
     parser.addoption("--schema-name", help="Name of schema for the test to operate on", default=None)
     parser.addoption("--table-name", help="Name of table for the test to operate on", default=None)
+    parser.addoption("--num-workers", help="Number of concurrent workers", default=1)
 
 
 @pytest.fixture(scope="session")
-def session(request):
-    return vastdb.connect(
+def session_kwargs(request, tabular_endpoint_urls):
+    return dict(
         access=request.config.getoption("--tabular-access-key"),
         secret=request.config.getoption("--tabular-secret-key"),
-        endpoint=request.config.getoption("--tabular-endpoint-url"),
+        endpoint=tabular_endpoint_urls[0],
     )
 
 
+@pytest.fixture(scope="session")
+def session(session_kwargs):
+    return vastdb.connect(**session_kwargs)
+
+
+@pytest.fixture(scope="session")
+def num_workers(request):
+    return int(request.config.getoption("--num-workers"))
+
+
 @pytest.fixture(scope="session")
 def test_bucket_name(request):
     return request.config.getoption("--tabular-bucket-name")
 
 
+@pytest.fixture(scope="session")
+def tabular_endpoint_urls(request):
+    return request.config.getoption("--tabular-endpoint-url") or ["http://localhost:9090"]
+
+
 def iter_schemas(s):
     """Recusively scan all schemas."""
     children = s.schemas()
@@ -55,12 +72,12 @@ def clean_bucket_name(request, test_bucket_name, session):
 
 
 @pytest.fixture(scope="session")
-def s3(request):
+def s3(request, tabular_endpoint_urls):
     return boto3.client(
         's3',
         aws_access_key_id=request.config.getoption("--tabular-access-key"),
         aws_secret_access_key=request.config.getoption("--tabular-secret-key"),
-        endpoint_url=request.config.getoption("--tabular-endpoint-url"))
+        endpoint_url=tabular_endpoint_urls[0])
 
 
 @pytest.fixture(scope="function")
@@ -81,3 +98,8 @@ def schema_name(request):
 @pytest.fixture(scope="function")
 def table_name(request):
     return request.config.getoption("--table-name")
+
+
+@pytest.fixture(scope="function")
+def perf_metrics_db(crater_path):
+    return sqlite3.connect(f"{crater_path}/metrics.sqlite")

vastdb/table.py

@@ -167,8 +167,13 @@ class Table:
         """Also, load columns' metadata."""
         self.arrow_schema = self.columns()
 
-        table_path = f'{self.schema.bucket.name}/{self.schema.name}/{self.name}'
-        self._ibis_table = ibis.table(ibis.Schema.from_pyarrow(self.arrow_schema), table_path)
+        self._table_path = f'{self.schema.bucket.name}/{self.schema.name}/{self.name}'
+        self._ibis_table = ibis.table(ibis.Schema.from_pyarrow(self.arrow_schema), self._table_path)
+
+    @property
+    def path(self):
+        """Return table's path."""
+        return self._table_path
 
     @property
     def tx(self):

vastdb/tests/metrics.py

@@ -0,0 +1,42 @@
+import dataclasses
+import sqlite3
+from typing import List
+
+_MAP_SQLITE_TYPES = {
+    str: "TEXT",
+    float: "REAL",
+    int: "INTEGER",
+}
+
+
+@dataclasses.dataclass
+class Row:
+    start: float
+    finish: float
+    table_path: str
+    op: str
+    nbytes: int
+    rows: int
+    cols: int
+    pid: int
+    tid: int
+
+
+class Table:
+    def __init__(self, conn: sqlite3.Connection, name: str):
+        self.fields = dataclasses.fields(Row)
+        self.conn = conn
+        self.name = name
+        columns = ", ".join(
+            f"{f.name} {_MAP_SQLITE_TYPES[f.type]}"
+            for f in self.fields
+        )
+        cmd = f"CREATE TABLE {self.name} ({columns})"
+        self.conn.execute(cmd).fetchall()
+
+    def insert(self, rows: List[Row]):
+        args = ", ".join(["?"] * len(self.fields))
+        cmd = f"INSERT INTO {self.name} VALUES ({args})"
+        data = [dataclasses.astuple(row) for row in rows]
+        self.conn.executemany(cmd, data).fetchall()
+        self.conn.commit()

vastdb/util.py

@@ -1,3 +1,4 @@
+import importlib
 import logging
 import re
 from typing import TYPE_CHECKING, Callable, List, Optional, Union
@@ -159,3 +160,7 @@ def check_supported_types(fields: pa.Schema):
         if isinstance(f.type, pa.TimestampType):
             if f.type.tz:
                 raise NotSupportedType(f)
+
+
+def version():
+    return importlib.metadata.distribution('vastdb').version

{vastdb-0.1.9.dist-info → vastdb-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vastdb
-Version: 0.1.9
+Version: 0.1.10
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA

{vastdb-0.1.9.dist-info → vastdb-0.1.10.dist-info}/RECORD

@@ -149,18 +149,20 @@ vast_flatbuf/tabular/S3File.py,sha256=KC9c2oS5-JXwTTriUVFdjOvRG0B54Cq9kviSDZY3NI
 vast_flatbuf/tabular/VipRange.py,sha256=_BJd1RRZAcK76T9vlsHzXKYVsPVaz6WTEAqStMQCAUQ,2069
 vast_flatbuf/tabular/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/__init__.py,sha256=8PLcZowy_vM0zuiYSQPXuxIEMcwHD7IRFpgcPK-03bk,386
-vastdb/_internal.py,sha256=bENTnMZFAXolkFaWa7op9bsPWMk3wVE6oKKU--0ukXk,89971
+vastdb/_internal.py,sha256=4vi6KgkfHnDOSZUrU3oQcNKKdeMH3alODLGk7Yt59Gk,90001
 vastdb/bucket.py,sha256=5KuKhPjZOevznZqWHDVVocejvAy7dcwobPuV6BJCfPc,2544
-vastdb/conftest.py,sha256=wnnPXjeLems2zpOBB1UGbZ_YW5S169NhGA5UZu7H5SM,2831
+vastdb/conftest.py,sha256=ePzQiEQmlNGcM2T4GZevE4XuvcnFWfnTSzr8IVZpVKk,3438
 vastdb/errors.py,sha256=jER5RQYsBRlQsjym1ItQYRukggMypATOo_sKvsJtMbo,4278
 vastdb/schema.py,sha256=yaueil92MSMYJf6bWseov_8fXTdW5zaKLXNjP5uuyzI,5963
 vastdb/session.py,sha256=3YHhG7IamFOKuy-Fkq_IVtPNriSfI6IN_4z4arBFbDU,3349
-vastdb/table.py,sha256=TGBiIp0pB7vHd-92q4_sDRDjd4klHDLFOeEgdn1ACQI,32880
+vastdb/table.py,sha256=C0kgV8CJVgoRxVx83SPTn75mgbTz9OWgYwK_RzLPZ5Q,32994
 vastdb/transaction.py,sha256=qu2rOlR7AS1ojMOzgWapQMpcorrutelZZLH1mLmTHxk,3186
-vastdb/util.py,sha256=2W5bBnlihIFvdV4im4HiDLArEhU8zjKMZB3Xw0lzgz0,5888
+vastdb/util.py,sha256=KQ2CjML-ipWxrJjwiaBbj4bxpTCtL24Pr2Co9woyw3Y,5983
 vastdb/bench/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/bench/test_perf.py,sha256=yn5gE7t_nzmJHBl9bCs1hxQOgzhvFphuYElsWGko8ts,1084
+vastdb/bench/test_sample.py,sha256=0qsKPj3i88J-YTrOrGvsP19xsyWGZy_-ptIt3oXBbSw,7181
 vastdb/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vastdb/tests/metrics.py,sha256=f1oOPKDsu-BzBLin0IQvjG-ueRDHTY-Hzl357TuoxCQ,989
 vastdb/tests/test_duckdb.py,sha256=STw_1PwTQR8Naz6s0p6lQTV1ZTKKhe3LPBUbhqzTCu0,1880
 vastdb/tests/test_imports.py,sha256=xKub3-bisFjH0BsZM8COfiUWuMrtoOoQKprF6VQT9RI,5669
 vastdb/tests/test_nested.py,sha256=22NAxBTm7Aq-Vn6AIYbi5Cb1ET8W0XeLK3pp4D8BYWI,3448
@@ -172,8 +174,8 @@ vastdb/tests/test_util.py,sha256=Ok_sAEBJsRGF5Voa_v5eu3eAd52GWu8jMjjQbadwW-s,126
 vastdb/tests/util.py,sha256=dpRJYbboDnlqL4qIdvScpp8--5fxRUBIcIYitrfcj9o,555
 vastdb/vast_tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vastdb/vast_tests/test_ha.py,sha256=744P4G6VJ09RIkHhMQL4wlipCBJWQVMhyvUrSc4k1HQ,975
-vastdb-0.1.9.dist-info/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
-vastdb-0.1.9.dist-info/METADATA,sha256=YRP3W_JzaDywneqZ0cFWJzLbLE_NGf9QYja6-CFwQl4,1350
-vastdb-0.1.9.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
-vastdb-0.1.9.dist-info/top_level.txt,sha256=Vsj2MKtlhPg0J4so64slQtnwjhgoPmJgcG-6YcVAwVc,20
-vastdb-0.1.9.dist-info/RECORD,,
+vastdb-0.1.10.dist-info/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
+vastdb-0.1.10.dist-info/METADATA,sha256=Pjw1EZvwnKhfEjuRVVMR0DBOSkmVo5wcHftqddxqRNY,1351
+vastdb-0.1.10.dist-info/WHEEL,sha256=y4mX-SOX4fYIkonsAGA5N0Oy-8_gI4FXw5HNI1xqvWg,91
+vastdb-0.1.10.dist-info/top_level.txt,sha256=Vsj2MKtlhPg0J4so64slQtnwjhgoPmJgcG-6YcVAwVc,20
+vastdb-0.1.10.dist-info/RECORD,,

{vastdb-0.1.9.dist-info → vastdb-0.1.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (70.1.1)
+Generator: setuptools (70.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any