feat: Vector Index & Columnar Engine Optimizations

docs/index.rst

@@ -7,6 +7,7 @@ API Reference
 
   langchain_google_alloydb_pg/engine
   langchain_google_alloydb_pg/vectorstore
+  langchain_google_alloydb_pg/indexes
   langchain_google_alloydb_pg/loader
   langchain_google_alloydb_pg/history
 

docs/langchain_google_alloydb_pg/indexes.rst

@@ -0,0 +1,7 @@
+Indexes
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. automodule:: langchain_google_alloydb_pg.indexes
+  :members:
+  :private-members:
+  :noindex:

src/langchain_google_alloydb_pg/async_vectorstore.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import base64
+import logging
 import re
 from typing import Any, Optional
 
@@ -26,6 +27,8 @@
 from langchain_postgres.v2.async_vectorstore import AsyncPGVectorStore
 from sqlalchemy import text
 
+logger = logging.getLogger(__name__)
+
 
 class AsyncAlloyDBVectorStore(AsyncPGVectorStore):
     """Google AlloyDB Vector Store class"""
@@ -134,8 +137,12 @@ async def asimilarity_search_image(
             embedding=embedding, k=k, filter=filter, **kwargs
         )
 
-    async def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> None:
+    async def aset_maintenance_work_mem(
+        self, num_leaves: Optional[int], vector_size: int
+    ) -> None:
         """Set database maintenance work memory (for ScaNN index creation)."""
+        if not num_leaves:
+            return
         # Required index memory in MB
         buffer = 1
         index_memory_required = (
@@ -146,6 +153,125 @@ async def set_maintenance_work_mem(self, num_leaves: int, vector_size: int) -> N
             await conn.execute(text(query))
             await conn.commit()
 
+    set_maintenance_work_mem = aset_maintenance_work_mem
+
+    async def ainitialize_auto_vector_embeddings(
+        self,
+        model_id: str,
+        content_column: Optional[str] = None,
+        embedding_column: Optional[str] = None,
+        schema_name: Optional[str] = None,
+    ) -> None:
+        """Asynchronously initialize auto vector embeddings.
+
+        Args:
+            model_id: The ID of the model to use for embeddings.
+            content_column: Optional name of the content column. Defaults to self.content_column.
+            embedding_column: Optional name of the embedding column. Defaults to self.embedding_column.
+            schema_name: Optional name of the database schema. Defaults to self.schema_name.
+        """
+        content_col = content_column or self.content_column
+        embedding_col = embedding_column or self.embedding_column
+        schema = schema_name or getattr(self, "schema_name", "public")
+
+        if not content_col:
+            raise ValueError(
+                "content_column must be provided or configured on the vector store."
+            )
+        if not embedding_col:
+            raise ValueError(
+                "embedding_column must be provided or configured on the vector store."
+            )
+
+        table_identifier = (
+            f'"{schema}"."{self.table_name}"' if schema else f'"{self.table_name}"'
+        )
+        query = "CALL ai.initialize_embeddings(:model_id, :table_name, :content_column, :embedding_column)"
+        async with self.engine.connect() as conn:
+            await conn.execute(
+                text(query),
+                {
+                    "model_id": model_id,
+                    "table_name": table_identifier,
+                    "content_column": content_col,
+                    "embedding_column": embedding_col,
+                },
+            )
+            await conn.commit()
+
+    async def aenable_columnar_engine(
+        self,
+        columns: Optional[list[str]] = None,
+    ) -> None:
+        """Asynchronously add the table and its columns to the columnar engine.
+
+        Args:
+            columns: Optional list of column names to add to the columnar engine.
+        """
+        if columns:
+            columns_str = ",".join(columns)
+            query = "SELECT google_columnar_engine_add(relation => :table_name, columns => :columns)"
+            params = {"table_name": self.table_name, "columns": columns_str}
+        else:
+            query = "SELECT google_columnar_engine_add(:table_name)"
+            params = {"table_name": self.table_name}
+
+        async with self.engine.connect() as conn:
+            await conn.execute(text(query), params)
+            await conn.commit()
+
+    async def aenable_auto_columnarization(self) -> None:
+        """Asynchronously trigger auto-columnarization recommendations."""
+        query = "SELECT google_columnar_engine_recommend('AUTO_COLUMNARIZATION')"
+        async with self.engine.connect() as conn:
+            await conn.execute(text(query))
+            await conn.commit()
+
+    async def adefine_vector_assist_spec(self) -> list[dict]:
+        """Asynchronously define a Vector Assist spec for the current table."""
+        query = "SELECT * FROM vector_assist.define_spec(table_name => :table_name, vector_column_name => :embedding_column)"
+        params = {
+            "table_name": self.table_name,
+            "embedding_column": self.embedding_column,
+        }
+        async with self.engine.connect() as conn:
+            result = await conn.execute(text(query), params)
+            return [dict(row) for row in result.mappings()]
+
+    async def aapply_vector_assist_spec(self) -> list[dict]:
+        """Asynchronously apply the Vector Assist spec for the current table."""
+        query = "SELECT * FROM vector_assist.apply_spec(table_name => :table_name, column_name => :embedding_column)"
+        params = {
+            "table_name": self.table_name,
+            "embedding_column": self.embedding_column,
+        }
+        async with self.engine.connect() as conn:
+            result = await conn.execute(text(query), params)
+            return [dict(row) for row in result.mappings()]
+
+    async def aget_vector_assist_recommendations(self) -> list[dict]:
+        """Asynchronously get Vector Assist recommendations for the current table."""
+        # First we need to get the spec ID for the current table
+        specs = await self.adefine_vector_assist_spec()
+        if not specs:
+            logger.warning(
+                "No vector assist spec found for table '%s'.", self.table_name
+            )
+            return []
+
+        spec_id = specs[0].get("vector_spec_id")
+        if spec_id is None:
+            logger.warning(
+                "Vector assist spec for table '%s' does not contain 'vector_spec_id'.",
+                self.table_name,
+            )
+            return []
+
+        query = "SELECT * FROM vector_assist.get_recommendations(:spec_id)"
+        async with self.engine.connect() as conn:
+            result = await conn.execute(text(query), {"spec_id": spec_id})
+            return [dict(row) for row in result.mappings()]
+
     def add_images(
         self,
         uris: list[str],

src/langchain_google_alloydb_pg/engine.py

@@ -621,6 +621,127 @@ def init_checkpoint_table(
         """
         self._run_as_sync(self._ainit_checkpoint_table(table_name, schema_name))
 
+    async def _aforecast(
+        self,
+        model_id: str,
+        source_table: str,
+        timestamp_col: str,
+        data_col: str,
+        horizon: int,
+        source_query: Optional[str] = None,
+        conf_level: Optional[float] = None,
+    ) -> list[dict]:
+        if not model_id:
+            raise ValueError("model_id must be provided.")
+        if not source_table:
+            raise ValueError("source_table must be provided.")
+        if not timestamp_col:
+            raise ValueError("timestamp_col must be provided.")
+        if not data_col:
+            raise ValueError("data_col must be provided.")
+        if horizon <= 0:
+            raise ValueError("horizon must be a positive integer.")
+        if conf_level is not None and not (0 < conf_level < 1):
+            raise ValueError("conf_level must be between 0 and 1.")
+
+        args = [
+            "model_id => :model_id",
+            "source_table => :source_table",
+            "timestamp_col => :timestamp_col",
+            "data_col => :data_col",
+            "horizon => :horizon",
+        ]
+        params: dict[str, Any] = {
+            "model_id": model_id,
+            "source_table": source_table,
+            "timestamp_col": timestamp_col,
+            "data_col": data_col,
+            "horizon": horizon,
+        }
+        if source_query is not None:
+            args.append("source_query => :source_query")
+            params["source_query"] = source_query
+        if conf_level is not None:
+            args.append("conf_level => :conf_level")
+            params["conf_level"] = conf_level
+
+        query = f"SELECT * FROM google_ml.forecast({', '.join(args)})"
+        async with self._pool.connect() as conn:
+            result = await conn.execute(text(query), params)
+            return [dict(row) for row in result.mappings()]
+
+    async def aforecast(
+        self,
+        model_id: str,
+        source_table: str,
+        timestamp_col: str,
+        data_col: str,
+        horizon: int,
+        source_query: Optional[str] = None,
+        conf_level: Optional[float] = None,
+    ) -> list[dict]:
+        """Asynchronously get forecasting from AlloyDB AI.
+
+        Args:
+            model_id: The ID of the time series forecasting model.
+            source_table: The table to read historical time series data from.
+            timestamp_col: The column containing the timestamp.
+            data_col: The column containing the data to forecast.
+            horizon: Number of future time steps to forecast.
+            source_query: Optional query to filter historical data.
+            conf_level: Optional confidence level for prediction intervals.
+
+        Returns:
+            A list of dictionaries with forecast_timestamp, forecast_value, and intervals.
+        """
+        return await self._run_as_async(
+            self._aforecast(
+                model_id,
+                source_table,
+                timestamp_col,
+                data_col,
+                horizon,
+                source_query,
+                conf_level,
+            )
+        )
+
+    def forecast(
+        self,
+        model_id: str,
+        source_table: str,
+        timestamp_col: str,
+        data_col: str,
+        horizon: int,
+        source_query: Optional[str] = None,
+        conf_level: Optional[float] = None,
+    ) -> list[dict]:
+        """Synchronously get forecasting from AlloyDB AI.
+
+        Args:
+            model_id: The ID of the time series forecasting model.
+            source_table: The table to read historical time series data from.
+            timestamp_col: The column containing the timestamp.
+            data_col: The column containing the data to forecast.
+            horizon: Number of future time steps to forecast.
+            source_query: Optional query to filter historical data.
+            conf_level: Optional confidence level for prediction intervals.
+
+        Returns:
+            A list of dictionaries with forecast_timestamp, forecast_value, and intervals.
+        """
+        return self._run_as_sync(
+            self._aforecast(
+                model_id,
+                source_table,
+                timestamp_col,
+                data_col,
+                horizon,
+                source_query,
+                conf_level,
+            )
+        )
+
     async def _aload_table_schema(
         self, table_name: str, schema_name: str = "public"
     ) -> Table:

src/langchain_google_alloydb_pg/indexes.py

@@ -14,6 +14,7 @@
 
 import warnings
 from dataclasses import dataclass, field
+from typing import Optional
 
 from langchain_postgres.v2.indexes import (
     DEFAULT_DISTANCE_STRATEGY,
@@ -62,15 +63,31 @@ def to_string(self) -> str:
 
 @dataclass
 class ScaNNIndex(BaseIndex):
+    """ScaNN index configuration for AlloyDB.
+
+    Args:
+        mode (Optional[str]): Index mode (e.g. 'AUTO' for auto-tuned indexing). Defaults to None.
+        num_leaves (Optional[int]): Number of leaves in index clusters. Defaults to 5.
+        quantizer (str): Quantizer type. Defaults to 'sq8'.
+        extension_name (str): Extension name. Defaults to 'alloydb_scann'.
+    """
+
     index_type: str = "ScaNN"
-    num_leaves: int = 5
+    mode: Optional[str] = None
+    num_leaves: Optional[int] = 5
     quantizer: str = field(
         default="sq8", init=False
     )  # Disable `quantizer` initialization currently only supports the value "sq8"
     extension_name: str = "alloydb_scann"
 
     def index_options(self) -> str:
         """Set index query options for vector store initialization."""
+        if self.mode is not None:
+            if self.mode.upper() != "AUTO":
+                raise ValueError(
+                    f"Invalid mode '{self.mode}'. Only mode='AUTO' is currently supported."
+                )
+            return "(mode = 'AUTO')"
         return f"(num_leaves = {self.num_leaves}, quantizer = {self.quantizer})"
 
     def get_index_function(self) -> str:
@@ -84,20 +101,41 @@ def get_index_function(self) -> str:
 
 @dataclass
 class ScaNNQueryOptions(QueryOptions):
-    num_leaves_to_search: int = 1
+    """Query options for ScaNN index.
+
+    Args:
+        num_leaves_to_search (Optional[int]): Absolute number of leaves to search. Defaults to 1.
+        pre_reordering_num_neighbors (int): Number of neighbors to consider before reordering. Defaults to -1.
+        pct_leaves_to_search (Optional[float]): Percentage of leaves to search (0.0 to 1.0 or proportion).
+            When specified, this takes precedence over `num_leaves_to_search`.
+    """
+
+    num_leaves_to_search: Optional[int] = 1
     pre_reordering_num_neighbors: int = -1
+    pct_leaves_to_search: Optional[float] = None
 
     def to_parameter(self) -> list[str]:
         """Convert index attributes to list of configurations."""
-        return [
-            f"scann.num_leaves_to_search = {self.num_leaves_to_search}",
-            f"scann.pre_reordering_num_neighbors = {self.pre_reordering_num_neighbors}",
-        ]
+        params = []
+        if self.pct_leaves_to_search is not None:
+            if self.num_leaves_to_search is not None and self.num_leaves_to_search != 1:
+                warnings.warn(
+                    "Both 'pct_leaves_to_search' and 'num_leaves_to_search' were provided. "
+                    "'pct_leaves_to_search' takes precedence.",
+                    UserWarning,
+                )
+            params.append(f"scann.pct_leaves_to_search = {self.pct_leaves_to_search}")
+        elif self.num_leaves_to_search is not None:
+            params.append(f"scann.num_leaves_to_search = {self.num_leaves_to_search}")
+        params.append(
+            f"scann.pre_reordering_num_neighbors = {self.pre_reordering_num_neighbors}"
+        )
+        return params
 
     def to_string(self) -> str:
         """Convert index attributes to string."""
         warnings.warn(
             "to_string is deprecated, use to_parameter instead.",
             DeprecationWarning,
         )
-        return f"scann.num_leaves_to_search = {self.num_leaves_to_search}, scann.pre_reordering_num_neighbors = {self.pre_reordering_num_neighbors}"
+        return ", ".join(self.to_parameter())