vastdb 1.4.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

vastdb/_table_interface.py

@@ -0,0 +1,136 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Iterable, Optional, Union
+
+import ibis
+import pyarrow as pa
+
+from .config import ImportConfig, QueryConfig
+from .table_metadata import TableRef
+
+if TYPE_CHECKING:
+    from .table import Projection
+
+
+class ITable(ABC):
+    """Interface for VAST Table operations."""
+
+    @property
+    @abstractmethod
+    def ref(self) -> TableRef:
+        """Return Table Ref."""
+        pass
+
+    @abstractmethod
+    def __eq__(self, other: object) -> bool:
+        """Table __eq__."""
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Table name."""
+        pass
+
+    @property
+    @abstractmethod
+    def arrow_schema(self) -> pa.Schema:
+        """Table arrow schema."""
+        pass
+
+    @property
+    @abstractmethod
+    def path(self) -> str:
+        """Return table's path."""
+        pass
+
+    @abstractmethod
+    def sorted_columns(self) -> list[str]:
+        """Return sorted columns' names."""
+        pass
+
+    @abstractmethod
+    def projection(self, name: str) -> "Projection":
+        """Get a specific semi-sorted projection of this table."""
+        pass
+
+    @abstractmethod
+    def projections(self, projection_name: str = "") -> Iterable["Projection"]:
+        """List semi-sorted projections."""
+        pass
+
+    @abstractmethod
+    def import_files(self, files_to_import: Iterable[str], config: Optional[ImportConfig] = None) -> None:
+        """Import files into table."""
+        pass
+
+    @abstractmethod
+    def import_partitioned_files(self, files_and_partitions: dict[str, pa.RecordBatch], config: Optional[ImportConfig] = None) -> None:
+        """Import partitioned files."""
+        pass
+
+    @abstractmethod
+    def select(self,
+               columns: Optional[list[str]] = None,
+               predicate: Union[ibis.expr.types.BooleanColumn,
+                                ibis.common.deferred.Deferred] = None,
+               config: Optional[QueryConfig] = None,
+               *,
+               internal_row_id: bool = False,
+               limit_rows: Optional[int] = None) -> pa.RecordBatchReader:
+        """Execute a query."""
+        pass
+
+    @abstractmethod
+    def insert(self, rows: Union[pa.RecordBatch, pa.Table]) -> pa.ChunkedArray:
+        """Insert rows into table."""
+        pass
+
+    @abstractmethod
+    def update(self,
+               rows: Union[pa.RecordBatch, pa.Table],
+               columns: Optional[list[str]] = None) -> None:
+        """Update rows in table."""
+        pass
+
+    @abstractmethod
+    def delete(self, rows: Union[pa.RecordBatch, pa.Table]) -> None:
+        """Delete rows from table."""
+        pass
+
+    @abstractmethod
+    def imports_table(self) -> Optional["ITable"]:
+        """Get imports table."""
+        pass
+
+    @abstractmethod
+    def sorting_done(self) -> bool:
+        """Check if sorting is done."""
+        pass
+
+    @abstractmethod
+    def sorting_score(self) -> int:
+        """Get sorting score."""
+        pass
+
+    @abstractmethod
+    def reload_schema(self) -> None:
+        """Reload Arrow Schema."""
+        pass
+
+    @abstractmethod
+    def reload_stats(self) -> None:
+        """Reload Table Stats."""
+        pass
+
+    @abstractmethod
+    def reload_sorted_columns(self) -> None:
+        """Reload Sorted Columns."""
+        pass
+
+    @abstractmethod
+    def __getitem__(self, col_name: str) -> ibis.Column:
+        """Allow constructing ibis-like column expressions from this table.
+
+        It is useful for constructing expressions for predicate pushdown in `ITable.select()` method.
+        """
+        pass

vastdb/bench/perf_bench/orchestrate/results_helpers.py

@@ -113,7 +113,7 @@ def calculate_aggregate_stats(
     )
     agg_df["duration_sec"] = (
         r_df.groupby(group_flds)
-        .apply(calc_total_time_coverage_seconds, include_groups=False)
+        .apply(calc_total_time_coverage_seconds, include_groups=False)  # type: ignore
         .sort_index()
     )
     agg_df["M_rows_per_sec"] = (agg_df["n_rows"] / agg_df["duration_sec"] / 1e6).astype(

vastdb/bucket.py

@@ -36,7 +36,7 @@ class Bucket:
         """Get a specific schema (a container of tables) under this bucket."""
         return self._root_schema.schema(name=name, fail_if_missing=fail_if_missing)
 
-    def schemas(self, batch_size=None):
+    def schemas(self, batch_size: Optional[int] = None) -> Iterable["Schema"]:
         """List bucket's schemas."""
         return self._root_schema.schemas(batch_size=batch_size)
 

vastdb/conftest.py

@@ -1,20 +1,42 @@
 import os
 import sqlite3
 from pathlib import Path
+from typing import Iterable
 
 import boto3
 import pytest
 
 import vastdb
 import vastdb.errors
+from vastdb.schema import Schema
+from vastdb.session import Session
+
+AWS_ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID"
+AWS_SECRET_ACCESS_KEY = "AWS_SECRET_ACCESS_KEY"
+
+
+def get_aws_cred_from_system(cred_name):
+    orion_dir = Path(os.path.expanduser(os.environ.get("ORION_DIR", "~/orion")))
+    orion_file_path = orion_dir / 'data' / cred_name
+
+    env_value = os.environ.get(cred_name)
+
+    if env_value is not None:
+        return env_value
+
+    if not orion_file_path.exists():
+        return None
+
+    with open(orion_file_path, "r") as f:
+        return f.read().strip()
 
 
 def pytest_addoption(parser):
     parser.addoption("--tabular-bucket-name", help="Name of the S3 bucket with Tabular enabled", default="vastdb")
     parser.addoption("--tabular-access-key", help="Access key with Tabular permissions (AWS_ACCESS_KEY_ID)",
-                     default=os.environ.get("AWS_ACCESS_KEY_ID", None))
+                     default=get_aws_cred_from_system(AWS_ACCESS_KEY_ID))
     parser.addoption("--tabular-secret-key", help="Secret key with Tabular permissions (AWS_SECRET_ACCESS_KEY)",
-                     default=os.environ.get("AWS_SECRET_ACCESS_KEY", None))
+                     default=get_aws_cred_from_system(AWS_SECRET_ACCESS_KEY))
     parser.addoption("--tabular-endpoint-url", help="Tabular server endpoint", default=[], action="append")
     parser.addoption("--data-path", help="Data files location", default=None)
     parser.addoption("--crater-path", help="Save benchmark results in a dedicated location", default=None)
@@ -24,7 +46,7 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope="session")
-def session_kwargs(request, tabular_endpoint_urls):
+def session_kwargs(request: pytest.FixtureRequest, tabular_endpoint_urls):
     return dict(
         access=request.config.getoption("--tabular-access-key"),
         secret=request.config.getoption("--tabular-secret-key"),
@@ -38,30 +60,30 @@ def session(session_kwargs):
 
 
 @pytest.fixture(scope="session")
-def num_workers(request):
+def num_workers(request: pytest.FixtureRequest):
     return int(request.config.getoption("--num-workers"))
 
 
 @pytest.fixture(scope="session")
-def test_bucket_name(request):
+def test_bucket_name(request: pytest.FixtureRequest):
     return request.config.getoption("--tabular-bucket-name")
 
 
 @pytest.fixture(scope="session")
-def tabular_endpoint_urls(request):
+def tabular_endpoint_urls(request: pytest.FixtureRequest):
     return request.config.getoption("--tabular-endpoint-url") or ["http://localhost:9090"]
 
 
-def iter_schemas(s):
-    """Recusively scan all schemas."""
-    children = s.schemas()
-    for c in children:
-        yield from iter_schemas(c)
-    yield s
+def iter_schemas(schema: Schema) -> Iterable[Schema]:
+    """Recursively scan all schemas."""
+    children = schema.schemas()
+    for child in children:
+        yield from iter_schemas(child)
+    yield schema
 
 
 @pytest.fixture(scope="function")
-def clean_bucket_name(request, test_bucket_name, session):
+def clean_bucket_name(request: pytest.FixtureRequest, test_bucket_name: str, session: Session) -> str:
     with session.transaction() as tx:
         b = tx.bucket(test_bucket_name)
         for top_schema in b.schemas():
@@ -69,7 +91,8 @@ def clean_bucket_name(request, test_bucket_name, session):
                 for t_name in s.tablenames():
                     try:
                         t = s.table(t_name)
-                        t.drop()
+                        if t is not None:
+                            t.drop()
                     except vastdb.errors.NotSupportedSchema:
                         # Use internal API to drop the table in case unsupported schema prevents creating a table
                         # object.
@@ -79,7 +102,7 @@ def clean_bucket_name(request, test_bucket_name, session):
 
 
 @pytest.fixture(scope="session")
-def s3(request, tabular_endpoint_urls):
+def s3(request: pytest.FixtureRequest, tabular_endpoint_urls):
     return boto3.client(
         's3',
         aws_access_key_id=request.config.getoption("--tabular-access-key"),
@@ -88,22 +111,22 @@ def s3(request, tabular_endpoint_urls):
 
 
 @pytest.fixture(scope="function")
-def parquets_path(request):
+def parquets_path(request: pytest.FixtureRequest):
     return Path(request.config.getoption("--data-path"))
 
 
 @pytest.fixture(scope="function")
-def crater_path(request):
+def crater_path(request: pytest.FixtureRequest):
     return request.config.getoption("--crater-path")
 
 
 @pytest.fixture(scope="function")
-def schema_name(request):
+def schema_name(request: pytest.FixtureRequest):
     return request.config.getoption("--schema-name")
 
 
 @pytest.fixture(scope="function")
-def table_name(request):
+def table_name(request: pytest.FixtureRequest):
     return request.config.getoption("--table-name")
 
 

vastdb/schema.py

@@ -10,7 +10,10 @@ from typing import TYPE_CHECKING, Iterable, List, Optional
 
 import pyarrow as pa
 
+from vastdb.table_metadata import TableMetadata, TableRef, TableType
+
 from . import bucket, errors, schema, table
+from ._ibis_support import validate_ibis_support_schema
 
 if TYPE_CHECKING:
     from .table import Table
@@ -91,7 +94,7 @@ class Schema:
         if use_external_row_ids_allocation:
             self.tx._rpc.features.check_external_row_ids_allocation()
 
-        table.Table.validate_ibis_support_schema(columns)
+        validate_ibis_support_schema(columns)
         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,
                                       sorting_key=sorting_key)
@@ -117,7 +120,7 @@ class Schema:
         while True:
             _bucket_name, _schema_name, curr_tables, next_key, is_truncated, _ = \
                 self.tx._rpc.api.list_tables(
-                    bucket=self.bucket.name, schema=self.name, next_key=next_key, max_keys=page_size, txid=self.tx.txid,
+                    bucket=self.bucket.name, schema=self.name, next_key=next_key, max_keys=page_size, txid=self.tx.active_txid,
                     exact_match=exact_match, name_prefix=name_prefix, include_list_stats=exact_match)
             if not curr_tables:
                 break
@@ -152,4 +155,13 @@ 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, sorted_table=table_info.sorting_key_enabled)
+    ref = TableRef(bucket=schema.bucket.name,
+                   schema=schema.name,
+                   table=table_info.name)
+
+    table_type = TableType.Elysium if table_info.sorting_key_enabled else TableType.Regular
+    table_metadata = TableMetadata(ref, table_type=table_type)
+
+    return table.Table(handle=int(table_info.handle),
+                       metadata=table_metadata,
+                       tx=schema.tx)

vastdb/session.py

@@ -10,6 +10,8 @@ For more details see:
 import os
 from typing import TYPE_CHECKING, Optional
 
+from vastdb.transaction import Transaction
+
 if TYPE_CHECKING:
     from .config import BackoffConfig
 
@@ -45,7 +47,7 @@ class Session:
         """Don't show the secret key."""
         return f'{self.__class__.__name__}(endpoint={self.api.url}, access={self.api.access_key})'
 
-    def transaction(self):
+    def transaction(self) -> Transaction:
         """Create a non-initialized transaction object.
 
         It should be used as a context manager: