vastdb 0.1.4__tar.gz → 0.1.5__tar.gz

{vastdb-0.1.4 → vastdb-0.1.5}/CHANGELOG.md

@@ -4,6 +4,19 @@ 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.5] (2024-05-16)
+[0.1.5]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.4...v0.1.5
+
+### Added
+ - Allow passing `ssl_verify` via `Session` c-tor
+ - Retry `requests.exceptions.ConnectionError`
+ - Document `QueryConfig` properties
+ - Support "deferred" `ibis` predicates using [Underscore (`_`) API](https://github.com/ibis-project/ibis/pull/3804)
+
+### Fixed
+ - Fix predicate pushdown when nested columns are present
+ - Attach original traceback to `MissingBucket` exception
+
 ## [0.1.4] (2024-05-13)
 [0.1.4]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.3...v0.1.4
 

{vastdb-0.1.4/vastdb.egg-info → vastdb-0.1.5}/PKG-INFO

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

{vastdb-0.1.4 → vastdb-0.1.5}/README.md

@@ -74,16 +74,18 @@ with session.transaction() as tx:
 Our SDK supports predicate and projection pushdown:
 
 ```python
+    from ibis import _
+
     # SELECT c1 FROM t WHERE (c2 > 2) AND (c3 IS NULL)
     table.select(columns=['c1'],
-                 predicate=(table['c2'] > 2) & table['c3'].isnull())
+                 predicate=(_.c2 > 2) & _.c3.isnull())
 
     # SELECT c2, c3 FROM t WHERE (c2 BETWEEN 0 AND 1) OR (c2 > 10)
     table.select(columns=['c2', 'c3'],
-                 predicate=table['c2'].between(0, 1) | (table['c2'] > 10))
+                 predicate=(_.c2.between(0, 1) | (table['c2'] > 10))
 
     # SELECT * FROM t WHERE c3 LIKE '%substring%'
-    table.select(predicate=table['c3'].contains('substring'))
+    table.select(predicate=_.c3.contains('substring'))
 ```
 
 See [here for more details](docs/predicate.md).
@@ -132,10 +134,12 @@ with contextlib.closing(pa.parquet.ParquetWriter('/path/to/file.parquet', batche
 We can use [DuckDB](https://duckdb.org/docs/guides/python/sql_on_arrow.html) to post-process the resulting stream of [PyArrow record batches](https://arrow.apache.org/docs/python/data.html#record-batches):
 
 ```python
+from ibis import _
+
 import duckdb
 conn = duckdb.connect()
 
-batches = table.select(columns=['c1'], predicate=(table['c2'] > 2))
+batches = table.select(columns=['c1'], predicate=(_.c2 > 2))
 print(conn.execute("SELECT sum(c1) FROM batches").arrow())
 ```
 

{vastdb-0.1.4 → vastdb-0.1.5}/setup.py

@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 
 from setuptools import find_packages, setup
 
@@ -28,13 +29,13 @@ setup(
     name='vastdb',
     python_requires='>=3.9.0',
     description='VAST Data SDK',
-    version='0.1.4' + suffix,
+    version='0.1.5' + suffix,
     url='https://github.com/vast-data/vastdb_sdk',
     author='VAST DATA',
     author_email='hello@vastdata.com',
     license='Copyright (C) VAST Data Ltd.',
     packages=find_packages(),
-    install_requires=open('requirements.txt').read().strip().split('\n'),
+    install_requires=Path('requirements.txt').read_text().strip().split(),
     long_description=long_description,
     long_description_content_type='text/markdown',
     classifiers=[

{vastdb-0.1.4 → vastdb-0.1.5}/vastdb/internal_commands.py

@@ -130,46 +130,13 @@ def get_unit_to_flatbuff_time_unit(type):
 class Predicate:
     def __init__(self, schema: 'pa.Schema', expr: ibis.expr.types.BooleanColumn):
         self.schema = schema
+        index = itertools.count()  # used to generate leaf column positions for VAST QueryData RPC
+        # Arrow schema contains the top-level columns, where each column may include multiple subfields
+        # We use DFS is used to enumerate all the sub-columns, using `index` as an ID allocator
+        nodes = [FieldNode(field, index) for field in schema]
+        self.nodes_map = {node.field.name: node for node in nodes}
         self.expr = expr
 
-    def get_field_indexes(self, field: 'pa.Field', field_name_per_index: list) -> None:
-        field_name_per_index.append(field.name)
-
-        if isinstance(field.type, pa.StructType):
-            flat_fields = field.flatten()
-        elif isinstance(field.type, pa.MapType):
-            flat_fields = [pa.field(f'{field.name}.entries', pa.struct([field.type.key_field, field.type.item_field]))]
-        elif isinstance(field.type, pa.ListType):
-            flat_fields = [pa.field(f'{field.name}.{field.type.value_field.name}', field.type.value_field.type)]
-        else:
-            return
-
-        for flat_field in flat_fields:
-            self.get_field_indexes(flat_field, field_name_per_index)
-
-    @property
-    def field_name_per_index(self):
-        if self._field_name_per_index is None:
-            _field_name_per_index = []
-            for field in self.schema:
-                self.get_field_indexes(field, _field_name_per_index)
-            self._field_name_per_index = {field: index for index, field in enumerate(_field_name_per_index)}
-        return self._field_name_per_index
-
-    def get_projections(self, builder: 'flatbuffers.builder.Builder', field_names: Optional[List[str]] = None):
-        if field_names is None:
-            field_names = self.field_name_per_index.keys()
-        projection_fields = []
-        for field_name in field_names:
-            fb_field_index.Start(builder)
-            fb_field_index.AddPosition(builder, self.field_name_per_index[field_name])
-            offset = fb_field_index.End(builder)
-            projection_fields.append(offset)
-        fb_source.StartProjectionVector(builder, len(projection_fields))
-        for offset in reversed(projection_fields):
-            builder.PrependUOffsetTRelative(offset)
-        return builder.EndVector()
-
     def serialize(self, builder: 'flatbuffers.builder.Builder'):
         from ibis.expr.operations.generic import (
             IsNull,
@@ -204,8 +171,6 @@ class Predicate:
             Between: self.build_between,
         }
 
-        positions_map = dict((f.name, index) for index, f in enumerate(self.schema))  # TODO: BFS
-
         self.builder = builder
 
         offsets = []
@@ -261,7 +226,11 @@ class Predicate:
                     elif prev_field_name != field_name:
                         raise NotImplementedError(self.expr)
 
-                    column_offset = self.build_column(position=positions_map[field_name])
+                    node = self.nodes_map[field_name]
+                    # 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]
@@ -839,12 +808,13 @@ class VastdbApi:
         return prefix
 
     def _fill_common_headers(self, txid=0, client_tags=[], version_id=1):
-        common_headers = {'tabular-txid': str(txid), 'tabular-api-version-id': str(version_id),
-                          'tabular-client-name': 'tabular-api'}
-        for tag in client_tags:
-            common_headers['tabular-client-tags-%d' % client_tags.index(tag)] = tag
+        common_headers = {
+            'tabular-txid': str(txid),
+            'tabular-api-version-id': str(version_id),
+            'tabular-client-name': 'tabular-api'
+        }
 
-        return common_headers
+        return common_headers | {f'tabular-client-tags-{index}': tag for index, tag in enumerate(client_tags)}
 
     def _check_res(self, res, cmd="", expected_retvals=[]):
         if exc := errors.from_response(res):
@@ -952,8 +922,7 @@ class VastdbApi:
             res_headers = res.headers
             next_key = int(res_headers['tabular-next-key'])
             is_truncated = res_headers['tabular-is-truncated'] == 'true'
-            flatbuf = b''.join(res.iter_content(chunk_size=128))
-            lists = list_schemas.GetRootAs(flatbuf)
+            lists = list_schemas.GetRootAs(res.content)
             bucket_name = lists.BucketName().decode()
             if not bucket.startswith(bucket_name):
                 raise ValueError(f'bucket: {bucket} did not start from {bucket_name}')
@@ -976,8 +945,7 @@ class VastdbApi:
         res = self.session.get(self._api_prefix(bucket=bucket, command="list", url_params=url_params), headers={}, stream=True)
         self._check_res(res, "list_snapshots")
 
-        out = b''.join(res.iter_content(chunk_size=128))
-        xml_str = out.decode()
+        xml_str = res.content.decode()
         xml_dict = xmltodict.parse(xml_str)
         list_res = xml_dict['ListBucketResult']
         is_truncated = list_res['IsTruncated'] == 'true'
@@ -1059,8 +1027,7 @@ class VastdbApi:
         res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=name, command="stats", url_params=url_params), headers=headers)
         self._check_res(res, "get_table_stats", expected_retvals)
 
-        flatbuf = b''.join(res.iter_content(chunk_size=128))
-        stats = get_table_stats.GetRootAs(flatbuf)
+        stats = get_table_stats.GetRootAs(res.content)
         num_rows = stats.NumRows()
         size_in_bytes = stats.SizeInBytes()
         is_external_rowid_alloc = stats.IsExternalRowidAlloc()
@@ -1159,8 +1126,7 @@ class VastdbApi:
             res_headers = res.headers
             next_key = int(res_headers['tabular-next-key'])
             is_truncated = res_headers['tabular-is-truncated'] == 'true'
-            flatbuf = b''.join(res.iter_content(chunk_size=128))
-            lists = list_tables.GetRootAs(flatbuf)
+            lists = list_tables.GetRootAs(res.content)
             bucket_name = lists.BucketName().decode()
             schema_name = lists.SchemaName().decode()
             if not bucket.startswith(bucket_name):  # ignore snapshot name
@@ -1288,11 +1254,7 @@ class VastdbApi:
             next_key = int(res_headers['tabular-next-key'])
             is_truncated = res_headers['tabular-is-truncated'] == 'true'
             count = int(res_headers['tabular-list-count'])
-            columns = []
-            if not count_only:
-                schema_buf = b''.join(res.iter_content(chunk_size=128))
-                schema_out = pa.ipc.open_stream(schema_buf).schema
-                columns = schema_out
+            columns = [] if count_only else pa.ipc.open_stream(res.content).schema
 
             return columns, next_key, is_truncated, count
 
@@ -1692,8 +1654,7 @@ class VastdbApi:
         res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=table, command="projection-stats", url_params=url_params),
                                headers=headers)
         if res.status_code == 200:
-            flatbuf = b''.join(res.iter_content(chunk_size=128))
-            stats = get_projection_table_stats.GetRootAs(flatbuf)
+            stats = get_projection_table_stats.GetRootAs(res.content)
             num_rows = stats.NumRows()
             size_in_bytes = stats.SizeInBytes()
             dirty_blocks_percentage = stats.DirtyBlocksPercentage()
@@ -1779,8 +1740,7 @@ class VastdbApi:
             next_key = int(res_headers['tabular-next-key'])
             is_truncated = res_headers['tabular-is-truncated'] == 'true'
             count = int(res_headers['tabular-list-count'])
-            flatbuf = b''.join(res.iter_content(chunk_size=128))
-            lists = list_projections.GetRootAs(flatbuf)
+            lists = list_projections.GetRootAs(res.content)
             bucket_name = lists.BucketName().decode()
             schema_name = lists.SchemaName().decode()
             table_name = lists.TableName().decode()
@@ -1827,13 +1787,8 @@ class VastdbApi:
             next_key = int(res_headers['tabular-next-key'])
             is_truncated = res_headers['tabular-is-truncated'] == 'true'
             count = int(res_headers['tabular-list-count'])
-            columns = []
-            if not count_only:
-                schema_buf = b''.join(res.iter_content(chunk_size=128))
-                schema_out = pa.ipc.open_stream(schema_buf).schema
-                for f in schema_out:
-                    columns.append([f.name, f.type, f.metadata])
-                #   sort_type = f.metadata[b'VAST:sort_type'].decode()
+            columns = [] if count_only else [[f.name, f.type, f.metadata] for f in
+                                             pa.ipc.open_stream(res.content).schema]
 
             return columns, next_key, is_truncated, count
 

{vastdb-0.1.4 → vastdb-0.1.5}/vastdb/session.py

@@ -35,7 +35,7 @@ class Features:
 class Session:
     """VAST database session."""
 
-    def __init__(self, access=None, secret=None, endpoint=None):
+    def __init__(self, access=None, secret=None, endpoint=None, ssl_verify=True):
         """Connect to a VAST Database endpoint, using specified credentials."""
         if access is None:
             access = os.environ['AWS_ACCESS_KEY_ID']
@@ -44,7 +44,7 @@ class Session:
         if endpoint is None:
             endpoint = os.environ['AWS_S3_ENDPOINT_URL']
 
-        self.api = internal_commands.VastdbApi(endpoint, access, secret)
+        self.api = internal_commands.VastdbApi(endpoint, access, secret, ssl_verify=ssl_verify)
         version_tuple = tuple(int(part) for part in self.api.vast_version.split('.'))
         self.features = Features(version_tuple)
         self.s3 = boto3.client('s3',

{vastdb-0.1.4 → vastdb-0.1.5}/vastdb/table.py

@@ -12,6 +12,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import backoff
 import ibis
 import pyarrow as pa
+import requests
 
 from . import errors, internal_commands, schema, util
 
@@ -39,20 +40,44 @@ class TableStats:
     endpoints: Tuple[str, ...] = ()
 
 
+RETRIABLE_ERRORS = (
+    errors.Slowdown,
+    requests.exceptions.ConnectionError,
+)
+
+
 @dataclass
 class QueryConfig:
     """Query execution configiration."""
 
+    # allows server-side parallel processing by issuing multiple reads concurrently for a single RPC
     num_sub_splits: int = 4
+
+    # used to split the table into disjoint subsets of rows, to be processed concurrently using multiple RPCs
     num_splits: int = 1
+
+    # each endpoint will be handled by a separate worker thread
+    # a single endpoint can be specified more than once to benefit from multithreaded execution
     data_endpoints: Optional[List[str]] = None
+
+    # a subsplit fiber will finish after sending this number of rows back to the client
     limit_rows_per_sub_split: int = 128 * 1024
+
+    # each fiber will read the following number of rowgroups coninuously before skipping
+    # in order to use semi-sorted projections this value must be 8
     num_row_groups_per_sub_split: int = 8
+
+    # can be disabled for benchmarking purposes
     use_semi_sorted_projections: bool = True
+
+    # used to estimate the number of splits, given the table rows' count
     rows_per_split: int = 4000000
+
+    # used for worker threads' naming
     query_id: str = ""
-    max_slowdown_retry: int = 10
-    backoff_func: Any = field(default=backoff.on_exception(backoff.expo, errors.Slowdown, max_tries=max_slowdown_retry))
+
+    # allows retrying QueryData when the server is overloaded
+    backoff_func: Any = field(default=backoff.on_exception(backoff.expo, RETRIABLE_ERRORS, max_tries=10))
 
 
 @dataclass
@@ -271,7 +296,7 @@ class Table:
         return TableStats(**stats_tuple._asdict())
 
     def select(self, columns: Optional[List[str]] = None,
-               predicate: ibis.expr.types.BooleanColumn = None,
+               predicate: Union[ibis.expr.types.BooleanColumn, ibis.common.deferred.Deferred] = None,
                config: Optional[QueryConfig] = None,
                *,
                internal_row_id: bool = False) -> pa.RecordBatchReader:
@@ -310,6 +335,9 @@ class Table:
             response_schema = internal_commands.get_response_schema(schema=query_schema, field_names=columns)
             return pa.RecordBatchReader.from_batches(response_schema, [])
 
+        if isinstance(predicate, ibis.common.deferred.Deferred):
+            predicate = predicate.resolve(self._ibis_table)  # may raise if the predicate is invalid (e.g. wrong types / missing column)
+
         query_data_request = internal_commands.build_query_data_request(
             schema=query_schema,
             predicate=predicate,

vastdb-0.1.5/vastdb/tests/test_nested.py

@@ -0,0 +1,102 @@
+import functools
+import itertools
+import operator
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+from .util import prepare_data
+
+
+def test_nested_select(session, clean_bucket_name):
+    columns = pa.schema([
+        ('l', pa.list_(pa.int8())),
+        ('m', pa.map_(pa.utf8(), pa.float64())),
+        ('s', pa.struct([('x', pa.int16()), ('y', pa.int32())])),
+    ])
+    expected = pa.table(schema=columns, data=[
+        [[1], [], [2, 3], None],
+        [None, {'a': 2.5}, {'b': 0.25, 'c': 0.025}, {}],
+        [{'x': 1, 'y': None}, None, {'x': 2, 'y': 3}, {'x': None, 'y': 4}],
+    ])
+
+    with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
+        actual = pa.Table.from_batches(t.select())
+        assert actual == expected
+
+        names = [f.name for f in columns]
+        for n in range(len(names) + 1):
+            for cols in itertools.permutations(names, n):
+                actual = pa.Table.from_batches(t.select(columns=cols))
+                assert actual == expected.select(cols)
+
+
+def test_nested_filter(session, clean_bucket_name):
+    columns = pa.schema([
+        ('x', pa.int64()),
+        ('l', pa.list_(pa.int8())),
+        ('y', pa.int64()),
+        ('m', pa.map_(pa.utf8(), pa.float64())),
+        ('z', pa.int64()),
+        ('s', pa.struct([('x', pa.int16()), ('y', pa.int32())])),
+        ('w', pa.int64()),
+    ])
+    expected = pa.table(schema=columns, data=[
+        [1, 2, 3, None],
+        [[1], [], [2, 3], None],
+        [1, 2, None, 3],
+        [None, {'a': 2.5}, {'b': 0.25, 'c': 0.025}, {}],
+        [1, None, 2, 3],
+        [{'x': 1, 'y': None}, None, {'x': 2, 'y': 3}, {'x': None, 'y': 4}],
+        [None, 1, 2, 3],
+    ])
+
+    with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
+        actual = pa.Table.from_batches(t.select())
+        assert actual == expected
+
+        names = list('xyzw')
+        for n in range(1, len(names) + 1):
+            for cols in itertools.permutations(names, n):
+                ibis_predicate = functools.reduce(
+                    operator.and_,
+                    (t[col] > 2 for col in cols))
+                actual = pa.Table.from_batches(t.select(predicate=ibis_predicate), t.arrow_schema)
+
+                arrow_predicate = functools.reduce(
+                    operator.and_,
+                    (pc.field(col) > 2 for col in cols))
+                assert actual == expected.filter(arrow_predicate)
+
+
+def test_nested_unsupported_filter(session, clean_bucket_name):
+    columns = pa.schema([
+        ('x', pa.int64()),
+        ('l', pa.list_(pa.int8())),
+        ('y', pa.int64()),
+        ('m', pa.map_(pa.utf8(), pa.float64())),
+        ('z', pa.int64()),
+        ('s', pa.struct([('x', pa.int16()), ('y', pa.int32())])),
+        ('w', pa.int64()),
+    ])
+    expected = pa.table(schema=columns, data=[
+        [1, 2, 3, None],
+        [[1], [], [2, 3], None],
+        [1, 2, None, 3],
+        [None, {'a': 2.5}, {'b': 0.25, 'c': 0.025}, {}],
+        [1, None, 2, 3],
+        [{'x': 1, 'y': None}, None, {'x': 2, 'y': 3}, {'x': None, 'y': 4}],
+        [None, 1, 2, 3],
+    ])
+
+    with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
+
+        with pytest.raises(NotImplementedError):
+            list(t.select(predicate=(t['l'].isnull())))
+
+        with pytest.raises(NotImplementedError):
+            list(t.select(predicate=(t['m'].isnull())))
+
+        with pytest.raises(NotImplementedError):
+            list(t.select(predicate=(t['s'].isnull())))