vastdb 0.1.9__tar.gz → 0.1.10__tar.gz

{vastdb-0.1.9 → vastdb-0.1.10}/CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [0.1.10] (2024-07-09)
+[0.1.10]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.9...v0.1.10
+
+## Fixed
+ - Support responses with no record batches
+ - Document supported Arrow data types
+
 ## [0.1.9] (2024-07-01)
 [0.1.9]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.8...v0.1.9
 

{vastdb-0.1.9/vastdb.egg-info → vastdb-0.1.10}/PKG-INFO

@@ -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 → vastdb-0.1.10}/README.md

@@ -19,6 +19,11 @@
 - [S3 access & secret keys on VAST cluster](https://support.vastdata.com/s/article/UUID-4d2e7e23-b2fb-7900-d98f-96c31a499626)
 - [Tabular identity policy with the proper permissions](https://support.vastdata.com/s/article/UUID-14322b60-d6a2-89ac-3df0-3dfbb6974182)
 
+### Required VAST release
+
+Currently, VAST DB Python SDK requires `5.0.0-sp10` or later VAST release.
+
+If the cluster is running an older VAST release, please contact customer.support@vastdata.com for more details.
 
 ### Installation
 
@@ -203,10 +208,3 @@ See the following blog posts for more examples:
 
 - https://vastdata.com/blog/the-vast-catalog-in-action-part-1
 - https://vastdata.com/blog/the-vast-catalog-in-action-part-2
-
-
-## Required VAST release
-
-Currently, VAST DB Python SDK requires `5.0.0-sp10` or later VAST release.
-
-If the cluster is running an older VAST release, please contact customer.support@vastdata.com for more details.

{vastdb-0.1.9 → vastdb-0.1.10}/setup.py

@@ -29,7 +29,7 @@ setup(
     name='vastdb',
     python_requires='>=3.9.0',
     description='VAST Data SDK',
-    version='0.1.9' + suffix,
+    version='0.1.10' + suffix,
     url='https://github.com/vast-data/vastdb_sdk',
     author='VAST DATA',
     author_email='hello@vastdata.com',

{vastdb-0.1.9 → vastdb-0.1.10}/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-0.1.10/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-0.1.9 → vastdb-0.1.10}/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-0.1.9 → vastdb-0.1.10}/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-0.1.10/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-0.1.9 → vastdb-0.1.10}/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 → vastdb-0.1.10/vastdb.egg-info}/PKG-INFO

@@ -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 → vastdb-0.1.10}/vastdb.egg-info/SOURCES.txt

@@ -171,7 +171,9 @@ vastdb.egg-info/requires.txt
 vastdb.egg-info/top_level.txt
 vastdb/bench/__init__.py
 vastdb/bench/test_perf.py
+vastdb/bench/test_sample.py
 vastdb/tests/__init__.py
+vastdb/tests/metrics.py
 vastdb/tests/test_duckdb.py
 vastdb/tests/test_imports.py
 vastdb/tests/test_nested.py