feat: add has_nodes_batch with UNWIND implementation for Neo4j

Author: ndcorder

lightrag/base.py

@@ -595,6 +595,25 @@ async def batch_upsert_edges(
         for src, tgt, edge_data in edges:
             await self.upsert_edge(src, tgt, edge_data)
 
+    async def has_nodes_batch(self, node_ids: list[str]) -> set[str]:
+        """Check which nodes exist. Returns set of existing node_ids.
+
+        Default implementation calls has_node for each ID.
+        Override this method for better performance in storage backends
+        that support batch operations.
+
+        Args:
+            node_ids: List of node IDs to check
+
+        Returns:
+            Set of node_ids that exist in the graph
+        """
+        result = set()
+        for node_id in node_ids:
+            if await self.has_node(node_id):
+                result.add(node_id)
+        return result
+
     @abstractmethod
     async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
         """Insert a new node or update an existing node in the graph.

lightrag/kg/neo4j_impl.py

@@ -1213,6 +1213,33 @@ async def _execute(tx, params=batch_params):
             logger.error(f"[{self.workspace}] Error during batch edge upsert: {str(e)}")
             raise
 
+    async def has_nodes_batch(self, node_ids: list[str]) -> set[str]:
+        """Batch check node existence using UNWIND."""
+        if not node_ids:
+            return set()
+
+        workspace_label = self._get_workspace_label()
+        result = set()
+
+        try:
+            async with self._driver.session(
+                database=self._DATABASE, default_access_mode="READ"
+            ) as session:
+                query = (
+                    "UNWIND $ids AS eid "
+                    f"MATCH (n:`{workspace_label}` {{entity_id: eid}}) "
+                    "RETURN n.entity_id AS entity_id"
+                )
+                records = await session.run(query, ids=node_ids)
+                async for record in records:
+                    result.add(record["entity_id"])
+                await records.consume()
+        except Exception as e:
+            logger.error(f"[{self.workspace}] Error during batch has_nodes: {str(e)}")
+            raise
+
+        return result
+
     async def get_knowledge_graph(
         self,
         node_label: str,

tests/test_batch_neo4j.py

@@ -101,3 +101,49 @@ async def test_batch_upsert_nodes_large_batch_chunked(neo4j_storage):
 
     # 1200 / 500 = 3 sub-batches
     assert mock_tx.run.call_count == 3
+
+
+@pytest.mark.asyncio
+async def test_has_nodes_batch(neo4j_storage):
+    storage, mock_tx = neo4j_storage
+
+    # Mock the session for read access
+    mock_records = [{"entity_id": "alice"}, {"entity_id": "charlie"}]
+
+    class MockResult:
+        def __aiter__(self):
+            return self
+
+        def __init__(self, records):
+            self._records = iter(records)
+
+        async def __anext__(self):
+            try:
+                return next(self._records)
+            except StopIteration:
+                raise StopAsyncIteration
+
+        async def consume(self):
+            pass
+
+    mock_result = MockResult(mock_records)
+
+    mock_session = AsyncMock()
+    mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+    mock_session.__aexit__ = AsyncMock(return_value=False)
+    mock_session.run = AsyncMock(return_value=mock_result)
+
+    storage._driver.session = MagicMock(return_value=mock_session)
+
+    result = await storage.has_nodes_batch(["alice", "bob", "charlie"])
+    assert result == {"alice", "charlie"}
+    # Verify UNWIND query was used
+    query = mock_session.run.call_args[0][0]
+    assert "UNWIND" in query
+
+
+@pytest.mark.asyncio
+async def test_has_nodes_batch_empty(neo4j_storage):
+    storage, mock_tx = neo4j_storage
+    result = await storage.has_nodes_batch([])
+    assert result == set()

tests/test_batch_upsert_base.py

@@ -71,3 +71,11 @@ async def test_batch_upsert_nodes_empty_list(storage):
 async def test_batch_upsert_edges_empty_list(storage):
     await storage.batch_upsert_edges([])
     storage.upsert_edge.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_has_nodes_batch_default(storage):
+    storage.has_node = AsyncMock(side_effect=[True, False, True])
+    result = await storage.has_nodes_batch(["a", "b", "c"])
+    assert result == {"a", "c"}
+    assert storage.has_node.call_count == 3