vastdb 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

vastdb/bench/test_perf.py

@@ -1,7 +1,6 @@
 import logging
 import time
 
-import pyarrow as pa
 import pytest
 
 from vastdb import util
@@ -20,7 +19,7 @@ def test_bench(session, clean_bucket_name, parquets_path, crater_path):
         t = util.create_table_from_files(s, 't1', files, config=ImportConfig(import_concurrency=8))
         config = QueryConfig(num_splits=8, num_sub_splits=4)
         s = time.time()
-        pa_table = pa.Table.from_batches(t.select(columns=['sid'], predicate=t['sid'] == 10033007, config=config))
+        pa_table = t.select(columns=['sid'], predicate=t['sid'] == 10033007, config=config).read_all()
         e = time.time()
         log.info("'SELECT sid from TABLE WHERE sid = 10033007' returned in %s seconds.", e - s)
         if crater_path:

vastdb/bucket.py

@@ -5,8 +5,8 @@ It is possible to list and access VAST snapshots generated over a bucket.
 """
 
 import logging
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Iterable, Optional
 
 from . import errors, schema, transaction
 
@@ -22,48 +22,23 @@ class Bucket:
 
     name: str
     tx: "transaction.Transaction"
+    _root_schema: "Schema" = field(init=False, compare=False, repr=False)
 
-    def create_schema(self, path: str, fail_if_exists=True) -> "Schema":
+    def __post_init__(self):
+        """Root schema is empty."""
+        self._root_schema = schema.Schema(name="", bucket=self)
+
+    def create_schema(self, name: str, fail_if_exists=True) -> "Schema":
         """Create a new schema (a container of tables) under this bucket."""
-        if current := self.schema(path, fail_if_missing=False):
-            if fail_if_exists:
-                raise errors.SchemaExists(self.name, path)
-            else:
-                return current
-        self.tx._rpc.api.create_schema(self.name, path, txid=self.tx.txid)
-        log.info("Created schema: %s", path)
-        return self.schema(path)  # type: ignore[return-value]
+        return self._root_schema.create_schema(name=name, fail_if_exists=fail_if_exists)
 
-    def schema(self, path: str, fail_if_missing=True) -> Optional["Schema"]:
+    def schema(self, name: str, fail_if_missing=True) -> Optional["Schema"]:
         """Get a specific schema (a container of tables) under this bucket."""
-        s = self.schemas(path)
-        log.debug("schema: %s", s)
-        if not s:
-            if fail_if_missing:
-                raise errors.MissingSchema(self.name, path)
-            else:
-                return None
-        assert len(s) == 1, f"Expected to receive only a single schema, but got: {len(s)}. ({s})"
-        log.debug("Found schema: %s", s[0].name)
-        return s[0]
+        return self._root_schema.schema(name=name, fail_if_missing=fail_if_missing)
 
-    def schemas(self, name: Optional[str] = None) -> List["Schema"]:
+    def schemas(self, batch_size=None):
         """List bucket's schemas."""
-        schemas = []
-        next_key = 0
-        exact_match = bool(name)
-        log.debug("list schemas param: schema=%s, exact_match=%s", name, exact_match)
-        while True:
-            _bucket_name, curr_schemas, next_key, is_truncated, _ = \
-                self.tx._rpc.api.list_schemas(bucket=self.name, next_key=next_key, txid=self.tx.txid,
-                                               name_prefix=name, exact_match=exact_match)
-            if not curr_schemas:
-                break
-            schemas.extend(curr_schemas)
-            if not is_truncated:
-                break
-
-        return [schema.Schema(name=name, bucket=self) for name, *_ in schemas]
+        return self._root_schema.schemas(batch_size=batch_size)
 
     def snapshot(self, name, fail_if_missing=True) -> Optional["Bucket"]:
         """Get snapshot by name (if exists)."""
@@ -80,7 +55,7 @@ class Bucket:
 
         return Bucket(name=f'{self.name}/{expected_name}', tx=self.tx)
 
-    def snapshots(self) -> List["Bucket"]:
+    def snapshots(self) -> Iterable["Bucket"]:
         """List bucket's snapshots."""
         snapshots = []
         next_key = 0

vastdb/conftest.py

@@ -30,14 +30,23 @@ def test_bucket_name(request):
     return request.config.getoption("--tabular-bucket-name")
 
 
+def iter_schemas(s):
+    """Recusively scan all schemas."""
+    children = s.schemas()
+    for c in children:
+        yield from iter_schemas(c)
+    yield s
+
+
 @pytest.fixture(scope="function")
 def clean_bucket_name(request, test_bucket_name, session):
     with session.transaction() as tx:
         b = tx.bucket(test_bucket_name)
-        for s in b.schemas():
-            for t in s.tables():
-                t.drop()
-            s.drop()
+        for top_schema in b.schemas():
+            for s in iter_schemas(top_schema):
+                for t in s.tables():
+                    t.drop()
+                s.drop()
     return test_bucket_name
 
 

vastdb/errors.py

@@ -89,7 +89,11 @@ class InvalidArgument(Exception):
     pass
 
 
-class TooWideRow(InvalidArgument):
+class TooLargeRequest(InvalidArgument):
+    pass
+
+
+class TooWideRow(TooLargeRequest):
     pass
 
 

vastdb/schema.py

@@ -6,7 +6,7 @@ It is possible to list and access VAST snapshots generated over a bucket.
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Iterable, List, Optional
 
 import pyarrow as pa
 
@@ -31,14 +31,63 @@ class Schema:
         """VAST transaction used for this schema."""
         return self.bucket.tx
 
-    def create_table(self, table_name: str, columns: pa.Schema, fail_if_exists=True) -> "Table":
+    def _subschema_full_name(self, name: str) -> str:
+        return f"{self.name}/{name}" if self.name else name
+
+    def create_schema(self, name: str, fail_if_exists=True) -> "Schema":
+        """Create a new schema (a container of tables) under this schema."""
+        if current := self.schema(name, fail_if_missing=False):
+            if fail_if_exists:
+                raise errors.SchemaExists(self.bucket.name, name)
+            else:
+                return current
+        full_name = self._subschema_full_name(name)
+        self.tx._rpc.api.create_schema(self.bucket.name, full_name, txid=self.tx.txid)
+        log.info("Created schema: %s", full_name)
+        return self.schema(name)  # type: ignore[return-value]
+
+    def schema(self, name: str, fail_if_missing=True) -> Optional["Schema"]:
+        """Get a specific schema (a container of tables) under this schema."""
+        _bucket_name, schemas, _next_key, _is_truncated, _ = \
+            self.tx._rpc.api.list_schemas(bucket=self.bucket.name, schema=self.name, next_key=0, txid=self.tx.txid,
+                                          name_prefix=name, exact_match=True, max_keys=1)
+        names = [name for name, *_ in schemas]
+        log.debug("Found schemas: %s", names)
+        if not names:
+            if fail_if_missing:
+                raise errors.MissingSchema(self.bucket.name, self._subschema_full_name(name))
+            else:
+                return None
+
+        assert len(names) == 1, f"Expected to receive only a single schema, but got {len(schemas)}: ({schemas})"
+        return schema.Schema(name=self._subschema_full_name(names[0]), bucket=self.bucket)
+
+    def schemas(self, batch_size=None) -> Iterable["Schema"]:
+        """List child schemas."""
+        next_key = 0
+        if not batch_size:
+            batch_size = 1000
+        result: List["Schema"] = []
+        while True:
+            _bucket_name, curr_schemas, next_key, is_truncated, _ = \
+                self.tx._rpc.api.list_schemas(bucket=self.bucket.name, schema=self.name, next_key=next_key, max_keys=batch_size, txid=self.tx.txid)
+            result.extend(schema.Schema(name=self._subschema_full_name(name), bucket=self.bucket) for name, *_ in curr_schemas)
+            if not is_truncated:
+                break
+        return result
+
+    def create_table(self, table_name: str, columns: pa.Schema, fail_if_exists=True, use_external_row_ids_allocation=False) -> "Table":
         """Create a new table under this schema."""
         if current := self.table(table_name, fail_if_missing=False):
             if fail_if_exists:
                 raise errors.TableExists(self.bucket.name, self.name, table_name)
             else:
                 return current
-        self.tx._rpc.api.create_table(self.bucket.name, self.name, table_name, columns, txid=self.tx.txid)
+        if use_external_row_ids_allocation:
+            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)
         log.info("Created table: %s", table_name)
         return self.table(table_name)  # type: ignore[return-value]
 

vastdb/session.py

@@ -7,11 +7,16 @@ For more details see:
 - [Tabular identity policy with the proper permissions](https://support.vastdata.com/s/article/UUID-14322b60-d6a2-89ac-3df0-3dfbb6974182)
 """
 
+import logging
 import os
+from typing import Optional
 
 import boto3
 
-from . import errors, internal_commands, transaction
+from . import _internal, errors, transaction
+from ._internal import BackoffConfig
+
+log = logging.getLogger()
 
 
 class Features:
@@ -21,21 +26,41 @@ class Features:
         """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)
+        self.check_imports_table = self._check(
+            "Imported objects' table feature requires 5.2+ VAST release",
+            vast_version >= (5, 2))
+
+        self.check_return_row_ids = self._check(
+            "Returning row IDs requires 5.1+ VAST release",
+            vast_version >= (5, 1))
+
+        self.check_enforce_semisorted_projection = self._check(
+            "Semi-sorted projection enforcement requires 5.1+ VAST release",
+            vast_version >= (5, 1))
+
+        self.check_external_row_ids_allocation = self._check(
+            "External row IDs allocation requires 5.1+ VAST release",
+            vast_version >= (5, 1))
+
+    def _check(self, msg, supported):
+        log.debug("%s (current version is %s): supported=%s", msg, self.vast_version, supported)
+        if not supported:
+            def fail():
+                raise errors.NotSupportedVersion(msg, self.vast_version)
+            return fail
 
-    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)
+        def noop():
+            pass
+        return noop
 
 
 class Session:
     """VAST database session."""
 
-    def __init__(self, access=None, secret=None, endpoint=None, ssl_verify=True):
+    def __init__(self, access=None, secret=None, endpoint=None,
+                 *,
+                 ssl_verify=True,
+                 backoff_config: Optional[BackoffConfig] = None):
         """Connect to a VAST Database endpoint, using specified credentials."""
         if access is None:
             access = os.environ['AWS_ACCESS_KEY_ID']
@@ -44,9 +69,13 @@ class Session:
         if endpoint is None:
             endpoint = os.environ['AWS_S3_ENDPOINT_URL']
 
-        self.api = internal_commands.VastdbApi(endpoint, access, secret, ssl_verify=ssl_verify)
-        version_tuple = tuple(int(part) for part in self.api.vast_version.split('.'))
-        self.features = Features(version_tuple)
+        self.api = _internal.VastdbApi(
+            endpoint=endpoint,
+            access_key=access,
+            secret_key=secret,
+            ssl_verify=ssl_verify,
+            backoff_config=backoff_config)
+        self.features = Features(self.api.vast_version)
         self.s3 = boto3.client('s3',
             aws_access_key_id=access,
             aws_secret_access_key=secret,

vastdb/table.py

@@ -7,14 +7,14 @@ import queue
 from dataclasses import dataclass, field
 from math import ceil
 from threading import Event
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import backoff
 import ibis
 import pyarrow as pa
 import requests
 
-from . import errors, internal_commands, schema, util
+from . import _internal, errors, schema, util
 
 log = logging.getLogger(__name__)
 
@@ -54,7 +54,8 @@ class QueryConfig:
     num_sub_splits: int = 4
 
     # used to split the table into disjoint subsets of rows, to be processed concurrently using multiple RPCs
-    num_splits: int = 1
+    # will be estimated from the table's row count, if not explicitly set
+    num_splits: Optional[int] = None
 
     # each endpoint will be handled by a separate worker thread
     # a single endpoint can be specified more than once to benefit from multithreaded execution
@@ -64,19 +65,27 @@ class QueryConfig:
     limit_rows_per_sub_split: int = 128 * 1024
 
     # each fiber will read the following number of rowgroups coninuously before skipping
-    # in order to use semi-sorted projections this value must be 8
+    # in order to use semi-sorted projections this value must be 8 (this is the hard coded size of a row groups per row block).
     num_row_groups_per_sub_split: int = 8
 
     # can be disabled for benchmarking purposes
     use_semi_sorted_projections: bool = True
 
+    # enforce using a specific semi-sorted projection (if enabled above)
+    semi_sorted_projection_name: Optional[str] = None
+
     # used to estimate the number of splits, given the table rows' count
     rows_per_split: int = 4000000
 
     # used for worker threads' naming
     query_id: str = ""
 
-    # allows retrying QueryData when the server is overloaded
+    # non-negative integer, used for server-side prioritization of queued requests:
+    # - requests with lower values will be served before requests with higher values.
+    # - if unset, the request will be added to the queue's end.
+    queue_priority: Optional[int] = None
+
+    # DEPRECATED: will be removed in a future release
     backoff_func: Any = field(default=backoff.on_exception(backoff.expo, RETRIABLE_ERRORS, max_tries=10))
 
 
@@ -98,14 +107,13 @@ class SelectSplitState:
         self.query_data_request = query_data_request
         self.table = table
 
-    def batches(self, api: internal_commands.VastdbApi):
+    def batches(self, api: _internal.VastdbApi):
         """Execute QueryData request, and yield parsed RecordBatch objects.
 
         Can be called repeatedly, to allow pagination.
         """
         while not self.done:
-            query_with_backoff = self.config.backoff_func(api.query_data)
-            response = query_with_backoff(
+            response = api.query_data(
                             bucket=self.table.bucket.name,
                             schema=self.table.schema.name,
                             table=self.table.name,
@@ -116,9 +124,11 @@ class SelectSplitState:
                             txid=self.table.tx.txid,
                             limit_rows=self.config.limit_rows_per_sub_split,
                             sub_split_start_row_ids=self.subsplits_state.items(),
+                            schedule_id=self.config.queue_priority,
                             enable_sorted_projections=self.config.use_semi_sorted_projections,
-                            query_imports_table=self.table._imports_table)
-            pages_iter = internal_commands.parse_query_data_response(
+                            query_imports_table=self.table._imports_table,
+                            projection=self.config.semi_sorted_projection_name)
+            pages_iter = _internal.parse_query_data_response(
                 conn=response.raw,
                 schema=self.query_data_request.response_schema,
                 start_row_ids=self.subsplits_state,
@@ -132,7 +142,7 @@ class SelectSplitState:
     @property
     def done(self):
         """Returns true iff the pagination over."""
-        return all(row_id == internal_commands.TABULAR_INVALID_ROW_ID for row_id in self.subsplits_state.values())
+        return all(row_id == _internal.TABULAR_INVALID_ROW_ID for row_id in self.subsplits_state.values())
 
 
 @dataclass
@@ -182,14 +192,14 @@ class Table:
         """Get a specific semi-sorted projection of this table."""
         if self._imports_table:
             raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
-        projs = self.projections(projection_name=name)
+        projs = tuple(self.projections(projection_name=name))
         if not projs:
             raise errors.MissingProjection(self.bucket.name, self.schema.name, self.name, name)
         assert len(projs) == 1, f"Expected to receive only a single projection, but got: {len(projs)}. projections: {projs}"
         log.debug("Found projection: %s", projs[0])
         return projs[0]
 
-    def projections(self, projection_name=None) -> List["Projection"]:
+    def projections(self, projection_name=None) -> Iterable["Projection"]:
         """List all semi-sorted projections of this table."""
         if self._imports_table:
             raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
@@ -209,7 +219,7 @@ class Table:
                 break
         return [_parse_projection_info(projection, self) for projection in projections]
 
-    def import_files(self, files_to_import: List[str], config: Optional[ImportConfig] = None) -> None:
+    def import_files(self, files_to_import: Iterable[str], config: Optional[ImportConfig] = None) -> None:
         """Import a list of Parquet files into this table.
 
         The files must be on VAST S3 server and be accessible using current credentials.
@@ -278,7 +288,7 @@ class Table:
                 max_workers=config.import_concurrency, thread_name_prefix='import_thread') as pool:
             try:
                 for endpoint in endpoints:
-                    session = internal_commands.VastdbApi(endpoint, self.tx._rpc.api.access_key, self.tx._rpc.api.secret_key)
+                    session = _internal.VastdbApi(endpoint, self.tx._rpc.api.access_key, self.tx._rpc.api.secret_key)
                     futures.append(pool.submit(import_worker, files_queue, session))
 
                 log.debug("Waiting for import workers to finish")
@@ -313,11 +323,16 @@ class Table:
 
         # Take a snapshot of enpoints
         stats = self.get_stats()
+        log.debug("stats: %s", stats)
         endpoints = stats.endpoints if config.data_endpoints is None else 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 stats.num_rows > config.rows_per_split and config.num_splits is None:
-            config.num_splits = stats.num_rows // config.rows_per_split
-        log.debug(f"num_rows={stats.num_rows} rows_per_splits={config.rows_per_split} num_splits={config.num_splits} ")
+        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]
@@ -332,16 +347,18 @@ class Table:
         if predicate is True:
             predicate = None
         if predicate is False:
-            response_schema = internal_commands.get_response_schema(schema=query_schema, field_names=columns)
+            response_schema = _internal.get_response_schema(schema=query_schema, field_names=columns)
             return pa.RecordBatchReader.from_batches(response_schema, [])
 
         if isinstance(predicate, ibis.common.deferred.Deferred):
             predicate = predicate.resolve(self._ibis_table)  # may raise if the predicate is invalid (e.g. wrong types / missing column)
 
-        query_data_request = internal_commands.build_query_data_request(
+        query_data_request = _internal.build_query_data_request(
             schema=query_schema,
             predicate=predicate,
             field_names=columns)
+        if len(query_data_request.serialized) > util.MAX_QUERY_DATA_REQUEST_SIZE:
+            raise errors.TooLargeRequest(f"{len(query_data_request.serialized)} bytes")
 
         splits_queue: queue.Queue[int] = queue.Queue()
 
@@ -364,7 +381,7 @@ class Table:
 
         def single_endpoint_worker(endpoint: str):
             try:
-                host_api = internal_commands.VastdbApi(endpoint=endpoint, access_key=self.tx._rpc.api.access_key, secret_key=self.tx._rpc.api.secret_key)
+                host_api = _internal.VastdbApi(endpoint=endpoint, access_key=self.tx._rpc.api.access_key, secret_key=self.tx._rpc.api.secret_key)
                 while True:
                     check_stop()
                     try:
@@ -461,7 +478,7 @@ class Table:
             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)
+                (batch,) = pa.RecordBatchStreamReader(res.content)
                 row_ids.append(batch[INTERNAL_ROW_ID])
             try:
                 self.tx._rpc.features.check_return_row_ids()
@@ -497,6 +514,8 @@ class Table:
         else:
             update_rows_rb = rows
 
+        update_rows_rb = util.sort_record_batch_if_needed(update_rows_rb, INTERNAL_ROW_ID)
+
         serialized_slices = util.iter_serialized_slices(update_rows_rb, MAX_ROWS_PER_BATCH)
         for slice in serialized_slices:
             self.tx._rpc.api.update_rows(self.bucket.name, self.schema.name, self.name, record_batch=slice,
@@ -516,6 +535,8 @@ class Table:
         delete_rows_rb = pa.record_batch(schema=pa.schema([(INTERNAL_ROW_ID, pa.uint64())]),
                                          data=[_combine_chunks(rows_chunk)])
 
+        delete_rows_rb = util.sort_record_batch_if_needed(delete_rows_rb, INTERNAL_ROW_ID)
+
         serialized_slices = util.iter_serialized_slices(delete_rows_rb, MAX_ROWS_PER_BATCH)
         for slice in serialized_slices:
             self.tx._rpc.api.delete_rows(self.bucket.name, self.schema.name, self.name, record_batch=slice,
@@ -581,7 +602,7 @@ class Table:
         return self.imports_table()  # type: ignore[return-value]
 
     def imports_table(self) -> Optional["Table"]:
-        """Get the imports table under of this 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), stats=self.stats, _imports_table=True)
 

vastdb/tests/test_duckdb.py

@@ -56,6 +56,6 @@ def test_closed_tx(session, clean_bucket_name):
         res = conn.execute('SELECT a FROM batches')
         log.debug("closing tx=%s after first batch=%s", t.tx, first)
 
-    # transaction is closed, collecting the result should fail
-    with pytest.raises(duckdb.InvalidInputException, match="Detail: Python exception: MissingTransaction"):
+    # transaction is closed, collecting the result should fail internally in DuckDB
+    with pytest.raises(duckdb.InvalidInputException):
         res.arrow()

vastdb/tests/test_imports.py

@@ -38,13 +38,13 @@ def test_parallel_imports(session, clean_bucket_name, s3):
             t.create_imports_table()
         log.info("Starting import of %d files", num_files)
         t.import_files(files)
-        arrow_table = pa.Table.from_batches(t.select(columns=['num']))
+        arrow_table = t.select(columns=['num']).read_all()
         assert arrow_table.num_rows == num_rows * num_files
-        arrow_table = pa.Table.from_batches(t.select(columns=['num'], predicate=t['num'] == 100))
+        arrow_table = t.select(columns=['num'], predicate=t['num'] == 100).read_all()
         assert arrow_table.num_rows == num_files
         import_table = t.imports_table()
         # checking all imports are on the imports table:
-        objects_name = pa.Table.from_batches(import_table.select(columns=["ObjectName"]))
+        objects_name = import_table.select(columns=["ObjectName"]).read_all()
         objects_name = objects_name.to_pydict()
         object_names = set(objects_name['ObjectName'])
         prefix = 'prq'

vastdb/tests/test_nested.py

@@ -22,13 +22,13 @@ def test_nested_select(session, clean_bucket_name):
     ])
 
     with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
-        actual = pa.Table.from_batches(t.select())
+        actual = t.select().read_all()
         assert actual == expected
 
         names = [f.name for f in columns]
         for n in range(len(names) + 1):
             for cols in itertools.permutations(names, n):
-                actual = pa.Table.from_batches(t.select(columns=cols))
+                actual = t.select(columns=cols).read_all()
                 assert actual == expected.select(cols)
 
 
@@ -53,7 +53,7 @@ def test_nested_filter(session, clean_bucket_name):
     ])
 
     with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
-        actual = pa.Table.from_batches(t.select())
+        actual = t.select().read_all()
         assert actual == expected
 
         names = list('xyzw')
@@ -62,7 +62,7 @@ def test_nested_filter(session, clean_bucket_name):
                 ibis_predicate = functools.reduce(
                     operator.and_,
                     (t[col] > 2 for col in cols))
-                actual = pa.Table.from_batches(t.select(predicate=ibis_predicate), t.arrow_schema)
+                actual = t.select(predicate=ibis_predicate).read_all()
 
                 arrow_predicate = functools.reduce(
                     operator.and_,

vastdb/tests/test_projections.py

@@ -1,7 +1,10 @@
 import logging
+import time
 
 import pyarrow as pa
 
+from vastdb.table import QueryConfig
+
 log = logging.getLogger(__name__)
 
 
@@ -41,3 +44,78 @@ def test_basic_projections(session, clean_bucket_name):
         projs = t.projections()
         assert len(projs) == 1
         assert projs[0].name == 'p_new'
+
+
+def test_query_data_with_projection(session, clean_bucket_name):
+    columns = pa.schema([
+        ('a', pa.int64()),
+        ('b', pa.int64()),
+        ('s', pa.utf8()),
+    ])
+    # need to be large enough in order to consider as projection
+
+    GROUP_SIZE = 128 * 1024
+    expected = pa.table(schema=columns, data=[
+        [i for i in range(GROUP_SIZE)],
+        [i for i in reversed(range(GROUP_SIZE))],
+        [f's{i}' for i in range(GROUP_SIZE)],
+    ])
+
+    expected_projection_p1 = pa.table(schema=columns, data=[
+        [i for i in reversed(range(GROUP_SIZE - 5, GROUP_SIZE))],
+        [i for i in range(5)],
+        [f's{i}' for i in reversed(range(GROUP_SIZE - 5, GROUP_SIZE))],
+    ])
+
+    expected_projection_p2 = pa.table(schema=columns, data=[
+        [i for i in range(GROUP_SIZE - 5, GROUP_SIZE)],
+        [i for i in reversed(range(5))],
+        [f's{i}' for i in range(GROUP_SIZE - 5, GROUP_SIZE)],
+    ])
+
+    schema_name = "schema"
+    table_name = "table"
+    with session.transaction() as tx:
+        s = tx.bucket(clean_bucket_name).create_schema(schema_name)
+        t = s.create_table(table_name, expected.schema)
+
+        sorted_columns = ['b']
+        unsorted_columns = ['a', 's']
+        t.create_projection('p1', sorted_columns, unsorted_columns)
+
+        sorted_columns = ['a']
+        unsorted_columns = ['b', 's']
+        t.create_projection('p2', sorted_columns, unsorted_columns)
+
+    with session.transaction() as tx:
+        s = tx.bucket(clean_bucket_name).schema(schema_name)
+        t = s.table(table_name)
+        t.insert(expected)
+        actual = pa.Table.from_batches(t.select(columns=['a', 'b', 's']))
+        assert actual == expected
+
+    time.sleep(3)
+
+    with session.transaction() as tx:
+        config = QueryConfig()
+        # in nfs mock server num row groups per row block is 1 so need to change this in the config
+        config.num_row_groups_per_sub_split = 1
+
+        s = tx.bucket(clean_bucket_name).schema(schema_name)
+        t = s.table(table_name)
+        projection_actual = pa.Table.from_batches(t.select(columns=['a', 'b', 's'], predicate=(t['b'] < 5), config=config))
+        # no projection supply - need to be with p1 projeciton
+        assert expected_projection_p1 == projection_actual
+
+        config.semi_sorted_projection_name = 'p1'
+        projection_actual = pa.Table.from_batches(t.select(columns=['a', 'b', 's'], predicate=(t['b'] < 5), config=config))
+        # expecting results of projection p1 since we asked it specificaly
+        assert expected_projection_p1 == projection_actual
+
+        config.semi_sorted_projection_name = 'p2'
+        projection_actual = pa.Table.from_batches(t.select(columns=['a', 'b', 's'], predicate=(t['b'] < 5), config=config))
+        # expecting results of projection p2 since we asked it specificaly
+        assert expected_projection_p2 == projection_actual
+
+        t.drop()
+        s.drop()