vastdb 0.1.2__tar.gz → 0.1.4__tar.gz

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

@@ -4,6 +4,29 @@ 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.4] (2024-05-13)
+[0.1.4]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.3...v0.1.4
+
+### Added
+ - Resume `Table.select()` on HTTP 503 code
+ - Allow pushing down True/False boolean literals
+ - Support `between` predicate pushdown
+ - Support inserting wide rows (larger than 5MB)
+
+## [0.1.3] (2024-05-05)
+[0.1.3]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.2...v0.1.3
+
+### Added
+ - Document predicate pushdown support
+ - Access imports' table (for VAST 5.2+)
+ - Support `is_in` predicate pushdown
+ - Document `table.py`
+
+### Fixed
+ - Freeze `ibis` dependency at 8.0.0
+ - Support snapshot-based access
+ - Optimize RecordBatch slicing
+
 ## [0.1.2] (2024-04-25)
 [0.1.2]: https://github.com/vast-data/vastdb_sdk/compare/v0.1.1...v0.1.2
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vastdb
-Version: 0.1.2
+Version: 0.1.4
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA
@@ -21,10 +21,11 @@ License-File: LICENSE
 Requires-Dist: aws-requests-auth
 Requires-Dist: boto3
 Requires-Dist: flatbuffers
-Requires-Dist: ibis-framework
+Requires-Dist: ibis-framework==8.0.0
 Requires-Dist: pyarrow
 Requires-Dist: requests
 Requires-Dist: xmltodict
+Requires-Dist: backoff==2.2.1
 
 
 `vastdb` is a Python-based SDK designed for interacting

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

@@ -86,6 +86,8 @@ Our SDK supports predicate and projection pushdown:
     table.select(predicate=table['c3'].contains('substring'))
 ```
 
+See [here for more details](docs/predicate.md).
+
 ## Import a single Parquet file via S3 protocol
 
 It is possible to efficiently create a table from a Parquet file (without copying it via the client):
@@ -150,10 +152,11 @@ p.drop()
 
 ## Snapshots
 
-It is possible to list [available snapshots](https://vastdata.com/blog/bringing-snapshots-to-vasts-element-store):
+It is possible to use [snapshots](https://vastdata.com/blog/bringing-snapshots-to-vasts-element-store) for accessing the Database:
 
 ```python
-print(bucket.list_snapshots())
+snaps = bucket.list_snapshots()
+batches = snaps[0].schema('schema-name').table('table-name').select()
 ```
 
 ## VAST Catalog

{vastdb-0.1.2 → vastdb-0.1.4}/requirements.txt

@@ -1,7 +1,9 @@
 aws-requests-auth
 boto3
 flatbuffers
-ibis-framework
+ibis-framework==8.0.0
 pyarrow
 requests
 xmltodict
+backoff==2.2.1
+

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

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

{vastdb-0.1.2 → vastdb-0.1.4}/vastdb/bucket.py

@@ -16,14 +16,6 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-@dataclass
-class Snapshot:
-    """VAST bucket-level snapshot."""
-
-    name: str
-    bucket: "Bucket"
-
-
 @dataclass
 class Bucket:
     """VAST bucket."""
@@ -73,7 +65,22 @@ class Bucket:
 
         return [schema.Schema(name=name, bucket=self) for name, *_ in schemas]
 
-    def snapshots(self) -> List[Snapshot]:
+    def snapshot(self, name, fail_if_missing=True) -> Optional["Bucket"]:
+        """Get snapshot by name (if exists)."""
+        snapshots, _is_truncated, _next_key = \
+            self.tx._rpc.api.list_snapshots(bucket=self.name, name_prefix=name, max_keys=1)
+
+        expected_name = f".snapshot/{name}"
+        exists = snapshots and snapshots[0] == expected_name + "/"
+        if not exists:
+            if fail_if_missing:
+                raise errors.MissingSnapshot(self.name, expected_name)
+            else:
+                return None
+
+        return Bucket(name=f'{self.name}/{expected_name}', tx=self.tx)
+
+    def snapshots(self) -> List["Bucket"]:
         """List bucket's snapshots."""
         snapshots = []
         next_key = 0
@@ -86,4 +93,7 @@ class Bucket:
             if not is_truncated:
                 break
 
-        return [Snapshot(name=snapshot, bucket=self) for snapshot in snapshots]
+        return [
+            Bucket(name=f'{self.name}/{snapshot.strip("/")}', tx=self.tx)
+            for snapshot in snapshots
+        ]

{vastdb-0.1.2 → vastdb-0.1.4}/vastdb/errors.py

@@ -71,6 +71,10 @@ class ServiceUnavailable(HttpError):
     pass
 
 
+class Slowdown(ServiceUnavailable):
+    pass
+
+
 class UnexpectedError(HttpError):
     pass
 
@@ -85,6 +89,10 @@ class InvalidArgument(Exception):
     pass
 
 
+class TooWideRow(InvalidArgument):
+    pass
+
+
 class Missing(Exception):
     pass
 
@@ -93,11 +101,25 @@ class MissingTransaction(Missing):
     pass
 
 
+class MissingRowIdColumn(Missing):
+    pass
+
+
+class NotSupported(Exception):
+    pass
+
+
 @dataclass
 class MissingBucket(Missing):
     bucket: str
 
 
+@dataclass
+class MissingSnapshot(Missing):
+    bucket: str
+    snapshot: str
+
+
 @dataclass
 class MissingSchema(Missing):
     bucket: str
@@ -136,6 +158,25 @@ class TableExists(Exists):
     table: str
 
 
+@dataclass
+class NotSupportedCommand(NotSupported):
+    bucket: str
+    schema: str
+    table: str
+
+
+@dataclass
+class NotSupportedVersion(NotSupported):
+    err_msg: str
+    version: str
+
+
+def handle_unavailable(**kwargs):
+    if kwargs['code'] == 'SlowDown':
+        raise Slowdown(**kwargs)
+    raise ServiceUnavailable(**kwargs)
+
+
 ERROR_TYPES_MAP = {
     HttpStatus.BAD_REQUEST: BadRequest,
     HttpStatus.FOBIDDEN: Forbidden,
@@ -145,7 +186,7 @@ ERROR_TYPES_MAP = {
     HttpStatus.CONFLICT: Conflict,
     HttpStatus.INTERNAL_SERVER_ERROR: InternalServerError,
     HttpStatus.NOT_IMPLEMENTED: NotImplemented,
-    HttpStatus.SERVICE_UNAVAILABLE: ServiceUnavailable,
+    HttpStatus.SERVICE_UNAVAILABLE: handle_unavailable,
 }
 
 
@@ -178,4 +219,4 @@ def from_response(res: requests.Response):
     log.warning("RPC failed: %s", kwargs)
     status = HttpStatus(res.status_code)
     error_type = ERROR_TYPES_MAP.get(status, UnexpectedError)
-    raise error_type(**kwargs)
+    return error_type(**kwargs)  # type: ignore

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

@@ -1,7 +1,6 @@
 import itertools
 import json
 import logging
-import math
 import re
 import struct
 import urllib.parse
@@ -179,9 +178,11 @@ class Predicate:
         )
         from ibis.expr.operations.logical import (
             And,
+            Between,
             Equals,
             Greater,
             GreaterEqual,
+            InValues,
             Less,
             LessEqual,
             Not,
@@ -200,6 +201,7 @@ class Predicate:
             IsNull: self.build_is_null,
             Not: self.build_is_not_null,
             StringContains: self.build_match_substring,
+            Between: self.build_between,
         }
 
         positions_map = dict((f.name, index) for index, f in enumerate(self.schema))  # TODO: BFS
@@ -219,40 +221,60 @@ class Predicate:
                 prev_field_name = None
                 for inner_op in or_args:
                     _logger.debug('inner_op %s', inner_op)
-                    builder_func: Any = builder_map.get(type(inner_op))
+                    op_type = type(inner_op)
+                    builder_func: Any = builder_map.get(op_type)
                     if not builder_func:
-                        raise NotImplementedError(inner_op.name)
+                        if op_type == InValues:
+                            builder_func = self.build_equal
+                        else:
+                            raise NotImplementedError(self.expr)
 
                     if builder_func == self.build_is_null:
                         column, = inner_op.args
-                        literal = None
+                        literals = (None,)
                     elif builder_func == self.build_is_not_null:
                         not_arg, = inner_op.args
                         # currently we only support not is_null, checking we really got is_null under the not:
                         if not builder_map.get(type(not_arg)) == self.build_is_null:
-                            raise NotImplementedError(not_arg.args[0].name)
+                            raise NotImplementedError(self.expr)
                         column, = not_arg.args
-                        literal = None
+                        literals = (None,)
+                    elif builder_func == self.build_between:
+                        column, lower, upper = inner_op.args
+                        literals = (None,)
                     else:
-                        column, literal = inner_op.args
-                        if not isinstance(literal, Literal):
-                            raise NotImplementedError(inner_op.name)
+                        column, arg = inner_op.args
+                        if isinstance(arg, tuple):
+                            literals = arg
+                        else:
+                            literals = (arg,)
+                        for literal in literals:
+                            if not isinstance(literal, Literal):
+                                raise NotImplementedError(self.expr)
 
                     if not isinstance(column, TableColumn):
-                        raise NotImplementedError(inner_op.name)
+                        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:
-                        raise NotImplementedError(op.name)
+                        raise NotImplementedError(self.expr)
 
-                    args_offsets = [self.build_column(position=positions_map[field_name])]
-                    if literal:
-                        field = self.schema.field(field_name)
-                        args_offsets.append(self.build_literal(field=field, value=literal.value))
+                    column_offset = self.build_column(position=positions_map[field_name])
+                    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))
+                        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))
 
-                    inner_offsets.append(builder_func(*args_offsets))
+                        inner_offsets.append(builder_func(*args_offsets))
+
+                if not inner_offsets:
+                    raise NotImplementedError(self.expr)  # an empty OR is equivalent to a 'FALSE' literal
 
                 domain_offset = self.build_or(inner_offsets)
                 offsets.append(domain_offset)
@@ -542,6 +564,13 @@ class Predicate:
     def build_match_substring(self, column: int, literal: int):
         return self.build_function('match_substring', column, literal)
 
+    def build_between(self, column: int, lower: int, upper: int):
+        offsets = [
+            self.build_greater_equal(column, lower),
+            self.build_less_equal(column, upper),
+        ]
+        return self.build_and(offsets)
+
 
 class FieldNodesState:
     def __init__(self) -> None:
@@ -719,20 +748,6 @@ def _parse_table_info(obj):
     return TableInfo(name, properties, handle, num_rows, used_bytes)
 
 
-def build_record_batch(column_info, column_values):
-    fields = [pa.field(column_name, column_type) for column_type, column_name in column_info]
-    schema = pa.schema(fields)
-    arrays = [pa.array(column_values[column_type], type=column_type) for column_type, _ in column_info]
-    batch = pa.record_batch(arrays, schema)
-    return serialize_record_batch(batch)
-
-
-def serialize_record_batch(batch):
-    sink = pa.BufferOutputStream()
-    with pa.ipc.new_stream(sink, batch.schema) as writer:
-        writer.write(batch)
-    return sink.getvalue()
-
 # Results that returns from tablestats
 
 
@@ -952,26 +967,27 @@ class VastdbApi:
 
             return bucket_name, schemas, next_key, is_truncated, count
 
-    def list_snapshots(self, bucket, max_keys=1000, next_token=None, expected_retvals=None):
+    def list_snapshots(self, bucket, max_keys=1000, next_token=None, name_prefix=''):
         next_token = next_token or ''
-        expected_retvals = expected_retvals or []
-        url_params = {'list_type': '2', 'prefix': '.snapshot/', 'delimiter': '/', 'max_keys': str(max_keys)}
+        url_params = {'list_type': '2', 'prefix': '.snapshot/' + name_prefix, 'delimiter': '/', 'max_keys': str(max_keys)}
         if next_token:
             url_params['continuation-token'] = next_token
 
         res = self.session.get(self._api_prefix(bucket=bucket, command="list", url_params=url_params), headers={}, stream=True)
-        self._check_res(res, "list_snapshots", expected_retvals)
-        if res.status_code == 200:
-            out = b''.join(res.iter_content(chunk_size=128))
-            xml_str = out.decode()
-            xml_dict = xmltodict.parse(xml_str)
-            list_res = xml_dict['ListBucketResult']
-            is_truncated = list_res['IsTruncated'] == 'true'
-            marker = list_res['Marker']
-            common_prefixes = list_res['CommonPrefixes'] if 'CommonPrefixes' in list_res else []
-            snapshots = [v['Prefix'] for v in common_prefixes]
+        self._check_res(res, "list_snapshots")
 
-            return snapshots, is_truncated, marker
+        out = b''.join(res.iter_content(chunk_size=128))
+        xml_str = out.decode()
+        xml_dict = xmltodict.parse(xml_str)
+        list_res = xml_dict['ListBucketResult']
+        is_truncated = list_res['IsTruncated'] == 'true'
+        marker = list_res['Marker']
+        common_prefixes = list_res.get('CommonPrefixes', [])
+        if isinstance(common_prefixes, dict):  # in case there is a single snapshot
+            common_prefixes = [common_prefixes]
+        snapshots = [v['Prefix'] for v in common_prefixes]
+
+        return snapshots, is_truncated, marker
 
     def create_table(self, bucket, schema, name, arrow_schema, txid=0, client_tags=[], expected_retvals=[],
                      topic_partitions=0, create_imports_table=False, use_external_row_ids_allocation=False):
@@ -1030,7 +1046,7 @@ class VastdbApi:
         # create the table
         return self.create_table(bucket, schema, name, arrow_schema, txid, client_tags, expected_retvals)
 
-    def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[]):
+    def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[], imports_table_stats=False):
         """
         GET /mybucket/myschema/mytable?stats HTTP/1.1
         tabular-txid: TransactionId
@@ -1039,7 +1055,8 @@ class VastdbApi:
         The Command will return the statistics in flatbuf format
         """
         headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
-        res = self.session.get(self._api_prefix(bucket=bucket, schema=schema, table=name, command="stats"), headers=headers)
+        url_params = {'sub-table': IMPORTED_OBJECTS_TABLE_NAME} if imports_table_stats else {}
+        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))
@@ -1527,11 +1544,18 @@ class VastdbApi:
             if response.status_code != 200:
                 return response
 
+            ALLOWED_IMPORT_STATES = {
+                'Success',
+                'TabularInProgress',
+                'TabularAlreadyImported',
+                'TabularImportNotStarted',
+            }
+
             chunk_size = 1024
             for chunk in response.iter_content(chunk_size=chunk_size):
                 chunk_dict = json.loads(chunk)
                 _logger.debug("import data chunk=%s, result: %s", chunk_dict, chunk_dict['res'])
-                if chunk_dict['res'] != 'Success' and chunk_dict['res'] != 'TabularInProgress' and chunk_dict['res'] != 'TabularAlreadyImported':
+                if chunk_dict['res'] not in ALLOWED_IMPORT_STATES:
                     raise errors.ImportFilesError(
                         f"Encountered an error during import_data. status: {chunk_dict['res']}, "
                         f"error message: {chunk_dict['err_msg'] or 'Unexpected error'} during import of "
@@ -1555,48 +1579,6 @@ class VastdbApi:
 
         return self._check_res(res, "import_data", expected_retvals)
 
-    def _record_batch_slices(self, batch, rows_per_slice=None):
-        max_slice_size_in_bytes = int(0.9 * 5 * 1024 * 1024)  # 0.9 * 5MB
-        batch_len = len(batch)
-        serialized_batch = serialize_record_batch(batch)
-        batch_size_in_bytes = len(serialized_batch)
-        _logger.debug('max_slice_size_in_bytes=%d batch_len=%d batch_size_in_bytes=%d',
-                      max_slice_size_in_bytes, batch_len, batch_size_in_bytes)
-
-        if not rows_per_slice:
-            if batch_size_in_bytes < max_slice_size_in_bytes:
-                rows_per_slice = batch_len
-            else:
-                rows_per_slice = int(0.9 * batch_len * max_slice_size_in_bytes / batch_size_in_bytes)
-
-        done_slicing = False
-        while not done_slicing:
-            # Attempt slicing according to the current rows_per_slice
-            offset = 0
-            serialized_slices = []
-            for i in range(math.ceil(batch_len / rows_per_slice)):
-                offset = rows_per_slice * i
-                if offset >= batch_len:
-                    done_slicing = True
-                    break
-                slice_batch = batch.slice(offset, rows_per_slice)
-                serialized_slice_batch = serialize_record_batch(slice_batch)
-                sizeof_serialized_slice_batch = len(serialized_slice_batch)
-
-                if sizeof_serialized_slice_batch <= max_slice_size_in_bytes:
-                    serialized_slices.append(serialized_slice_batch)
-                else:
-                    _logger.info(f'Using rows_per_slice {rows_per_slice} slice {i} size {sizeof_serialized_slice_batch} exceeds {max_slice_size_in_bytes} bytes, trying smaller rows_per_slice')
-                    # We have a slice that is too large
-                    rows_per_slice = int(rows_per_slice / 2)
-                    if rows_per_slice < 1:
-                        raise ValueError('cannot decrease batch size below 1 row')
-                    break
-            else:
-                done_slicing = True
-
-        return serialized_slices
-
     def insert_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
         """
         POST /mybucket/myschema/mytable?rows HTTP/1.1
@@ -1611,8 +1593,7 @@ class VastdbApi:
         headers['Content-Length'] = str(len(record_batch))
         res = self.session.post(self._api_prefix(bucket=bucket, schema=schema, table=table, command="rows"),
                                 data=record_batch, headers=headers, stream=True)
-        self._check_res(res, "insert_rows", expected_retvals)
-        res.raw.read()  # flush the response
+        return self._check_res(res, "insert_rows", expected_retvals)
 
     def update_rows(self, bucket, schema, table, record_batch, txid=0, client_tags=[], expected_retvals=[]):
         """
@@ -2140,6 +2121,13 @@ class QueryDataRequest:
         self.response_parser = response_parser
 
 
+def get_response_schema(schema: 'pa.Schema' = pa.schema([]), field_names: Optional[List[str]] = None):
+    if field_names is None:
+        field_names = [field.name for field in schema]
+
+    return pa.schema([schema.field(name) for name in field_names])
+
+
 def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibis.expr.types.BooleanColumn = None, field_names: Optional[List[str]] = None):
     builder = flatbuffers.Builder(1024)
 
@@ -2160,13 +2148,11 @@ def build_query_data_request(schema: 'pa.Schema' = pa.schema([]), predicate: ibi
     filter_obj = predicate.serialize(builder)
 
     parser = QueryDataParser(schema)
-    fields_map = {node.field.name: node.field for node in parser.nodes}
     leaves_map = {node.field.name: [leaf.index for leaf in node._iter_leaves()] for node in parser.nodes}
 
-    if field_names is None:
-        field_names = [field.name for field in schema]
+    response_schema = get_response_schema(schema, field_names)
+    field_names = [field.name for field in response_schema]
 
-    response_schema = pa.schema([fields_map[name] for name in field_names])
     projection_fields = []
     for field_name in field_names:
         # TODO: only root-level projection pushdown is supported (i.e. no support for SELECT s.x FROM t)

{vastdb-0.1.2 → vastdb-0.1.4}/vastdb/schema.py

@@ -87,4 +87,4 @@ class Schema:
 
 def _parse_table_info(table_info, schema: "schema.Schema"):
     stats = table.TableStats(num_rows=table_info.num_rows, size_in_bytes=table_info.size_in_bytes)
-    return table.Table(name=table_info.name, schema=schema, handle=int(table_info.handle), stats=stats)
+    return table.Table(name=table_info.name, schema=schema, handle=int(table_info.handle), stats=stats, _imports_table=False)

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

@@ -11,7 +11,25 @@ import os
 
 import boto3
 
-from . import internal_commands, transaction
+from . import errors, internal_commands, transaction
+
+
+class Features:
+    """VAST database features - check if server is already support a feature."""
+
+    def __init__(self, vast_version):
+        """Save the server version."""
+        self.vast_version = vast_version
+
+    def check_imports_table(self):
+        """Check if the feature that support imports table is supported."""
+        if self.vast_version < (5, 2):
+            raise errors.NotSupportedVersion("import_table requires 5.2+", self.vast_version)
+
+    def check_return_row_ids(self):
+        """Check if insert/update/delete can return the row_ids."""
+        if self.vast_version < (5, 1):
+            raise errors.NotSupportedVersion("return_row_ids requires 5.1+", self.vast_version)
 
 
 class Session:
@@ -27,6 +45,8 @@ class Session:
             endpoint = os.environ['AWS_S3_ENDPOINT_URL']
 
         self.api = internal_commands.VastdbApi(endpoint, access, secret)
+        version_tuple = tuple(int(part) for part in self.api.vast_version.split('.'))
+        self.features = Features(version_tuple)
         self.s3 = boto3.client('s3',
             aws_access_key_id=access,
             aws_secret_access_key=secret,