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

vastdb/tests/test_vector_index.py

@@ -0,0 +1,162 @@
+"""Tests for vector index functionality."""
+
+import logging
+from typing import Optional
+
+import pyarrow as pa
+import pytest
+
+from vastdb import errors
+from vastdb._internal import VectorIndexSpec
+from vastdb.session import Session
+
+log = logging.getLogger(__name__)
+
+
+@pytest.mark.parametrize("table_name,vector_index", [
+    # Test 1: Table without vector index
+    ("table_without_index", None),
+    # Test 2: Table with L2 vector index
+    ("table_with_l2_index", VectorIndexSpec("embedding", "l2sq")),
+    # Test 3: Table with inner product vector index
+    ("table_with_ip_index", VectorIndexSpec("embedding", "ip")),
+])
+def test_create_table_with_vector_index_metadata(session: Session,
+                                                 clean_bucket_name: str,
+                                                 table_name: str,
+                                                 vector_index: Optional[VectorIndexSpec]):
+    """Test that table creation and stats retrieval work correctly with vector index metadata."""
+    schema_name = "schema1"
+
+    with session.transaction() as tx:
+        log.info(f"Testing table '{table_name}' with {vector_index}")
+
+        # Create schema
+        bucket = tx.bucket(clean_bucket_name)
+        schema = bucket.create_schema(schema_name)
+
+        # Create the appropriate schema based on whether vector index is needed
+        if vector_index is None:
+            # Simple table without vector index
+            arrow_schema = pa.schema([
+                ('id', pa.int64()),
+                ('data', pa.string())
+            ])
+        else:
+            # Table with vector column
+            vector_dimension = 128  # Fixed-size vector dimension
+            vec_type = pa.list_(pa.field('', pa.float32(), False), vector_dimension)
+            arrow_schema = pa.schema([
+                ('id', pa.int64()),
+                ('embedding', vec_type)  # Fixed-size vector column
+            ])
+
+        # Create table with or without vector index
+        log.info(f"Creating table: {table_name}")
+        table = schema.create_table(
+            table_name=table_name,
+            columns=arrow_schema,
+            vector_index=vector_index
+        )
+
+        # Reload stats to ensure we get the vector index metadata
+        table.reload_stats()
+
+        # Get vector index metadata
+        result_vector_index = table._metadata._vector_index
+
+        log.info(f"Vector index metadata: {result_vector_index}")
+
+        # Assert expected values (should match input parameters)
+        result_vector_index_spec = (
+            None
+            if result_vector_index is None
+            else result_vector_index.to_vector_index_spec()
+        )
+        assert result_vector_index_spec == vector_index
+
+        log.info(f"✓ Test passed for table '{table_name}'")
+
+
+@pytest.mark.parametrize("table_name,vector_index,expected_error", [
+    # Test 1: Invalid column name (column doesn't exist in schema)
+    ("table_invalid_column", VectorIndexSpec("nonexistent_column", "l2sq"), "invalid vector indexed column name nonexistent_column"),
+    # Test 2: Invalid distance metric
+    ("table_invalid_metric", VectorIndexSpec("embedding", "invalid_metric"), "invalid vector index distance metric invalid_metric, supported metrics: 'l2sq', 'ip'"),
+])
+def test_create_table_with_invalid_vector_index(session: Session,
+                                                clean_bucket_name: str,
+                                                table_name: str,
+                                                vector_index: VectorIndexSpec,
+                                                expected_error: str):
+    """Test that table creation fails with appropriate error messages for invalid vector index parameters."""
+    schema_name = "schema1"
+
+    with session.transaction() as tx:
+        log.info(f"Testing invalid table '{table_name}' with vector_index={vector_index}, expected_error={expected_error}")
+
+        # Create schema
+        bucket = tx.bucket(clean_bucket_name)
+        schema = bucket.create_schema(schema_name)
+
+        # Table with vector column
+        vector_dimension = 128  # Fixed-size vector dimension
+        vec_type = pa.list_(pa.field('', pa.float32(), False), vector_dimension)
+        arrow_schema = pa.schema([
+            ('id', pa.int64()),
+            ('embedding', vec_type)  # Fixed-size vector column
+        ])
+
+        # Attempt to create table with invalid parameters - should raise an error
+        log.info(f"Attempting to create invalid table: {table_name}")
+        with pytest.raises((errors.BadRequest)) as exc_info:
+            schema.create_table(
+                table_name=table_name,
+                columns=arrow_schema,
+                vector_index=vector_index
+            )
+
+        # Verify the error message contains the expected error text
+        assert expected_error in str(exc_info.value), \
+            f"Expected error message to contain '{expected_error}', got '{str(exc_info.value)}'"
+
+        log.info(f"✓ Test passed for invalid table '{table_name}'")
+
+
+def test_vector_index_metadata_from_stats(session: Session, clean_bucket_name: str):
+    """Test that vector index metadata is correctly retrieved from table stats."""
+    schema_name = "schema1"
+    table_name = "vector_table"
+
+    with session.transaction() as tx:
+        # Create schema
+        bucket = tx.bucket(clean_bucket_name)
+        schema = bucket.create_schema(schema_name)
+
+        # Create table with vector index
+        vector_dimension = 128
+        vec_type = pa.list_(pa.field('', pa.float32(), False), vector_dimension)
+        arrow_schema = pa.schema([
+            ('id', pa.int64()),
+            ('embedding', vec_type)
+        ])
+
+        table = schema.create_table(
+            table_name=table_name,
+            columns=arrow_schema,
+            vector_index=VectorIndexSpec("embedding", "l2sq")
+        )
+
+        # Check stats object directly
+        stats = table.stats
+        assert stats is not None
+        assert stats.vector_index is not None
+        assert stats.vector_index.column == "embedding"
+        assert stats.vector_index.distance_metric == "l2sq"
+
+        # Check via the table methods
+        assert table._metadata._vector_index is not None
+        assert table._metadata._vector_index.column == "embedding"
+        assert table._metadata._vector_index.distance_metric == "l2sq"
+
+        log.info("✓ Vector index metadata correctly retrieved from stats")

vastdb/tests/test_vector_search.py

@@ -0,0 +1,210 @@
+from collections import defaultdict
+
+import numpy as np
+import pyarrow as pa
+import pytest
+
+from vastdb._adbc import _ibis_to_qe_predicates
+from vastdb._internal import VectorIndex, VectorIndexSpec
+from vastdb._table_interface import IbisPredicate
+from vastdb.conftest import SessionFactory
+from vastdb.table_metadata import TableMetadata, TableRef, TableType
+
+DIM = 8
+TEST_DISTANCE_FUNC = 'array_distance'
+TEST_DISTANCE_METRIC = 'l2sq'
+
+VectorColumnArrowType = pa.list_(
+    pa.field(name='item', type=pa.float32(), nullable=False), DIM)
+
+query_vector: np.ndarray = np.ones(DIM)
+first_closest_vector = (query_vector * 1.1).tolist()
+second_closest_vector = (query_vector * 1.2).tolist()
+third_closest_vector = (query_vector * 1.3).tolist()
+fourth_closest_vector = (query_vector * 1.4).tolist()
+fifth_closest_vector = (query_vector * 1.5).tolist()
+furthest_vector = (np.ones(DIM) * 5).tolist()
+
+vector_column_name = 'vector_col'
+
+data = [
+    {
+        "id": 0,
+        "n1": 1,
+        "n2": 100,
+        vector_column_name: first_closest_vector,
+    },
+    {
+        "id": 1,
+        "n1": 2,
+        "n2": 100,
+        vector_column_name: second_closest_vector,
+    },
+    {
+        "id": 2,
+        "n1": 3,
+        "n2": 200,
+        vector_column_name: third_closest_vector,
+    },
+    {
+        "id": 3,
+        "n1": 4,
+        "n2": 200,
+        vector_column_name: fourth_closest_vector,
+    },
+    {
+        "id": 4,
+        "n1": 5,
+        "n2": 200,
+        vector_column_name: fifth_closest_vector,
+    },
+    {
+        "id": 5,
+        "n1": 6,
+        "n2": 200,
+        vector_column_name: furthest_vector,
+    },
+
+]
+
+
+def into_arrow_arrays(data: list[dict]) -> list[list]:
+    agg = defaultdict(list)
+    for d in data:
+        for k, v in d.items():
+            agg[k].append(v)
+
+    return [v for v in agg.values()]
+
+
+def test_sanity(session_factory: SessionFactory, clean_bucket_name: str):
+    session = session_factory(with_adbc=True)
+
+    arrow_schema = pa.schema([('id', pa.int32()), ('n1', pa.int32(
+    )), ('n2', pa.int32()), (vector_column_name, VectorColumnArrowType),])
+
+    limit = 3
+
+    ref = TableRef(clean_bucket_name, 's', 't')
+    table_md = TableMetadata(ref,
+                             arrow_schema,
+                             TableType.Regular,
+                             vector_index=VectorIndex(vector_column_name,
+                                                      TEST_DISTANCE_METRIC,
+                                                      TEST_DISTANCE_FUNC))
+    data_table = pa.table(schema=arrow_schema, data=into_arrow_arrays(data))
+
+    with session.transaction() as tx:
+        table = tx.bucket(clean_bucket_name).create_schema(
+            's').create_table('t', arrow_schema)
+        table.insert(data_table)
+
+    # TODO merge in same tx
+    with session.transaction() as tx:
+        table = tx.table_from_metadata(table_md)
+
+        reader = table.vector_search(vec=query_vector.tolist(),
+                                     columns=['id', 'n1', 'n2'],
+                                     limit=limit)
+
+        result_table = reader.read_all()
+
+        assert set([v.as_py() for v in result_table['n1']]) == {1, 2, 3}
+
+
+def test_with_predicates(session_factory: SessionFactory, clean_bucket_name: str):
+    session = session_factory(with_adbc=True)
+
+    vector_column_name = 'vector_column'
+    arrow_schema = pa.schema([('id', pa.int32()), ('n1', pa.int32(
+    )), ('n2', pa.int32()), (vector_column_name, VectorColumnArrowType),])
+    limit = 3
+
+    ref = TableRef(clean_bucket_name, 's', 't')
+    table_md = TableMetadata(ref, arrow_schema, TableType.Regular,
+                             vector_index=VectorIndex(vector_column_name,
+                                                      TEST_DISTANCE_METRIC,
+                                                      TEST_DISTANCE_FUNC))
+    data_table = pa.table(schema=arrow_schema, data=into_arrow_arrays(data))
+
+    with session.transaction() as tx:
+        table = tx.bucket(clean_bucket_name).create_schema(
+            's').create_table('t', arrow_schema)
+        table.insert(data_table)
+
+    with session.transaction() as tx:
+        table = tx.table_from_metadata(table_md)
+
+        pred = table_md.ibis_table['n2'] == 200
+        reader = table.vector_search(vec=query_vector.tolist(),
+                                     columns=['id', 'n1', 'n2'],
+                                     limit=limit,
+                                     predicate=pred)
+
+        result_table = reader.read_all()
+
+        assert set([v.as_py() for v in result_table['n1']]) == {3, 4, 5}
+
+
+arrow_schema = pa.schema([('id', pa.int32()), ('n1', pa.int32(
+)), ('n2', pa.int32()), (vector_column_name, VectorColumnArrowType),])
+
+ref = TableRef('b', 's', 't')
+table_md = TableMetadata(ref, arrow_schema, TableType.Regular)
+
+
+@pytest.mark.parametrize('ibis_predicate, expected', [
+    ((table_md.ibis_table['n1'] == 1) & (table_md.ibis_table['n2'] == 2),
+     '("n1" = 1) AND ("n2" = 2)'),
+    ((table_md.ibis_table['n1'] == 1),
+     '"n1" = 1')
+])
+def test_ibis_to_query_engine_predicates(ibis_predicate: IbisPredicate, expected: str):
+    assert _ibis_to_qe_predicates(ibis_predicate) == expected
+
+
+def test_with_predicates_get_vector_index_properties_from_server(
+    session_factory: SessionFactory,
+    clean_bucket_name: str
+    ):
+    session = session_factory(with_adbc=True)
+
+    vector_column_name = 'vector_column'
+    vector_index_distance_metric = 'l2sq'
+    vector_index_sql_distance_function = "array_distance"
+
+    arrow_schema = pa.schema([('id', pa.int32()), ('n1', pa.int32(
+    )), ('n2', pa.int32()), (vector_column_name, VectorColumnArrowType),])
+    limit = 3
+
+    ref = TableRef(clean_bucket_name, 's', 't')
+    table_md = TableMetadata(ref, arrow_schema, TableType.Regular)
+    data_table = pa.table(schema=arrow_schema, data=into_arrow_arrays(data))
+
+    with session.transaction() as tx:
+        table = (tx.bucket(clean_bucket_name)
+                 .create_schema('s')
+                 .create_table('t',
+                               arrow_schema,
+                               vector_index=VectorIndexSpec(vector_column_name,
+                                                            vector_index_distance_metric)))
+        table.insert(data_table)
+
+    with session.transaction() as tx:
+        table = tx.table_from_metadata(table_md)
+
+        table.reload_stats()
+        assert table.vector_index == VectorIndex(
+            column=vector_column_name,
+            distance_metric=vector_index_distance_metric,
+            sql_distance_function=vector_index_sql_distance_function)
+
+        pred = table_md.ibis_table['n2'] == 200
+        reader = table.vector_search(vec=query_vector.tolist(),
+                                     columns=['id', 'n1', 'n2'],
+                                     limit=limit,
+                                     predicate=pred)
+
+        result_table = reader.read_all()
+
+        assert set([v.as_py() for v in result_table['n1']]) == {3, 4, 5}

vastdb/tests/util.py

@@ -25,11 +25,12 @@ def assert_row_ids_ascending_on_first_insertion_to_table(row_ids, expected_num_r
 @contextmanager
 def prepare_data(session: Session,
                  clean_bucket_name: str, schema_name: str, table_name: str,
-                 arrow_table: pa.Table, sorting_key: List[str] = []) -> Iterator[Table]:
+                 arrow_table: pa.Table, sorting_key: List[str] = [],
+                 insert_by_columns: bool = True) -> Iterator[Table]:
     with session.transaction() as tx:
         s = tx.bucket(clean_bucket_name).create_schema(schema_name)
         t = s.create_table(table_name, arrow_table.schema, sorting_key=sorting_key)
-        row_ids_array = t.insert(arrow_table)
+        row_ids_array = t.insert(arrow_table, by_columns=insert_by_columns)
         row_ids = row_ids_array.to_pylist()
         assert_row_ids_ascending_on_first_insertion_to_table(row_ids, arrow_table.num_rows, t.sorted_table)
         yield t

vastdb/transaction.py

@@ -11,6 +11,7 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Iterable, Optional
 
+from vastdb._adbc import AdbcConnection, AdbcDriver
 from vastdb._table_interface import ITable
 from vastdb.table import TableInTransaction
 from vastdb.table_metadata import TableMetadata
@@ -39,17 +40,35 @@ class TransactionNotActiveError(Exception):
     pass
 
 
+class NoAdbcConnectionError(Exception):
+    """No Adbc Connection Error."""
+
+    pass
+
+
 @dataclass
 class Transaction:
     """A holder of a single VAST transaction."""
 
     _rpc: "session.Session"
     txid: Optional[int] = None
+    _adbc_driver: Optional[AdbcDriver] = None
+    _adbc_conn: Optional[AdbcConnection] = None
 
     def __enter__(self):
         """Create a transaction and store its ID."""
         response = self._rpc.api.begin_transaction()
         self.txid = int(response.headers['tabular-txid'])
+
+        if self._adbc_driver is not None:
+            self._adbc_conn = AdbcConnection(
+                self._adbc_driver,
+                self._rpc.endpoint,
+                self._rpc.access,
+                self._rpc.secret,
+                self.txid,
+            )
+
         log.debug("opened txid=%016x", self.txid)
         return self
 
@@ -57,6 +76,10 @@ class Transaction:
         """On success, the transaction is committed. Otherwise, it is rolled back."""
         txid = self.txid
         self.txid = None
+        if self._adbc_conn is not None:
+            self.adbc_conn.close()
+            self._adbc_conn = None
+
         if (exc_type, exc_value, exc_traceback) == (None, None, None):
             log.debug("committing txid=%016x", txid)
             self._rpc.api.commit_transaction(txid)
@@ -110,3 +133,10 @@ class Transaction:
     def table_from_metadata(self, metadata: TableMetadata) -> ITable:
         """Create Table from TableMetadata."""
         return TableInTransaction(deepcopy(metadata), tx=self)
+
+    @property
+    def adbc_conn(self) -> AdbcConnection:
+        """ADBC connection in transaction."""
+        if self._adbc_conn is None:
+            raise NoAdbcConnectionError("Adbc Driver may not have been supplied")
+        return self._adbc_conn

vastdb/vast_flatbuf/tabular/GetTableStatsResponse.py

@@ -105,74 +105,66 @@ class GetTableStatsResponse(object):
             return self._tab.Get(flatbuffers.number_types.Int64Flags, o + self._tab.Pos)
         return 0
 
-def GetTableStatsResponseStart(builder):
-    builder.StartObject(9)
-
-def Start(builder):
-    GetTableStatsResponseStart(builder)
+    # GetTableStatsResponse
+    def VectorIndexMetadata(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(22))
+        if o != 0:
+            x = self._tab.Indirect(o + self._tab.Pos)
+            from vastdb.vast_flatbuf.tabular.VectorIndexMetadata import VectorIndexMetadata
+            obj = VectorIndexMetadata()
+            obj.Init(self._tab.Bytes, x)
+            return obj
+        return None
 
+def Start(builder): builder.StartObject(10)
+def GetTableStatsResponseStart(builder):
+    """This method is deprecated. Please switch to Start."""
+    return Start(builder)
+def AddNumRows(builder, numRows): builder.PrependInt64Slot(0, numRows, 0)
 def GetTableStatsResponseAddNumRows(builder, numRows):
-    builder.PrependInt64Slot(0, numRows, 0)
-
-def AddNumRows(builder, numRows):
-    GetTableStatsResponseAddNumRows(builder, numRows)
-
+    """This method is deprecated. Please switch to AddNumRows."""
+    return AddNumRows(builder, numRows)
+def AddSizeInBytes(builder, sizeInBytes): builder.PrependInt64Slot(1, sizeInBytes, 0)
 def GetTableStatsResponseAddSizeInBytes(builder, sizeInBytes):
-    builder.PrependInt64Slot(1, sizeInBytes, 0)
-
-def AddSizeInBytes(builder, sizeInBytes):
-    GetTableStatsResponseAddSizeInBytes(builder, sizeInBytes)
-
+    """This method is deprecated. Please switch to AddSizeInBytes."""
+    return AddSizeInBytes(builder, sizeInBytes)
+def AddIsExternalRowidAlloc(builder, isExternalRowidAlloc): builder.PrependBoolSlot(2, isExternalRowidAlloc, 0)
 def GetTableStatsResponseAddIsExternalRowidAlloc(builder, isExternalRowidAlloc):
-    builder.PrependBoolSlot(2, isExternalRowidAlloc, 0)
-
-def AddIsExternalRowidAlloc(builder, isExternalRowidAlloc):
-    GetTableStatsResponseAddIsExternalRowidAlloc(builder, isExternalRowidAlloc)
-
+    """This method is deprecated. Please switch to AddIsExternalRowidAlloc."""
+    return AddIsExternalRowidAlloc(builder, isExternalRowidAlloc)
+def AddAddressType(builder, addressType): builder.PrependUOffsetTRelativeSlot(3, flatbuffers.number_types.UOffsetTFlags.py_type(addressType), 0)
 def GetTableStatsResponseAddAddressType(builder, addressType):
-    builder.PrependUOffsetTRelativeSlot(3, flatbuffers.number_types.UOffsetTFlags.py_type(addressType), 0)
-
-def AddAddressType(builder, addressType):
-    GetTableStatsResponseAddAddressType(builder, addressType)
-
+    """This method is deprecated. Please switch to AddAddressType."""
+    return AddAddressType(builder, addressType)
+def AddVips(builder, vips): builder.PrependUOffsetTRelativeSlot(4, flatbuffers.number_types.UOffsetTFlags.py_type(vips), 0)
 def GetTableStatsResponseAddVips(builder, vips):
-    builder.PrependUOffsetTRelativeSlot(4, flatbuffers.number_types.UOffsetTFlags.py_type(vips), 0)
-
-def AddVips(builder, vips):
-    GetTableStatsResponseAddVips(builder, vips)
-
+    """This method is deprecated. Please switch to AddVips."""
+    return AddVips(builder, vips)
+def StartVipsVector(builder, numElems): return builder.StartVector(4, numElems, 4)
 def GetTableStatsResponseStartVipsVector(builder, numElems):
-    return builder.StartVector(4, numElems, 4)
-
-def StartVipsVector(builder, numElems):
-    return GetTableStatsResponseStartVipsVector(builder, numElems)
-
+    """This method is deprecated. Please switch to Start."""
+    return StartVipsVector(builder, numElems)
+def AddSortingKeyEnabled(builder, sortingKeyEnabled): builder.PrependBoolSlot(5, sortingKeyEnabled, 0)
 def GetTableStatsResponseAddSortingKeyEnabled(builder, sortingKeyEnabled):
-    builder.PrependBoolSlot(5, sortingKeyEnabled, 0)
-
-def AddSortingKeyEnabled(builder, sortingKeyEnabled):
-    GetTableStatsResponseAddSortingKeyEnabled(builder, sortingKeyEnabled)
-
+    """This method is deprecated. Please switch to AddSortingKeyEnabled."""
+    return AddSortingKeyEnabled(builder, sortingKeyEnabled)
+def AddSortingScore(builder, sortingScore): builder.PrependInt64Slot(6, sortingScore, 0)
 def GetTableStatsResponseAddSortingScore(builder, sortingScore):
-    builder.PrependInt64Slot(6, sortingScore, 0)
-
-def AddSortingScore(builder, sortingScore):
-    GetTableStatsResponseAddSortingScore(builder, sortingScore)
-
+    """This method is deprecated. Please switch to AddSortingScore."""
+    return AddSortingScore(builder, sortingScore)
+def AddWriteAmplification(builder, writeAmplification): builder.PrependInt64Slot(7, writeAmplification, 0)
 def GetTableStatsResponseAddWriteAmplification(builder, writeAmplification):
-    builder.PrependInt64Slot(7, writeAmplification, 0)
-
-def AddWriteAmplification(builder, writeAmplification):
-    GetTableStatsResponseAddWriteAmplification(builder, writeAmplification)
-
+    """This method is deprecated. Please switch to AddWriteAmplification."""
+    return AddWriteAmplification(builder, writeAmplification)
+def AddAcummulativeRowInseritionCount(builder, acummulativeRowInseritionCount): builder.PrependInt64Slot(8, acummulativeRowInseritionCount, 0)
 def GetTableStatsResponseAddAcummulativeRowInseritionCount(builder, acummulativeRowInseritionCount):
-    builder.PrependInt64Slot(8, acummulativeRowInseritionCount, 0)
-
-def AddAcummulativeRowInseritionCount(builder, acummulativeRowInseritionCount):
-    GetTableStatsResponseAddAcummulativeRowInseritionCount(builder, acummulativeRowInseritionCount)
-
+    """This method is deprecated. Please switch to AddAcummulativeRowInseritionCount."""
+    return AddAcummulativeRowInseritionCount(builder, acummulativeRowInseritionCount)
+def AddVectorIndexMetadata(builder, vectorIndexMetadata): builder.PrependUOffsetTRelativeSlot(9, flatbuffers.number_types.UOffsetTFlags.py_type(vectorIndexMetadata), 0)
+def GetTableStatsResponseAddVectorIndexMetadata(builder, vectorIndexMetadata):
+    """This method is deprecated. Please switch to AddVectorIndexMetadata."""
+    return AddVectorIndexMetadata(builder, vectorIndexMetadata)
+def End(builder): return builder.EndObject()
 def GetTableStatsResponseEnd(builder):
-    return builder.EndObject()
-
-def End(builder):
-    return GetTableStatsResponseEnd(builder)
+    """This method is deprecated. Please switch to End."""
+    return End(builder)