search method

Implementation

List<VectorSearchHit> search(
  Vector query,
  int k, {
  VectorMetric? metric,
  int? nprobe,
}) {
  if (!_trained) {
    throw StateError('IvfPqIndex.search: call train() first');
  }
  if (query.dim != dim) {
    throw StateError(
      'IvfPqIndex.search: query dim ${query.dim} != index dim $dim',
    );
  }
  if (k <= 0 || _n == 0) return const [];
  final cells = _probeCells(query, nprobe ?? this.nprobe);
  final effK = math.min(k, _n);
  final scores = List<double>.filled(effK, 0.0);
  final ids = List<Object?>.filled(effK, null);
  var filled = 0;
  final residual = Float32List(dim);

  for (final cell in cells) {
    final cOff = cell * dim;
    for (var j = 0; j < dim; j++) {
      residual[j] = query.values[j] - _coarseQuantizer.centroids[cOff + j];
    }
    final lut = _pq.buildDistanceTable(residual);
    final codes = _cellCodes[cell];
    final cellIds = _cellIds[cell];
    final count = _cellCounts[cell];
    for (var r = 0; r < count; r++) {
      final di = _pq.lookupDistance(lut, codes, r * m);
      if (filled < effK) {
        var j = filled;
        while (j > 0 && di < scores[j - 1]) {
          scores[j] = scores[j - 1];
          ids[j] = ids[j - 1];
          j--;
        }
        scores[j] = di;
        ids[j] = cellIds[r];
        filled++;
      } else if (di < scores[effK - 1]) {
        var j = effK - 1;
        while (j > 0 && di < scores[j - 1]) {
          scores[j] = scores[j - 1];
          ids[j] = ids[j - 1];
          j--;
        }
        scores[j] = di;
        ids[j] = cellIds[r];
      }
    }
  }
  return [
    for (var i = 0; i < filled; i++) VectorSearchHit(ids[i], scores[i]),
  ];
}