fix: FalkorDB changing pagination to standard method (#3688)

Author: davidsbatista

integrations/falkordb/src/haystack_integrations/document_stores/falkordb/document_store.py

@@ -607,25 +607,24 @@ def get_metadata_field_unique_values(
         self,
         metadata_field: str,
         search_term: str | None = None,
-        size: int | None = 10,
-        after: dict[str, Any] | None = None,
-    ) -> tuple[list[Any], dict[str, Any] | None]:
+        from_: int = 0,
+        size: int = 10,
+    ) -> tuple[list[str], int]:
         """
         Return distinct values for the given metadata field with optional filtering and pagination.
 
         :param metadata_field: Metadata field name. May include or omit the `meta.` prefix.
         :param search_term: Optional case-insensitive substring filter applied to the metadata
             field's own value.
-        :param size: Maximum number of values to return per page. Defaults to 10 000.
-        :param after: Pagination cursor returned by a previous call. Pass `None` for the first page.
-        :returns: Tuple of `(values, next_cursor)`. `next_cursor` is `None` on the last page.
+        :param from_: The offset for pagination (0-based).
+        :param size: Maximum number of values to return per page. Defaults to 10.
+        :returns: Tuple of `(values, total_count)`. `total_count` is the number of distinct
+            values matching the filter, independent of pagination.
         """
         self._ensure_connected()
         field = metadata_field[5:] if metadata_field.startswith("meta.") else metadata_field
-        offset = after.get("offset", 0) if after else 0
-        limit = size if size is not None else 10000
 
-        query_params: dict[str, Any] = {}
+        query_params: dict[str, Any] = {"from_": from_, "size": size}
         where_parts = [f"d.{field} IS NOT NULL"]
         if search_term:
             where_parts.append(f"toLower(toString(d.{field})) CONTAINS toLower($search_term)")
@@ -634,15 +633,16 @@ def get_metadata_field_unique_values(
         where = " AND ".join(where_parts)
         cypher = (
             f"MATCH (d:{self.node_label}) WHERE {where} "
-            f"RETURN DISTINCT d.{field} AS val "
+            f"WITH DISTINCT d.{field} AS val "
             f"ORDER BY val "
-            f"SKIP {offset} LIMIT {limit + 1}"
+            f"WITH collect(val) AS vals "
+            f"RETURN vals[$from_..$from_ + $size] AS page, size(vals) AS total"
         )
         result = self.graph.query(cypher, query_params)
-        rows = result.result_set
-        values = [row[0] for row in rows[:limit]]
-        next_cursor: dict[str, Any] | None = {"offset": offset + limit} if len(rows) > limit else None
-        return values, next_cursor
+        if not result.result_set:
+            return [], 0
+        page, total = result.result_set[0]
+        return list(page), total
 
     # ------------------------------------------------------------------
     # Internal retrieval helpers (called by retriever components)

integrations/falkordb/tests/test_document_store.py

@@ -371,26 +371,39 @@ def test_get_metadata_field_min_max(self, mock_falkordb, rows, expected):
 
     def test_get_metadata_field_unique_values(self, mock_falkordb):
         _, _, graph = mock_falkordb
-        graph.query.side_effect = [_result([]), _result([]), _result([["A"], ["B"], ["C"]])]
-        values, cursor = FalkorDBDocumentStore().get_metadata_field_unique_values("category", size=10)
+        graph.query.side_effect = [_result([]), _result([]), _result([[["A", "B", "C"], 3]])]
+        values, total = FalkorDBDocumentStore().get_metadata_field_unique_values("category", size=10)
         assert values == ["A", "B", "C"]
-        assert cursor is None
+        assert total == 3
 
     def test_get_metadata_field_unique_values_pagination(self, mock_falkordb):
         _, _, graph = mock_falkordb
-        graph.query.side_effect = [_result([]), _result([]), _result([["A"], ["B"], ["C"]])]
-        values, cursor = FalkorDBDocumentStore().get_metadata_field_unique_values("category", size=2)
-        assert values == ["A", "B"]
-        assert cursor == {"offset": 2}
+        all_values = ["A", "B", "C", "D", "E"]
+        total = len(all_values)
+        pages = [all_values[i : i + 2] for i in range(0, total, 2)]
+        graph.query.side_effect = [_result([]), _result([]), *(_result([[page, total]]) for page in pages)]
+        store = FalkorDBDocumentStore()
+
+        for page_index, expected_page in enumerate(pages):
+            from_ = page_index * 2
+            values, returned_total = store.get_metadata_field_unique_values("category", from_=from_, size=2)
+            assert values == expected_page
+            assert returned_total == total
+            _, params = graph.query.call_args[0]
+            assert params["from_"] == from_
+            assert params["size"] == 2
+
+        # last page is a partial page, smaller than the requested size
+        assert len(pages[-1]) == 1
 
     def test_get_metadata_field_unique_values_search_term_case_insensitive(self, mock_falkordb):
         _, _, graph = mock_falkordb
-        graph.query.side_effect = [_result([]), _result([]), _result([["Apple"]])]
+        graph.query.side_effect = [_result([]), _result([]), _result([[["Apple"], 1]])]
         values, _ = FalkorDBDocumentStore().get_metadata_field_unique_values("category", search_term="APP")
         assert values == ["Apple"]
         cypher, params = graph.query.call_args[0]
         assert "toLower(toString(d.category)) CONTAINS toLower($search_term)" in cypher
-        assert params == {"search_term": "APP"}
+        assert params == {"from_": 0, "size": 10, "search_term": "APP"}
 
     def test_close(self):
         store = FalkorDBDocumentStore()