Merge pull request #1182 from weaviate/batch_ref_fix

Dont add references if the .to object has not been added yet

Author: dirkkul

integration/test_batch_v4.py

@@ -6,6 +6,7 @@
 from _pytest.fixtures import SubRequest
 
 import weaviate
+import weaviate.classes as wvc
 from integration.conftest import _sanitize_collection_name
 from weaviate.collections.classes.batch import Shard
 from weaviate.collections.classes.config import (
@@ -587,3 +588,37 @@ def test_uuids_keys_and_original_index(client_factory: ClientFactory) -> None:
     assert [objs[k][0] for k in client.batch.results.objs.uuids.keys()] == list(
         client.batch.results.objs.uuids.values()
     )
+
+
+def test_references_with_to_uuids(client_factory: ClientFactory) -> None:
+    """Test that batch waits until the to object is created."""
+    client, name = client_factory()
+
+    client.collections.delete(["target", "source"])
+    target = client.collections.create(
+        "target", multi_tenancy_config=wvc.config.Configure.multi_tenancy(enabled=True)
+    )
+    source = client.collections.create(
+        "source",
+        references=[wvc.config.ReferenceProperty(name="to", target_collection="target")],
+        multi_tenancy_config=wvc.config.Configure.multi_tenancy(enabled=True),
+    )
+
+    target.tenants.create("tenant-1")
+    source.tenants.create("tenant-1")
+    from_uuid = source.with_tenant("tenant-1").data.insert(properties={})
+    objs = 20
+
+    with client.batch.fixed_size(batch_size=10, concurrent_requests=1) as batch:
+        for _ in range(objs):
+            to = batch.add_object(collection="target", properties={}, tenant="tenant-1")
+            batch.add_reference(
+                from_uuid=from_uuid,
+                from_property="to",
+                to=to,
+                from_collection="source",
+                tenant="tenant-1",
+            )
+
+    assert len(client.batch.failed_references) == 0, client.batch.failed_references
+    client.collections.delete(["target", "source"])

weaviate/collections/batch/base.py

@@ -100,7 +100,9 @@ def pop_items(self, pop_amount: int, uuid_lookup: Set[str]) -> List[_BatchRefere
         i = 0
         self._lock.acquire()
         while len(ret) < pop_amount and len(self._items) > 0 and i < len(self._items):
-            if self._items[i].from_uuid not in uuid_lookup:
+            if self._items[i].from_uuid not in uuid_lookup and (
+                self._items[i].to_uuid is None or self._items[i].to_uuid not in uuid_lookup
+            ):
                 ret.append(self._items.pop(i))
             else:
                 i += 1

weaviate/collections/classes/batch.py

@@ -31,6 +31,7 @@ class _BatchReference:
     to: str
     tenant: Optional[str]
     from_uuid: str
+    to_uuid: Optional[str] = None
 
 
 class BatchObject(BaseModel):
@@ -131,6 +132,7 @@ def _to_internal(self) -> _BatchReference:
             from_uuid=str(self.from_object_uuid),
             from_=f"{BEACON}{self.from_object_collection}/{self.from_object_uuid}/{self.from_property_name}",
             to=f"{BEACON}{self.to_object_collection}{str(self.to_object_uuid)}",
+            to_uuid=str(self.to_object_uuid),
             tenant=self.tenant,
         )
 

weaviate/collections/data.py

@@ -215,6 +215,7 @@ def _reference_add_many(self, refs: List[DataReferences]) -> BatchReferenceRetur
                 to=beacon,
                 tenant=self._tenant,
                 from_uuid=str(ref.from_uuid),
+                to_uuid=None,  # not relevant here, this entry is only needed for the batch module
             )
             for ref in refs
             for beacon in ref._to_beacons()