vastdb 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

vastdb/errors.py

@@ -71,6 +71,10 @@ class ServiceUnavailable(HttpError):
     pass
 
 
+class Slowdown(ServiceUnavailable):
+    pass
+
+
 class UnexpectedError(HttpError):
     pass
 
@@ -97,6 +101,10 @@ class MissingTransaction(Missing):
     pass
 
 
+class MissingRowIdColumn(Missing):
+    pass
+
+
 class NotSupported(Exception):
     pass
 
@@ -163,6 +171,12 @@ class NotSupportedVersion(NotSupported):
     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,
@@ -172,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,
 }
 
 
@@ -205,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)
-    return error_type(**kwargs)
+    return error_type(**kwargs)  # type: ignore

vastdb/internal_commands.py

@@ -178,6 +178,7 @@ class Predicate:
         )
         from ibis.expr.operations.logical import (
             And,
+            Between,
             Equals,
             Greater,
             GreaterEqual,
@@ -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
@@ -237,6 +239,9 @@ class Predicate:
                             raise NotImplementedError(self.expr)
                         column, = not_arg.args
                         literals = (None,)
+                    elif builder_func == self.build_between:
+                        column, lower, upper = inner_op.args
+                        literals = (None,)
                     else:
                         column, arg = inner_op.args
                         if isinstance(arg, tuple):
@@ -262,6 +267,9 @@ class Predicate:
                         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))
 
@@ -556,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:
@@ -1578,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=[]):
         """
@@ -2107,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)
 
@@ -2127,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/session.py

@@ -26,6 +26,11 @@ class Features:
         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:
     """VAST database session."""

vastdb/table.py

@@ -7,8 +7,9 @@ import queue
 from dataclasses import dataclass, field
 from math import ceil
 from threading import Event
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
+import backoff
 import ibis
 import pyarrow as pa
 
@@ -18,10 +19,14 @@ log = logging.getLogger(__name__)
 
 
 INTERNAL_ROW_ID = "$row_id"
+INTERNAL_ROW_ID_FIELD = pa.field(INTERNAL_ROW_ID, pa.uint64())
+
 MAX_ROWS_PER_BATCH = 512 * 1024
 # for insert we need a smaller limit due to response amplification
 # for example insert of 512k uint8 result in 512k*8bytes response since row_ids are uint64
 MAX_INSERT_ROWS_PER_PATCH = 512 * 1024
+# in case insert has TooWideRow - need to insert in smaller batches - each cell could contain up to 128K, and our wire is limited to 5MB
+MAX_COLUMN_IN_BATCH = int(5 * 1024 / 128)
 
 
 @dataclass
@@ -46,6 +51,8 @@ class QueryConfig:
     use_semi_sorted_projections: bool = True
     rows_per_split: int = 4000000
     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))
 
 
 @dataclass
@@ -72,7 +79,8 @@ class SelectSplitState:
         Can be called repeatedly, to allow pagination.
         """
         while not self.done:
-            response = api.query_data(
+            query_with_backoff = self.config.backoff_func(api.query_data)
+            response = query_with_backoff(
                             bucket=self.table.bucket.name,
                             schema=self.table.schema.name,
                             table=self.table.name,
@@ -291,11 +299,17 @@ class Table:
 
         query_schema = self.arrow_schema
         if internal_row_id:
-            queried_fields = [pa.field(INTERNAL_ROW_ID, pa.uint64())]
+            queried_fields = [INTERNAL_ROW_ID_FIELD]
             queried_fields.extend(column for column in self.arrow_schema)
             query_schema = pa.schema(queried_fields)
             columns.append(INTERNAL_ROW_ID)
 
+        if predicate is True:
+            predicate = None
+        if predicate is False:
+            response_schema = internal_commands.get_response_schema(schema=query_schema, field_names=columns)
+            return pa.RecordBatchReader.from_batches(response_schema, [])
+
         query_data_request = internal_commands.build_query_data_request(
             schema=query_schema,
             predicate=predicate,
@@ -385,27 +399,68 @@ class Table:
 
         return pa.RecordBatchReader.from_batches(query_data_request.response_schema, batches_iterator())
 
-    def insert(self, rows: pa.RecordBatch) -> pa.RecordBatch:
+    def insert_in_column_batches(self, rows: pa.RecordBatch):
+        """Split the RecordBatch into max_columns that can be inserted in single RPC.
+
+        Insert first MAX_COLUMN_IN_BATCH columns and get the row_ids. Then loop on the rest of the columns and
+        update in groups of MAX_COLUMN_IN_BATCH.
+        """
+        column_record_batch = pa.RecordBatch.from_arrays([_combine_chunks(rows.column(i)) for i in range(0, MAX_COLUMN_IN_BATCH)],
+                                                         schema=pa.schema([rows.schema.field(i) for i in range(0, MAX_COLUMN_IN_BATCH)]))
+        row_ids = self.insert(rows=column_record_batch)  # type: ignore
+
+        columns_names = [field.name for field in rows.schema]
+        columns = list(rows.schema)
+        arrays = [_combine_chunks(rows.column(i)) for i in range(len(rows.schema))]
+        for start in range(MAX_COLUMN_IN_BATCH, len(rows.schema), MAX_COLUMN_IN_BATCH):
+            end = start + MAX_COLUMN_IN_BATCH if start + MAX_COLUMN_IN_BATCH < len(rows.schema) else len(rows.schema)
+            columns_name_chunk = columns_names[start:end]
+            columns_chunks = columns[start:end]
+            arrays_chunks = arrays[start:end]
+            columns_chunks.append(INTERNAL_ROW_ID_FIELD)
+            arrays_chunks.append(row_ids.to_pylist())
+            column_record_batch = pa.RecordBatch.from_arrays(arrays_chunks, schema=pa.schema(columns_chunks))
+            self.update(rows=column_record_batch, columns=columns_name_chunk)
+        return row_ids
+
+    def insert(self, rows: pa.RecordBatch):
         """Insert a RecordBatch into this table."""
         if self._imports_table:
             raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
-        serialized_slices = util.iter_serialized_slices(rows, MAX_INSERT_ROWS_PER_PATCH)
-        for slice in serialized_slices:
-            self.tx._rpc.api.insert_rows(self.bucket.name, self.schema.name, self.name, record_batch=slice,
-                                               txid=self.tx.txid)
+        try:
+            row_ids = []
+            serialized_slices = util.iter_serialized_slices(rows, MAX_INSERT_ROWS_PER_PATCH)
+            for slice in serialized_slices:
+                res = self.tx._rpc.api.insert_rows(self.bucket.name, self.schema.name, self.name, record_batch=slice,
+                                                   txid=self.tx.txid)
+                (batch,) = pa.RecordBatchStreamReader(res.raw)
+                row_ids.append(batch[INTERNAL_ROW_ID])
+            try:
+                self.tx._rpc.features.check_return_row_ids()
+            except errors.NotSupportedVersion:
+                return  # type: ignore
+            return pa.chunked_array(row_ids)
+        except errors.TooWideRow:
+            self.tx._rpc.features.check_return_row_ids()
+            return self.insert_in_column_batches(rows)
 
     def update(self, rows: Union[pa.RecordBatch, pa.Table], columns: Optional[List[str]] = None) -> None:
         """Update a subset of cells in this table.
 
-        Row IDs are specified using a special field (named "$row_id" of uint64 type).
+        Row IDs are specified using a special field (named "$row_id" of uint64 type) - this function assume that this
+        special field is part of arguments.
 
         A subset of columns to be updated can be specified via the `columns` argument.
         """
         if self._imports_table:
             raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
+        try:
+            rows_chunk = rows[INTERNAL_ROW_ID]
+        except KeyError:
+            raise errors.MissingRowIdColumn
         if columns is not None:
             update_fields = [(INTERNAL_ROW_ID, pa.uint64())]
-            update_values = [_combine_chunks(rows[INTERNAL_ROW_ID])]
+            update_values = [_combine_chunks(rows_chunk)]
             for col in columns:
                 update_fields.append(rows.field(col))
                 update_values.append(_combine_chunks(rows[col]))
@@ -424,8 +479,14 @@ class Table:
 
         Row IDs are specified using a special field (named "$row_id" of uint64 type).
         """
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
+        try:
+            rows_chunk = rows[INTERNAL_ROW_ID]
+        except KeyError:
+            raise errors.MissingRowIdColumn
         delete_rows_rb = pa.record_batch(schema=pa.schema([(INTERNAL_ROW_ID, pa.uint64())]),
-                                         data=[_combine_chunks(rows[INTERNAL_ROW_ID])])
+                                         data=[_combine_chunks(rows_chunk)])
 
         serialized_slices = util.iter_serialized_slices(delete_rows_rb, MAX_ROWS_PER_BATCH)
         for slice in serialized_slices:

vastdb/tests/test_tables.py

@@ -71,6 +71,16 @@ def test_tables(session, clean_bucket_name):
         }
 
 
+def test_insert_wide_row(session, clean_bucket_name):
+    columns = pa.schema([pa.field(f's{i}', pa.utf8()) for i in range(500)])
+    data = [['a' * 10**4] for i in range(500)]
+    expected = pa.table(schema=columns, data=data)
+
+    with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
+        actual = pa.Table.from_batches(t.select())
+        assert actual == expected
+
+
 def test_exists(session, clean_bucket_name):
     with session.transaction() as tx:
         s = tx.bucket(clean_bucket_name).create_schema('s1')
@@ -265,6 +275,11 @@ def test_filters(session, clean_bucket_name):
             return pa.Table.from_batches(t.select(predicate=predicate), t.arrow_schema)
 
         assert select(None) == expected
+        assert select(True) == expected
+        assert select(False) == pa.Table.from_batches([], schema=columns)
+
+        assert select(t['a'].between(222, 444)) == expected.filter((pc.field('a') >= 222) & (pc.field('a') <= 444))
+        assert select((t['a'].between(222, 444)) & (t['b'] > 2.5)) == expected.filter((pc.field('a') >= 222) & (pc.field('a') <= 444) & (pc.field('b') > 2.5))
 
         assert select(t['a'] > 222) == expected.filter(pc.field('a') > 222)
         assert select(t['a'] < 222) == expected.filter(pc.field('a') < 222)
@@ -331,7 +346,8 @@ def test_parquet_export(session, clean_bucket_name):
             ['a', 'b'],
         ])
         expected = pa.Table.from_batches([rb])
-        t.insert(rb)
+        rb = t.insert(rb)
+        assert rb.to_pylist() == [0, 1]
         actual = pa.Table.from_batches(t.select())
         assert actual == expected
 

vastdb/tests/util.py

@@ -9,7 +9,9 @@ def prepare_data(session, clean_bucket_name, schema_name, table_name, arrow_tabl
     with session.transaction() as tx:
         s = tx.bucket(clean_bucket_name).create_schema(schema_name)
         t = s.create_table(table_name, arrow_table.schema)
-        t.insert(arrow_table)
+        row_ids_array = t.insert(arrow_table)
+        row_ids = row_ids_array.to_pylist()
+        assert row_ids == list(range(arrow_table.num_rows))
         yield t
         t.drop()
         s.drop()

{vastdb-0.1.3.dist-info → vastdb-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vastdb
-Version: 0.1.3
+Version: 0.1.4
 Summary: VAST Data SDK
 Home-page: https://github.com/vast-data/vastdb_sdk
 Author: VAST DATA
@@ -25,6 +25,7 @@ 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.3.dist-info → vastdb-0.1.4.dist-info}/RECORD

@@ -151,11 +151,11 @@ vast_flatbuf/tabular/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hS
 vastdb/__init__.py,sha256=cMJtZuJ0IL9aKyM3DUWqTCzuP1H1MXXVivKKE1-q0DY,292
 vastdb/bucket.py,sha256=4rPEm9qlPTg7ccWO6VGmd4LKb8w-BDhJYwzXGjn03sc,3566
 vastdb/conftest.py,sha256=pKpo_46Vq4QHzTDQAFxasrVhnZ2V2L-y6IMLxojxaFM,2132
-vastdb/errors.py,sha256=vKWoq1yXrHyafMWwJgW_sQkSxQYxlI1JbTVCLz5Xi9Y,3793
-vastdb/internal_commands.py,sha256=ZD2YXYvZ3lJWYzZU0oHtv8G3lNtDQUF0e8yg8813Xt4,99575
+vastdb/errors.py,sha256=fj8IlPnGi1lbJWIl1-8MSjLavL9bYQ-YUoboWbXCo54,4047
+vastdb/internal_commands.py,sha256=yS6ylyuJjaAwAm4OqVGX4tq-Un5cvM-LXp7F4eYOUDw,100414
 vastdb/schema.py,sha256=ql4TPB1W_FQ_BHov3CKHI8JX3krXMlcKWz7dTrjpQ1w,3346
-vastdb/session.py,sha256=ciYS8Je2cRpuaAEE6Wjk79VsW0KAPdnRB2cqfxFCjis,2323
-vastdb/table.py,sha256=xnSTWUUa0QHzXC5MUQWsGT1fsG8yAgMLy3nrgSH4j5Q,25661
+vastdb/session.py,sha256=2tu5cp7xG28ynyQfEl9_HM2dtNcLM2AoJmm3bfNLC0o,2563
+vastdb/table.py,sha256=apRXCrglg6_glozJXu8D7q6du5seP7NMi42PNjyGcTM,28891
 vastdb/transaction.py,sha256=g8YTcYnsNPIhB2udbHyT5RIFB5kHnBLJcvV2CWRICwI,2845
 vastdb/util.py,sha256=rs7nLL2Qz-OVEZDSVIqAvS-uETMq-zxQs5jBksB5-JA,4276
 vastdb/bench/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -167,11 +167,11 @@ vastdb/tests/test_nested.py,sha256=3kejEvtSqV0LrUgb1QglRjrlxnKI4_AXTFw2nE7Q520,9
 vastdb/tests/test_projections.py,sha256=_cDNfD5zTwbCXLk6uGpPUWGN0P-4HElu5OjubWu-Jg0,1255
 vastdb/tests/test_sanity.py,sha256=ixx0QPo73hLHjAa7bByFXjS1XST0WvmSwLEpgnHh_JY,2960
 vastdb/tests/test_schemas.py,sha256=qoHTLX51D-0S4bMxdCpRh9gaYQd-BkZdT_agGOwFwTM,1739
-vastdb/tests/test_tables.py,sha256=joeEQ30TwKBQc-2N_qGIdviZVnQr4rs6thlNsy5s_og,26672
+vastdb/tests/test_tables.py,sha256=pfQx0OZm6oVJj1-CziPWUoEn3l2-OET4Bpl8M9Z4mws,27499
 vastdb/tests/test_util.py,sha256=owRAU3TCKMq-kz54NRdA5wX2O_bZIHqG5ucUR77jm5k,1046
-vastdb/tests/util.py,sha256=NaCzKymEGy1xuiyMxyt2_0frKVfVk9iGrFwLf3GHjTI,435
-vastdb-0.1.3.dist-info/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
-vastdb-0.1.3.dist-info/METADATA,sha256=3h3JttUxw9oMMsxV_CVG_LMYwhgegsS9-b4gZkihrM0,1319
-vastdb-0.1.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-vastdb-0.1.3.dist-info/top_level.txt,sha256=Vsj2MKtlhPg0J4so64slQtnwjhgoPmJgcG-6YcVAwVc,20
-vastdb-0.1.3.dist-info/RECORD,,
+vastdb/tests/util.py,sha256=dpRJYbboDnlqL4qIdvScpp8--5fxRUBIcIYitrfcj9o,555
+vastdb-0.1.4.dist-info/LICENSE,sha256=obffan7LYrq7hLHNrY7vHcn2pKUTBUYXMKu-VOAvDxU,11333
+vastdb-0.1.4.dist-info/METADATA,sha256=SyZkyjQSwklzsq3oub8m8w9lY-HuI4XOG72y8trKvf4,1350
+vastdb-0.1.4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+vastdb-0.1.4.dist-info/top_level.txt,sha256=Vsj2MKtlhPg0J4so64slQtnwjhgoPmJgcG-6YcVAwVc,20
+vastdb-0.1.4.dist-info/RECORD,,