Introduce experimental batching to async client

integration/test_batch_v4.py

@@ -4,6 +4,7 @@
 from typing import Callable, Generator, List, Optional, Protocol, Tuple
 
 import pytest
+import pytest_asyncio
 from _pytest.fixtures import SubRequest
 
 import weaviate
@@ -119,6 +120,53 @@ def _factory(
             client_fixture.close()
 
 
+class AsyncClientFactory(Protocol):
+    """Typing for fixture."""
+
+    async def __call__(
+        self, name: str = "", ports: Tuple[int, int] = (8080, 50051), multi_tenant: bool = False
+    ) -> Tuple[weaviate.WeaviateAsyncClient, str]:
+        """Typing for fixture."""
+        ...
+
+
+@pytest_asyncio.fixture
+async def async_client_factory(request: SubRequest):
+    name_fixtures: List[str] = []
+    client_fixture: Optional[weaviate.WeaviateAsyncClient] = None
+
+    async def _factory(
+        name: str = "", ports: Tuple[int, int] = (8080, 50051), multi_tenant: bool = False
+    ):
+        nonlocal client_fixture, name_fixtures  # noqa: F824
+        name_fixture = _sanitize_collection_name(request.node.name) + name
+        name_fixtures.append(name_fixture)
+        if client_fixture is None:
+            client_fixture = weaviate.use_async_with_local(grpc_port=ports[1], port=ports[0])
+            await client_fixture.connect()
+
+        if await client_fixture.collections.exists(name_fixture):
+            await client_fixture.collections.delete(name_fixture)
+
+        await client_fixture.collections.create(
+            name=name_fixture,
+            properties=[
+                Property(name="name", data_type=DataType.TEXT),
+                Property(name="age", data_type=DataType.INT),
+            ],
+            references=[ReferenceProperty(name="test", target_collection=name_fixture)],
+            multi_tenancy_config=Configure.multi_tenancy(multi_tenant),
+            vectorizer_config=Configure.Vectorizer.none(),
+        )
+        return client_fixture, name_fixture
+
+    try:
+        yield _factory
+    finally:
+        if client_fixture is not None:
+            await client_fixture.close()
+
+
 def test_add_objects_in_multiple_batches(client_factory: ClientFactory) -> None:
     client, name = client_factory()
     with client.batch.rate_limit(50) as batch:
@@ -367,13 +415,13 @@ def test_add_ref_batch_with_tenant(client_factory: ClientFactory) -> None:
     [
         lambda client: client.batch.dynamic(),
         lambda client: client.batch.fixed_size(),
-        # lambda client: client.batch.rate_limit(9999),
+        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_rate_limit",
         "test_add_ten_thousand_data_objects_experimental",
     ],
 )
@@ -767,3 +815,34 @@ def test_references_with_to_uuids(client_factory: ClientFactory) -> None:
 
     assert len(client.batch.failed_references) == 0, client.batch.failed_references
     client.collections.delete(["target", "source"])
+
+
+@pytest.mark.asyncio
+async def test_add_one_hundred_thousand_objects_async_client(
+    async_client_factory: AsyncClientFactory,
+) -> None:
+    """Test adding one hundred thousand data objects."""
+    client, name = await async_client_factory()
+    if client._connection._weaviate_version.is_lower_than(1, 36, 0):
+        pytest.skip("Server-side batching not supported in Weaviate < 1.36.0")
+    nr_objects = 100000
+    import time
+
+    start = time.time()
+    async with client.batch.experimental(concurrency=1) as batch:
+        for i in range(nr_objects):
+            await batch.add_object(
+                collection=name,
+                properties={"name": "test" + str(i)},
+            )
+    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 await client.collections.use(name).length() == 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
+    ]
+    await client.collections.delete(name)

integration/test_collection_batch.py

@@ -1,10 +1,10 @@
 import uuid
 from dataclasses import dataclass
-from typing import Any, Generator, Optional, Protocol, Union
+from typing import Any, Awaitable, Generator, Optional, Protocol, Union
 
 import pytest
 
-from integration.conftest import CollectionFactory, CollectionFactoryGet
+from integration.conftest import AsyncCollectionFactory, CollectionFactory, CollectionFactoryGet
 from weaviate.collections import Collection
 from weaviate.collections.classes.config import (
     Configure,
@@ -17,6 +17,8 @@
 from weaviate.collections.classes.tenants import Tenant
 from weaviate.types import VECTORS
 
+from weaviate.collections.collection.async_ import CollectionAsync
+
 UUID = Union[str, uuid.UUID]
 
 
@@ -55,11 +57,21 @@ def __call__(self, name: str = "", multi_tenancy: bool = False) -> Collection[An
         ...
 
 
+class BatchCollectionAsync(Protocol):
+    """Typing for fixture."""
+
+    def __call__(
+        self, name: str = "", multi_tenancy: bool = False
+    ) -> Awaitable[CollectionAsync[Any, Any]]:
+        """Typing for fixture."""
+        ...
+
+
 @pytest.fixture
 def batch_collection(
     collection_factory: CollectionFactory,
 ) -> Generator[BatchCollection, None, None]:
-    def _factory(name: str = "", multi_tenancy: bool = False) -> Collection[Any, Any]:
+    def _factory(name: str = "", multi_tenancy: bool = False):
         collection = collection_factory(
             name=name,
             vectorizer_config=Configure.Vectorizer.none(),
@@ -78,6 +90,29 @@ def _factory(name: str = "", multi_tenancy: bool = False) -> Collection[Any, Any
     yield _factory
 
 
+@pytest.fixture
+def batch_collection_async(
+    async_collection_factory: AsyncCollectionFactory,
+) -> Generator[BatchCollectionAsync, None, None]:
+    async def _factory(name: str = "", multi_tenancy: bool = False):
+        collection = await async_collection_factory(
+            name=name,
+            vectorizer_config=Configure.Vectorizer.none(),
+            properties=[
+                Property(name="name", data_type=DataType.TEXT),
+                Property(name="age", data_type=DataType.INT),
+            ],
+            multi_tenancy_config=Configure.multi_tenancy(multi_tenancy),
+        )
+        await collection.config.add_reference(
+            ReferenceProperty(name="test", target_collection=collection.name)
+        )
+
+        return collection
+
+    yield _factory
+
+
 @pytest.mark.parametrize(
     "vector",
     [None, [1, 2, 3], MockNumpyTorch([1, 2, 3]), MockTensorFlow([1, 2, 3])],
@@ -233,3 +268,30 @@ def test_non_existant_collection(collection_factory_get: CollectionFactoryGet) -
 
     # above should not throw - depending on the autoschema config this might create an error or
     # not, so we do not check for errors here
+
+
+@pytest.mark.asyncio
+async def test_add_one_hundred_thousand_objects_async_collection(
+    batch_collection_async: BatchCollectionAsync,
+) -> None:
+    """Test adding one hundred thousand data objects."""
+    col = await batch_collection_async()
+    if col._connection._weaviate_version.is_lower_than(1, 36, 0):
+        pytest.skip("Server-side batching not supported in Weaviate < 1.36.0")
+    nr_objects = 100000
+    import time
+
+    start = time.time()
+    async with col.batch.experimental() as batch:
+        for i in range(nr_objects):
+            await batch.add_object(
+                properties={"name": "test" + str(i)},
+            )
+    end = time.time()
+    print(f"Time taken to add {nr_objects} objects: {end - start} seconds")
+    assert len(col.batch.results.objs.errors) == 0
+    assert len(col.batch.results.objs.all_responses) == nr_objects
+    assert len(col.batch.results.objs.uuids) == nr_objects
+    assert await col.length() == nr_objects
+    assert col.batch.results.objs.has_errors is False
+    assert len(col.batch.failed_objects) == 0, [obj.message for obj in col.batch.failed_objects]

weaviate/client.py

@@ -10,7 +10,7 @@
 from .auth import AuthCredentials
 from .backup import _Backup, _BackupAsync
 from .cluster import _Cluster, _ClusterAsync
-from .collections.batch.client import _BatchClientWrapper
+from .collections.batch.client import _BatchClientWrapper, _BatchClientWrapperAsync
 from .collections.collections import _Collections, _CollectionsAsync
 from .config import AdditionalConfig
 from .connect import executor
@@ -76,6 +76,7 @@ def __init__(
         )
         self.alias = _AliasAsync(self._connection)
         self.backup = _BackupAsync(self._connection)
+        self.batch = _BatchClientWrapperAsync(self._connection)
         self.cluster = _ClusterAsync(self._connection)
         self.collections = _CollectionsAsync(self._connection)
         self.debug = _DebugAsync(self._connection)

weaviate/client.pyi

@@ -18,7 +18,7 @@ from weaviate.users.sync import _Users
 
 from .backup import _Backup, _BackupAsync
 from .cluster import _Cluster, _ClusterAsync
-from .collections.batch.client import _BatchClientWrapper
+from .collections.batch.client import _BatchClientWrapper, _BatchClientWrapperAsync
 from .debug import _Debug, _DebugAsync
 from .rbac import _Roles, _RolesAsync
 from .types import NUMBER
@@ -29,6 +29,7 @@ class WeaviateAsyncClient(_WeaviateClientExecutor[ConnectionAsync]):
     _connection: ConnectionAsync
     alias: _AliasAsync
     backup: _BackupAsync
+    batch: _BatchClientWrapperAsync
     collections: _CollectionsAsync
     cluster: _ClusterAsync
     debug: _DebugAsync