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

vastdb/table.py

@@ -1,3 +1,5 @@
+"""VAST Database table."""
+
 import concurrent.futures
 import logging
 import os
@@ -5,25 +7,32 @@ 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
 
-from . import errors, internal_commands, schema
+from . import errors, internal_commands, schema, util
 
 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
 class TableStats:
+    """Table-related information."""
+
     num_rows: int
     size_in_bytes: int
     is_external_rowid_alloc: bool = False
@@ -32,6 +41,8 @@ class TableStats:
 
 @dataclass
 class QueryConfig:
+    """Query execution configiration."""
+
     num_sub_splits: int = 4
     num_splits: int = 1
     data_endpoints: Optional[List[str]] = None
@@ -40,15 +51,22 @@ 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
 class ImportConfig:
+    """Import execution configiration."""
+
     import_concurrency: int = 2
 
 
-class SelectSplitState():
+class SelectSplitState:
+    """State of a specific query split execution."""
+
     def __init__(self, query_data_request, table: "Table", split_id: int, config: QueryConfig) -> None:
+        """Initialize query split state."""
         self.split_id = split_id
         self.subsplits_state = {i: 0 for i in range(config.num_sub_splits)}
         self.config = config
@@ -56,8 +74,13 @@ class SelectSplitState():
         self.table = table
 
     def batches(self, api: internal_commands.VastdbApi):
+        """Execute QueryData request, and yield parsed RecordBatch objects.
+
+        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,
@@ -68,7 +91,8 @@ 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(),
-                            enable_sorted_projections=self.config.use_semi_sorted_projections)
+                            enable_sorted_projections=self.config.use_semi_sorted_projections,
+                            query_imports_table=self.table._imports_table)
             pages_iter = internal_commands.parse_query_data_response(
                 conn=response.raw,
                 schema=self.query_data_request.response_schema,
@@ -82,19 +106,24 @@ 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())
 
 
 @dataclass
 class Table:
+    """VAST Table."""
+
     name: str
     schema: "schema.Schema"
     handle: int
     stats: TableStats
-    arrow_schema: pa.Schema = field(init=False, compare=False)
-    _ibis_table: ibis.Schema = field(init=False, compare=False)
+    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
 
     def __post_init__(self):
+        """Also, load columns' metadata."""
         self.arrow_schema = self.columns()
 
         table_path = f'{self.schema.bucket.name}/{self.schema.name}/{self.name}'
@@ -102,21 +131,21 @@ class Table:
 
     @property
     def tx(self):
+        """Return transaction."""
         return self.schema.tx
 
     @property
     def bucket(self):
+        """Return bucket."""
         return self.schema.bucket
 
-    def __repr__(self):
-        return f"{type(self).__name__}(name={self.name})"
-
     def columns(self) -> pa.Schema:
+        """Return columns' metadata."""
         fields = []
         next_key = 0
         while True:
             cur_columns, next_key, is_truncated, _count = self.tx._rpc.api.list_columns(
-                bucket=self.bucket.name, schema=self.schema.name, table=self.name, next_key=next_key, txid=self.tx.txid)
+                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
@@ -125,6 +154,9 @@ class Table:
         return self.arrow_schema
 
     def projection(self, name: str) -> "Projection":
+        """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)
         if not projs:
             raise errors.MissingProjection(self.bucket.name, self.schema.name, self.name, name)
@@ -133,6 +165,9 @@ class Table:
         return projs[0]
 
     def projections(self, projection_name=None) -> List["Projection"]:
+        """List all semi-sorted projections of this table."""
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         projections = []
         next_key = 0
         name_prefix = projection_name if projection_name else ""
@@ -150,6 +185,12 @@ class Table:
         return [_parse_projection_info(projection, self) for projection in projections]
 
     def import_files(self, files_to_import: List[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.
+        """
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         source_files = {}
         for f in files_to_import:
             bucket_name, object_path = _parse_bucket_and_object_names(f)
@@ -158,6 +199,13 @@ class Table:
         self._execute_import(source_files, config=config)
 
     def import_partitioned_files(self, files_and_partitions: Dict[str, pa.RecordBatch], 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.
+        Each file must have its own partition values defined as an Arrow RecordBatch.
+        """
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         source_files = {}
         for f, record_batch in files_and_partitions.items():
             bucket_name, object_path = _parse_bucket_and_object_names(f)
@@ -216,8 +264,10 @@ class Table:
                 # ThreadPoolExecutor will be joined at the end of the context
 
     def get_stats(self) -> TableStats:
+        """Get the statistics of this table."""
         stats_tuple = self.tx._rpc.api.get_table_stats(
-            bucket=self.bucket.name, schema=self.schema.name, name=self.name, txid=self.tx.txid)
+            bucket=self.bucket.name, schema=self.schema.name, name=self.name, txid=self.tx.txid,
+            imports_table_stats=self._imports_table)
         return TableStats(**stats_tuple._asdict())
 
     def select(self, columns: Optional[List[str]] = None,
@@ -225,6 +275,14 @@ class Table:
                config: Optional[QueryConfig] = None,
                *,
                internal_row_id: bool = False) -> pa.RecordBatchReader:
+        """Execute a query over this table.
+
+        To read a subset of the columns, specify their names via `columns` argument. Otherwise, all columns will be read.
+
+        In order to apply a filter, a predicate can be specified. See https://github.com/vast-data/vastdb_sdk/blob/main/README.md#filters-and-projections for more details.
+
+        Query-execution configuration options can be specified via the optional `config` argument.
+        """
         if config is None:
             config = QueryConfig()
 
@@ -241,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,
@@ -335,82 +399,176 @@ class Table:
 
         return pa.RecordBatchReader.from_batches(query_data_request.response_schema, batches_iterator())
 
-    def _combine_chunks(self, col):
-        if hasattr(col, "combine_chunks"):
-            return col.combine_chunks()
-        else:
-            return col
-
-    def insert(self, rows: pa.RecordBatch) -> pa.RecordBatch:
-        serialized_slices = self.tx._rpc.api._record_batch_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)
+    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)
+        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) - 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 = [self._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(self._combine_chunks(rows[col]))
+                update_values.append(_combine_chunks(rows[col]))
 
             update_rows_rb = pa.record_batch(schema=pa.schema(update_fields), data=update_values)
         else:
             update_rows_rb = rows
 
-        serialized_slices = self.tx._rpc.api._record_batch_slices(update_rows_rb, MAX_ROWS_PER_BATCH)
+        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,
                                          txid=self.tx.txid)
 
     def delete(self, rows: Union[pa.RecordBatch, pa.Table]) -> None:
+        """Delete a subset of rows in this 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=[self._combine_chunks(rows[INTERNAL_ROW_ID])])
+                                         data=[_combine_chunks(rows_chunk)])
 
-        serialized_slices = self.tx._rpc.api._record_batch_slices(delete_rows_rb, MAX_ROWS_PER_BATCH)
+        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,
-                                         txid=self.tx.txid)
+                                         txid=self.tx.txid, delete_from_imports_table=self._imports_table)
 
     def drop(self) -> None:
-        self.tx._rpc.api.drop_table(self.bucket.name, self.schema.name, self.name, txid=self.tx.txid)
+        """Drop this table."""
+        self.tx._rpc.api.drop_table(self.bucket.name, self.schema.name, self.name, txid=self.tx.txid, remove_imports_table=self._imports_table)
         log.info("Dropped table: %s", self.name)
 
     def rename(self, new_name) -> None:
+        """Rename this table."""
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         self.tx._rpc.api.alter_table(
             self.bucket.name, self.schema.name, self.name, txid=self.tx.txid, new_name=new_name)
         log.info("Renamed table from %s to %s ", self.name, new_name)
         self.name = new_name
 
     def add_column(self, new_column: pa.Schema) -> None:
+        """Add a new column."""
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         self.tx._rpc.api.add_columns(self.bucket.name, self.schema.name, self.name, new_column, txid=self.tx.txid)
         log.info("Added column(s): %s", new_column)
         self.arrow_schema = self.columns()
 
     def drop_column(self, column_to_drop: pa.Schema) -> None:
+        """Drop an existing column."""
+        if self._imports_table:
+            raise errors.NotSupported(self.bucket.name, self.schema.name, self.name)
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         self.tx._rpc.api.drop_columns(self.bucket.name, self.schema.name, self.name, column_to_drop, txid=self.tx.txid)
         log.info("Dropped column(s): %s", column_to_drop)
         self.arrow_schema = self.columns()
 
     def rename_column(self, current_column_name: str, new_column_name: str) -> None:
+        """Rename an existing column."""
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         self.tx._rpc.api.alter_column(self.bucket.name, self.schema.name, self.name, name=current_column_name,
                                        new_name=new_column_name, txid=self.tx.txid)
         log.info("Renamed column: %s to %s", current_column_name, new_column_name)
         self.arrow_schema = self.columns()
 
     def create_projection(self, projection_name: str, sorted_columns: List[str], unsorted_columns: List[str]) -> "Projection":
+        """Create a new semi-sorted projection."""
+        if self._imports_table:
+            raise errors.NotSupportedCommand(self.bucket.name, self.schema.name, self.name)
         columns = [(sorted_column, "Sorted") for sorted_column in sorted_columns] + [(unsorted_column, "Unorted") for unsorted_column in unsorted_columns]
         self.tx._rpc.api.create_projection(self.bucket.name, self.schema.name, self.name, projection_name, columns=columns, txid=self.tx.txid)
         log.info("Created projection: %s", projection_name)
         return self.projection(projection_name)
 
+    def create_imports_table(self, fail_if_exists=True) -> "Table":
+        """Create imports table."""
+        self.tx._rpc.features.check_imports_table()
+        empty_schema = pa.schema([])
+        self.tx._rpc.api.create_table(self.bucket.name, self.schema.name, self.name, empty_schema, txid=self.tx.txid,
+                                        create_imports_table=True)
+        log.info("Created imports table for table: %s", self.name)
+        return self.imports_table()  # type: ignore[return-value]
+
+    def imports_table(self) -> Optional["Table"]:
+        """Get the imports table under 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)
+
     def __getitem__(self, col_name):
+        """Allow constructing ibis-like column expressions from this table.
+
+        It is useful for constructing expressions for predicate pushdown in `Table.select()` method.
+        """
         return self._ibis_table[col_name]
 
 
 @dataclass
 class Projection:
+    """VAST semi-sorted projection."""
+
     name: str
     table: Table
     handle: int
@@ -418,20 +576,21 @@ class Projection:
 
     @property
     def bucket(self):
+        """Return bucket."""
         return self.table.schema.bucket
 
     @property
     def schema(self):
+        """Return schema."""
         return self.table.schema
 
     @property
     def tx(self):
+        """Return transaction."""
         return self.table.schema.tx
 
-    def __repr__(self):
-        return f"{type(self).__name__}(name={self.name})"
-
     def columns(self) -> pa.Schema:
+        """Return this projections' columns as an Arrow schema."""
         columns = []
         next_key = 0
         while True:
@@ -447,12 +606,14 @@ class Projection:
         return self.arrow_schema
 
     def rename(self, new_name) -> None:
+        """Rename this projection."""
         self.tx._rpc.api.alter_projection(self.bucket.name, self.schema.name,
                                                 self.table.name, self.name, txid=self.tx.txid, new_name=new_name)
         log.info("Renamed projection from %s to %s ", self.name, new_name)
         self.name = new_name
 
     def drop(self) -> None:
+        """Drop this projection."""
         self.tx._rpc.api.drop_projection(self.bucket.name, self.schema.name, self.table.name,
                                          self.name, txid=self.tx.txid)
         log.info("Dropped projection: %s", self.name)
@@ -478,3 +639,10 @@ def _serialize_record_batch(record_batch: pa.RecordBatch) -> pa.lib.Buffer:
     with pa.ipc.new_stream(sink, record_batch.schema) as writer:
         writer.write(record_batch)
     return sink.getvalue()
+
+
+def _combine_chunks(col):
+    if hasattr(col, "combine_chunks"):
+        return col.combine_chunks()
+    else:
+        return col

vastdb/tests/test_imports.py

@@ -6,7 +6,7 @@ import pyarrow.parquet as pq
 import pytest
 
 from vastdb import util
-from vastdb.errors import ImportFilesError, InvalidArgument
+from vastdb.errors import ImportFilesError, InternalServerError, InvalidArgument
 
 log = logging.getLogger(__name__)
 
@@ -34,12 +34,24 @@ def test_parallel_imports(session, clean_bucket_name, s3):
         b = tx.bucket(clean_bucket_name)
         s = b.create_schema('s1')
         t = s.create_table('t1', pa.schema([('num', pa.int64())]))
+        with pytest.raises(InternalServerError):
+            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']))
         assert arrow_table.num_rows == num_rows * num_files
         arrow_table = pa.Table.from_batches(t.select(columns=['num'], predicate=t['num'] == 100))
         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 = objects_name.to_pydict()
+        object_names = set(objects_name['ObjectName'])
+        prefix = 'prq'
+        numbers = set(range(53))
+        assert all(name.startswith(prefix) for name in object_names)
+        numbers.issubset(int(name.replace(prefix, '')) for name in object_names)
+        assert len(object_names) == len(objects_name['ObjectName'])
 
 
 def test_create_table_from_files(session, clean_bucket_name, s3):

vastdb/tests/test_schemas.py

@@ -60,5 +60,4 @@ def test_commits_and_rollbacks(session, clean_bucket_name):
 def test_list_snapshots(session, clean_bucket_name):
     with session.transaction() as tx:
         b = tx.bucket(clean_bucket_name)
-        s = b.snapshots()
-        assert s == []
+        b.snapshots()  # VAST Catalog may create some snapshots

vastdb/tests/test_tables.py

@@ -3,6 +3,7 @@ import decimal
 import logging
 import random
 import threading
+import time
 from contextlib import closing
 from tempfile import NamedTemporaryFile
 
@@ -70,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')
@@ -261,9 +272,14 @@ def test_filters(session, clean_bucket_name):
 
     with prepare_data(session, clean_bucket_name, 's', 't', expected) as t:
         def select(predicate):
-            return pa.Table.from_batches(t.select(predicate=predicate))
+            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)
@@ -304,6 +320,13 @@ def test_filters(session, clean_bucket_name):
         assert select(t['s'].contains('b')) == expected.filter(pc.field('s') == 'bb')
         assert select(t['s'].contains('y')) == expected.filter(pc.field('s') == 'xyz')
 
+        assert select(t['a'].isin([555])) == expected.filter(pc.field('a').isin([555]))
+        assert select(t['a'].isin([111, 222, 999])) == expected.filter(pc.field('a').isin([111, 222, 999]))
+        assert select((t['a'] == 111) | t['a'].isin([333, 444]) | (t['a'] > 600)) == expected.filter((pc.field('a') == 111) | pc.field('a').isin([333, 444]) | (pc.field('a') > 600))
+
+        with pytest.raises(NotImplementedError):
+            select(t['a'].isin([]))
+
 
 def test_parquet_export(session, clean_bucket_name):
     with session.transaction() as tx:
@@ -323,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
 
@@ -638,3 +662,20 @@ def test_select_stop(session, clean_bucket_name):
 
     # validate that all query threads were killed.
     assert active_threads() == 0
+
+
+def test_big_catalog_select(session, clean_bucket_name):
+    with session.transaction() as tx:
+        bc = tx.catalog()
+        actual = pa.Table.from_batches(bc.select(['name']))
+        assert actual
+        log.info("actual=%s", actual)
+
+
+def test_audit_log_select(session, clean_bucket_name):
+    with session.transaction() as tx:
+        a = tx.audit_log()
+        a.columns()
+        time.sleep(1)
+        actual = pa.Table.from_batches(a.select(), a.arrow_schema)
+        log.info("actual=%s", actual)

vastdb/tests/test_util.py

@@ -0,0 +1,39 @@
+import pyarrow as pa
+import pytest
+
+from .. import errors, util
+
+
+def test_slices():
+    ROWS = 1 << 20
+    t = pa.table({"x": range(ROWS), "y": [i / 1000 for i in range(ROWS)]})
+
+    chunks = list(util.iter_serialized_slices(t))
+    assert len(chunks) > 1
+    sizes = [len(c) for c in chunks]
+
+    assert max(sizes) < util.MAX_RECORD_BATCH_SLICE_SIZE
+    assert t == pa.Table.from_batches(_parse(chunks))
+
+    chunks = list(util.iter_serialized_slices(t, 1000))
+    assert len(chunks) > 1
+    sizes = [len(c) for c in chunks]
+
+    assert max(sizes) < util.MAX_RECORD_BATCH_SLICE_SIZE
+    assert t == pa.Table.from_batches(_parse(chunks))
+
+
+def test_wide_row():
+    cols = [pa.field(f"x{i}", pa.utf8()) for i in range(1000)]
+    values = [['a' * 10000]] * len(cols)
+    t = pa.table(values, schema=pa.schema(cols))
+    assert len(t) == 1
+
+    with pytest.raises(errors.TooWideRow):
+        list(util.iter_serialized_slices(t))
+
+
+def _parse(bufs):
+    for buf in bufs:
+        with pa.ipc.open_stream(buf) as reader:
+            yield from reader

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()