velesdb_server/handlers/search/

batch.rs

//! Batch search handler: multiple vector queries in a single request.

use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::IntoResponse,
    Json,
};
use std::sync::Arc;

use crate::types::{BatchSearchRequest, BatchSearchResponse, ErrorResponse, SearchResponse};
use crate::AppState;

use super::pipeline::{
    actionable_search_error, build_search_response, record_circuit_breaker,
    validate_query_dimension,
};
use crate::handlers::helpers::{
    apply_pre_check, extract_client_id, get_vector_collection_or_404, notify_query_timing,
};

/// Batch search for multiple vectors.
#[utoipa::path(
    post,
    path = "/collections/{name}/search/batch",
    tag = "search",
    params(
        ("name" = String, Path, description = "Collection name")
    ),
    request_body = BatchSearchRequest,
    responses(
        (status = 200, description = "Batch search results", body = BatchSearchResponse),
        (status = 404, description = "Collection not found", body = ErrorResponse),
        (status = 400, description = "Invalid request", body = ErrorResponse)
    )
)]
pub async fn batch_search(
    State(state): State<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Path(name): Path<String>,
    Json(req): Json<BatchSearchRequest>,
) -> impl IntoResponse {
    let start = std::time::Instant::now();
    state.onboarding_metrics.record_search_request();

    let collection = match get_vector_collection_or_404(&state, &name) {
        Ok(c) => c,
        Err(resp) => return resp,
    };

    // Record query type only after confirming the collection exists, so
    // 404s do not inflate queries_total or vector_queries.
    state.operational_metrics.record_vector_query();

    let client_id = extract_client_id(&headers);
    if let Err(resp) = apply_pre_check(collection.guard_rails(), &client_id) {
        state.operational_metrics.inc_rate_limited();
        return resp;
    }

    if let Err(resp) = validate_batch_dimensions(&state, &name, &collection, &req) {
        state.operational_metrics.inc_errors();
        return resp;
    }

    let filters = match parse_batch_filters(&state, &req) {
        Ok(f) => f,
        Err(resp) => {
            state.operational_metrics.inc_errors();
            return resp;
        }
    };

    // Gate the read (CORE-2) before dispatching the off-thread batch. A denied
    // decision — or a scope narrowing, which batch has no per-query channel to
    // apply — refuses the batch (fail closed).
    match state.db.authorize_read(
        &name,
        velesdb_core::observer::QueryOperationKind::VectorSearch,
        None,
        None,
    ) {
        Ok(None) => {}
        Ok(Some(_)) | Err(_) => {
            state.operational_metrics.inc_errors();
            return (
                StatusCode::FORBIDDEN,
                Json(ErrorResponse {
                    error: "Read denied by governance policy".to_string(),
                    code: None,
                }),
            )
                .into_response();
        }
    }

    // F-01 sweep: batch search is CPU-bound (one HNSW pass per query) and must
    // not run on the async runtime thread — a large batch would otherwise block
    // a Tokio worker and starve concurrent requests. Move it to a blocking
    // worker like the other search endpoints. `spawn_blocking` needs a 'static
    // closure, so owned query vectors + filters are moved in and the `&[f32]`
    // views are rebuilt inside.
    let collection_for_work = collection.clone();
    let query_vectors: Vec<Vec<f32>> = req.searches.iter().map(|s| s.vector.clone()).collect();
    let max_top_k = req.searches.iter().map(|s| s.top_k).max().unwrap_or(10);

    let batch_result = match tokio::task::spawn_blocking(move || {
        let queries: Vec<&[f32]> = query_vectors.iter().map(Vec::as_slice).collect();
        run_batch_search(&collection_for_work, &queries, max_top_k, &filters)
    })
    .await
    {
        Ok(inner) => inner,
        Err(_) => {
            state.operational_metrics.inc_errors();
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(ErrorResponse {
                    error: "Batch search worker failed".to_string(),
                    code: None,
                }),
            )
                .into_response();
        }
    };
    record_circuit_breaker(&collection, &batch_result);

    let all_results = match batch_result {
        Ok(batch_results) => build_batch_responses(&state, batch_results, &req),
        Err(e) => {
            state.operational_metrics.inc_errors();
            return (StatusCode::BAD_REQUEST, Json(actionable_search_error(&e))).into_response();
        }
    };

    finish_batch_search(&state, &name, start, all_results)
}

/// Dispatch a batch search to the throughput-optimized parallel kernel when no
/// query carries a filter, falling back to the per-query filtered kernel
/// otherwise. Both paths share the same HNSW traversal, so results are
/// identical for the unfiltered case — only the rayon parallelism differs.
fn run_batch_search(
    collection: &velesdb_core::collection::VectorCollection,
    queries: &[&[f32]],
    max_top_k: usize,
    filters: &[Option<velesdb_core::Filter>],
) -> velesdb_core::Result<Vec<Vec<velesdb_core::SearchResult>>> {
    if filters.iter().all(Option::is_none) {
        collection.search_batch_parallel(queries, max_top_k)
    } else {
        collection.search_batch_with_filters(queries, max_top_k, filters)
    }
}

/// Record timing metrics and build the final batch response envelope.
fn finish_batch_search(
    state: &AppState,
    name: &str,
    start: std::time::Instant,
    results: Vec<SearchResponse>,
) -> axum::response::Response {
    let elapsed = start.elapsed();
    let timing_ms = elapsed.as_secs_f64() * 1000.0;
    notify_query_timing(state, name, start);
    state
        .query_duration_histogram
        .observe(elapsed.as_secs_f64());

    Json(BatchSearchResponse { results, timing_ms }).into_response()
}

/// Validate that every query vector in a batch request matches the collection dimension.
#[allow(clippy::result_large_err)]
fn validate_batch_dimensions(
    state: &AppState,
    name: &str,
    collection: &velesdb_core::collection::VectorCollection,
    req: &BatchSearchRequest,
) -> Result<(), axum::response::Response> {
    let expected_dimension = collection.config().dimension;
    for (idx, search) in req.searches.iter().enumerate() {
        if let Err(error) =
            validate_query_dimension(state, name, expected_dimension, &search.vector)
        {
            return Err((
                StatusCode::BAD_REQUEST,
                Json(ErrorResponse {
                    error: format!("Invalid query at index {idx}: {}", error.error),
                    code: error.code.clone(),
                }),
            )
                .into_response());
        }
    }
    Ok(())
}

/// Parse filters from each search in a batch request.
#[allow(clippy::result_large_err)]
fn parse_batch_filters(
    state: &AppState,
    req: &BatchSearchRequest,
) -> Result<Vec<Option<velesdb_core::Filter>>, axum::response::Response> {
    let mut filters: Vec<Option<velesdb_core::Filter>> = Vec::with_capacity(req.searches.len());
    for (idx, search) in req.searches.iter().enumerate() {
        if let Some(filter_json) = &search.filter {
            match serde_json::from_value(filter_json.clone()) {
                Ok(filter) => filters.push(Some(filter)),
                Err(e) => {
                    state.onboarding_metrics.record_filter_parse_error();
                    return Err((
                        StatusCode::BAD_REQUEST,
                        Json(ErrorResponse {
                            error: format!(
                                "Invalid filter at index {idx}: {e}. Hint: validate filter syntax and start with a broader query before reintroducing strict filters."
                            ),
                            code: None,
                        }),
                    )
                        .into_response());
                }
            }
        } else {
            filters.push(None);
        }
    }
    Ok(filters)
}

/// Convert batch search results into response objects, recording metrics for empty results.
fn build_batch_responses(
    state: &AppState,
    batch_results: Vec<Vec<velesdb_core::SearchResult>>,
    req: &BatchSearchRequest,
) -> Vec<SearchResponse> {
    let empty_count = batch_results
        .iter()
        .filter(|results| results.is_empty())
        .count();
    for _ in 0..empty_count {
        state.onboarding_metrics.record_empty_search_results();
    }
    debug_assert_eq!(
        batch_results.len(),
        req.searches.len(),
        "search_batch_with_filters must return one result-vec per query"
    );
    batch_results
        .into_iter()
        .zip(req.searches.iter())
        .map(|(results, search)| {
            let truncated: Vec<_> = results.into_iter().take(search.top_k).collect();
            build_search_response(truncated)
        })
        .collect()
}