vastdb 1.3.6__py3-none-any.whl → 1.3.8__py3-none-any.whl

vastdb/_internal.py

@@ -3,6 +3,7 @@ import json
 import logging
 import re
 import struct
+import time
 import urllib.parse
 from collections import defaultdict, namedtuple
 from enum import Enum
@@ -37,75 +38,82 @@ from ibis.expr.operations.relations import Field
 from ibis.expr.operations.strings import StartsWith, StringContains
 from ibis.expr.operations.structs import StructField
 
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.BinaryLiteral as fb_binary_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.BooleanLiteral as fb_bool_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Call as fb_call
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.DateLiteral as fb_date32_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.DecimalLiteral as fb_decimal_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Expression as fb_expression
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldIndex as fb_field_index
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldRef as fb_field_ref
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Float32Literal as fb_float32_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Float64Literal as fb_float64_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int8Literal as fb_int8_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int16Literal as fb_int16_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int32Literal as fb_int32_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int64Literal as fb_int64_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Literal as fb_literal
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Relation as fb_relation
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.RelationImpl as rel_impl
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.Source as fb_source
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.StringLiteral as fb_string_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.TimeLiteral as fb_time_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.TimestampLiteral as fb_timestamp_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt8Literal as fb_uint8_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt16Literal as fb_uint16_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt32Literal as fb_uint32_lit
-import vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt64Literal as fb_uint64_lit
-import vast_flatbuf.org.apache.arrow.flatbuf.Binary as fb_binary
-import vast_flatbuf.org.apache.arrow.flatbuf.Bool as fb_bool
-import vast_flatbuf.org.apache.arrow.flatbuf.Date as fb_date
-import vast_flatbuf.org.apache.arrow.flatbuf.Decimal as fb_decimal
-import vast_flatbuf.org.apache.arrow.flatbuf.Field as fb_field
-import vast_flatbuf.org.apache.arrow.flatbuf.FixedSizeBinary as fb_fixed_size_binary
-import vast_flatbuf.org.apache.arrow.flatbuf.FloatingPoint as fb_floating_point
-import vast_flatbuf.org.apache.arrow.flatbuf.Int as fb_int
-import vast_flatbuf.org.apache.arrow.flatbuf.List as fb_list
-import vast_flatbuf.org.apache.arrow.flatbuf.Map as fb_map
-import vast_flatbuf.org.apache.arrow.flatbuf.Schema as fb_schema
-import vast_flatbuf.org.apache.arrow.flatbuf.Struct_ as fb_struct
-import vast_flatbuf.org.apache.arrow.flatbuf.Time as fb_time
-import vast_flatbuf.org.apache.arrow.flatbuf.Timestamp as fb_timestamp
-import vast_flatbuf.org.apache.arrow.flatbuf.Utf8 as fb_utf8
-import vast_flatbuf.tabular.AlterColumnRequest as tabular_alter_column
-import vast_flatbuf.tabular.AlterProjectionTableRequest as tabular_alter_projection
-import vast_flatbuf.tabular.AlterSchemaRequest as tabular_alter_schema
-import vast_flatbuf.tabular.AlterTableRequest as tabular_alter_table
-import vast_flatbuf.tabular.Column as tabular_projecion_column
-import vast_flatbuf.tabular.ColumnType as tabular_proj_column_type
-import vast_flatbuf.tabular.CreateProjectionRequest as tabular_create_projection
-import vast_flatbuf.tabular.CreateSchemaRequest as tabular_create_schema
-import vast_flatbuf.tabular.ImportDataRequest as tabular_import_data
-import vast_flatbuf.tabular.S3File as tabular_s3_file
-from vast_flatbuf.org.apache.arrow.computeir.flatbuf.Deref import Deref
-from vast_flatbuf.org.apache.arrow.computeir.flatbuf.ExpressionImpl import (
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.BinaryLiteral as fb_binary_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.BooleanLiteral as fb_bool_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Call as fb_call
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.DateLiteral as fb_date32_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.DecimalLiteral as fb_decimal_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Expression as fb_expression
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldIndex as fb_field_index
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.FieldRef as fb_field_ref
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Float32Literal as fb_float32_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Float64Literal as fb_float64_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int8Literal as fb_int8_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int16Literal as fb_int16_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int32Literal as fb_int32_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Int64Literal as fb_int64_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Literal as fb_literal
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Relation as fb_relation
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.RelationImpl as rel_impl
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Source as fb_source
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.StringLiteral as fb_string_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.TimeLiteral as fb_time_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.TimestampLiteral as fb_timestamp_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt8Literal as fb_uint8_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt16Literal as fb_uint16_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt32Literal as fb_uint32_lit
+import vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.UInt64Literal as fb_uint64_lit
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Binary as fb_binary
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Bool as fb_bool
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Date as fb_date
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Decimal as fb_decimal
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Field as fb_field
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.FixedSizeBinary as fb_fixed_size_binary
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.FloatingPoint as fb_floating_point
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Int as fb_int
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.List as fb_list
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Map as fb_map
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Schema as fb_schema
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Struct_ as fb_struct
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Time as fb_time
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Timestamp as fb_timestamp
+import vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Utf8 as fb_utf8
+import vastdb.vast_flatbuf.tabular.AlterColumnRequest as tabular_alter_column
+import vastdb.vast_flatbuf.tabular.AlterProjectionTableRequest as tabular_alter_projection
+import vastdb.vast_flatbuf.tabular.AlterSchemaRequest as tabular_alter_schema
+import vastdb.vast_flatbuf.tabular.AlterTableRequest as tabular_alter_table
+import vastdb.vast_flatbuf.tabular.Column as tabular_projecion_column
+import vastdb.vast_flatbuf.tabular.ColumnType as tabular_proj_column_type
+import vastdb.vast_flatbuf.tabular.CreateProjectionRequest as tabular_create_projection
+import vastdb.vast_flatbuf.tabular.CreateSchemaRequest as tabular_create_schema
+import vastdb.vast_flatbuf.tabular.ImportDataRequest as tabular_import_data
+import vastdb.vast_flatbuf.tabular.KeyName as import_key_name
+import vastdb.vast_flatbuf.tabular.S3File as tabular_s3_file
+from vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.Deref import Deref
+from vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.ExpressionImpl import (
     ExpressionImpl,
 )
-from vast_flatbuf.org.apache.arrow.computeir.flatbuf.LiteralImpl import LiteralImpl
-from vast_flatbuf.org.apache.arrow.flatbuf.DateUnit import DateUnit
-from vast_flatbuf.org.apache.arrow.flatbuf.TimeUnit import TimeUnit
-from vast_flatbuf.org.apache.arrow.flatbuf.Type import Type
-from vast_flatbuf.tabular.GetProjectionTableStatsResponse import (
+from vastdb.vast_flatbuf.org.apache.arrow.computeir.flatbuf.LiteralImpl import (
+    LiteralImpl,
+)
+from vastdb.vast_flatbuf.org.apache.arrow.flatbuf.DateUnit import DateUnit
+from vastdb.vast_flatbuf.org.apache.arrow.flatbuf.TimeUnit import TimeUnit
+from vastdb.vast_flatbuf.org.apache.arrow.flatbuf.Type import Type
+from vastdb.vast_flatbuf.tabular.GetProjectionTableStatsResponse import (
     GetProjectionTableStatsResponse as get_projection_table_stats,
 )
-from vast_flatbuf.tabular.GetTableStatsResponse import (
+from vastdb.vast_flatbuf.tabular.GetTableStatsResponse import (
     GetTableStatsResponse as get_table_stats,
 )
-from vast_flatbuf.tabular.ListProjectionsResponse import (
+from vastdb.vast_flatbuf.tabular.ListProjectionsResponse import (
     ListProjectionsResponse as list_projections,
 )
-from vast_flatbuf.tabular.ListSchemasResponse import ListSchemasResponse as list_schemas
-from vast_flatbuf.tabular.ListTablesResponse import ListTablesResponse as list_tables
+from vastdb.vast_flatbuf.tabular.ListSchemasResponse import (
+    ListSchemasResponse as list_schemas,
+)
+from vastdb.vast_flatbuf.tabular.ListTablesResponse import (
+    ListTablesResponse as list_tables,
+)
 
 from . import errors, util
 from .config import BackoffConfig
@@ -238,7 +246,7 @@ class Predicate:
                         field_path.append(subfield_name)
 
                     if not isinstance(column, Field):
-                        raise NotImplementedError(self.expr)
+                        raise NotImplementedError("predicates not supported for column types ", column, "expr=", self.expr)
 
                     field_path.append(column.name)
                     field_path.reverse()  # first entry should be the top-level column name
@@ -786,7 +794,7 @@ def _decode_table_props(s):
     return {y: _prop_coding[x][1](z) for x, y, z in triplets if z != ''}
 
 
-TableInfo = namedtuple('TableInfo', 'name properties handle num_rows size_in_bytes num_partitions')
+TableInfo = namedtuple('TableInfo', 'name properties handle num_rows size_in_bytes num_partitions sorting_key_enabled')
 
 
 def _parse_table_info(obj, parse_properties):
@@ -797,7 +805,8 @@ def _parse_table_info(obj, parse_properties):
     used_bytes = obj.SizeInBytes()
     num_partitions = obj.NumPartitions()
     properties = parse_properties(properties)
-    return TableInfo(name, properties, handle, num_rows, used_bytes, num_partitions)
+    sorting_key_enabled = obj.SortingKeyEnabled()
+    return TableInfo(name, properties, handle, num_rows, used_bytes, num_partitions, sorting_key_enabled)
 
 
 # Results that returns from tablestats
@@ -961,12 +970,12 @@ class VastdbApi:
         prefix += '&'.join(params_list)
         return prefix
 
-    def _fill_common_headers(self, txid=0, client_tags=[], version_id=1):
+    def _fill_common_headers(self, txid=0, client_tags=[], version_id=1, sorting_key=[]):
         common_headers = {
             'tabular-txid': str(txid),
             'tabular-api-version-id': str(version_id),
             'tabular-client-name': 'tabular-api'
-        }
+        } | {f"tabular-sorted-column-{i}": str(k) for i, k in enumerate(sorting_key)}
 
         return common_headers | {f'tabular-client-tags-{index}': tag for index, tag in enumerate(client_tags)}
 
@@ -1110,17 +1119,28 @@ class VastdbApi:
 
         return snapshots, is_truncated, marker
 
+    def _make_sure_schema_exists(self, bucket, schema, timeout_sec=3):
+        start_time = time.time()
+        while time.time() - start_time < timeout_sec:
+            _, schemas, _, _, _ = self.list_schemas(bucket)
+            if schema in (x[0] for x in schemas):
+                return
+            self.create_schema(bucket, schema)
+        raise Exception(f"Failed to find or create schema {schema} in {timeout_sec} seconds.")
+
     def create_table(self, bucket, schema, name, arrow_schema=None,
                      txid=0, client_tags=[], expected_retvals=[],
-                     create_imports_table=False, use_external_row_ids_allocation=False, table_props=None):
+                     create_imports_table=False, use_external_row_ids_allocation=False, table_props=None,
+                     sorting_key=[]):
         self._create_table_internal(bucket=bucket, schema=schema, name=name, arrow_schema=arrow_schema,
                                     txid=txid, client_tags=client_tags, expected_retvals=expected_retvals,
                                     create_imports_table=create_imports_table, use_external_row_ids_allocation=use_external_row_ids_allocation,
-                                    table_props=table_props)
+                                    table_props=table_props, sorting_key=sorting_key)
 
     def create_topic(self, bucket, name, topic_partitions, expected_retvals=[],
                      message_timestamp_type=None, retention_ms=None, message_timestamp_after_max_ms=None,
                      message_timestamp_before_max_ms=None):
+        self._make_sure_schema_exists(bucket, KAFKA_TOPICS_SCHEMA_NAME)
         table_props = _encode_table_props(message_timestamp_type=message_timestamp_type,
                                           retention_ms=retention_ms,
                                           message_timestamp_after_max_ms=message_timestamp_after_max_ms,
@@ -1132,7 +1152,8 @@ class VastdbApi:
 
     def _create_table_internal(self, bucket, schema, name, arrow_schema=None,
                                txid=0, client_tags=[], expected_retvals=[], topic_partitions=0,
-                               create_imports_table=False, use_external_row_ids_allocation=False, table_props=None):
+                               create_imports_table=False, use_external_row_ids_allocation=False, table_props=None,
+                               sorting_key=[]):
         """
         Create a table, use the following request
         POST /bucket/schema/table?table HTTP/1.1
@@ -1149,8 +1170,7 @@ class VastdbApi:
         The request will look like:
         POST /bucket/schema/table?table&sub-table=vastdb-imported-objects HTTP/1.1
         """
-        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
-
+        headers = self._fill_common_headers(txid=txid, client_tags=client_tags, sorting_key=sorting_key)
         if arrow_schema is None:
             arrow_schema = pa.schema([])
 
@@ -1211,7 +1231,7 @@ class VastdbApi:
                          table_properties=table_properties, new_name=new_name, expected_retvals=expected_retvals)
 
     def alter_table(self, bucket, schema, name, txid=0, client_tags=[], table_properties="",
-                    new_name="", expected_retvals=[]):
+                    new_name="", expected_retvals=[], sorting_key=[]):
         """
         PUT /mybucket/myschema/mytable?table HTTP/1.1
         Content-Length: ContentLength
@@ -1235,7 +1255,7 @@ class VastdbApi:
         builder.Finish(params)
         alter_table_req = builder.Output()
 
-        headers = self._fill_common_headers(txid=txid, client_tags=client_tags)
+        headers = self._fill_common_headers(txid=txid, client_tags=client_tags, sorting_key=sorting_key)
         headers['Content-Length'] = str(len(alter_table_req))
         url_params = {'tabular-new-table-name': schema + "/" + new_name} if len(new_name) else {}
 
@@ -1404,9 +1424,9 @@ class VastdbApi:
             url=self._url(bucket=bucket, schema=schema, table=table, command="column"),
             data=serialized_schema, headers=headers)
 
-    def list_columns(self, bucket, schema, table, *, txid=0, client_tags=None, max_keys=None, next_key=0,
-                     count_only=False, name_prefix="", exact_match=False,
-                     expected_retvals=None, bc_list_internals=False, list_imports_table=False):
+    def _list_columns_internal(self, command, bucket, schema, table, txid, client_tags, max_keys, next_key,
+                               count_only, name_prefix, exact_match, expected_retvals, bc_list_internals,
+                               list_imports_table):
         """
         GET /mybucket/myschema/mytable?columns HTTP/1.1
         tabular-txid: TransactionId
@@ -1436,7 +1456,7 @@ class VastdbApi:
         url_params = {'sub-table': IMPORTED_OBJECTS_TABLE_NAME} if list_imports_table else {}
         res = self._request(
             method="GET",
-            url=self._url(bucket=bucket, schema=schema, table=table, command="column", url_params=url_params),
+            url=self._url(bucket=bucket, schema=schema, table=table, command=command, url_params=url_params),
             headers=headers)
 
         res_headers = res.headers
@@ -1447,6 +1467,20 @@ class VastdbApi:
 
         return columns, next_key, is_truncated, count
 
+    def list_columns(self, bucket, schema, table, *, txid=0, client_tags=None, max_keys=None, next_key=0,
+                     count_only=False, name_prefix="", exact_match=False,
+                     expected_retvals=None, bc_list_internals=False, list_imports_table=False):
+        return self._list_columns_internal('column', bucket, schema, table, txid, client_tags, max_keys, next_key,
+                                           count_only, name_prefix, exact_match, expected_retvals, bc_list_internals,
+                                           list_imports_table)
+
+    def list_sorted_columns(self, bucket, schema, table, *, txid=0, client_tags=None, max_keys=None, next_key=0,
+                            count_only=False, name_prefix="", exact_match=False,
+                            expected_retvals=None, bc_list_internals=False, list_imports_table=False):
+        return self._list_columns_internal('sorted-columns', bucket, schema, table, txid, client_tags, max_keys, next_key,
+                                           count_only, name_prefix, exact_match, expected_retvals, bc_list_internals,
+                                           list_imports_table)
+
     def head_bucket(self, bucket_name):
         """
         Reimplemented, instead of depending on boto3 for checking the existence of a bucket.
@@ -1585,7 +1619,7 @@ class VastdbApi:
     source_files: list of (bucket_name, file_name)
     """
     def import_data(self, bucket, schema, table, source_files, txid=0, client_tags=[], expected_retvals=[], case_sensitive=True,
-                    schedule_id=None, retry_count=0, blocking=True):
+                    schedule_id=None, retry_count=0, blocking=True, key_names=[]):
         """
         POST /mybucket/myschema/mytable?data HTTP/1.1
         Content-Length: ContentLength
@@ -1627,8 +1661,23 @@ class VastdbApi:
             builder.PrependUOffsetTRelative(f)
 
         files = builder.EndVector()
+
+        key_names_arr = []
+        for key in key_names:
+            kname = builder.CreateString(key)
+            import_key_name.Start(builder)
+            import_key_name.AddName(builder, kname)
+            key_names_arr.append(import_key_name.End(builder))
+
+        tabular_import_data.StartKeyNamesVector(builder, len(key_names_arr))
+        for key in reversed(key_names_arr):
+            builder.PrependUOffsetTRelative(key)
+
+        key_names_vec = builder.EndVector()
+
         tabular_import_data.Start(builder)
         tabular_import_data.AddS3Files(builder, files)
+        tabular_import_data.AddKeyNames(builder, key_names_vec)
         params = tabular_import_data.End(builder)
         builder.Finish(params)
         import_req = builder.Output()
@@ -1942,6 +1991,49 @@ class QueryDataInternalError(Exception):
     pass
 
 
+def read_first_batch(fileobj):
+    readers = {}  # {stream_id: pa.ipc.RecordBatchStreamReader}
+    while True:
+        stream_id_bytes = fileobj.read(4)
+        if not stream_id_bytes:
+            if readers:
+                raise EOFError(f'no readers ({readers}) should be open at EOF')
+            break
+
+        stream_id, = struct.unpack('<L', stream_id_bytes)
+        if stream_id == TABULAR_KEEP_ALIVE_STREAM_ID:
+            continue
+
+        if stream_id == TABULAR_QUERY_DATA_COMPLETED_STREAM_ID:
+            # read the terminating end chunk from socket
+            res = fileobj.read()
+            _logger.debug("stream_id=%d res=%s (finish)", stream_id, res)
+            return None
+
+        if stream_id == TABULAR_QUERY_DATA_FAILED_STREAM_ID:
+            # read the terminating end chunk from socket
+            res = fileobj.read()
+            _logger.debug("stream_id=%d res=%s (failed)", stream_id, res)
+            raise QueryDataInternalError()  # connection closed by server due to an internal error
+
+        next_row_id_bytes = fileobj.read(8)
+        next_row_id, = struct.unpack('<Q', next_row_id_bytes)
+        _logger.debug("stream_id=%d next_row_id=%d", stream_id, next_row_id)
+
+        if stream_id not in readers:
+            # we implicitly read 1st message (Arrow schema) when constructing RecordBatchStreamReader
+            reader = pa.ipc.RecordBatchStreamReader(fileobj)
+            _logger.debug("stream_id=%d schema=%s", stream_id, reader.schema)
+            readers[stream_id] = reader
+            continue
+
+        reader = readers[stream_id]
+        try:
+            return reader.read_next_batch()  # read single-column chunk data
+        except StopIteration:  # we got an end-of-stream IPC message for a given stream ID
+            return None
+
+
 def _iter_query_data_response_columns(fileobj, stream_ids=None):
     readers = {}  # {stream_id: pa.ipc.RecordBatchStreamReader}
     while True:

vastdb/bench/test_perf.py

@@ -10,11 +10,11 @@ log = logging.getLogger(__name__)
 
 
 @pytest.mark.benchmark
-def test_bench(session, bucket_name, parquets_path, crater_path):
+def test_bench(session, test_bucket_name, parquets_path, crater_path):
     files = [str(parquets_path / f) for f in (parquets_path.glob('**/*.pq'))]
 
     with session.transaction() as tx:
-        b = tx.bucket(bucket_name)
+        b = tx.bucket(test_bucket_name)
         s = b.create_schema('s1')
         t = util.create_table_from_files(s, 't1', files, config=ImportConfig(import_concurrency=8))
         config = QueryConfig(num_splits=8, num_sub_splits=4)

vastdb/config.py

@@ -63,3 +63,6 @@ class ImportConfig:
     """Import execution configiration."""
 
     import_concurrency: int = 2
+
+    # import key column names
+    key_names: Optional[List[str]] = None

vastdb/errors.py

@@ -17,6 +17,7 @@ class HttpStatus(Enum):
     INTERNAL_SERVER_ERROR = 500
     NOT_IMPLEMENTED = 501
     SERVICE_UNAVAILABLE = 503
+    INSUFFICIENT_CAPACITY = 507
 
 
 log = logging.getLogger(__name__)
@@ -79,6 +80,10 @@ class UnexpectedError(HttpError):
     pass
 
 
+class InsufficientCapacity(HttpError):
+    pass
+
+
 @dataclass
 class ImportFilesError(Exception):
     message: str
@@ -197,6 +202,7 @@ ERROR_TYPES_MAP = {
     HttpStatus.INTERNAL_SERVER_ERROR: InternalServerError,
     HttpStatus.NOT_IMPLEMENTED: NotImplemented,
     HttpStatus.SERVICE_UNAVAILABLE: handle_unavailable,
+    HttpStatus.INSUFFICIENT_CAPACITY: InsufficientCapacity,
 }
 
 

vastdb/features.py

@@ -13,6 +13,7 @@ class Features:
     def __init__(self, vast_version):
         """Save the server version."""
         self.vast_version = vast_version
+        log.info("VAST version: %s", self.vast_version)
 
         self.check_imports_table = self._check(
             "Imported objects' table feature requires 5.2+ VAST release",
@@ -30,6 +31,14 @@ class Features:
             "External row IDs allocation requires 5.1+ VAST release",
             vast_version >= (5, 1))
 
+        self.check_elysium = self._check(
+            "Elysium requires 5.3.5+ VAST release",
+            vast_version >= (5, 3, 5))
+
+        self.check_zip_import = self._check(
+            "Zip import requires 5.3.1+ VAST release",
+            vast_version >= (5, 3, 1))
+
     def _check(self, msg, supported):
         log.debug("%s (current version is %s): supported=%s", msg, self.vast_version, supported)
         if not supported:

vastdb/schema.py

@@ -76,7 +76,8 @@ class Schema:
                 break
         return result
 
-    def create_table(self, table_name: str, columns: pa.Schema, fail_if_exists=True, use_external_row_ids_allocation=False) -> "Table":
+    def create_table(self, table_name: str, columns: pa.Schema, fail_if_exists=True,
+                     use_external_row_ids_allocation=False, sorting_key=[]) -> "Table":
         """Create a new table under this schema.
 
         A virtual `vastdb_rowid` column (of `int64` type) can be created to access and filter by internal VAST row IDs.
@@ -91,7 +92,8 @@ class Schema:
             self.tx._rpc.features.check_external_row_ids_allocation()
 
         self.tx._rpc.api.create_table(self.bucket.name, self.name, table_name, columns, txid=self.tx.txid,
-                                      use_external_row_ids_allocation=use_external_row_ids_allocation)
+                                      use_external_row_ids_allocation=use_external_row_ids_allocation,
+                                      sorting_key=sorting_key)
         log.info("Created table: %s", table_name)
         return self.table(table_name)  # type: ignore[return-value]
 
@@ -149,4 +151,4 @@ class Schema:
 
 
 def _parse_table_info(table_info, schema: "schema.Schema"):
-    return table.Table(name=table_info.name, schema=schema, handle=int(table_info.handle), _imports_table=False)
+    return table.Table(name=table_info.name, schema=schema, handle=int(table_info.handle), _imports_table=False, sorted_table=table_info.sorting_key_enabled)