feat(v2): add validate_schema option and skip embedding column in similarity search

langchain_postgres/v2/async_vectorstore.py

@@ -155,6 +155,7 @@ async def create(
         lambda_mult: float = 0.5,
         index_query_options: Optional[QueryOptions] = None,
         hybrid_search_config: Optional[HybridSearchConfig] = None,
+        validate_schema: bool = True,
     ) -> AsyncPGVectorStore:
         """Create an AsyncPGVectorStore instance.
 
@@ -175,6 +176,14 @@ async def create(
             lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
             index_query_options (QueryOptions): Index query option.
             hybrid_search_config (HybridSearchConfig): Hybrid search configuration. Defaults to None.
+            validate_schema (bool): Whether to query `information_schema.columns` to validate
+                that the configured columns exist and have compatible types. Defaults to True.
+                Set to False to skip that round trip (e.g. when constructing a short-lived store
+                per request, such as in a serverless/edge function, against a table whose schema
+                is already known to be valid). Can not be used together with
+                `ignore_metadata_columns`, since resolving it requires knowing the full column
+                list. When False, `metadata_json_column` and `hybrid_search_config.tsv_column`
+                are trusted as given instead of being probed for existence.
 
         Returns:
             AsyncPGVectorStore
@@ -187,64 +196,77 @@ async def create(
             raise ValueError(
                 "Can not use both metadata_columns and ignore_metadata_columns."
             )
-        # Get field type information
-        stmt = "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = :table_name AND table_schema = :schema_name"
-        async with engine._pool.connect() as conn:
-            result = await conn.execute(
-                text(stmt),
-                {"table_name": table_name, "schema_name": schema_name},
-            )
-            result_map = result.mappings()
-            results = result_map.fetchall()
-        columns = {}
-        for field in results:
-            columns[field["column_name"]] = field["data_type"]
-
-        # Check columns
-        if id_column not in columns:
-            raise ValueError(f"Id column, {id_column}, does not exist.")
-        if content_column not in columns:
-            raise ValueError(f"Content column, {content_column}, does not exist.")
-        content_type = columns[content_column]
-        if content_type != "text" and "char" not in content_type:
+        if not validate_schema and ignore_metadata_columns:
             raise ValueError(
-                f"Content column, {content_column}, is type, {content_type}. It must be a type of character string."
-            )
-        if hybrid_search_config:
-            tsv_column_name = (
-                hybrid_search_config.tsv_column
-                if hybrid_search_config.tsv_column
-                else content_column + "_tsv"
-            )
-            if tsv_column_name not in columns or columns[tsv_column_name] != "tsvector":
-                # mark tsv_column as empty because there is no TSV column in table
-                hybrid_search_config.tsv_column = ""
-        if embedding_column not in columns:
-            raise ValueError(f"Embedding column, {embedding_column}, does not exist.")
-        if columns[embedding_column] not in ["USER-DEFINED", "vector"]:
-            raise ValueError(
-                f"Embedding column, {embedding_column}, is not type Vector."
+                "Can not use both validate_schema=False and ignore_metadata_columns, "
+                "since resolving ignore_metadata_columns requires querying the full "
+                "column list. Pass metadata_columns explicitly instead."
             )
 
-        metadata_json_column = (
-            None if metadata_json_column not in columns else metadata_json_column
-        )
+        if validate_schema:
+            # Get field type information
+            stmt = "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = :table_name AND table_schema = :schema_name"
+            async with engine._pool.connect() as conn:
+                result = await conn.execute(
+                    text(stmt),
+                    {"table_name": table_name, "schema_name": schema_name},
+                )
+                result_map = result.mappings()
+                results = result_map.fetchall()
+            columns = {}
+            for field in results:
+                columns[field["column_name"]] = field["data_type"]
+
+            # Check columns
+            if id_column not in columns:
+                raise ValueError(f"Id column, {id_column}, does not exist.")
+            if content_column not in columns:
+                raise ValueError(f"Content column, {content_column}, does not exist.")
+            content_type = columns[content_column]
+            if content_type != "text" and "char" not in content_type:
+                raise ValueError(
+                    f"Content column, {content_column}, is type, {content_type}. It must be a type of character string."
+                )
+            if hybrid_search_config:
+                tsv_column_name = (
+                    hybrid_search_config.tsv_column
+                    if hybrid_search_config.tsv_column
+                    else content_column + "_tsv"
+                )
+                if (
+                    tsv_column_name not in columns
+                    or columns[tsv_column_name] != "tsvector"
+                ):
+                    # mark tsv_column as empty because there is no TSV column in table
+                    hybrid_search_config.tsv_column = ""
+            if embedding_column not in columns:
+                raise ValueError(
+                    f"Embedding column, {embedding_column}, does not exist."
+                )
+            if columns[embedding_column] not in ["USER-DEFINED", "vector"]:
+                raise ValueError(
+                    f"Embedding column, {embedding_column}, is not type Vector."
+                )
+
+            metadata_json_column = (
+                None if metadata_json_column not in columns else metadata_json_column
+            )
 
-        # If using metadata_columns check to make sure column exists
-        for column in metadata_columns:
-            if column not in columns:
-                raise ValueError(f"Metadata column, {column}, does not exist.")
+            # If using metadata_columns check to make sure column exists
+            for column in metadata_columns:
+                if column not in columns:
+                    raise ValueError(f"Metadata column, {column}, does not exist.")
 
-        # If using ignore_metadata_columns, filter out known columns and set known metadata columns
-        all_columns = columns
-        if ignore_metadata_columns:
-            for column in ignore_metadata_columns:
-                del all_columns[column]
+            # If using ignore_metadata_columns, filter out known columns and set known metadata columns
+            all_columns = columns
+            if ignore_metadata_columns:
+                for column in ignore_metadata_columns:
+                    del all_columns[column]
 
-            del all_columns[id_column]
-            del all_columns[content_column]
-            del all_columns[embedding_column]
-            metadata_columns = [k for k in all_columns.keys()]
+                del all_columns[id_column]
+                del all_columns[content_column]
+                del all_columns[embedding_column]
+                metadata_columns = [k for k in all_columns.keys()]
 
         return cls(
             cls.__create_key,
@@ -649,11 +671,15 @@ async def __query_collection(
         operator = self.distance_strategy.operator
         search_function = self.distance_strategy.search_function
 
-        columns = [
-            self.id_column,
-            self.content_column,
-            self.embedding_column,
-        ] + self.metadata_columns
+        # The embedding column itself is only read back by callers that need the
+        # matched rows' raw vectors (e.g. MMR); plain similarity search only needs
+        # `distance`, which is computed via `search_function` below regardless of
+        # whether the column is also included in the SELECT list.
+        include_embedding = kwargs.get("include_embedding", True)
+        columns = [self.id_column, self.content_column]
+        if include_embedding:
+            columns.append(self.embedding_column)
+        columns += self.metadata_columns
         if self.metadata_json_column:
             columns.append(self.metadata_json_column)
 
@@ -851,6 +877,9 @@ async def asimilarity_search_with_score_by_vector(
         **kwargs: Any,
     ) -> list[tuple[Document, float]]:
         """Return docs and distance scores selected by vector similarity search."""
+        # The embedding vector of each matched row is never read below, so skip
+        # selecting it (a caller can still opt back in via include_embedding=True).
+        kwargs.setdefault("include_embedding", False)
         results = await self.__query_collection(
             embedding=embedding, k=k, filter=filter, **kwargs
         )

langchain_postgres/v2/vectorstores.py

@@ -61,6 +61,7 @@ async def create(
         lambda_mult: float = 0.5,
         index_query_options: Optional[QueryOptions] = None,
         hybrid_search_config: Optional[HybridSearchConfig] = None,
+        validate_schema: bool = True,
     ) -> PGVectorStore:
         """Create an PGVectorStore instance.
 
@@ -81,6 +82,9 @@ async def create(
             lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
             index_query_options (QueryOptions): Index query option.
             hybrid_search_config (HybridSearchConfig): Hybrid search configuration. Defaults to None.
+            validate_schema (bool): Whether to query `information_schema.columns` to validate
+                the configured columns. Defaults to True. See
+                `AsyncPGVectorStore.create` for details on when to disable it.
 
         Returns:
             PGVectorStore
@@ -102,6 +106,7 @@ async def create(
             lambda_mult=lambda_mult,
             index_query_options=index_query_options,
             hybrid_search_config=hybrid_search_config,
+            validate_schema=validate_schema,
         )
         vs = await engine._run_as_async(coro)
         return cls(cls.__create_key, engine, vs)
@@ -125,6 +130,7 @@ def create_sync(
         lambda_mult: float = 0.5,
         index_query_options: Optional[QueryOptions] = None,
         hybrid_search_config: Optional[HybridSearchConfig] = None,
+        validate_schema: bool = True,
     ) -> PGVectorStore:
         """Create an PGVectorStore instance.
 
@@ -146,6 +152,9 @@ def create_sync(
             lambda_mult (float, optional): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
             index_query_options (Optional[QueryOptions], optional): Index query option. Defaults to None.
             hybrid_search_config (HybridSearchConfig): Hybrid search configuration. Defaults to None.
+            validate_schema (bool): Whether to query `information_schema.columns` to validate
+                the configured columns. Defaults to True. See
+                `AsyncPGVectorStore.create` for details on when to disable it.
 
         Returns:
             PGVectorStore
@@ -167,6 +176,7 @@ def create_sync(
             lambda_mult=lambda_mult,
             index_query_options=index_query_options,
             hybrid_search_config=hybrid_search_config,
+            validate_schema=validate_schema,
         )
         vs = engine._run_as_sync(coro)
         return cls(cls.__create_key, engine, vs)

tests/unit_tests/v2/test_async_pg_vectorstore.py

@@ -528,6 +528,43 @@ async def test_create_vectorstore_with_invalid_parameters_5(
                 ],  # invalid use of metadata_columns and ignore columns
             )
 
+    async def test_create_with_validate_schema_false(self, engine: PGEngine) -> None:
+        # Skips the information_schema round trip entirely; the store must
+        # still be fully functional against a table whose schema is valid.
+        vs = await AsyncPGVectorStore.create(
+            engine,
+            embedding_service=embeddings_service,
+            table_name=CUSTOM_TABLE,
+            id_column="myid",
+            content_column="mycontent",
+            embedding_column="myembedding",
+            metadata_columns=["page", "source"],
+            metadata_json_column="mymeta",
+            validate_schema=False,
+        )
+        await vs.aadd_texts(
+            ["hello validate_schema=False"],
+            metadatas=[{"page": "9", "source": "test"}],
+        )
+        docs = await vs.asimilarity_search("hello validate_schema=False")
+        assert docs[0].page_content == "hello validate_schema=False"
+        await aexecute(engine, f'TRUNCATE TABLE "{CUSTOM_TABLE}"')
+
+    async def test_create_with_validate_schema_false_and_ignore_metadata_columns(
+        self, engine: PGEngine
+    ) -> None:
+        with pytest.raises(ValueError):
+            await AsyncPGVectorStore.create(
+                engine,
+                embedding_service=embeddings_service,
+                table_name=CUSTOM_TABLE,
+                id_column="myid",
+                content_column="mycontent",
+                embedding_column="myembedding",
+                ignore_metadata_columns=["source"],
+                validate_schema=False,  # can not be combined with ignore_metadata_columns
+            )
+
     async def test_create_vectorstore_with_init(self, engine: PGEngine) -> None:
         with pytest.raises(Exception):
             AsyncPGVectorStore(

tests/unit_tests/v2/test_async_pg_vectorstore_search.py

@@ -253,6 +253,35 @@ async def test_asimilarity_search_by_vector(self, vs: AsyncPGVectorStore) -> Non
         assert result[0][0] == Document(page_content="foo", id=ids[0])
         assert result[0][1] == 0
 
+    async def test_query_collection_include_embedding(
+        self, vs: AsyncPGVectorStore
+    ) -> None:
+        """`__query_collection` defaults to including the embedding column
+        (matching prior behavior), but plain similarity search opts out of it
+        since only MMR ever reads the embedding value back from the row."""
+        embedding = embeddings_service.embed_query("foo")
+
+        rows_with_embedding = await vs._AsyncPGVectorStore__query_collection(  # type: ignore[attr-defined]
+            embedding, k=4
+        )
+        assert vs.embedding_column in rows_with_embedding[0]
+
+        rows_without_embedding = await vs._AsyncPGVectorStore__query_collection(  # type: ignore[attr-defined]
+            embedding, k=4, include_embedding=False
+        )
+        assert vs.embedding_column not in rows_without_embedding[0]
+
+        # asimilarity_search_with_score_by_vector never reads the embedding
+        # column back, so it opts out of selecting it by default.
+        result = await vs.asimilarity_search_with_score_by_vector(embedding)
+        assert result[0][0] == Document(page_content="foo", id=ids[0])
+
+        # ...but a caller can still opt back in via include_embedding=True.
+        result_with_embedding = await vs.asimilarity_search_with_score_by_vector(
+            embedding, include_embedding=True
+        )
+        assert result_with_embedding[0][0] == Document(page_content="foo", id=ids[0])
+
     async def test_similarity_search_with_relevance_scores_threshold_cosine(
         self, vs: AsyncPGVectorStore
     ) -> None: