Introduce batch.experimental() while server-side batching is in beta

.github/workflows/main.yaml

@@ -20,11 +20,12 @@ concurrency:
 env:
   WEAVIATE_127: 1.27.27
   WEAVIATE_128: 1.28.16
-  WEAVIATE_129: 1.29.9
-  WEAVIATE_130: 1.30.12
-  WEAVIATE_131: 1.31.5
+  WEAVIATE_129: 1.29.11
+  WEAVIATE_130: 1.30.18
+  WEAVIATE_131: 1.31.16
   WEAVIATE_132: 1.32.5
-  WEAVIATE_133: 1.33.0-rc.1
+  WEAVIATE_133: 1.33.0
+  WEAVIATE_134: 1.34.0-rc.0-ba35f50.amd64
 
 jobs:
   lint-and-format:
@@ -152,11 +153,11 @@ jobs:
       fail-fast: false
       matrix:
         versions: [
-          { py: "3.9", weaviate: $WEAVIATE_131, grpc: "1.59.0"},
-          { py: "3.10", weaviate: $WEAVIATE_131, grpc: "1.66.0"},
-          { py: "3.11", weaviate: $WEAVIATE_131, grpc: "1.70.0"},
-          { py: "3.12", weaviate: $WEAVIATE_131, grpc: "1.72.1"},
-          { py: "3.13", weaviate: $WEAVIATE_131, grpc: "1.74.0"}
+          { py: "3.9", weaviate: $WEAVIATE_132, grpc: "1.59.0"},
+          { py: "3.10", weaviate: $WEAVIATE_132, grpc: "1.66.0"},
+          { py: "3.11", weaviate: $WEAVIATE_132, grpc: "1.70.0"},
+          { py: "3.12", weaviate: $WEAVIATE_132, grpc: "1.72.1"},
+          { py: "3.13", weaviate: $WEAVIATE_132, grpc: "1.74.0"}
         ]
         optional_dependencies: [false]
     steps:
@@ -207,11 +208,11 @@ jobs:
       fail-fast: false
       matrix:
         versions: [
-          { py: "3.9", weaviate: $WEAVIATE_131},
-          { py: "3.10", weaviate: $WEAVIATE_131},
-          { py: "3.11", weaviate: $WEAVIATE_131},
-          { py: "3.12", weaviate: $WEAVIATE_131},
-          { py: "3.13", weaviate: $WEAVIATE_131}
+          { py: "3.9", weaviate: $WEAVIATE_132},
+          { py: "3.10", weaviate: $WEAVIATE_132},
+          { py: "3.11", weaviate: $WEAVIATE_132},
+          { py: "3.12", weaviate: $WEAVIATE_132},
+          { py: "3.13", weaviate: $WEAVIATE_132}
         ]
         optional_dependencies: [false]
     steps:
@@ -302,7 +303,8 @@ jobs:
           $WEAVIATE_130,
           $WEAVIATE_131,
           $WEAVIATE_132,
-          $WEAVIATE_133
+          $WEAVIATE_133,
+          $WEAVIATE_134
         ]
     steps:
       - name: Checkout

integration/test_batch_v4.py

@@ -9,7 +9,7 @@
 import weaviate
 import weaviate.classes as wvc
 from integration.conftest import _sanitize_collection_name
-from weaviate import BatchClient, ClientBatchingContextManager
+from weaviate import ClientBatchingContextManager
 from weaviate.collections.classes.batch import Shard
 from weaviate.collections.classes.config import (
     Configure,
@@ -175,11 +175,11 @@ def test_add_reference(
     client_factory: ClientFactory,
     from_object_uuid: UUID,
     to_object_uuid: UUID,
-    to_object_collection: Optional[bool],
+    to_object_collection: bool,
 ) -> None:
     """Test the `add_reference` method."""
     client, name = client_factory()
-    with client.batch.fixed_size() as batch:
+    with client.batch.dynamic() as batch:
         batch.add_object(
             properties={},
             collection=name,
@@ -194,8 +194,16 @@ def test_add_reference(
             from_uuid=from_object_uuid,
             from_collection=name,
             from_property="test",
-            to=to_object_uuid,
+            to=ReferenceToMulti(target_collection=name, uuids=to_object_uuid)
+            if to_object_collection
+            else to_object_uuid,
         )
+    assert len(client.batch.failed_objects) == 0, [
+        obj.message for obj in client.batch.failed_objects
+    ]
+    assert len(client.batch.failed_references) == 0, [
+        ref.message for ref in client.batch.failed_references
+    ]
     objs = (
         client.collections.use(name)
         .query.fetch_objects(return_references=QueryReference(link_on="test"))
@@ -357,14 +365,16 @@ def test_add_ref_batch_with_tenant(client_factory: ClientFactory) -> None:
 @pytest.mark.parametrize(
     "batching_method",
     [
-        lambda client: client.batch.dynamic(),
-        lambda client: client.batch.fixed_size(),
-        lambda client: client.batch.rate_limit(9999),
+        # lambda client: client.batch.dynamic(),
+        # lambda client: client.batch.fixed_size(),
+        # lambda client: client.batch.rate_limit(9999),
+        lambda client: client.batch.experimental(concurrency=1),
     ],
     ids=[
-        "test_add_ten_thousand_data_objects_dynamic",
-        "test_add_ten_thousand_data_objects_fixed_size",
-        "test_add_ten_thousand_data_objects_rate_limit",
+        # "test_add_ten_thousand_data_objects_dynamic",
+        # "test_add_ten_thousand_data_objects_fixed_size",
+        # "test_add_ten_thousand_data_objects_rate_limit",
+        "test_add_ten_thousand_data_objects_experimental",
     ],
 )
 def test_add_ten_thousand_data_objects(
@@ -374,16 +384,31 @@ def test_add_ten_thousand_data_objects(
 ) -> None:
     """Test adding ten thousand data objects."""
     client, name = client_factory()
-
-    nr_objects = 10000
+    if (
+        request.node.callspec.id == "test_add_ten_thousand_data_objects_experimental"
+        and client._connection._weaviate_version.is_lower_than(1, 34, 0)
+    ):
+        pytest.skip("Server-side batching not supported in Weaviate < 1.34.0")
+    nr_objects = 100000
+    import time
+
+    start = time.time()
     with batching_method(client) as batch:
         for i in range(nr_objects):
             batch.add_object(
                 collection=name,
                 properties={"name": "test" + str(i)},
             )
-    objs = client.collections.use(name).query.fetch_objects(limit=nr_objects).objects
-    assert len(objs) == nr_objects
+    end = time.time()
+    print(f"Time taken to add {nr_objects} objects: {end - start} seconds")
+    assert len(client.batch.results.objs.errors) == 0
+    assert len(client.batch.results.objs.all_responses) == nr_objects
+    assert len(client.batch.results.objs.uuids) == nr_objects
+    assert len(client.collections.use(name)) == nr_objects
+    assert client.batch.results.objs.has_errors is False
+    assert len(client.batch.failed_objects) == 0, [
+        obj.message for obj in client.batch.failed_objects
+    ]
     client.collections.delete(name)
 
 
@@ -551,19 +576,28 @@ def test_add_1000_tenant_objects_with_async_indexing_and_wait_for_only_one(
         lambda client: client.batch.dynamic(),
         lambda client: client.batch.fixed_size(),
         lambda client: client.batch.rate_limit(1000),
+        lambda client: client.batch.experimental(),
     ],
     ids=[
         "test_add_one_hundred_objects_and_references_between_all_dynamic",
         "test_add_one_hundred_objects_and_references_between_all_fixed_size",
         "test_add_one_hundred_objects_and_references_between_all_rate_limit",
+        "test_add_one_hundred_objects_and_references_between_all_experimental",
     ],
 )
 def test_add_one_object_and_a_self_reference(
     client_factory: ClientFactory,
     batching_method: Callable[[weaviate.WeaviateClient], ClientBatchingContextManager],
+    request: SubRequest,
 ) -> None:
     """Test adding one object and a self reference."""
     client, name = client_factory()
+    if (
+        request.node.callspec.id
+        == "test_add_one_hundred_objects_and_references_between_all_experimental"
+        and client._connection._weaviate_version.is_lower_than(1, 34, 0)
+    ):
+        pytest.skip("Server-side batching not supported in Weaviate < 1.34.0")
     with batching_method(client) as batch:
         uuid = batch.add_object(collection=name, properties={})
         batch.add_reference(
@@ -586,7 +620,7 @@ def test_multi_threaded_batching(
     nr_objects = 1000
     nr_threads = 10
 
-    def batch_insert(batch: BatchClient) -> None:
+    def batch_insert(batch) -> None:
         for i in range(nr_objects):
             batch.add_object(
                 collection=name,
@@ -683,6 +717,12 @@ def test_batching_error_logs(
     client_factory: ClientFactory, caplog: pytest.LogCaptureFixture
 ) -> None:
     client, name = client_factory()
+    if client._connection._weaviate_version.is_at_least(
+        1, 32, 1
+    ):  # TODO: change to 1.33.0 when released
+        pytest.skip(
+            "Batching error logs do not get emitted by the new server-side batching functionality."
+        )
     with client.batch.fixed_size() as batch:
         for obj in [{"name": i} for i in range(100)]:
             batch.add_object(properties=obj, collection=name)

integration/test_client.py

@@ -345,7 +345,12 @@ def test_client_cluster_with_lazy_shard_loading(
         assert len(nodes[0].shards) == 1
         assert nodes[0].shards[0].collection == collection.name
         assert nodes[0].shards[0].object_count == 0
-        assert nodes[0].shards[0].vector_indexing_status == "READY"
+        assert nodes[0].shards[0].vector_indexing_status in [
+            "READONLY",
+            "INDEXING",
+            "READY",
+            "LAZY_LOADING",
+        ]
         assert nodes[0].shards[0].vector_queue_length == 0
         assert nodes[0].shards[0].compressed is False
         assert nodes[0].shards[0].loaded is True
@@ -358,12 +363,6 @@ def test_client_cluster_without_lazy_shard_loading(
 ) -> None:
     client = client_factory(8090, 50061)
 
-    # Lazy-loading behaviour was changed in 1.32.4:
-    # https://github.com/weaviate/weaviate/pull/8829
-    #
-    # We also accept LOADING/READY because it may vary
-    # based on the machine running the tests.
-
     try:
         collection = client.collections.create(
             name=request.node.name, vectorizer_config=Configure.Vectorizer.none()
@@ -374,10 +373,12 @@ def test_client_cluster_without_lazy_shard_loading(
         assert len(nodes[0].shards) == 1
         assert nodes[0].shards[0].collection == collection.name
         assert nodes[0].shards[0].object_count == 0
-        if collection._connection._weaviate_version.is_lower_than(1, 32, 0):
-            assert nodes[0].shards[0].vector_indexing_status == "READY"
-        else:
-            assert nodes[0].shards[0].vector_indexing_status == "LAZY_LOADING"
+        assert nodes[0].shards[0].vector_indexing_status in [
+            "READONLY",
+            "INDEXING",
+            "READY",
+            "LAZY_LOADING",
+        ]
         assert nodes[0].shards[0].vector_queue_length == 0
         assert nodes[0].shards[0].compressed is False
         if collection._connection._weaviate_version.is_lower_than(1, 25, 0):

integration/test_rbac.py

@@ -6,6 +6,7 @@
 from integration.conftest import ClientFactory, _sanitize_collection_name
 from weaviate.auth import Auth
 from weaviate.classes.rbac import Actions, Permissions, RoleScope
+from weaviate.connect.helpers import connect_to_local
 from weaviate.rbac.models import (
     AliasPermissionOutput,
     BackupsPermissionOutput,
@@ -734,3 +735,21 @@ def test_permission_joining(client_factory: ClientFactory) -> None:
 
         finally:
             client.roles.delete(role_name)
+
+
+def test_server_side_batching_with_auth() -> None:
+    collection_name = "TestSSBAuth"
+    with connect_to_local(
+        port=RBAC_PORTS[0], grpc_port=RBAC_PORTS[1], auth_credentials=RBAC_AUTH_CREDS
+    ) as client:
+        if client._connection._weaviate_version.is_lower_than(1, 34, 0):
+            pytest.skip("Server-side batching not supported in Weaviate < 1.34.0")
+        collection = client.collections.create(collection_name)
+        with client.batch.experimental() as batch:
+            batch.add_object(collection_name)
+            batch.add_object(collection_name)
+            batch.add_object(collection_name)
+        try:
+            assert len(collection) == 3
+        finally:
+            client.collections.delete(collection_name)

profiling/test_shutdown.py

@@ -0,0 +1,76 @@
+import json
+import random
+import time
+import weaviate
+import weaviate.classes.config as wvcc
+
+
+def setup(client: weaviate.WeaviateClient, collection: str) -> weaviate.collections.Collection:
+    if client.collections.exists(collection):
+        client.collections.delete(collection)
+    return client.collections.create(
+        name=collection,
+        properties=[
+            wvcc.Property(
+                name="title",
+                data_type=wvcc.DataType.TEXT,
+            ),
+            wvcc.Property(
+                name="content",
+                data_type=wvcc.DataType.TEXT,
+            ),
+        ],
+        replication_config=wvcc.Configure.replication(factor=3, async_enabled=True),
+        vector_config=wvcc.Configure.Vectors.self_provided(),
+    )
+
+
+def import_(client: weaviate.WeaviateClient, collection: str, how_many: int = 1_000_000) -> None:
+    uuids: dict[str, int] = {}
+    with client.batch.experimental(concurrency=1) as batch:
+        for i in range(how_many):
+            uuid = batch.add_object(
+                collection=collection,
+                properties={
+                    "title": f"Title {i}",
+                    "content": f"Content {i}",
+                },
+                vector=random_vector(),
+            )
+            uuids[str(uuid)] = i
+            if batch.number_errors > 0:
+                print(f"There are some errors {batch.number_errors}")
+
+    for err in client.batch.failed_objects:
+        print(err.message)
+    assert len(client.batch.failed_objects) == 0, "Expected there to be no errors when importing"
+    client.batch.wait_for_vector_indexing()
+    with open("uuids.json", "w") as f:
+        json.dump(uuids, f)
+
+
+def verify(client: weaviate.WeaviateClient, collection: str, expected: int = 1_000_000) -> None:
+    actual = 0
+    count = 0
+    c = client.collections.use(collection)
+    while actual != expected:
+        actual = len(c)
+        print(f"Found {actual} objects, waiting for async repl to reach {expected}...")
+        time.sleep(1)
+        count += 1
+        if count == 120:
+            break
+    assert actual == expected, f"Expected {expected} objects, found {actual}"
+
+
+def random_vector() -> list[float]:
+    return [random.uniform(0, 1) for _ in range(128)]
+
+
+def test_main() -> None:
+    collection = "BatchImportShutdownJourney"
+    how_many = 500000
+    with weaviate.connect_to_local() as client:
+        collection = setup(client, collection)
+        import_(client, collection.name, how_many)
+        verify(client, collection.name, how_many)

weaviate/classes/data.py

@@ -1,9 +1,11 @@
 from weaviate.collections.classes.data import DataObject, DataReference
+from weaviate.collections.classes.internal import ReferenceToMulti
 from weaviate.collections.classes.types import GeoCoordinate, PhoneNumber
 
 __all__ = [
     "DataObject",
     "DataReference",
     "GeoCoordinate",
     "PhoneNumber",
+    "ReferenceToMulti",
 ]

weaviate/client.py

@@ -145,7 +145,11 @@ def __init__(
         self.alias = _Alias(
             self._connection,
         )
-        self.batch = _BatchClientWrapper(self._connection, config=collections)
+        self.batch = _BatchClientWrapper(
+            self._connection,
+            config=collections,
+            consistency_level=None,
+        )
         self.backup = _Backup(self._connection)
         self.cluster = _Cluster(self._connection)
         self.collections = collections

weaviate/collections/batch/__init__.py

@@ -1,6 +1,7 @@
 __all__ = ["_BatchClient", "_BatchCollection", "_BatchGRPC", "_BatchREST"]
 
+from weaviate.collections.batch.grpc_batch import _BatchGRPC
+
 from .client import _BatchClient
 from .collection import _BatchCollection
-from .grpc_batch_objects import _BatchGRPC
 from .rest import _BatchREST