search method

Implementation

List<VectorSearchHit> search(
  Vector query,
  int k, {
  VectorMetric? metric,
}) {
  if (query.dim != dim) {
    throw StateError(
      'FlatIndex.search: query dim ${query.dim} != index dim $dim',
    );
  }
  if (k <= 0 || _ids.isEmpty) return const [];
  final m = metric ?? defaultMetric;
  final larger = m == VectorMetric.innerProduct;
  final n = _ids.length;
  final effK = math.min(k, n);

  // Precompute query-side scalars once.
  double qNorm = 0.0;
  if (m == VectorMetric.cosine) {
    for (final x in query.values) {
      qNorm += x * x;
    }
    qNorm = math.sqrt(qNorm);
    if (qNorm == 0.0) qNorm = 1.0; // avoid div-by-zero
  }

  // Simple partial selection: keep the running "worst kept" score in
  // `_scores`/`_hitIds`. For small k this is fine; for very large k
  // a proper heap would help. We size for effK.
  final scores = List<double>.filled(effK, 0.0);
  final ids = List<Object?>.filled(effK, null);
  var filled = 0;

  for (var row = 0; row < n; row++) {
    final base = row * dim;
    double score;
    switch (m) {
      case VectorMetric.l2sq:
      case VectorMetric.l2:
        var acc = 0.0;
        for (var i = 0; i < dim; i++) {
          final d = _data[base + i] - query.values[i];
          acc += d * d;
        }
        score = acc;
        break;
      case VectorMetric.innerProduct:
        var acc = 0.0;
        for (var i = 0; i < dim; i++) {
          acc += _data[base + i] * query.values[i];
        }
        score = acc;
        break;
      case VectorMetric.cosine:
        var dot = 0.0;
        var norm = 0.0;
        for (var i = 0; i < dim; i++) {
          final a = _data[base + i];
          final b = query.values[i];
          dot += a * b;
          norm += a * a;
        }
        norm = math.sqrt(norm);
        score = norm == 0.0 ? 1.0 : 1.0 - dot / (norm * qNorm);
        break;
    }

    // Insert into the top-k list. `larger` flips the comparison.
    if (filled < effK) {
      // Insert in sorted position.
      var j = filled;
      while (j > 0 && _better(score, scores[j - 1], larger)) {
        scores[j] = scores[j - 1];
        ids[j] = ids[j - 1];
        j--;
      }
      scores[j] = score;
      ids[j] = _ids[row];
      filled++;
    } else if (_better(score, scores[effK - 1], larger)) {
      // Displace the current worst-kept.
      var j = effK - 1;
      while (j > 0 && _better(score, scores[j - 1], larger)) {
        scores[j] = scores[j - 1];
        ids[j] = ids[j - 1];
        j--;
      }
      scores[j] = score;
      ids[j] = _ids[row];
    }
  }

  // Convert stored score to the caller-visible distance:
  //   - l2sq: pass through
  //   - l2  : sqrt at the end
  //   - ip  : pass through (larger = better)
  //   - cos : pass through
  final out = <VectorSearchHit>[];
  for (var i = 0; i < filled; i++) {
    final s = m == VectorMetric.l2 ? math.sqrt(scores[i]) : scores[i];
    out.add(VectorSearchHit(ids[i], s));
  }
  return out;
}