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

vastdb/_internal.py

@@ -38,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
@@ -787,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):
@@ -798,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
@@ -962,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)}
 
@@ -1122,11 +1130,12 @@ class VastdbApi:
 
     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,
@@ -1143,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
@@ -1160,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([])
 
@@ -1222,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
@@ -1246,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 {}
 
@@ -1415,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
@@ -1447,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
@@ -1458,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.
@@ -1596,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
@@ -1638,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()
@@ -1953,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)

vastdb/table.py

@@ -21,6 +21,7 @@ log = logging.getLogger(__name__)
 
 INTERNAL_ROW_ID = "$row_id"
 INTERNAL_ROW_ID_FIELD = pa.field(INTERNAL_ROW_ID, pa.uint64())
+INTERNAL_ROW_ID_SORTED_FIELD = pa.field(INTERNAL_ROW_ID, pa.decimal128(38, 0))  # Sorted tables have longer row ids
 
 MAX_ROWS_PER_BATCH = 512 * 1024
 # for insert we need a smaller limit due to response amplification
@@ -115,6 +116,7 @@ class Table:
     arrow_schema: pa.Schema = field(init=False, compare=False, repr=False)
     _ibis_table: ibis.Schema = field(init=False, compare=False, repr=False)
     _imports_table: bool
+    sorted_table: bool
 
     def __post_init__(self):
         """Also, load columns' metadata."""
@@ -157,6 +159,29 @@ class Table:
         self.arrow_schema = pa.schema(fields)
         return self.arrow_schema
 
+    def sorted_columns(self) -> list:
+        """Return sorted columns' metadata."""
+        fields = []
+        try:
+            self.tx._rpc.features.check_elysium()
+            next_key = 0
+            while True:
+                cur_columns, next_key, is_truncated, _count = self.tx._rpc.api.list_sorted_columns(
+                    bucket=self.bucket.name, schema=self.schema.name, table=self.name, next_key=next_key, txid=self.tx.txid, list_imports_table=self._imports_table)
+                fields.extend(cur_columns)
+                if not is_truncated:
+                    break
+        except errors.BadRequest:
+            pass
+        except errors.InternalServerError as ise:
+            log.warning("Failed to get the sorted columns Elysium might not be supported: %s", ise)
+            pass
+        except errors.NotSupportedVersion:
+            log.warning("Failed to get the sorted columns, Elysium not supported")
+            pass
+
+        return fields
+
     def projection(self, name: str) -> "Projection":
         """Get a specific semi-sorted projection of this table."""
         if self._imports_table:
@@ -228,6 +253,10 @@ class Table:
         endpoints = [self.tx._rpc.api.url for _ in range(config.import_concurrency)]  # TODO: use valid endpoints...
         files_queue = queue.Queue()
 
+        key_names = config.key_names or []
+        if key_names:
+            self.tx._rpc.features.check_zip_import()
+
         for source_file in source_files.items():
             files_queue.put(source_file)
 
@@ -249,8 +278,10 @@ class Table:
                             pass
                         if files_batch:
                             log.debug("Starting import batch of %s files", len(files_batch))
+                            log.info(f"starting import of {files_batch}")
                             session.import_data(
-                                self.bucket.name, self.schema.name, self.name, files_batch, txid=self.tx.txid)
+                                self.bucket.name, self.schema.name, self.name, files_batch, txid=self.tx.txid,
+                                key_names=key_names)
             except (Exception, KeyboardInterrupt) as e:
                 stop_event.set()
                 log.error("Got exception inside import_worker. exception: %s", e)
@@ -277,6 +308,21 @@ class Table:
             imports_table_stats=self._imports_table)
         return TableStats(**stats_tuple._asdict())
 
+    def _get_row_estimate(self, columns: List[str], predicate: ibis.expr.types.BooleanColumn, arrow_schema: pa.Schema):
+        query_data_request = _internal.build_query_data_request(
+            schema=arrow_schema,
+            predicate=predicate,
+            field_names=columns)
+        response = self.tx._rpc.api.query_data(
+            bucket=self.bucket.name,
+            schema=self.schema.name,
+            table=self.name,
+            params=query_data_request.serialized,
+            split=(0xffffffff - 3, 1, 1),
+            txid=self.tx.txid)
+        batch = _internal.read_first_batch(response.raw)
+        return batch.num_rows * 2**16 if batch is not None else 0
+
     def select(self, columns: Optional[List[str]] = None,
                predicate: Union[ibis.expr.types.BooleanColumn, ibis.common.deferred.Deferred] = None,
                config: Optional[QueryConfig] = None,
@@ -293,30 +339,22 @@ class Table:
         if config is None:
             config = QueryConfig()
 
+        stats = None
         # Retrieve snapshots only if needed
-        if config.data_endpoints is None or config.num_splits is None:
+        if config.data_endpoints is None:
             stats = self.get_stats()
             log.debug("stats: %s", stats)
-
-        if config.data_endpoints is None:
             endpoints = stats.endpoints
         else:
             endpoints = tuple(config.data_endpoints)
         log.debug("endpoints: %s", endpoints)
 
-        if config.num_splits is None:
-            config.num_splits = max(1, stats.num_rows // config.rows_per_split)
-        log.debug("config: %s", config)
-
-        if config.semi_sorted_projection_name:
-            self.tx._rpc.features.check_enforce_semisorted_projection()
-
         if columns is None:
             columns = [f.name for f in self.arrow_schema]
 
         query_schema = self.arrow_schema
         if internal_row_id:
-            queried_fields = [INTERNAL_ROW_ID_FIELD]
+            queried_fields = [INTERNAL_ROW_ID_SORTED_FIELD if self.sorted_table else 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)
@@ -330,6 +368,22 @@ class Table:
         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)
 
+        if config.num_splits is None:
+            num_rows = 0
+            if self.sorted_table:
+                num_rows = self._get_row_estimate(columns, predicate, query_schema)
+                log.info(f'sorted estimate: {num_rows}')
+            if num_rows == 0:
+                if stats is None:
+                    stats = self.get_stats()
+                num_rows = stats.num_rows
+
+            config.num_splits = max(1, num_rows // config.rows_per_split)
+        log.debug("config: %s", config)
+
+        if config.semi_sorted_projection_name:
+            self.tx._rpc.features.check_enforce_semisorted_projection()
+
         query_data_request = _internal.build_query_data_request(
             schema=query_schema,
             predicate=predicate,
@@ -485,7 +539,7 @@ class Table:
         if columns is None:
             columns = [name for name in rows.schema.names if name != INTERNAL_ROW_ID]
 
-        update_fields = [(INTERNAL_ROW_ID, pa.uint64())]
+        update_fields = [INTERNAL_ROW_ID_SORTED_FIELD if self.sorted_table else INTERNAL_ROW_ID_FIELD]
         update_values = [_combine_chunks(rows_chunk)]
         for col in columns:
             update_fields.append(rows.field(col))
@@ -511,7 +565,7 @@ class Table:
             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())]),
+        delete_rows_rb = pa.record_batch(schema=pa.schema([INTERNAL_ROW_ID_SORTED_FIELD if self.sorted_table else INTERNAL_ROW_ID_FIELD]),
                                          data=[_combine_chunks(rows_chunk)])
 
         delete_rows_rb = util.sort_record_batch_if_needed(delete_rows_rb, INTERNAL_ROW_ID)
@@ -535,6 +589,13 @@ class Table:
         log.info("Renamed table from %s to %s ", self.name, new_name)
         self.name = new_name
 
+    def add_sorting_key(self, sorting_key: list) -> None:
+        """Ads a sorting key to a table that doesn't have any."""
+        self.tx._rpc.features.check_elysium()
+        self.tx._rpc.api.alter_table(
+            self.bucket.name, self.schema.name, self.name, txid=self.tx.txid, sorting_key=sorting_key)
+        log.info("Enabled Elysium for table %s with sorting key %s ", self.name, str(sorting_key))
+
     def add_column(self, new_column: pa.Schema) -> None:
         """Add a new column."""
         if self._imports_table:
@@ -583,7 +644,7 @@ class Table:
     def imports_table(self) -> Optional["Table"]:
         """Get the imports table of this table."""
         self.tx._rpc.features.check_imports_table()
-        return Table(name=self.name, schema=self.schema, handle=int(self.handle), _imports_table=True)
+        return Table(name=self.name, schema=self.schema, handle=int(self.handle), _imports_table=True, sorted_table=self.sorted_table)
 
     def __getitem__(self, col_name: str):
         """Allow constructing ibis-like column expressions from this table.