diff --git a/docs/concepts/search-and-indexing.md b/docs/concepts/search-and-indexing.md index 5312d7df..058c950c 100644 --- a/docs/concepts/search-and-indexing.md +++ b/docs/concepts/search-and-indexing.md @@ -84,6 +84,51 @@ index = SearchIndex.from_existing("my-index", redis_url="redis://localhost:6379" **Clearing data** with `clear()` removes all documents from the index without deleting the index itself. The schema remains intact, ready for new data. +## Bulk Delete and Update + +Redis has no server-side "delete/update by query", so mutating many documents traditionally means scanning for keys and issuing writes yourself. RedisVL wraps that pattern behind two filter-driven methods (available on both `SearchIndex` and `AsyncSearchIndex`). + +**Deleting by filter** resolves every document matching a filter expression and removes it with non-blocking `UNLINK`, in batches: + +```python +from redisvl.query.filter import Tag, Num + +# Remove every archived document from before 2020 +index.drop_by_filter((Tag("status") == "archived") & (Num("year") < 2020)) +``` + +**Updating by filter** applies a partial field update to every match. Fields you don't mention are left untouched—hash fields are written with `HSET`, and JSON documents are merged at the root with `JSON.MERGE` (RFC 7396: nested objects merge recursively, arrays are replaced wholesale, and a `None` value deletes that path): + +```python +# Mark all draft documents as published, leaving other fields intact +index.update_by_filter(Tag("status") == "draft", {"status": "published"}) +``` + +Two caveats for `update_by_filter` values: + +- **No schema validation.** Unlike `load()`, values are written as-is—pre-encode vectors/bytes and format numerics as your schema expects, or you may corrupt query results. +- **JSON keys must match the document layout, not the schema field name.** Values merge at the document root (`$`). If a field is indexed at a nested path (e.g. `$.metadata.status`), pass the correspondingly nested mapping (`{"metadata": {"status": "published"}}`); passing the flat field name writes the wrong path and leaves the indexed field unchanged. Only static values are supported (no callable/expression transforms). + +Both methods share the same safety-oriented options and return a `BulkResult`: + +- The return value carries `matched` (documents matching the filter), `processed` (documents actually affected), `completed` (False if a delete stopped early at its runaway backstop; always True for update), and `dry_run`. Read the fields explicitly—`result.processed`, `result.completed`, etc. +- `dry_run=True` reports how many documents *would* be affected—via a count query—without changing anything. +- A match-all filter (empty, `"*"`, or `None`) is refused unless you pass `allow_all=True`; use `clear()` when you intentionally want to empty the index. This guard catches the canonical match-all forms only—it is a convenience backstop, not a security control. +- `on_progress(processed, matched)` is called after each write batch for observability on large operations (it runs synchronously—don't pass a coroutine—and raising from it aborts the run). Note `update_by_filter` resolves all matching keys before it writes, so progress callbacks begin only once the write phase starts. `batch_size` tunes how many documents are processed per round-trip. + +The related `drop_documents()` and `drop_keys()` helpers delete by document ID or full Redis key and batch large inputs. On Redis Cluster, `drop_keys()` unlinks per-key so it works across hash slots, while `drop_documents()` requires the target keys to share a hash tag and raises `ValueError` otherwise. + +**Durability and partial failure.** Because Redis has no server-side delete/update-by-query, these operations run as a series of batched writes rather than a single transaction—they are **not atomic across the match set** and there is no rollback. The unit of atomicity is a single document: each key is deleted with one `UNLINK`, and each document is updated with one `HSET` or `JSON.MERGE`, so you never get a half-deleted key or a half-updated document. But batches are applied incrementally, so a crash or connection error mid-run leaves some documents changed and the rest untouched. + +The intended recovery is simply to **re-run the same call**: both operations are idempotent, so a repeat pass removes (or re-sets) only what still needs it and converges on the desired state. There is no built-in checkpoint or resume token—`on_progress` reports live progress but is not a restart point. `update_by_filter` resolves keys before it writes, so a document can be deleted by another client in between; each write is applied only if the key still exists (atomically), so such a document is **skipped, not recreated** as a partial document—it just isn't counted in `processed`, which is why `processed` can be less than `matched` under concurrent deletion. Note this guard is *existence*-based, not *identity*-based: if a key is deleted and a different document is recreated at the same key mid-run, the update can land on that new document. Running against a quiescent partition (below) avoids all of these concurrency cases. + +**Operational considerations at scale.** These are single-threaded, foreground operations that issue one round-trip per batch; a very large match set is a long-running job. Keep the following in mind for production use: + +- **Live-traffic impact.** Every delete/update also updates the search index synchronously. Running a large bulk job against a node that is serving queries competes for CPU and reindex bandwidth and can raise query latency—prefer off-peak windows, or narrow the filter and run in waves. +- **Memory.** `drop_by_filter` holds only one batch of keys at a time (`O(batch_size)`). `update_by_filter` must resolve *all* matching keys before it can write (an open aggregation cursor can't be read while the index is being written), so its client memory grows with the match count—roughly the total size of the matched keys. For very large match sets, partition the filter (see below) rather than updating everything in one call. +- **Redis Cluster.** Cross-slot multi-key commands aren't allowed, so writes are issued per key (no pipelining across slots)—expect this to be slower on Cluster for large match sets. +- **Resumability.** There is no checkpoint; recovery is a full re-run. For very large corpora, **partition the filter** (e.g. by a tag or numeric range) and process partition-by-partition. This is the recommended pattern at scale: it bounds memory and per-run time, keeps each call independently retryable, and lets you spread load across off-peak windows. + ## Data Validation RedisVL can validate data against your schema before loading it to Redis. This catches type mismatches, missing required fields, and invalid values early—before they cause problems in production. diff --git a/redisvl/index/index.py b/redisvl/index/index.py index ac5e3982..b92486d7 100644 --- a/redisvl/index/index.py +++ b/redisvl/index/index.py @@ -4,6 +4,7 @@ import time import warnings import weakref +from dataclasses import dataclass from math import ceil from typing import ( TYPE_CHECKING, @@ -12,6 +13,7 @@ Callable, Generator, Iterable, + Iterator, Sequence, cast, ) @@ -47,6 +49,7 @@ from redis import __version__ as redis_version from redis.client import NEVER_DECODE +from redis.commands.search.aggregation import AggregateRequest, Cursor from redisvl.utils.redis_protocol import get_protocol_version @@ -126,6 +129,114 @@ def _sql_executor_cache_key(sql_redis_options: dict[str, Any]) -> str: return json.dumps(sql_redis_options, sort_keys=True, default=repr) +# Default number of documents processed per round-trip in bulk operations. +DEFAULT_BULK_BATCH_SIZE = 500 + +# Idle timeout (seconds) applied to the FT.AGGREGATE cursor used by bulk +# updates, so an abandoned cursor is reclaimed server-side rather than +# lingering. redis-py's AggregateRequest.cursor(max_idle=...) takes seconds. +BULK_CURSOR_MAX_IDLE_SECONDS = 300 + +# update_by_filter resolves keys before it writes, so a key can be deleted by +# another client in between. These scripts make each write conditional on the +# key still existing, atomically — so a concurrently-deleted document is +# skipped rather than resurrected as a partial document. Each returns 1 if the +# document was written, 0 if it no longer existed. +_UPDATE_HASH_IF_EXISTS_LUA = """ +if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end +redis.call('HSET', KEYS[1], unpack(ARGV)) +return 1 +""" +_UPDATE_JSON_IF_EXISTS_LUA = """ +if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end +redis.call('JSON.MERGE', KEYS[1], '$', ARGV[1]) +return 1 +""" + + +@dataclass(frozen=True) +class BulkResult: + """Outcome of a bulk operation (:meth:`SearchIndex.drop_by_filter` / + :meth:`SearchIndex.update_by_filter`). + + Because bulk operations are not atomic across the match set, the result + distinguishes how many documents matched from how many were actually + processed, and whether the run ran to completion. Read the fields + explicitly (``result.processed`` / ``result.matched`` / ``result.completed``). + + Attributes: + matched (int): Documents matching the filter when the run started. + processed (int): Documents actually deleted/updated (for ``dry_run``, + the number that *would* be processed). + completed (bool): False if the run stopped early (e.g. a delete hit its + runaway backstop under heavy concurrent inserts); re-run to finish. + dry_run (bool): True if this was a preview and nothing was mutated. + """ + + matched: int + processed: int + completed: bool = True + dry_run: bool = False + + +def _is_match_all_filter(filter_expression: str | FilterExpression | None) -> bool: + """Return True if the filter would match every document in the index. + + Guards the bulk ``*_by_filter`` methods against an accidental full-index + wipe/update. ``None`` is treated as match-all because it defaults to + ``FilterExpression("*")`` downstream; a default/empty ``FilterExpression`` + (whose ``str()`` raises) is likewise treated as match-all rather than + surfacing an opaque error. + """ + if filter_expression is None: + return True + try: + rendered = str(filter_expression).strip() + except ValueError: + # Improperly initialized FilterExpression() - treat as the match-all sentinel + return True + return rendered in ("", "*") + + +def _require_specific_filter( + filter_expression: str | FilterExpression | None, allow_all: bool +) -> None: + """Raise unless the filter is specific or the caller opted into match-all.""" + if not allow_all and _is_match_all_filter(filter_expression): + raise ValueError( + "Refusing to run a bulk operation that matches all documents. " + "Pass a specific filter_expression, set allow_all=True to override, " + "or use clear() to intentionally remove every document." + ) + + +def _agg_row_to_key(row: Any) -> str: + """Extract the document key from a ``LOAD 1 @__key`` aggregation row. + + Aggregation rows are a flat ``[field, value, field, value, ...]`` list, so + the key is the element immediately after the ``__key`` field name. + """ + decoded = convert_bytes(row) + key_field_index = decoded.index("__key") + return decoded[key_field_index + 1] + + +def _update_script_and_args( + is_json: bool, values: dict[str, Any] +) -> tuple[str, list[Any]]: + """Return the (Lua script, args) for an existence-guarded partial update. + + Hash values are flattened into ``HSET`` field/value pairs; JSON values are + serialized into a single ``JSON.MERGE`` payload. + """ + if is_json: + return _UPDATE_JSON_IF_EXISTS_LUA, [json.dumps(values)] + flattened: list[Any] = [] + for field, value in values.items(): + flattened.extend((field, value)) + return _UPDATE_HASH_IF_EXISTS_LUA, flattened + + def process_results( results: "Result", query: BaseQuery, schema: IndexSchema ) -> list[dict[str, Any]]: @@ -823,7 +934,35 @@ def invalidate_sql_schema_cache(self) -> None: """Clear cached sql-redis executors and schema state for this index.""" self._sql_executors.clear() - def drop_keys(self, keys: str | list[str]) -> int: + def _unlink_batch(self, batch_keys: list[str]) -> int: + """Unlink a batch of keys from Redis. + + Mirrors ``_delete_batch`` but uses non-blocking ``UNLINK``. For Redis + Cluster, keys are unlinked individually to avoid cross-slot errors; + for standalone Redis a single variadic ``UNLINK`` is used. + + Args: + batch_keys (List[str]): List of Redis keys to unlink. + + Returns: + int: Count of records unlinked from Redis. + """ + client = cast(SyncRedisClient, self._redis_client) + + if isinstance(client, RedisCluster): + unlinked = 0 + for key_to_unlink in batch_keys: + try: + unlinked += cast(int, client.unlink(key_to_unlink)) + except redis.exceptions.RedisError as e: + logger.warning(f"Failed to unlink key {key_to_unlink}: {e}") + return unlinked + + return cast(int, client.unlink(*batch_keys)) + + def drop_keys( + self, keys: str | list[str], batch_size: int = DEFAULT_BULK_BATCH_SIZE + ) -> int: """Remove a specific entry or entries from the index by it's key ID. Uses ``UNLINK`` rather than ``DEL`` so memory reclamation runs on a @@ -831,18 +970,31 @@ def drop_keys(self, keys: str | list[str]) -> int: number of keys are dropped at once (for example, scope-targeted ``SemanticCache`` invalidation). The returned count is unchanged. + Large key lists are unlinked in chunks of ``batch_size``. On Redis + Cluster, keys are unlinked individually so a chunk that spans hash + slots does not raise ``CROSSSLOT``. + Args: keys (Union[str, List[str]]): The document ID or IDs to remove from the index. + batch_size (int): Number of keys to unlink per round-trip. Returns: int: Count of records deleted from Redis. """ - if isinstance(keys, list): - return self._redis_client.unlink(*keys) # type: ignore - else: + if not isinstance(keys, list): return self._redis_client.unlink(keys) # type: ignore - def drop_documents(self, ids: str | list[str]) -> int: + if not keys: + return 0 + + total = 0 + for i in range(0, len(keys), batch_size): + total += self._unlink_batch(keys[i : i + batch_size]) + return total + + def drop_documents( + self, ids: str | list[str], batch_size: int = DEFAULT_BULK_BATCH_SIZE + ) -> int: """Remove documents from the index by their document IDs. This method converts document IDs to Redis keys automatically by applying @@ -854,25 +1006,41 @@ def drop_documents(self, ids: str | list[str]) -> int: Args: ids (Union[str, List[str]]): The document ID or IDs to remove from the index. + batch_size (int): Number of documents to delete per round-trip + (standalone Redis only; cluster deletes in a single call after + the shared-hash-tag check). Returns: int: Count of documents deleted from Redis. + + Raises: + ValueError: On Redis Cluster, if the resolved keys do not all share + a hash tag (a cross-slot ``DELETE`` is not permitted). """ - if isinstance(ids, list): - if not ids: - return 0 - keys = [self.key(id) for id in ids] - # Check for cluster compatibility - if isinstance( - self._redis_client, RedisCluster - ) and not _keys_share_hash_tag(keys): + if not isinstance(ids, list): + key = self.key(ids) + return self._redis_client.delete(key) # type: ignore + + if not ids: + return 0 + + keys = [self.key(id) for id in ids] + + # On cluster, a multi-key DELETE must stay within one hash slot, so + # require the keys to share a hash tag and delete them in a single call. + if isinstance(self._redis_client, RedisCluster): + if not _keys_share_hash_tag(keys): raise ValueError( "All keys must share a hash tag when using Redis Cluster." ) return self._redis_client.delete(*keys) # type: ignore - else: - key = self.key(ids) - return self._redis_client.delete(key) # type: ignore + + # Standalone: delete in chunks to bound the size of any single command. + total = 0 + for i in range(0, len(keys), batch_size): + batch = keys[i : i + batch_size] + total += cast(int, self._redis_client.delete(*batch)) + return total def expire_keys(self, keys: str | list[str], ttl: int) -> int | list[int]: """Set the expiration time for a specific entry or entries in Redis. @@ -889,6 +1057,269 @@ def expire_keys(self, keys: str | list[str], ttl: int) -> int | list[int]: else: return self._redis_client.expire(keys, ttl) # type: ignore + def drop_by_filter( + self, + filter_expression: str | FilterExpression, + *, + batch_size: int = DEFAULT_BULK_BATCH_SIZE, + dry_run: bool = False, + allow_all: bool = False, + on_progress: Callable[[int, int], None] | None = None, + ) -> BulkResult: + """Delete every document matching a filter expression. + + Redis has no server-side "delete by query", so RedisVL resolves the + matching document keys and removes them with non-blocking ``UNLINK`` in + batches. Matching documents leave the result set as they are deleted, so + this re-queries from offset 0 each round (the same strategy as + :meth:`clear`) and is not subject to the ``MAXSEARCHRESULTS`` limit + (provided ``batch_size`` itself does not exceed it). + + Args: + filter_expression (Union[str, FilterExpression]): Selects the + documents to delete. Prefer the escaping builders (``Tag``, + ``Num``, ``Text``...) over raw filter strings; **never + string-concatenate untrusted input into a filter** — unlike a + read query, an injected predicate here deletes data. + batch_size (int): Number of documents to resolve and unlink per + round-trip. Defaults to 500. + dry_run (bool): If True, report how many documents *would* be + deleted (via a count query) without deleting anything. + allow_all (bool): Must be True to run against a match-all filter + (empty/``"*"``/``None``). Prefer :meth:`clear` to intentionally + empty the index. + on_progress (Optional[Callable[[int, int], None]]): Called after each + batch with ``(processed, matched)`` — the cumulative documents + deleted and the total matched at the start. Invoked + synchronously (do not pass a coroutine); raising from it aborts + the run (already-deleted documents stay deleted). + + Returns: + BulkResult: ``matched``/``processed`` counts, plus ``completed`` + (False if the runaway backstop tripped) and ``dry_run``. Read the + fields explicitly (``result.processed``, ``result.completed``). + + See Also: + :meth:`update_by_filter` (bulk partial update), :meth:`drop_documents` + / :meth:`drop_keys` (delete by id/key), :meth:`clear` (delete all). + + Note: + This operation is **not atomic** across the match set. Each key is + unlinked atomically, but batches are applied incrementally with no + rollback, so a crash or connection error mid-run leaves the already + -deleted documents gone and the rest in place. Deletes are + idempotent: re-running the same call after a failure removes only + whatever still matches, converging on the intended state. + """ + _require_specific_filter(filter_expression, allow_all) + + matched = cast(int, self.query(CountQuery(filter_expression))) + if dry_run: + return BulkResult( + matched=matched, processed=matched, completed=True, dry_run=True + ) + + # Runaway backstop sized to the matched count (plus slack for concurrent + # inserts); normal termination is an empty offset-0 re-query. + max_records = ceil(matched * 1.5) + batch_size + total_deleted = 0 + completed = True + query = FilterQuery(filter_expression, return_fields=["id"]) + query.paging(0, batch_size) + + while True: + if total_deleted > max_records: + logger.warning( + "drop_by_filter hit its runaway backstop (%d) with documents " + "possibly still matching; returning an incomplete result. " + "Re-run to continue.", + max_records, + ) + completed = False + break + batch = self._query(query) + if not batch: + break + batch_keys = [record["id"] for record in batch] + total_deleted += self._unlink_batch(batch_keys) + if on_progress is not None: + on_progress(total_deleted, matched) + + self.invalidate_sql_schema_cache() + return BulkResult(matched=matched, processed=total_deleted, completed=completed) + + def update_by_filter( + self, + filter_expression: str | FilterExpression, + values: dict[str, Any], + *, + batch_size: int = DEFAULT_BULK_BATCH_SIZE, + dry_run: bool = False, + allow_all: bool = False, + on_progress: Callable[[int, int], None] | None = None, + ) -> BulkResult: + """Set ``values`` on every document matching a filter expression. + + This is a partial update: fields not present in ``values`` are left + untouched. For hash indexes the fields are written with ``HSET``; for + JSON indexes they are merged at the document root (``$``) with + ``JSON.MERGE`` (RFC 7396), so nested objects merge recursively, arrays + are replaced wholesale, and a ``None`` value deletes that path. + + Because the read phase (an ``FT.AGGREGATE`` cursor) cannot safely run + while the index is being written, all matching keys are resolved into + memory *before* any write, then updated in batches. Memory use is + therefore proportional to the match count; for very large match sets, + narrow the filter and run in partitions (see the user guide). + + Args: + filter_expression (Union[str, FilterExpression]): Selects the + documents to update. Prefer the escaping builders (``Tag``, + ``Num``...) over raw filter strings; **never string-concatenate + untrusted input into a filter** — an injected predicate here + mutates data. + values (Dict[str, Any]): Field/value pairs to set on each match. + Values are written **as-is with no schema validation** (unlike + :meth:`load`): callers must pre-encode vectors/bytes and format + numerics as the schema expects. For JSON indexes, keys must + match the document's JSON **layout**, not the schema field name + — a field indexed at a nested path (e.g. ``$.metadata.status``) + must be passed nested (``{"metadata": {"status": ...}}``); + passing the flat field name writes the wrong path and leaves the + indexed field unchanged. Only static values are supported (no + callable/expression transforms). + batch_size (int): Number of documents to update per round-trip. + dry_run (bool): If True, report how many documents *would* be + updated without writing anything. + allow_all (bool): Must be True to run against a match-all filter. + on_progress (Optional[Callable[[int, int], None]]): Called after each + write batch with ``(processed, matched)``. Invoked synchronously + (do not pass a coroutine); raising from it aborts the run. + + Returns: + BulkResult: ``matched``/``processed`` counts, plus ``completed`` + (always True for update — it runs to completion or raises) and + ``dry_run``. + + See Also: + :meth:`drop_by_filter`, :meth:`load` (validated whole-document + upsert by key). + + Note: + This operation is **not atomic** across the match set. Each + document is updated atomically (one ``HSET``/``JSON.MERGE``), but + batches use a non-transactional pipeline and are applied + incrementally with no rollback, so a crash or connection error + mid-run can leave some documents updated and others not. Because + the update is a fixed field set, it is idempotent: re-running the + same call after a failure converges on the intended state. + + Keys are resolved before writing, so a document may be deleted by + another client in between. Each write is conditional on the key + still existing (applied atomically), so such a document is + **skipped rather than recreated** as a partial document; it simply + isn't counted in ``processed``. ``processed`` therefore reflects the + documents actually written, which can be less than ``matched`` under + concurrent deletion. + """ + _require_specific_filter(filter_expression, allow_all) + if not values: + raise ValueError("values must be a non-empty mapping of field to value.") + + matched = cast(int, self.query(CountQuery(filter_expression))) + if dry_run: + return BulkResult( + matched=matched, processed=matched, completed=True, dry_run=True + ) + + # Phase 1 (read): fully resolve matching keys before mutating. The + # aggregation cursor cannot be read while the index is written, so the + # read phase must finish first. on_progress fires only in phase 2. + keys = [ + key + for batch in self._iter_keys_by_filter(filter_expression, batch_size) + for key in batch + ] + + # Phase 2 (write): cursor closed, safe to mutate the index. + total_updated = 0 + for i in range(0, len(keys), batch_size): + total_updated += self._apply_update_batch(keys[i : i + batch_size], values) + if on_progress is not None: + on_progress(total_updated, matched) + + return BulkResult(matched=matched, processed=total_updated, completed=True) + + def _iter_keys_by_filter( + self, filter_expression: str | FilterExpression, batch_size: int + ) -> Iterator[list[str]]: + """Yield batches of document keys matching a filter using a cursor. + + Uses ``FT.AGGREGATE ... LOAD 1 @__key WITHCURSOR`` for resumable, + unbounded iteration (unlike ``FT.SEARCH`` + ``LIMIT``, which is capped + by ``MAXSEARCHRESULTS`` and non-deterministic without a unique sort). + The server-side cursor is always released on exit, even if iteration is + abandoned early or errors. + """ + request = ( + AggregateRequest(str(filter_expression)) + .load("__key") + .cursor(count=batch_size, max_idle=BULK_CURSOR_MAX_IDLE_SECONDS) + .dialect(2) + ) + name = self.schema.index.name + ft = self._redis_client.ft(name) # type: ignore[union-attr] + cid = 0 + try: + result = ft.aggregate(request) + while True: + # Capture the cursor id BEFORE parsing rows, so the finally + # block can release the cursor even if row parsing raises. + # A cursor id of 0 signals the server has released the cursor. + cid = result.cursor.cid if result.cursor else 0 + keys = [_agg_row_to_key(row) for row in result.rows] + if keys: + yield keys + if not cid: + break + result = ft.aggregate(Cursor(cid)) + except redis.exceptions.RedisError as e: + raise RedisSearchError( + f"Error while iterating documents by filter: {str(e)}" + ) from e + finally: + # Release the cursor if it is still open (early break/error/abandon). + if cid: + try: + self._redis_client.execute_command( # type: ignore[union-attr] + "FT.CURSOR", "DEL", name, cid + ) + except redis.exceptions.RedisError: + pass + + def _apply_update_batch(self, batch_keys: list[str], values: dict[str, Any]) -> int: + """Apply a partial update to a batch of keys, skipping any that no longer + exist (guards against recreating a concurrently-deleted document). + + Returns the number of documents actually written. + """ + lua, args = _update_script_and_args( + self.storage_type == StorageType.JSON, values + ) + client = self._redis_client + + if isinstance(client, RedisCluster): + # Per-key to stay within a single slot per command. + written = 0 + for key in batch_keys: + written += int(cast(int, client.eval(lua, 1, key, *args))) + return written + + pipe = client.pipeline(transaction=False) # type: ignore[union-attr] + for key in batch_keys: + pipe.eval(lua, 1, key, *args) + return sum(int(r) for r in pipe.execute()) + def load( self, data: Iterable[Any], @@ -1781,7 +2212,28 @@ def invalidate_sql_schema_cache(self) -> None: """Clear cached sql-redis executors and schema state for this index.""" self._sql_executors.clear() - async def drop_keys(self, keys: str | list[str]) -> int: + async def _unlink_batch(self, batch_keys: list[str]) -> int: + """Unlink a batch of keys from Redis (async). + + Mirrors ``_delete_batch`` but uses non-blocking ``UNLINK``. For Redis + Cluster, keys are unlinked individually to avoid cross-slot errors. + """ + client = await self._get_client() + + if isinstance(client, AsyncRedisCluster): + unlinked = 0 + for key_to_unlink in batch_keys: + try: + unlinked += cast(int, await client.unlink(key_to_unlink)) + except redis.exceptions.RedisError as e: + logger.warning(f"Failed to unlink key {key_to_unlink}: {e}") + return unlinked + + return cast(int, await client.unlink(*batch_keys)) + + async def drop_keys( + self, keys: str | list[str], batch_size: int = DEFAULT_BULK_BATCH_SIZE + ) -> int: """Remove a specific entry or entries from the index by it's key ID. Uses ``UNLINK`` rather than ``DEL`` so memory reclamation runs on a @@ -1789,19 +2241,33 @@ async def drop_keys(self, keys: str | list[str]) -> int: number of keys are dropped at once (for example, scope-targeted ``SemanticCache`` invalidation). The returned count is unchanged. + Large key lists are unlinked in chunks of ``batch_size``. On Redis + Cluster, keys are unlinked individually so a chunk that spans hash + slots does not raise ``CROSSSLOT``. + Args: keys (Union[str, List[str]]): The document ID or IDs to remove from the index. + batch_size (int): Number of keys to unlink per round-trip. Returns: int: Count of records deleted from Redis. """ client = await self._get_client() - if isinstance(keys, list): - return await client.unlink(*keys) - else: + + if not isinstance(keys, list): return await client.unlink(keys) - async def drop_documents(self, ids: str | list[str]) -> int: + if not keys: + return 0 + + total = 0 + for i in range(0, len(keys), batch_size): + total += await self._unlink_batch(keys[i : i + batch_size]) + return total + + async def drop_documents( + self, ids: str | list[str], batch_size: int = DEFAULT_BULK_BATCH_SIZE + ) -> int: """Remove documents from the index by their document IDs. This method converts document IDs to Redis keys automatically by applying @@ -1813,24 +2279,43 @@ async def drop_documents(self, ids: str | list[str]) -> int: Args: ids (Union[str, List[str]]): The document ID or IDs to remove from the index. + batch_size (int): Number of documents to delete per round-trip + (standalone Redis only; cluster deletes in a single call after + the shared-hash-tag check). Returns: int: Count of documents deleted from Redis. + + Raises: + ValueError: On Redis Cluster, if the resolved keys do not all share + a hash tag (a cross-slot ``DELETE`` is not permitted). """ client = await self._get_client() - if isinstance(ids, list): - if not ids: - return 0 - keys = [self.key(id) for id in ids] - # Check for cluster compatibility - if isinstance(client, AsyncRedisCluster) and not _keys_share_hash_tag(keys): + + if not isinstance(ids, list): + key = self.key(ids) + return await client.delete(key) + + if not ids: + return 0 + + keys = [self.key(id) for id in ids] + + # On cluster, a multi-key DELETE must stay within one hash slot, so + # require the keys to share a hash tag and delete them in a single call. + if isinstance(client, AsyncRedisCluster): + if not _keys_share_hash_tag(keys): raise ValueError( "All keys must share a hash tag when using Redis Cluster." ) return await client.delete(*keys) - else: - key = self.key(ids) - return await client.delete(key) + + # Standalone: delete in chunks to bound the size of any single command. + total = 0 + for i in range(0, len(keys), batch_size): + batch = keys[i : i + batch_size] + total += cast(int, await client.delete(*batch)) + return total async def expire_keys(self, keys: str | list[str], ttl: int) -> int | list[int]: """Set the expiration time for a specific entry or entries in Redis. @@ -1848,6 +2333,160 @@ async def expire_keys(self, keys: str | list[str], ttl: int) -> int | list[int]: else: return await client.expire(keys, ttl) + async def drop_by_filter( + self, + filter_expression: str | FilterExpression, + *, + batch_size: int = DEFAULT_BULK_BATCH_SIZE, + dry_run: bool = False, + allow_all: bool = False, + on_progress: Callable[[int, int], None] | None = None, + ) -> BulkResult: + """Delete every document matching a filter expression (async). + + See :meth:`SearchIndex.drop_by_filter` for full semantics. + """ + _require_specific_filter(filter_expression, allow_all) + + matched = cast(int, await self.query(CountQuery(filter_expression))) + if dry_run: + return BulkResult( + matched=matched, processed=matched, completed=True, dry_run=True + ) + + max_records = ceil(matched * 1.5) + batch_size + total_deleted = 0 + completed = True + query = FilterQuery(filter_expression, return_fields=["id"]) + query.paging(0, batch_size) + + while True: + if total_deleted > max_records: + logger.warning( + "drop_by_filter hit its runaway backstop (%d) with documents " + "possibly still matching; returning an incomplete result. " + "Re-run to continue.", + max_records, + ) + completed = False + break + batch = await self._query(query) + if not batch: + break + batch_keys = [record["id"] for record in batch] + total_deleted += await self._unlink_batch(batch_keys) + if on_progress is not None: + on_progress(total_deleted, matched) + + self.invalidate_sql_schema_cache() + return BulkResult(matched=matched, processed=total_deleted, completed=completed) + + async def update_by_filter( + self, + filter_expression: str | FilterExpression, + values: dict[str, Any], + *, + batch_size: int = DEFAULT_BULK_BATCH_SIZE, + dry_run: bool = False, + allow_all: bool = False, + on_progress: Callable[[int, int], None] | None = None, + ) -> BulkResult: + """Set ``values`` on every document matching a filter expression (async). + + See :meth:`SearchIndex.update_by_filter` for full semantics. + """ + _require_specific_filter(filter_expression, allow_all) + if not values: + raise ValueError("values must be a non-empty mapping of field to value.") + + matched = cast(int, await self.query(CountQuery(filter_expression))) + if dry_run: + return BulkResult( + matched=matched, processed=matched, completed=True, dry_run=True + ) + + # Phase 1 (read): fully resolve matching keys before mutating (the + # cursor cannot be read while the index is written). on_progress fires + # only in phase 2. + keys: list[str] = [] + async for batch in self._iter_keys_by_filter(filter_expression, batch_size): + keys.extend(batch) + + # Phase 2 (write): cursor closed, safe to mutate the index. + total_updated = 0 + for i in range(0, len(keys), batch_size): + total_updated += await self._apply_update_batch( + keys[i : i + batch_size], values + ) + if on_progress is not None: + on_progress(total_updated, matched) + + return BulkResult(matched=matched, processed=total_updated, completed=True) + + async def _iter_keys_by_filter( + self, filter_expression: str | FilterExpression, batch_size: int + ) -> AsyncGenerator[list[str], None]: + """Yield batches of document keys matching a filter using a cursor (async). + + Always releases the server-side cursor on exit (see sync counterpart). + """ + client = await self._get_client() + request = ( + AggregateRequest(str(filter_expression)) + .load("__key") + .cursor(count=batch_size, max_idle=BULK_CURSOR_MAX_IDLE_SECONDS) + .dialect(2) + ) + name = self.schema.index.name + ft = client.ft(name) + cid = 0 + try: + result = await ft.aggregate(request) # type: ignore[arg-type] + while True: + # Capture the cursor id BEFORE parsing rows, so the finally + # block can release the cursor even if row parsing raises. + # A cursor id of 0 signals the server has released the cursor. + cid = result.cursor.cid if result.cursor else 0 + keys = [_agg_row_to_key(row) for row in result.rows] + if keys: + yield keys + if not cid: + break + result = await ft.aggregate(Cursor(cid)) + except redis.exceptions.RedisError as e: + raise RedisSearchError( + f"Error while iterating documents by filter: {str(e)}" + ) from e + finally: + if cid: + try: + await client.execute_command("FT.CURSOR", "DEL", name, cid) + except redis.exceptions.RedisError: + pass + + async def _apply_update_batch( + self, batch_keys: list[str], values: dict[str, Any] + ) -> int: + """Apply a partial update to a batch of keys, skipping any that no longer + exist (async; see sync counterpart). Returns documents actually written. + """ + client = await self._get_client() + lua, args = _update_script_and_args( + self.storage_type == StorageType.JSON, values + ) + + if isinstance(client, AsyncRedisCluster): + written = 0 + for key in batch_keys: + written += int(cast(int, await client.eval(lua, 1, key, *args))) # type: ignore[misc] + return written + + pipe = client.pipeline(transaction=False) + for key in batch_keys: + pipe.eval(lua, 1, key, *args) + results = await pipe.execute() + return sum(int(r) for r in results) + @deprecated_argument("concurrency", "Use batch_size instead.") async def load( self, diff --git a/tests/integration/test_bulk_operations.py b/tests/integration/test_bulk_operations.py new file mode 100644 index 00000000..dc584b84 --- /dev/null +++ b/tests/integration/test_bulk_operations.py @@ -0,0 +1,321 @@ +"""Integration tests for bulk delete/update by filter (RAAE-1326).""" + +import pytest + +from redisvl.index import AsyncSearchIndex, SearchIndex +from redisvl.index.index import BulkResult +from redisvl.query import CountQuery +from redisvl.query.filter import Num, Tag +from redisvl.schema import IndexSchema + +BULK_FIELDS = [ + {"name": "cat", "type": "tag"}, + {"name": "status", "type": "tag"}, + {"name": "n", "type": "numeric", "attrs": {"sortable": True}}, +] + + +def _schema(name, storage_type): + return IndexSchema.from_dict( + { + "index": {"name": name, "prefix": name, "storage_type": storage_type}, + "fields": BULK_FIELDS, + } + ) + + +def _data(n=40): + return [ + { + "id": str(i), + "cat": "a" if i % 2 else "b", + "status": "draft", + "n": i, + "keep": "orig", + } + for i in range(n) + ] + + +@pytest.fixture +def hash_index(worker_id, client): + index = SearchIndex( + schema=_schema(f"bulk_h_{worker_id}", "hash"), redis_client=client + ) + index.create(overwrite=True, drop=True) + yield index + index.delete(drop=True) + + +@pytest.fixture +def json_index(worker_id, client): + index = SearchIndex( + schema=_schema(f"bulk_j_{worker_id}", "json"), redis_client=client + ) + index.create(overwrite=True, drop=True) + yield index + index.delete(drop=True) + + +# --------------------------------------------------------------------------- # +# BulkResult contract +# --------------------------------------------------------------------------- # +def test_bulk_result_fields(hash_index): + hash_index.load(_data(), id_field="id") + result = hash_index.drop_by_filter(Tag("cat") == "a") + assert isinstance(result, BulkResult) + assert result.matched == 20 + assert result.processed == 20 + assert result.completed is True + assert result.dry_run is False + + +# --------------------------------------------------------------------------- # +# drop_by_filter +# --------------------------------------------------------------------------- # +def test_drop_by_filter_removes_only_matches(hash_index): + hash_index.load(_data(), id_field="id") + result = hash_index.drop_by_filter(Tag("cat") == "a", batch_size=10) + assert result.processed == 20 + assert hash_index.query(CountQuery(Tag("cat") == "a")) == 0 + assert hash_index.query(CountQuery(Tag("cat") == "b")) == 20 + + +def test_drop_by_filter_json(json_index): + json_index.load(_data(30), id_field="id") # 15 'a' + 15 'b' + result = json_index.drop_by_filter(Tag("cat") == "a", batch_size=10) + assert result.processed == 15 + assert json_index.query(CountQuery(Tag("cat") == "a")) == 0 + assert json_index.query(CountQuery(Tag("cat") == "b")) == 15 + + +def test_drop_by_filter_dry_run_does_not_mutate(hash_index): + hash_index.load(_data(), id_field="id") + result = hash_index.drop_by_filter(Num("n") < 10, dry_run=True) + assert result.dry_run is True + assert result.matched == 10 + assert result.processed == 10 + assert hash_index.query(CountQuery(Num("n") < 10)) == 10 # nothing removed + + +def test_drop_by_filter_reports_progress(hash_index): + hash_index.load(_data(), id_field="id") + progress = [] + hash_index.drop_by_filter( + Tag("cat") == "a", + batch_size=5, + on_progress=lambda p, t: progress.append((p, t)), + ) + assert progress + processed = [p for p, _ in progress] + assert processed == sorted(processed) + assert processed[-1] == 20 + assert all(total == 20 for _, total in progress) # matched total passed through + + +@pytest.mark.parametrize("bad_filter", [None, "*", ""]) +def test_drop_by_filter_guards_match_all(hash_index, bad_filter): + hash_index.load(_data(), id_field="id") + with pytest.raises(ValueError): + hash_index.drop_by_filter(bad_filter) + assert hash_index.query(CountQuery(Num("n") >= 0)) == 40 + + +def test_drop_by_filter_allow_all_override(hash_index): + hash_index.load(_data(), id_field="id") + result = hash_index.drop_by_filter("*", allow_all=True, batch_size=10) + assert result.processed == 40 + assert hash_index.query(CountQuery(Num("n") >= 0)) == 0 + + +# --------------------------------------------------------------------------- # +# update_by_filter +# --------------------------------------------------------------------------- # +def test_update_by_filter_hash_is_partial(hash_index): + hash_index.load(_data(), id_field="id") + result = hash_index.update_by_filter( + Tag("cat") == "b", {"status": "published"}, batch_size=10 + ) + assert result.processed == 20 + doc = hash_index.fetch("0") # id 0 -> cat "b" + assert doc["status"] == "published" + assert doc["keep"] == "orig" # untouched field preserved + assert doc["cat"] == "b" + assert hash_index.fetch("1")["status"] == "draft" # non-matching untouched + + +def test_update_by_filter_json_merges_partially(json_index): + json_index.load( + [{"id": "x", "cat": "a", "status": "draft", "n": 1, "obj": {"p": 1, "q": 2}}], + id_field="id", + ) + json_index.update_by_filter( + Tag("cat") == "a", {"status": "done", "obj": {"q": 9, "r": 3}} + ) + doc = json_index.fetch("x") + assert doc["status"] == "done" + # nested object merged recursively, not replaced + assert doc["obj"] == {"p": 1, "q": 9, "r": 3} + + +def test_update_by_filter_json_none_deletes_path(json_index): + json_index.load( + [{"id": "x", "cat": "a", "status": "draft", "n": 1, "drop_me": "gone"}], + id_field="id", + ) + json_index.update_by_filter(Tag("cat") == "a", {"drop_me": None}) + doc = json_index.fetch("x") + assert "drop_me" not in doc # None deletes the path (RFC 7396) + + +def test_update_by_filter_spans_multiple_batches_hash(hash_index): + # more docs than batch_size to exercise cursor staging + multi-batch writes + hash_index.load(_data(120), id_field="id") + progress = [] + result = hash_index.update_by_filter( + Tag("cat") == "a", + {"status": "x"}, + batch_size=25, + on_progress=lambda p, t: progress.append((p, t)), + ) + assert result.processed == 60 + assert len([p for p, _ in progress]) >= 2 # spanned multiple batches + assert progress[-1] == (60, 60) + + +def test_update_by_filter_spans_multiple_batches_json(json_index): + json_index.load(_data(120), id_field="id") + result = json_index.update_by_filter( + Tag("cat") == "b", {"status": "live"}, batch_size=25 + ) + assert result.processed == 60 + # spot-check a couple of docs across batches + assert json_index.fetch("0")["status"] == "live" + assert json_index.fetch("118")["status"] == "live" + + +def test_update_by_filter_skips_concurrently_deleted_key(hash_index, client): + # Simulate the resolve->write race: a matched key is deleted before the + # write lands. The guard must skip it (not recreate a partial doc) and not + # count it in processed. + hash_index.load(_data(4), id_field="id") # ids 1,3 -> 'a'; 0,2 -> 'b' + a_keys = [hash_index.key("1"), hash_index.key("3")] + client.unlink(a_keys[0]) # id "1" vanishes mid-flight + + written = hash_index._apply_update_batch(a_keys, {"status": "x"}) + + assert written == 1 # deleted key skipped, not counted + assert not client.exists(a_keys[0]) # NOT recreated as a partial doc + assert hash_index.fetch("3")["status"] == "x" # survivor updated + + +def test_update_by_filter_json_skips_missing_key(json_index, client): + json_index.load([{"id": "x", "cat": "a", "status": "draft", "n": 1}], id_field="id") + missing = json_index.key("does-not-exist") + written = json_index._apply_update_batch( + [json_index.key("x"), missing], {"status": "done"} + ) + assert written == 1 + assert not client.exists(missing) # not recreated + assert json_index.fetch("x")["status"] == "done" + + +def test_update_by_filter_processed_less_than_matched_on_delete(hash_index, client): + # End-to-end through the public API: delete a still-pending matching doc + # mid-run (via on_progress) and confirm it's skipped, so processed < matched. + hash_index.load(_data(6), id_field="id") # 'a' ids: 1, 3, 5 + state = {"deleted": False} + + def on_progress(processed, matched): + if state["deleted"]: + return + for i in ("1", "3", "5"): + doc = hash_index.fetch(i) + if doc and doc.get("status") != "x": # not yet written -> still pending + client.unlink(hash_index.key(i)) + state["deleted"] = True + return + + result = hash_index.update_by_filter( + Tag("cat") == "a", {"status": "x"}, batch_size=1, on_progress=on_progress + ) + assert result.matched == 3 + assert result.processed == 2 # one pending match deleted mid-run -> skipped + + +def test_update_by_filter_dry_run_and_empty_values(hash_index): + hash_index.load(_data(), id_field="id") + result = hash_index.update_by_filter( + Tag("cat") == "a", {"status": "z"}, dry_run=True + ) + assert result.dry_run is True and result.processed == 20 + assert hash_index.fetch("1")["status"] == "draft" # unchanged + with pytest.raises(ValueError): + hash_index.update_by_filter(Tag("cat") == "a", {}) + + +# --------------------------------------------------------------------------- # +# batched drop_documents / drop_keys +# --------------------------------------------------------------------------- # +def test_drop_documents_batched(hash_index): + hash_index.load(_data(30), id_field="id") + dropped = hash_index.drop_documents([str(i) for i in range(20)], batch_size=7) + assert dropped == 20 + assert hash_index.query(CountQuery(Num("n") >= 0)) == 10 + + +def test_drop_keys_batched(hash_index): + keys = hash_index.load(_data(30), id_field="id") + dropped = hash_index.drop_keys(keys[:20], batch_size=7) + assert dropped == 20 + assert hash_index.query(CountQuery(Num("n") >= 0)) == 10 + + +# --------------------------------------------------------------------------- # +# async coverage +# --------------------------------------------------------------------------- # +@pytest.mark.asyncio +async def test_async_delete_and_update_by_filter(worker_id, async_client): + index = AsyncSearchIndex( + schema=_schema(f"bulk_ah_{worker_id}", "hash"), redis_client=async_client + ) + await index.create(overwrite=True, drop=True) + try: + await index.load(_data(120), id_field="id") + + with pytest.raises(ValueError): + await index.drop_by_filter(None) + + result = await index.update_by_filter( + Tag("cat") == "b", {"status": "live"}, batch_size=25 + ) + assert result.processed == 60 + doc = await index.fetch("0") + assert doc["status"] == "live" and doc["keep"] == "orig" + + result = await index.drop_by_filter(Num("n") < 60, batch_size=25) + assert result.processed == 60 + assert await index.query(CountQuery(Num("n") < 60)) == 0 + finally: + await index.delete(drop=True) + + +@pytest.mark.asyncio +async def test_async_update_skips_missing_key(worker_id, async_client): + # Cover the async existence-guard path: a missing key is skipped, not recreated. + index = AsyncSearchIndex( + schema=_schema(f"bulk_askip_{worker_id}", "hash"), redis_client=async_client + ) + await index.create(overwrite=True, drop=True) + try: + await index.load(_data(4), id_field="id") # 'a' ids: 1, 3 + keys = [index.key("1"), index.key("3")] + await async_client.unlink(keys[0]) + + written = await index._apply_update_batch(keys, {"status": "x"}) + + assert written == 1 + assert not await async_client.exists(keys[0]) # not recreated + assert (await index.fetch("3"))["status"] == "x" + finally: + await index.delete(drop=True)