Improved deeplake.get_nodes() performance (#14920)

Co-authored-by: Nathan Voxland <nathan@voxland.net>
run-llama · Jul 24, 2024 · 3eea9f9 · 3eea9f9
1 parent 2dca097
commit 3eea9f9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/...ctor_stores/llama-index-vector-stores-deeplake/llama_index/vector_stores/deeplake/base.py b/...ctor_stores/llama-index-vector-stores-deeplake/llama_index/vector_stores/deeplake/base.py
@@ -8,7 +8,7 @@
 from typing import Any, List, Optional, cast
 
 from llama_index.core.bridge.pydantic import PrivateAttr
-from llama_index.core.schema import BaseNode, MetadataMode
+from llama_index.core.schema import BaseNode, MetadataMode, TextNode
 from llama_index.core.vector_stores.types import (
     BasePydanticVectorStore,
     VectorStoreQuery,
@@ -149,10 +149,10 @@ def get_nodes(
         filters: Optional[MetadataFilters] = None,
     ) -> List[BaseNode]:
         """Get nodes from vector store."""
-        if not node_ids:
-            data = self._vectorstore.search(filter=lambda x: True)
+        if node_ids:
+            data = self._vectorstore.search(filter={"id": node_ids})
         else:
-            data = self._vectorstore.search(filter=lambda x: x.id.text() in node_ids)
+            data = self._vectorstore.search(filter={})
 
         nodes = []
         for metadata in data["metadata"]:
@@ -196,7 +196,10 @@ def filter_func(doc):
 
             return found_one
 
-        return [x for x in nodes if filter_func(x)]
+        if filters:
+            return [x for x in nodes if filter_func(x)]
+        else:
+            return nodes
 
     def delete_nodes(
         self,
@@ -294,6 +297,8 @@ def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResul
         metadatas = data["metadata"]
         nodes = []
         for metadata in metadatas:
+            if "_node_type" not in metadata:
+                metadata["_node_type"] = TextNode.class_name()
             nodes.append(metadata_dict_to_node(metadata))
 
         return VectorStoreQueryResult(nodes=nodes, similarities=similarities, ids=ids)
diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-deeplake/pyproject.toml
@@ -27,12 +27,12 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-deeplake"
 readme = "README.md"
-version = "0.1.4"
+version = "0.1.5"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
 llama-index-core = "^0.10.1"
-deeplake = ">=3.9.0"
+deeplake = ">=3.9.12"
 
 [tool.poetry.group.dev.dependencies]
 ipython = "8.10.0"