diff --git a/redisvl/index/storage.py b/redisvl/index/storage.py index da9290b4..43ecdc1b 100644 --- a/redisvl/index/storage.py +++ b/redisvl/index/storage.py @@ -375,7 +375,9 @@ def write( ValueError: If the length of provided keys does not match the length of objects, or if validation fails. """ - if keys and len(keys) != len(objects): # type: ignore + objects = list(objects) + keys = list(keys) if keys is not None else None + if keys is not None and len(keys) != len(objects): raise ValueError("Length of keys does not match the length of objects") if batch_size is None: @@ -386,7 +388,7 @@ def write( # Pass 1: Preprocess and validate all objects prepared_objects = self._preprocess_and_validate_objects( - list(objects), # Convert Iterable to List + objects, id_field=id_field, keys=keys, preprocess=preprocess, @@ -454,7 +456,9 @@ async def awrite( ValueError: If the length of provided keys does not match the length of objects, or if validation fails. """ - if keys and len(keys) != len(objects): # type: ignore + objects = list(objects) + keys = list(keys) if keys is not None else None + if keys is not None and len(keys) != len(objects): raise ValueError("Length of keys does not match the length of objects") if batch_size is None: @@ -465,7 +469,7 @@ async def awrite( # Pass 1: Preprocess and validate all objects prepared_objects = self._preprocess_and_validate_objects( - list(objects), # Convert Iterable to List + objects, id_field=id_field, keys=keys, preprocess=preprocess, diff --git a/tests/integration/test_load_iterables.py b/tests/integration/test_load_iterables.py new file mode 100644 index 00000000..433cb70e --- /dev/null +++ b/tests/integration/test_load_iterables.py @@ -0,0 +1,97 @@ +"""Explicit load keys and data may be single-pass iterables.""" + +import pytest + +from redisvl.exceptions import RedisVLError +from redisvl.index import AsyncSearchIndex, SearchIndex +from redisvl.schema import IndexSchema + + +@pytest.fixture(params=["hash", "json"]) +def load_schema(request, redis_test_name): + return IndexSchema.from_dict( + { + "index": { + "name": redis_test_name("iterable_load"), + "prefix": redis_test_name("iterable_load"), + "storage_type": request.param, + }, + "fields": [{"name": "title", "type": "text"}], + } + ) + + +@pytest.mark.parametrize("data_iterator", [False, True]) +@pytest.mark.parametrize("keys_iterator", [False, True]) +def test_load_with_iterable_keys(client, load_schema, data_iterator, keys_iterator): + index = SearchIndex(load_schema, redis_client=client) + records = [{"title": "first"}, {"title": "second"}] + keys = [index.key(str(i)) for i in range(len(records))] + try: + loaded = index.load( + iter(records) if data_iterator else records, + keys=iter(keys) if keys_iterator else keys, + batch_size=1, + ) + assert loaded == keys + assert [index.fetch(str(i)) for i in range(len(records))] == records + finally: + client.delete(*keys) + + +@pytest.mark.parametrize("data_iterator", [False, True]) +@pytest.mark.parametrize("keys_iterator", [False, True]) +async def test_async_load_with_iterable_keys( + async_client, load_schema, data_iterator, keys_iterator +): + index = AsyncSearchIndex(load_schema, redis_client=async_client) + records = [{"title": "first"}, {"title": "second"}] + keys = [index.key(str(i)) for i in range(len(records))] + try: + loaded = await index.load( + iter(records) if data_iterator else records, + keys=iter(keys) if keys_iterator else keys, + batch_size=1, + ) + assert loaded == keys + assert [await index.fetch(str(i)) for i in range(len(records))] == records + finally: + await async_client.delete(*keys) + + +@pytest.mark.parametrize("key_count", [0, 1, 3]) +def test_load_rejects_mismatched_keys_before_writing(client, load_schema, key_count): + index = SearchIndex(load_schema, redis_client=client) + keys = [index.key(str(i)) for i in range(key_count)] + try: + with pytest.raises(RedisVLError, match="Length of keys does not match"): + index.load( + iter([{"title": "first"}, {"title": "second"}]), + keys=iter(keys) if keys else [], + batch_size=1, + ) + assert list(client.scan_iter(f"{index.prefix}:*")) == [] + finally: + stored_keys = list(client.scan_iter(f"{index.prefix}:*")) + if stored_keys: + client.delete(*stored_keys) + + +@pytest.mark.parametrize("key_count", [0, 1, 3]) +async def test_async_load_rejects_mismatched_keys_before_writing( + async_client, load_schema, key_count +): + index = AsyncSearchIndex(load_schema, redis_client=async_client) + keys = [index.key(str(i)) for i in range(key_count)] + try: + with pytest.raises(RedisVLError, match="Length of keys does not match"): + await index.load( + iter([{"title": "first"}, {"title": "second"}]), + keys=iter(keys) if keys else [], + batch_size=1, + ) + assert [key async for key in async_client.scan_iter(f"{index.prefix}:*")] == [] + finally: + stored_keys = [key async for key in async_client.scan_iter(f"{index.prefix}:*")] + if stored_keys: + await async_client.delete(*stored_keys)