query_runner.py

# Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

"""Query execution and ground truth computation."""

import asyncio
import logging
import time

import numpy as np
from vespa.application import Vespa

from config import Config
from recall import recall_at_k
from tuning_request import TuningRequest

log = logging.getLogger(__name__)

# Per-query ground truth cache: (filter_key, query_index) -> list of doc IDs
_exact_gt_cache: dict[tuple[float, int], list[int]] = {}


async def compute_exact_ground_truth(
    app: Vespa,
    cfg: Config,
    query_vectors: np.ndarray,
    filter_threshold: float,
    num_queries: int,
    concurrency: int = 4,
) -> list[list[int]]:
    """Compute exact ground truth for queries at a given filter threshold using Vespa."""
    filter_key = round(filter_threshold, 5)

    cached = []
    to_compute = []
    for idx in range(num_queries):
        key = (filter_key, idx)
        if key in _exact_gt_cache:
            cached.append((idx, _exact_gt_cache[key]))
        else:
            to_compute.append(idx)

    if not to_compute:
        return [r for _, r in sorted(cached)]

    log.info(
        "Computing exact ground truth for %d/%d queries (filter=%.5f, %d cached)",
        len(to_compute),
        num_queries,
        filter_threshold,
        len(cached),
    )

    filter_clause = "" if filter_threshold < 0 else f" and random >= {filter_threshold}"
    sem = asyncio.Semaphore(concurrency)

    async def run_exact(idx: int) -> tuple[int, list[int]]:
        async with sem:
            vec = query_vectors[idx]
            yql = (
                f"select id from {cfg.schema_name} where "
                f"{{targetHits:{cfg.target_hits},approximate:false}}"
                f"nearestNeighbor(embedding,q){filter_clause}"
            )
            params = {
                "yql": yql,
                "ranking": "closeness",
                "input.query(q)": str(vec.tolist()),
                "timeout": "30s",
            }
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(None, lambda: app.query(body=params))
            if not response.is_successful():
                return idx, []
            return idx, [hit["fields"]["id"] for hit in response.hits]

    computed = await asyncio.gather(*[run_exact(idx) for idx in to_compute])

    for idx, ids in computed:
        result = ids[: cfg.target_hits]
        _exact_gt_cache[(filter_key, idx)] = result
        cached.append((idx, result))

    return [r for _, r in sorted(cached)]


async def run_single_query(
    app: Vespa,
    cfg: Config,
    idx: int,
    query_vec: np.ndarray,
    gt_neighbors: list[int],
    explore_additional_hits: int,
    approximate_threshold: float,
    filter_threshold: float,
    post_filter_threshold: float,
    use_exact: bool,
    filter_first_threshold: float,
    filter_first_exploration: float,
    target_hits_max_adjustment_factor: float,
    exploration_slack: float,
    max_latency_ms: int,
) -> dict | None:
    """Execute a single query and return latency/recall."""
    filter_clause = "" if filter_threshold < 0 else f" and random >= {filter_threshold}"

    if use_exact:
        yql = (
            f"select id from {cfg.schema_name} where "
            f"{{targetHits:{cfg.target_hits},approximate:false}}"
            f"nearestNeighbor(embedding,q){filter_clause}"
        )
    else:
        yql = (
            f"select id from {cfg.schema_name} where "
            f"{{targetHits:{cfg.target_hits},"
            f"hnsw.exploreAdditionalHits:{explore_additional_hits}}}"
            f"nearestNeighbor(embedding,q){filter_clause}"
        )

    query_params = {
        "yql": yql,
        "ranking": "closeness",
        "input.query(q)": str(query_vec.tolist()),
        "ranking.matching.approximateThreshold": approximate_threshold,
        "ranking.matching.postFilterThreshold": post_filter_threshold,
        "ranking.matching.filterFirstThreshold": filter_first_threshold,
        "ranking.matching.filterFirstExploration": filter_first_exploration,
        "ranking.matching.targetHitsMaxAdjustmentFactor": target_hits_max_adjustment_factor,
        "ranking.matching.explorationSlack": exploration_slack,
        "trace.level": 5,
        "presentation.timing": "true",
        "timeout": f"{max_latency_ms}ms",
    }

    loop = asyncio.get_event_loop()
    start = time.time()
    try:
        response = await loop.run_in_executor(None, lambda: app.query(body=query_params))
        latency_ms = (time.time() - start) * 1000

        # Check for timeout / degraded coverage
        if hasattr(response, "json") and "root" in response.json:
            root = response.json["root"]
            for err in root.get("errors", []):
                if err.get("code") == 12 or "Timed out" in err.get("summary", ""):
                    return {"timeout": True}
            coverage = root.get("coverage", {})
            if coverage.get("degraded", {}).get("timeout", False):
                return {"timeout": True}

        if not response.is_successful():
            log.warning("Query %d failed: %s", idx, getattr(response, "json", "unknown"))
            return None

        retrieved_ids = [hit["fields"]["id"] for hit in response.hits]

        # Detect exact vs approximate from trace
        was_exact = False
        if hasattr(response, "json") and "trace" in response.json:
            trace_str = str(response.json["trace"])
            if "approximate=false" in trace_str or "approximate = false" in trace_str:
                was_exact = True
            elif "Skip calculate global filter" in trace_str and "estimated_hit_ratio" in trace_str:
                # Vespa auto-switched to exact because estimated hit ratio fell
                # below the approximateThreshold (global_filter.lower_limit)
                was_exact = True

        # Internal Vespa timing
        internal_latency_ms = None
        if hasattr(response, "json") and "timing" in response.json:
            qt = response.json["timing"].get("querytime")
            if qt is not None:
                internal_latency_ms = qt * 1000

        recall = recall_at_k(retrieved_ids, gt_neighbors, k=cfg.target_hits)

        result = {
            "latency_ms": latency_ms,
            "internal_latency_ms": internal_latency_ms,
            "recall": recall,
            "exact": was_exact,
        }
        if idx == 0:
            result["query_params"] = query_params
        return result

    except Exception as e:
        error_str = str(e).lower()
        if any(kw in error_str for kw in ("timed out", "timeout", "code': 12")):
            return {"timeout": True}
        log.warning("Query %d exception: %s", idx, e)
        return None


async def run_queries(
    app: Vespa,
    cfg: Config,
    query_vectors: np.ndarray,
    ground_truth: list[list[int]],
    req: TuningRequest,
) -> dict:
    """Run a batch of queries and return aggregated results."""
    num_available = min(len(query_vectors), len(ground_truth))
    num_queries = min(req.num_queries, num_available)

    sem = asyncio.Semaphore(req.concurrency)

    async def run_with_sem(idx: int):
        async with sem:
            return await run_single_query(
                app,
                cfg,
                idx,
                query_vectors[idx],
                ground_truth[idx],
                req.explore_additional_hits,
                req.approximate_threshold,
                req.filter_threshold,
                req.post_filter_threshold,
                req.use_exact,
                req.filter_first_threshold,
                req.filter_first_exploration,
                req.target_hits_max_adjustment_factor,
                req.exploration_slack,
                req.max_latency_ms,
            )

    start = time.time()
    results = await asyncio.gather(*[run_with_sem(i) for i in range(num_queries)])
    total_time_ms = (time.time() - start) * 1000

    successful = []
    timeout_count = 0
    failed_count = 0
    for r in results:
        if r is None:
            failed_count += 1
        elif r.get("timeout"):
            timeout_count += 1
        else:
            successful.append(r)

    log.info(
        "Query results: %d successful, %d timeouts, %d failed / %d total",
        len(successful),
        timeout_count,
        failed_count,
        num_queries,
    )

    if not successful:
        return {
            "error": "All queries failed",
            "timeout_count": timeout_count,
            "failed_count": failed_count,
            "total_requested": num_queries,
        }

    latencies = [r["latency_ms"] for r in successful]
    recalls = [r["recall"] for r in successful]

    example_query_params = None
    if successful and "query_params" in successful[0]:
        example_query_params = successful[0].pop("query_params")

    return {
        "points": successful,
        "stats": {
            "min_latency_ms": float(np.min(latencies)),
            "mean_latency_ms": float(np.mean(latencies)),
            "max_latency_ms": float(np.max(latencies)),
            "p50_latency_ms": float(np.percentile(latencies, 50)),
            "p95_latency_ms": float(np.percentile(latencies, 95)),
            "p99_latency_ms": float(np.percentile(latencies, 99)),
            "mean_recall": float(np.mean(recalls)),
            "num_queries": len(successful),
            "total_requested": num_queries,
            "total_time_ms": total_time_ms,
            "concurrency": req.concurrency,
            "timeout_count": timeout_count,
            "failed_count": failed_count,
        },
        "example_query_params": example_query_params,
    }