Implement count_by_index in redis store.

Author: dkostyrev

nativelink-store/src/redis_store.rs

@@ -65,6 +65,7 @@ use uuid::Uuid;
 use crate::cas_utils::is_zero_digest;
 use crate::redis_utils::{
     FtAggregateCursor, FtAggregateOptions, FtCreateOptions, SearchSchema, ft_aggregate, ft_create,
+    ft_search_count,
 };
 
 /// The default size of the read chunk when reading data from Redis.
@@ -1632,7 +1633,80 @@ where
     where
         K: SchedulerIndexProvider + Send,
     {
-        Err(make_err!(Code::Unimplemented, "Not implemented"))
+        if index.is_empty() {
+            return Ok(Vec::new());
+        }
+        let index_name = format!(
+            "{}",
+            get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY)
+        );
+        let mut counts = Vec::with_capacity(index.len());
+        for idx in index {
+            let index_value = idx.index_value();
+            let sanitized_field = try_sanitize(index_value.as_ref()).err_tip(|| {
+                format!("In RedisStore::count_by_index::try_sanitize - {index_value:?}")
+            })?;
+            let query = if sanitized_field.is_empty() {
+                "*".to_string()
+            } else {
+                format!("@{}:{{ {} }}", K::INDEX_NAME, sanitized_field)
+            };
+            let run_ft_search = || {
+                Ok::<_, Error>(async {
+                    let mut client = self.get_client().await?;
+                    ft_search_count(
+                        &mut client.connection_manager,
+                        index_name.as_str(),
+                        query.as_str(),
+                    )
+                    .await
+                })
+            };
+            let count = run_ft_search()?
+                .or_else(|_| async {
+                    let mut schema = vec![SearchSchema {
+                        field_name: K::INDEX_NAME.into(),
+                        sortable: false,
+                    }];
+                    if let Some(sort_key) = K::MAYBE_SORT_KEY {
+                        schema.push(SearchSchema {
+                            field_name: sort_key.into(),
+                            sortable: true,
+                        });
+                    }
+                    let create_result = ft_create(
+                        self.connection_manager.get_connection().await?.0,
+                        format!(
+                            "{}",
+                            get_index_name!(K::KEY_PREFIX, K::INDEX_NAME, K::MAYBE_SORT_KEY)
+                        ),
+                        FtCreateOptions {
+                            prefixes: vec![K::KEY_PREFIX.into()],
+                            nohl: true,
+                            nofields: true,
+                            nofreqs: true,
+                            nooffsets: true,
+                            temporary: Some(INDEX_TTL_S),
+                        },
+                        schema,
+                    )
+                    .await
+                    .err_tip(|| {
+                        format!(
+                            "Error with ft_create in RedisStore::count_by_index({index_name})",
+                        )
+                    });
+                    let count_result = run_ft_search()?.await.err_tip(|| {
+                        format!(
+                            "Error with second FT.SEARCH count in RedisStore::count_by_index({index_name})",
+                        )
+                    });
+                    count_result.or_else(|e| create_result.merge(Err(e)))
+                })
+                .await?;
+            counts.push(count);
+        }
+        Ok(counts)
     }
 
     async fn search_by_index_prefix<K>(

nativelink-store/src/redis_utils/ft_search_count.rs

@@ -0,0 +1,77 @@
+// Copyright 2024-2026 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    See LICENSE file for details
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use nativelink_error::{Code, Error, ResultExt, make_err};
+use redis::aio::ConnectionLike;
+use redis::Value;
+
+/// Returns the number of documents matching `query` on RediSearch index `index`, without loading
+/// document payloads (`LIMIT 0 0`).
+pub(crate) async fn ft_search_count<C: ConnectionLike + Send>(
+    connection_manager: &mut C,
+    index: &str,
+    query: &str,
+) -> Result<usize, Error> {
+    let res: Value = redis::cmd("FT.SEARCH")
+        .arg(index)
+        .arg(query)
+        .arg("LIMIT")
+        .arg(0_i64)
+        .arg(0_i64)
+        .query_async(connection_manager)
+        .await
+        .err_tip(|| format!("FT.SEARCH count index={index} query={query:?}"))?;
+
+    parse_ft_search_total(res).err_tip(|| format!("parse FT.SEARCH total index={index}"))
+}
+
+fn parse_ft_search_total(value: Value) -> Result<usize, Error> {
+    match value {
+        Value::Array(arr) if !arr.is_empty() => int_to_document_count(&arr[0]),
+        Value::Map(ref entries) => {
+            for (k, v) in entries.iter() {
+                let Value::SimpleString(key) = k else {
+                    continue;
+                };
+                if key == "total_results" {
+                    return int_to_document_count(v);
+                }
+            }
+            Err(make_err!(
+                Code::Internal,
+                "FT.SEARCH map missing total_results: {value:?}"
+            ))
+        }
+        other => Err(make_err!(
+            Code::Internal,
+            "Unexpected FT.SEARCH response: {other:?}"
+        )),
+    }
+}
+
+fn int_to_document_count(v: &Value) -> Result<usize, Error> {
+    let Value::Int(n) = v else {
+        return Err(make_err!(
+            Code::Internal,
+            "FT.SEARCH count field not integer: {v:?}"
+        ));
+    };
+    if *n < 0 {
+        return Err(make_err!(
+            Code::Internal,
+            "Negative document count from FT.SEARCH: {n}"
+        ));
+    }
+    Ok(*n as usize)
+}

nativelink-store/src/redis_utils/mod.rs

@@ -16,5 +16,7 @@ mod aggregate_types;
 mod ft_aggregate;
 mod ft_create;
 mod ft_cursor_read;
+mod ft_search_count;
 pub(crate) use ft_aggregate::{FtAggregateCursor, FtAggregateOptions, ft_aggregate};
 pub(crate) use ft_create::{FtCreateOptions, SearchSchema, ft_create};
+pub(crate) use ft_search_count::ft_search_count;