Repository navigation
Expand file tree
/
Copy pathVectorSearchMethod.cs
More file actions
172 lines (144 loc) · 6.44 KB
/
Copy pathVectorSearchMethod.cs
File metadata and controls
172 lines (144 loc) · 6.44 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
using System.Diagnostics.CodeAnalysis;
namespace NRedisStack.Search;
public abstract class VectorSearchMethod
{
private protected VectorSearchMethod()
{
}
private protected abstract string Method { get; }
internal abstract int GetOwnArgsCount();
internal abstract void AddOwnArgs(List<object> args);
/// <inheritdoc />
public override string ToString() => Method;
public static VectorSearchMethod Range(double radius, double? epsilon = null) => RangeVectorSearchMethod.Create(radius, epsilon, null);
internal static VectorSearchMethod Range(double radius, double? epsilon, string? distanceAlias)
=> RangeVectorSearchMethod.Create(radius, epsilon, distanceAlias);
public static VectorSearchMethod NearestNeighbour(
int count, int? maxCandidates) // retained for binary compat
=> NearestNeighbourVectorSearchMethod.Create(count, maxCandidates, null, null);
public static VectorSearchMethod NearestNeighbour(
int? count = NearestNeighbourVectorSearchMethod.DEFAULT_NEAREST_NEIGHBOUR_COUNT, int? maxTopCandidates = null, string? distanceAlias = null, double? shardRatio = null)
=> NearestNeighbourVectorSearchMethod.Create(count ?? NearestNeighbourVectorSearchMethod.DEFAULT_NEAREST_NEIGHBOUR_COUNT, maxTopCandidates, distanceAlias, shardRatio);
private sealed class NearestNeighbourVectorSearchMethod : VectorSearchMethod
{
private static NearestNeighbourVectorSearchMethod? s_Default;
internal static NearestNeighbourVectorSearchMethod Create(int count, int? maxTopCandidates,
string? distanceAlias, double? shardRatio)
=> count == DEFAULT_NEAREST_NEIGHBOUR_COUNT & maxTopCandidates == null & distanceAlias == null & !shardRatio.HasValue
? (s_Default ??= new NearestNeighbourVectorSearchMethod(DEFAULT_NEAREST_NEIGHBOUR_COUNT, null, null, null))
: new(count, maxTopCandidates, distanceAlias, shardRatio);
private NearestNeighbourVectorSearchMethod(int nearestNeighbourCount, int? maxTopCandidates,
string? distanceAlias, double? shardRatio)
{
NearestNeighbourCount = nearestNeighbourCount;
MaxTopCandidates = maxTopCandidates;
DistanceAlias = distanceAlias;
ShardRatio = shardRatio;
}
internal const int DEFAULT_NEAREST_NEIGHBOUR_COUNT = 10;
private protected override string Method => "KNN";
/// <summary>
/// The number of nearest neighbors to find. This is the K in KNN.
/// </summary>
public int NearestNeighbourCount { get; }
/// <summary>
/// Max top candidates during KNN search. Higher values increase accuracy, but also increase search latency.
/// This corresponds to the HNSW "EF_RUNTIME" parameter.
/// </summary>
public int? MaxTopCandidates { get; }
/// <summary>
/// Include the distance from the query vector in the results.
/// </summary>
public string? DistanceAlias { get; }
/// <summary>
/// Limits the number of documents processed per shard. Only relevant for cluster scenarios. This corresponds
/// to the "SHARD_K_RATIO" parameter.
/// </summary>
public double? ShardRatio { get; }
internal override int GetOwnArgsCount()
{
int count = 4;
if (MaxTopCandidates.HasValue) count += 2;
if (DistanceAlias != null) count += 2;
if (ShardRatio.HasValue) count += 2;
return count;
}
internal override void AddOwnArgs(List<object> args)
{
args.Add(Method);
int tokens = 2;
if (MaxTopCandidates.HasValue) tokens += 2;
if (DistanceAlias != null) tokens += 2;
if (ShardRatio.HasValue) tokens += 2;
args.Add(tokens);
args.Add("K");
args.Add(NearestNeighbourCount);
if (MaxTopCandidates.HasValue)
{
args.Add("EF_RUNTIME");
args.Add(MaxTopCandidates.GetValueOrDefault());
}
if (DistanceAlias != null)
{
args.Add("YIELD_DISTANCE_AS");
args.Add(DistanceAlias);
}
if (ShardRatio.HasValue)
{
args.Add("SHARD_K_RATIO");
args.Add(ShardRatio.GetValueOrDefault());
}
}
}
private sealed class RangeVectorSearchMethod : VectorSearchMethod
{
internal static RangeVectorSearchMethod Create(double radius, double? epsilon, string? distanceAlias)
=> new(radius, epsilon, distanceAlias);
private RangeVectorSearchMethod(double radius, double? epsilon, string? distanceAlias)
{
Radius = radius;
Epsilon = epsilon;
DistanceAlias = distanceAlias;
}
private protected override string Method => "RANGE";
/// <summary>
/// The search radius/threshold. Finds all vectors within this distance.
/// </summary>
public double Radius { get; }
/// <summary>
/// Relative factor that sets the boundaries in which a range query may search for candidates. That is, vector candidates whose distance from the query vector is radius * (1 + EPSILON) are potentially scanned, allowing more extensive search and more accurate results, at the expense of run time.
/// </summary>
public double? Epsilon { get; }
/// <summary>
/// Include the distance from the query vector in the results.
/// </summary>
public string? DistanceAlias { get; }
internal override int GetOwnArgsCount()
{
int count = 4;
if (Epsilon != null) count += 2;
if (DistanceAlias != null) count += 2;
return count;
}
internal override void AddOwnArgs(List<object> args)
{
args.Add(Method);
int tokens = 2;
if (Epsilon != null) tokens += 2;
if (DistanceAlias != null) tokens += 2;
args.Add(tokens);
args.Add("RADIUS");
args.Add(Radius);
if (Epsilon != null)
{
args.Add("EPSILON");
args.Add(Epsilon);
}
if (DistanceAlias != null)
{
args.Add("YIELD_DISTANCE_AS");
args.Add(DistanceAlias);
}
}
}
}