From b6181b200c04bacd6c6134bec5e99e02014b01c7 Mon Sep 17 00:00:00 2001 From: fei <204683769+feiiiiii5@users.noreply.github.com> Date: Sat, 19 Sep 2026 02:23:19 +0800 Subject: [PATCH] fix(batch): key wholesale-failure object errors by original index __send_batch built the wholesale-failure error map with chunk-local enumerate() keys while BatchObjectReturn.__add__ merges with dict.update, so keys from separate chunks collided and failures were dropped or attributed to the wrong input object. --- mock_tests/test_batch_error_indices.py | 118 +++++++++++++++++++++++++ weaviate/collections/batch/base.py | 4 +- 2 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 mock_tests/test_batch_error_indices.py diff --git a/mock_tests/test_batch_error_indices.py b/mock_tests/test_batch_error_indices.py new file mode 100644 index 000000000..aef3ea3f5 --- /dev/null +++ b/mock_tests/test_batch_error_indices.py @@ -0,0 +1,118 @@ +"""`BatchObjectReturn.errors` must be keyed by the original batch index on every path. + +`BatchObjectReturn.__add__` merges chunk results with a plain `dict.update`, so the +wholesale-failure path in `_BatchBase.__send_batch` has to key its errors by +`obj.index` like the success path does. Keying them by the position inside the chunk +makes the keys of different chunks collide. +""" + +from typing import List + +import grpc +import pytest +import weaviate +from weaviate.proto.v1 import batch_pb2, weaviate_pb2_grpc + +from .conftest import MOCK_IP, MOCK_PORT, MOCK_PORT_GRPC, mock_class + + +class MockChunkedService(weaviate_pb2_grpc.WeaviateServicer): + """Answers each `BatchObjects` call according to `per_chunk`, a list of bools.""" + + def __init__(self, per_chunk: List[bool]) -> None: + self.per_chunk = per_chunk + self.calls: List[List[str]] = [] + + def BatchObjects( + self, request: batch_pb2.BatchObjectsRequest, context: grpc.ServicerContext + ) -> batch_pb2.BatchObjectsReply: + index = len(self.calls) + self.calls.append([str(obj.uuid) for obj in request.objects]) + if self.per_chunk[index]: + return batch_pb2.BatchObjectsReply() + return batch_pb2.BatchObjectsReply( + errors=[ + {"index": i, "error": f"mock rejection {i}"} for i in range(len(request.objects)) + ] + ) + + +def _client_with(service, weaviate_mock, start_grpc_server): + weaviate_mock.expect_request(f"/v1/schema/{mock_class['class']}").respond_with_json(mock_class) + weaviate_pb2_grpc.add_WeaviateServicer_to_server(service, start_grpc_server) + return weaviate.connect_to_local(port=MOCK_PORT, host=MOCK_IP, grpc_port=MOCK_PORT_GRPC) + + +@pytest.fixture +def make_client(weaviate_mock, start_grpc_server): + clients = [] + + def _make(per_chunk: List[bool]): + service = MockChunkedService(per_chunk) + client = _client_with(service, weaviate_mock, start_grpc_server) + clients.append(client) + return client, service + + yield _make + for client in clients: + client.close() + + +def _add_objects(collection, count: int) -> List: + uuids = [] + with collection.batch.fixed_size(batch_size=2, concurrent_requests=1) as batch: + for i in range(count): + uuids.append(batch.add_object({"name": f"Object {i}"})) + batch.flush() + return uuids + + +def test_wholesale_failed_chunks_report_every_object(make_client): + client, service = make_client([False, False]) + collection = client.collections.use(mock_class["class"]) + + _add_objects(collection, 4) + + assert len(service.calls) == 2, "the 4 objects must be split into two chunks" + + errors = collection.batch.results.objs.errors + assert sorted(errors) == [0, 1, 2, 3], ( + f"every failed object must be reported once, got {sorted(errors)}" + ) + for index, error in errors.items(): + assert error.object_.index == index, ( + f"errors[{index}] describes the object at index {error.object_.index}" + ) + assert len(collection.batch.failed_objects) == len(errors) + + +def test_failed_and_successful_chunks_do_not_share_keys(make_client): + client, service = make_client([True, False]) + collection = client.collections.use(mock_class["class"]) + + uuids = _add_objects(collection, 4) + + assert len(service.calls) == 2 + result = collection.batch.results.objs + assert sorted(result.uuids) == [0, 1] + assert sorted(result.errors) == [2, 3], ( + f"the failed chunk must be reported under its own indices, got {sorted(result.errors)}" + ) + assert not set(result.uuids) & set(result.errors), ( + "a key cannot describe both an inserted and a failed object" + ) + for index in result.errors: + assert str(result.errors[index].object_.uuid) == str(uuids[index]) + + +def test_successful_chunks_keep_original_indices(make_client): + client, service = make_client([True, True, True]) + collection = client.collections.use(mock_class["class"]) + + uuids = _add_objects(collection, 6) + + assert len(service.calls) == 3 + result = collection.batch.results.objs + assert sorted(result.uuids) == [0, 1, 2, 3, 4, 5] + for index, uuid in enumerate(uuids): + assert str(result.uuids[index]) == str(uuid) diff --git a/weaviate/collections/batch/base.py b/weaviate/collections/batch/base.py index 1d736303c..45ff9edfb 100644 --- a/weaviate/collections/batch/base.py +++ b/weaviate/collections/batch/base.py @@ -645,9 +645,7 @@ def __send_batch( } ) except Exception as e: - errors_obj = { - idx: ErrorObject(message=repr(e), object_=obj) for idx, obj in enumerate(objs) - } + errors_obj = {obj.index: ErrorObject(message=repr(e), object_=obj) for obj in objs} logger.error( { "message": f"Failed to send all objects in a batch of {len(objs)}",