feat: Add support for ANY/ALL postgres operators (#19553)

Author: rodcanales

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/llama_index/vector_stores/postgres/base.py

@@ -644,6 +644,10 @@ def _to_postgres_operator(self, operator: FilterOperator) -> str:
             return "ILIKE"
         elif operator == FilterOperator.IS_EMPTY:
             return "IS NULL"
+        elif operator == FilterOperator.ANY:
+            return "?|"
+        elif operator == FilterOperator.ALL:
+            return "?&"
         else:
             _logger.warning(f"Unknown operator: {operator}, fallback to '='")
             return "="
@@ -663,6 +667,17 @@ def _build_filter_clause(self, filter_: MetadataFilter) -> Any:
                 f"{self._to_postgres_operator(filter_.operator)} "
                 f"({filter_value})"
             )
+        elif filter_.operator in [FilterOperator.ANY, FilterOperator.ALL]:
+            # Expects a list stored in the metadata, and a single value to compare
+
+            # We apply same logic as above, but as an array
+            filter_value = ", ".join(f"'{e}'" for e in filter_.value)
+
+            return text(
+                f"metadata_::jsonb->'{filter_.key}' "
+                f"{self._to_postgres_operator(filter_.operator)} "
+                f"array[{filter_value}]"
+            )
         elif filter_.operator == FilterOperator.CONTAINS:
             # Expects a list stored in the metadata, and a single value to compare
             return text(

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/pyproject.toml

@@ -27,7 +27,7 @@ dev = [
 
 [project]
 name = "llama-index-vector-stores-postgres"
-version = "0.5.4"
+version = "0.5.5"
 description = "llama-index vector_stores postgres integration"
 authors = [{name = "Your Name", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/tests/test_postgres.py

@@ -222,7 +222,7 @@ def node_embeddings() -> List[TextNode]:
             text="consectetur adipiscing elit",
             id_="ccc",
             relationships={NodeRelationship.SOURCE: RelatedNodeInfo(node_id="ccc")},
-            extra_info={"test_key_list": ["test_value"]},
+            extra_info={"test_key_list": ["test_value_1", "test_value_2"]},
             embedding=_get_sample_vector(0.1),
         ),
         TextNode(
@@ -518,6 +518,91 @@ async def test_add_to_db_and_query_with_metadata_filters_with_in_operator_and_si
     assert res.nodes[0].node_id == "bbb"
 
 
+@pytest.mark.skipif(postgres_not_available, reason="postgres db is not available")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("pg_fixture", ["pg", "pg_halfvec"], indirect=True)
+@pytest.mark.parametrize("use_async", [True, False])
+async def test_add_to_db_and_query_with_metadata_filters_with_any_operator(
+    pg_fixture: PGVectorStore, node_embeddings: List[TextNode], use_async: bool
+) -> None:
+    if use_async:
+        await pg_fixture.async_add(node_embeddings)
+    else:
+        pg_fixture.add(node_embeddings)
+    assert isinstance(pg_fixture, PGVectorStore)
+    assert hasattr(pg_fixture, "_engine")
+    filters = MetadataFilters(
+        filters=[
+            MetadataFilter(
+                key="test_key_list",
+                value=["test_value_1", "test_value_new"],
+                operator=FilterOperator.ANY,
+            )
+        ]
+    )
+    q = VectorStoreQuery(
+        query_embedding=_get_sample_vector(0.5), similarity_top_k=10, filters=filters
+    )
+    if use_async:
+        res = await pg_fixture.aquery(q)
+    else:
+        res = pg_fixture.query(q)
+    assert res.nodes
+    assert len(res.nodes) == 1
+    assert res.nodes[0].node_id == "ccc"
+
+
+@pytest.mark.skipif(postgres_not_available, reason="postgres db is not available")
+@pytest.mark.asyncio
+@pytest.mark.parametrize("pg_fixture", ["pg", "pg_halfvec"], indirect=True)
+@pytest.mark.parametrize("use_async", [True, False])
+async def test_add_to_db_and_query_with_metadata_filters_with_all_operator(
+    pg_fixture: PGVectorStore, node_embeddings: List[TextNode], use_async: bool
+) -> None:
+    if use_async:
+        await pg_fixture.async_add(node_embeddings)
+    else:
+        pg_fixture.add(node_embeddings)
+    assert isinstance(pg_fixture, PGVectorStore)
+    assert hasattr(pg_fixture, "_engine")
+    filters = MetadataFilters(
+        filters=[
+            MetadataFilter(
+                key="test_key_list",
+                value=["test_value_1", "test_value_2"],
+                operator=FilterOperator.ALL,
+            )
+        ]
+    )
+    filters_no_all_match = MetadataFilters(
+        filters=[
+            MetadataFilter(
+                key="test_key_list",
+                value=["test_value_1", "test_value_3"],
+                operator=FilterOperator.ALL,
+            )
+        ]
+    )
+    q = VectorStoreQuery(
+        query_embedding=_get_sample_vector(0.5), similarity_top_k=10, filters=filters
+    )
+    q2 = VectorStoreQuery(
+        query_embedding=_get_sample_vector(0.5),
+        similarity_top_k=10,
+        filters=filters_no_all_match,
+    )
+    if use_async:
+        res = await pg_fixture.aquery(q)
+        res_no_match = await pg_fixture.aquery(q2)
+    else:
+        res = pg_fixture.query(q)
+        res_no_match = pg_fixture.query(q2)
+    assert res.nodes
+    assert len(res.nodes) == 1
+    assert res.nodes[0].node_id == "ccc"
+    assert not res_no_match.nodes
+
+
 @pytest.mark.skipif(postgres_not_available, reason="postgres db is not available")
 @pytest.mark.asyncio
 @pytest.mark.parametrize("pg_fixture", ["pg", "pg_halfvec"], indirect=True)
@@ -535,7 +620,7 @@ async def test_add_to_db_and_query_with_metadata_filters_with_contains_operator(
         filters=[
             MetadataFilter(
                 key="test_key_list",
-                value="test_value",
+                value="test_value_1",
                 operator=FilterOperator.CONTAINS,
             )
         ]