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

vastdb/_adbc.py

@@ -0,0 +1,194 @@
+import hashlib
+import logging
+import os
+import urllib.request
+from typing import Optional
+
+import pyarrow as pa
+import sqlglot
+from adbc_driver_manager.dbapi import Connection, Cursor, connect
+from sqlglot import exp
+
+from vastdb._internal import VectorIndex
+from vastdb._table_interface import IbisPredicate
+from vastdb.table_metadata import TableRef
+
+log = logging.getLogger(__name__)
+
+
+TXID_OVERRIDE_PROPERTY: str = "vast.db.external_txid"
+VAST_DIST_ALIAS = "vast_pysdk_vector_dist"
+DEFAULT_ADBC_DRIVER_CACHE_DIR: str = "~/.vast/adbc_drivers_cache"
+DEFAULT_ADBC_DRIVER_CACHE_BY_URL_DIR: str = f"{DEFAULT_ADBC_DRIVER_CACHE_DIR}/by_url"
+
+
+class LocalAdbcDriverNotFound(Exception):
+    """LocalAdbcDriverNotFound."""
+
+    pass
+
+
+class RemoteAdbcDriverDownloadFailed(Exception):
+    """RemoteAdbcDriverDownloadFailed."""
+
+    pass
+
+
+class AdbcDriver:
+    _local_path: str
+
+    def __init__(self, local_path: str):
+        self._local_path = local_path
+
+    @staticmethod
+    def from_local_path(local_path: str) -> "AdbcDriver":
+        """AdbcDriver from a local_path to shared-library."""
+        if not os.path.exists(local_path):
+            raise LocalAdbcDriverNotFound(local_path)
+
+        return AdbcDriver(local_path)
+
+    @staticmethod
+    def from_url(url: str) -> "AdbcDriver":
+        """AdbcDriver to be downloaded by url to shared-library (uses cache if exists)."""
+        expected_local_path = AdbcDriver._url_to_local_path(url)
+
+        if os.path.exists(expected_local_path):
+            return AdbcDriver(expected_local_path)
+
+        AdbcDriver._download_driver(url, expected_local_path)
+        return AdbcDriver(expected_local_path)
+
+    @staticmethod
+    def _url_to_local_path(url: str) -> str:
+        url_hash = hashlib.sha256(url.encode("utf-8")).hexdigest()
+        return os.path.join(DEFAULT_ADBC_DRIVER_CACHE_BY_URL_DIR, url_hash)
+
+    @staticmethod
+    def _download_driver(url: str, target_path: str):
+        os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+        try:
+            log.info(f"Downloading ADBC driver from {url} to {target_path}...")
+            urllib.request.urlretrieve(url, target_path)
+            log.info(f"Successfully downloaded driver to {target_path}.")
+        except Exception as e:
+            raise RemoteAdbcDriverDownloadFailed(
+                f"Failed to download ADBC driver from {url}: {e}"
+            ) from e
+
+    @property
+    def local_path(self) -> str:
+        return self._local_path
+
+
+def _get_adbc_connection(
+    adbc_driver_path: str, endpoint: str, access_key: str, secret_key: str, txid: int
+) -> Connection:
+    """Get an adbc connection in transaction."""
+    return connect(
+        driver=adbc_driver_path,
+        db_kwargs={
+            "vast.db.endpoint": endpoint,
+            "vast.db.access_key": access_key,
+            "vast.db.secret_key": secret_key,
+        },
+        conn_kwargs={TXID_OVERRIDE_PROPERTY: str(txid)},
+    )
+
+
+def _remove_table_qualification_from_columns(expression: exp.Expression):
+    """Goes over all columns which are fully qualified with "t0" table reference (ibis default table qualification for unbound tables.
+
+    Note: use only if one table is involved - if two tables exist in the expression columns might become ambiguous.
+    """
+    for col in expression.find_all(exp.Column):
+        col.set("table", None)
+    return expression
+
+
+def _ibis_to_qe_predicates(predicate: IbisPredicate) -> str:
+    ibis_sql = predicate.to_sql()
+    parsed = sqlglot.parse_one(ibis_sql)
+
+    # currently there is a single table
+    # removing the
+    without_table_qualification = _remove_table_qualification_from_columns(
+        parsed.expressions[0].this
+    )
+
+    return without_table_qualification.sql()
+
+
+def _vector_search_sql(
+    query_vector: list[float],
+    vector_index: VectorIndex,
+    table_ref: TableRef,
+    columns: list[str],
+    limit: int,
+    predicate: Optional[IbisPredicate] = None,
+) -> str:
+    query_vector_dim = len(query_vector)
+
+    query_vector_literal = f"{query_vector}::FLOAT[{query_vector_dim}]"
+    dist_func = f"{vector_index.sql_distance_function}({vector_index.column}::FLOAT[{query_vector_dim}], {query_vector_literal})"
+    dist_alias = f"{dist_func} as {VAST_DIST_ALIAS}"
+
+    projection_str = ",".join(columns + [dist_alias])
+
+    if predicate is not None:
+        where = f"WHERE {_ibis_to_qe_predicates(predicate)}"
+    else:
+        where = ""
+
+    return f"""
+            SELECT {projection_str}
+            FROM {table_ref.query_engine_full_path}
+            {where}
+            ORDER BY {VAST_DIST_ALIAS}
+            LIMIT {limit}"""
+
+
+class AdbcConnection:
+    def __init__(
+        self,
+        adbc_driver: AdbcDriver,
+        endpoint: str,
+        access_key: str,
+        secret_key: str,
+        txid: int,
+    ):
+        self._adbc_conn = _get_adbc_connection(
+            adbc_driver.local_path, endpoint, access_key, secret_key, txid
+        )
+
+        self._cursor = self._adbc_conn.cursor()
+
+    @property
+    def cursor(self) -> Cursor:
+        return self._cursor
+
+    def close(self):
+        self._cursor.close()
+
+    def vector_search(
+        self,
+        query_vector: list[float],
+        vector_index: VectorIndex,
+        table_ref: TableRef,
+        columns: list[str],
+        limit: int,
+        predicate: Optional[IbisPredicate] = None,
+    ) -> pa.RecordBatchReader:
+        """Top-n on vector-column."""
+        sql = _vector_search_sql(
+            query_vector=query_vector,
+            vector_index=vector_index,
+            table_ref=table_ref,
+            columns=columns,
+            limit=limit,
+            predicate=predicate,
+        )
+
+        self._cursor.execute(sql)
+        return self._cursor.fetch_record_batch()

vastdb/_internal.py

@@ -6,6 +6,7 @@ import struct
 import time
 import urllib.parse
 from collections import defaultdict, namedtuple
+from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast
 
@@ -116,6 +117,7 @@ from vastdb.vast_flatbuf.tabular.ListSchemasResponse import (
 from vastdb.vast_flatbuf.tabular.ListTablesResponse import (
     ListTablesResponse as list_tables,
 )
+from vastdb.vast_flatbuf.tabular.VectorIndexMetadata import VectorIndexMetadata
 
 from . import errors, util
 from .config import BackoffConfig
@@ -803,10 +805,40 @@ def _parse_table_info(obj, parse_properties):
                      sorting_score, write_amplification, acummulative_row_insertion_count, sorting_done)
 
 
-# Results that returns from tablestats
+@dataclass
+class VectorIndexSpec:
+    """
+    Vector Index Specification when creating a table.
+    """
+    column: str
+    distance_metric: str
+
+
+@dataclass
+class VectorIndex:
+    column: str
+    distance_metric: str
+    sql_distance_function: str
 
+    def to_vector_index_spec(self) -> VectorIndexSpec:
+        return VectorIndexSpec(self.column,
+                               self.distance_metric)
 
-TableStatsResult = namedtuple("TableStatsResult", 'num_rows size_in_bytes is_external_rowid_alloc sorting_key_enabled sorting_score write_amplification acummulative_row_inserition_count sorting_done endpoints')
+
+@dataclass
+class TableStats:
+    """Table-related information."""
+
+    num_rows: int
+    size_in_bytes: int
+    sorting_score: int
+    write_amplification: int
+    acummulative_row_inserition_count: int
+    is_external_rowid_alloc: bool = False
+    sorting_key_enabled: bool = False
+    sorting_done: bool = False
+    endpoints: Tuple[str, ...] = ()
+    vector_index: Optional[VectorIndex] = None
 
 
 _RETRIABLE_EXCEPTIONS = (
@@ -1128,11 +1160,31 @@ 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,
-                     sorting_key=[]):
+                     sorting_key=[], vector_index: Optional[VectorIndexSpec] = None):
+        """
+        Create a table in the specified bucket and schema.
+
+        Args:
+            bucket: Name of the bucket
+            schema: Name of the schema
+            name: Name of the table
+            arrow_schema: PyArrow schema defining the table columns
+            txid: Transaction ID
+            client_tags: Client tags for the request
+            expected_retvals: Expected return values
+            create_imports_table: Whether this is an imports table
+            use_external_row_ids_allocation: Whether to use external row ID allocation
+            table_props: Table properties
+            sorting_key: List of column indices to sort by (for Elysium tables)
+            vector_index: Optional vector index
+        """
         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, sorting_key=sorting_key)
+                                    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, sorting_key=sorting_key,
+                                    vector_index=vector_index)
 
     def create_topic(self, bucket, name, topic_partitions, expected_retvals=[],
                      message_timestamp_type=None, retention_ms=None, message_timestamp_after_max_ms=None,
@@ -1149,8 +1201,9 @@ 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,
-                               sorting_key=[]):
+                               create_imports_table=False, use_external_row_ids_allocation=False,
+                               table_props=None, sorting_key=[],
+                               vector_index: Optional[VectorIndexSpec] = None):
         """
         Create a table, use the following request
         POST /bucket/schema/table?table HTTP/1.1
@@ -1176,6 +1229,10 @@ class VastdbApi:
         if use_external_row_ids_allocation:
             headers['use-external-row-ids-alloc'] = str(use_external_row_ids_allocation)
 
+        if vector_index is not None:
+            headers['tabular-vector-index-column'] = vector_index.column
+            headers['tabular-vector-index-distance-metric'] = vector_index.distance_metric
+
         url_params = {'topic_partitions': str(topic_partitions)} if topic_partitions else {}
         if create_imports_table:
             url_params['sub-table'] = IMPORTED_OBJECTS_TABLE_NAME
@@ -1188,10 +1245,10 @@ class VastdbApi:
             url=self._url(bucket=bucket, schema=schema, table=name, command="table", url_params=url_params),
             data=serialized_schema, headers=headers)
 
-    def get_topic_stats(self, bucket, name, expected_retvals=[]):
+    def get_topic_stats(self, bucket, name, expected_retvals=[]) -> TableStats:
         return self.get_table_stats(bucket=bucket, schema=KAFKA_TOPICS_SCHEMA_NAME, name=name, expected_retvals=expected_retvals)
 
-    def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[], imports_table_stats=False):
+    def get_table_stats(self, bucket, schema, name, txid=0, client_tags=[], expected_retvals=[], imports_table_stats=False) -> TableStats:
         """
         GET /mybucket/myschema/mytable?stats HTTP/1.1
         tabular-txid: TransactionId
@@ -1218,8 +1275,40 @@ class VastdbApi:
         sorting_score = sorting_score_raw & ((1 << 63) - 1)
         sorting_done = bool(sorting_score_raw >> 63)
 
+        vector_index_metadata: Optional[VectorIndexMetadata] = stats.VectorIndexMetadata()
+
+        if vector_index_metadata is not None:
+            column_name = vector_index_metadata.ColumnName()
+            distance_metric = vector_index_metadata.DistanceMetric()
+            sql_distance_function = vector_index_metadata.SqlFunctionName()
+
+            if (column_name is None or
+                distance_metric is None or
+                sql_distance_function is None):
+                raise errors.ApiResponseError(
+                    "VectorIndexMetadata properties (column_name, distance_metric, sql_function_name) must all be set."
+                )
+
+            vector_index = VectorIndex(
+                column=column_name.decode('utf-8'),
+                distance_metric=distance_metric.decode('utf-8'),
+                sql_distance_function=sql_distance_function.decode('utf-8'))
+        else:
+            vector_index = None
+
         endpoints = [self.url]  # we cannot replace the host by a VIP address in HTTPS-based URLs
-        return TableStatsResult(num_rows, size_in_bytes, is_external_rowid_alloc, sorting_key_enabled, sorting_score, write_amplification, acummulative_row_inserition_count, sorting_done, tuple(endpoints))
+
+        return TableStats(
+          num_rows=num_rows,
+          size_in_bytes=size_in_bytes,
+          sorting_score=sorting_score,
+          write_amplification=write_amplification,
+          acummulative_row_inserition_count=acummulative_row_inserition_count,
+          is_external_rowid_alloc=is_external_rowid_alloc,
+          sorting_key_enabled=sorting_key_enabled,
+          sorting_done=sorting_done,
+          endpoints=tuple(endpoints),
+          vector_index=vector_index)
 
     def alter_topic(self, bucket, name,
                     new_name="", expected_retvals=[],
@@ -2339,7 +2428,7 @@ def build_field(builder: flatbuffers.Builder, f: pa.Field, include_name=True):
 class QueryDataRequest:
     def __init__(self, serialized, response_schema, response_parser):
         self.serialized = serialized
-        self.response_schema = response_schema
+        self.response_schema: pa.Schema = response_schema
         self.response_parser = response_parser
 
 

vastdb/_table_interface.py

@@ -1,15 +1,18 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Iterable, Optional, TypeAlias, Union
 
 import ibis
 import pyarrow as pa
 
+from ._internal import VectorIndex
 from .config import ImportConfig, QueryConfig
 from .table_metadata import TableRef
 
 if TYPE_CHECKING:
     from .table import Projection
 
+IbisPredicate: TypeAlias = Union[ibis.expr.types.BooleanColumn, ibis.common.deferred.Deferred]
+
 
 class ITable(ABC):
     """Interface for VAST Table operations."""
@@ -71,8 +74,7 @@ class ITable(ABC):
     @abstractmethod
     def select(self,
                columns: Optional[list[str]] = None,
-               predicate: Union[ibis.expr.types.BooleanColumn,
-                                ibis.common.deferred.Deferred] = None,
+               predicate: Optional[IbisPredicate] = None,
                config: Optional[QueryConfig] = None,
                *,
                internal_row_id: bool = False,
@@ -134,3 +136,18 @@ class ITable(ABC):
         It is useful for constructing expressions for predicate pushdown in `ITable.select()` method.
         """
         pass
+
+    @property
+    @abstractmethod
+    def vector_index(self) -> Optional[VectorIndex]:
+        """Table's Vector Index if exists."""
+        pass
+
+    @abstractmethod
+    def vector_search(self,
+                      vec: list[float],
+                      columns: list[str],
+                      limit: int,
+                      predicate: Optional[IbisPredicate] = None) -> pa.RecordBatchReader:
+        """Top-n on vector-column."""
+        pass

vastdb/conftest.py

@@ -1,13 +1,14 @@
 import os
 import sqlite3
 from pathlib import Path
-from typing import Iterable
+from typing import Iterable, Protocol
 
 import boto3
 import pytest
 
 import vastdb
 import vastdb.errors
+from vastdb._adbc import AdbcDriver
 from vastdb.schema import Schema
 from vastdb.session import Session
 
@@ -45,6 +46,10 @@ def pytest_addoption(parser):
     parser.addoption("--num-workers", help="Number of concurrent workers", default=1)
 
 
+def _get_adbc_driver_url(pipeline: str) -> str:
+    return f"https://artifactory.vastdata.com/artifactory/files/vastdb-native-client/{pipeline}/libadbc_driver_vastdb.so"
+
+
 @pytest.fixture(scope="session")
 def session_kwargs(request: pytest.FixtureRequest, tabular_endpoint_urls):
     return dict(
@@ -59,6 +64,23 @@ def session(session_kwargs):
     return vastdb.connect(**session_kwargs)
 
 
+class SessionFactory(Protocol):
+    def __call__(self, *, with_adbc: bool) -> Session: ...
+
+
+@pytest.fixture(scope="session")
+def session_factory(session_kwargs) -> SessionFactory:
+    def create_session(with_adbc: bool = False) -> Session:
+
+        if with_adbc:
+            # TODO use other not hard coded driver
+            session_kwargs['adbc_driver'] = AdbcDriver.from_url(url=_get_adbc_driver_url("2103686"))
+
+        return vastdb.connect(**session_kwargs)
+
+    return create_session
+
+
 @pytest.fixture(scope="session")
 def num_workers(request: pytest.FixtureRequest):
     return int(request.config.getoption("--num-workers"))

vastdb/errors.py

@@ -231,6 +231,11 @@ class ConnectionError(Exception):
         self.args = [vars(self)]
 
 
+class ApiResponseError(Exception):
+    """Indicates a logically invalid or inconsistent server response."""
+    pass
+
+
 def handle_unavailable(**kwargs):
     if kwargs['code'] == 'SlowDown':
         raise Slowdown(**kwargs)

vastdb/schema.py

@@ -14,6 +14,7 @@ from vastdb.table_metadata import TableMetadata, TableRef, TableType
 
 from . import bucket, errors, schema, table
 from ._ibis_support import validate_ibis_support_schema
+from ._internal import VectorIndexSpec
 
 if TYPE_CHECKING:
     from .table import Table
@@ -80,11 +81,24 @@ class Schema:
         return result
 
     def create_table(self, table_name: str, columns: pa.Schema, fail_if_exists=True,
-                     use_external_row_ids_allocation=False, sorting_key=[]) -> "Table":
+                     use_external_row_ids_allocation=False, sorting_key=[],
+                     vector_index: Optional[VectorIndexSpec] = None) -> "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.
         See https://support.vastdata.com/s/article/UUID-48d0a8cf-5786-5ef3-3fa3-9c64e63a0967 for more details.
+
+        Args:
+            table_name: Name of the table to create
+            columns: PyArrow schema defining the table columns
+            fail_if_exists: Whether to fail if the table already exists
+            use_external_row_ids_allocation: Whether to use external row ID allocation
+            sorting_key: List of column names to use as sorting key (for Elysium tables)
+            vector_index: Optional vector index.
+
+        Returns:
+            The created table
+
         """
         if current := self.table(table_name, fail_if_missing=False):
             if fail_if_exists:
@@ -97,7 +111,8 @@ class Schema:
         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)
+                                      sorting_key=sorting_key,
+                                      vector_index=vector_index)
         log.info("Created table: %s", table_name)
         return self.table(table_name)  # type: ignore[return-value]
 

vastdb/session.py

@@ -10,6 +10,7 @@ For more details see:
 import os
 from typing import TYPE_CHECKING, Optional
 
+from vastdb._adbc import AdbcDriver
 from vastdb.transaction import Transaction
 
 if TYPE_CHECKING:
@@ -23,7 +24,8 @@ class Session:
                  *,
                  ssl_verify=True,
                  timeout=None,
-                 backoff_config: Optional["BackoffConfig"] = None):
+                 backoff_config: Optional["BackoffConfig"] = None,
+                 adbc_driver: Optional[AdbcDriver] = None):
         """Connect to a VAST Database endpoint, using specified credentials."""
         from . import _internal, features
 
@@ -34,14 +36,19 @@ class Session:
         if endpoint is None:
             endpoint = os.environ['AWS_S3_ENDPOINT_URL']
 
+        self.endpoint = endpoint
+        self.access = access
+        self.secret = secret
+
         self.api = _internal.VastdbApi(
-            endpoint=endpoint,
-            access_key=access,
-            secret_key=secret,
+            endpoint=self.endpoint,
+            access_key=self.access,
+            secret_key=self.secret,
             ssl_verify=ssl_verify,
             timeout=timeout,
             backoff_config=backoff_config)
         self.features = features.Features(self.api.vast_version)
+        self.adbc_driver: Optional[AdbcDriver] = adbc_driver
 
     def __repr__(self):
         """Don't show the secret key."""
@@ -56,4 +63,4 @@ class Session:
                 tx.bucket("bucket").create_schema("schema")
         """
         from . import transaction
-        return transaction.Transaction(self)
+        return transaction.Transaction(self, _adbc_driver=self.adbc_driver)