Skip to content

Commit 6e0d430

Browse files
ggivoggivo
authored
[vector sets] Support for VSIM WITHATTRIBS (CAE-1421) (#4260)
* [vector sets] Support for VSIM WITHATTRIBS * [vector sets] Add misisng file - * [vector sets] format * [vector sets] fix tests - mark tests for WITHATTRIBS with SINCE 8.2.0 - Use JedisByteMap for binary methods * [vector sets] fix tests - mark tests for WITHATTRIBS with SINCE 8.2.0 * [vector sets] address review comments - remove unused import - remove unused helper methods
1 parent 1528e65 commit 6e0d430

File tree

11 files changed

+520
-14
lines changed

11 files changed

+520
-14
lines changed

‎pom.xml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,7 @@
352352
<include>**/VectorTestUtils.java</include>
353353
<include>**/VAddParams.java</include>
354354
<include>**/VSimParams.java</include>
355+
<include>**/VSimScoreAttribs.java</include>
355356
<include>**/*FunctionCommandsTest*</include>
356357
</includes>
357358
</configuration>

‎src/main/java/redis/clients/jedis/BuilderFactory.java

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -514,15 +514,14 @@ public Map<byte[], Double> build(Object data) {
514514
final List<Object> list = (List<Object>) data;
515515
if (list.isEmpty()) return Collections.emptyMap();
516516

517+
final JedisByteMap<Double> map = new JedisByteMap<>();
517518
if (list.get(0) instanceof KeyValue) {
518-
final Map<byte[], Double> map = new LinkedHashMap<>(list.size(), 1f);
519519
for (Object o : list) {
520520
KeyValue<?, ?> kv = (KeyValue<?, ?>) o;
521521
map.put(BINARY.build(kv.getKey()), DOUBLE.build(kv.getValue()));
522522
}
523523
return map;
524524
} else {
525-
final Map<byte[], Double> map = new LinkedHashMap<>(list.size() / 2, 1f);
526525
final Iterator iterator = list.iterator();
527526
while (iterator.hasNext()) {
528527
map.put(BINARY.build(iterator.next()), DOUBLE.build(iterator.next()));
@@ -2249,7 +2248,73 @@ public String toString() {
22492248
};
22502249

22512250
// Vector Set builders
2251+
public static final Builder<Map<String, VSimScoreAttribs>> VSIM_SCORE_ATTRIBS_MAP = new Builder<Map<String, VSimScoreAttribs>>() {
22522252

2253+
@Override
2254+
@SuppressWarnings("unchecked")
2255+
public Map<String, VSimScoreAttribs> build(Object data) {
2256+
if (data == null) return null;
2257+
List<Object> list = (List<Object>) data;
2258+
if (list.isEmpty()) return Collections.emptyMap();
2259+
2260+
if (list.get(0) instanceof KeyValue) {
2261+
final Map<String, VSimScoreAttribs> result = new LinkedHashMap<>(list.size(), 1f);
2262+
for (Object o : list) {
2263+
KeyValue<?, ?> kv = (KeyValue<?, ?>) o;
2264+
List<Object> scoreAndAttribs = (List<Object>) kv.getValue();
2265+
result.put(STRING.build(kv.getKey()),
2266+
new VSimScoreAttribs(DOUBLE.build(scoreAndAttribs.get(0)),
2267+
STRING.build(scoreAndAttribs.get(1))));
2268+
}
2269+
return result;
2270+
} else {
2271+
final Map<String, VSimScoreAttribs> result = new LinkedHashMap<>(list.size() / 3, 1f);
2272+
for (int i = 0; i < list.size(); i += 3) {
2273+
result.put(STRING.build(list.get(i)),
2274+
new VSimScoreAttribs(DOUBLE.build(list.get(i + 1)), STRING.build(list.get(i + 2))));
2275+
}
2276+
return result;
2277+
}
2278+
}
2279+
2280+
@Override
2281+
public String toString() {
2282+
return "Map<String, VSimScoreAttribs>";
2283+
}
2284+
};
2285+
2286+
public static final Builder<Map<byte[], VSimScoreAttribs>> VSIM_SCORE_ATTRIBS_BINARY_MAP = new Builder<Map<byte[], VSimScoreAttribs>>() {
2287+
2288+
@Override
2289+
@SuppressWarnings("unchecked")
2290+
public Map<byte[], VSimScoreAttribs> build(Object data) {
2291+
if (data == null) return null;
2292+
List<Object> list = (List<Object>) data;
2293+
if (list.isEmpty()) return Collections.emptyMap();
2294+
2295+
JedisByteMap<VSimScoreAttribs> result = new JedisByteMap<>();
2296+
if (list.get(0) instanceof KeyValue) {
2297+
for (Object o : list) {
2298+
KeyValue<?, ?> kv = (KeyValue<?, ?>) o;
2299+
List<Object> scoreAndAttribs = (List<Object>) kv.getValue();
2300+
result.put(BINARY.build(kv.getKey()),
2301+
new VSimScoreAttribs(DOUBLE.build(scoreAndAttribs.get(0)),
2302+
STRING.build(scoreAndAttribs.get(1))));
2303+
}
2304+
} else {
2305+
for (int i = 0; i < list.size(); i += 3) {
2306+
result.put(BINARY.build(list.get(i)),
2307+
new VSimScoreAttribs(DOUBLE.build(list.get(i + 1)), STRING.build(list.get(i + 2))));
2308+
}
2309+
}
2310+
return result;
2311+
}
2312+
2313+
@Override
2314+
public String toString() {
2315+
return "Map<byte[], VSimScoreAttribs>";
2316+
}
2317+
};
22532318

22542319
public static final Builder<RawVector> VEMB_RAW_RESULT = new Builder<RawVector>() {
22552320
@Override

‎src/main/java/redis/clients/jedis/CommandObjects.java

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4883,6 +4883,15 @@ public final CommandObject<Map<String, Double>> vsimWithScores(String key, float
48834883
return new CommandObject<>(args, BuilderFactory.STRING_DOUBLE_MAP);
48844884
}
48854885

4886+
public final CommandObject<Map<String, VSimScoreAttribs>> vsimWithScoresAndAttribs(String key, float[] vector, VSimParams params) {
4887+
CommandArguments args = commandArguments(Command.VSIM).key(key);
4888+
addVectors(vector, args);
4889+
args.add(Keyword.WITHSCORES);
4890+
args.add(Keyword.WITHATTRIBS);
4891+
addOptionalParams(params, args);
4892+
return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_MAP);
4893+
}
4894+
48864895
public final CommandObject<List<String>> vsimByElement(String key, String element) {
48874896
return vsimByElement(key, element, null);
48884897
}
@@ -4902,6 +4911,15 @@ public final CommandObject<Map<String, Double>> vsimByElementWithScores(String k
49024911
return new CommandObject<>(args, BuilderFactory.STRING_DOUBLE_MAP);
49034912
}
49044913

4914+
public final CommandObject<Map<String, VSimScoreAttribs>> vsimByElementWithScoresAndAttribs(String key, String element, VSimParams params) {
4915+
CommandArguments args = commandArguments(Command.VSIM).key(key);
4916+
args.add(Keyword.ELE).add(element);
4917+
args.add(Keyword.WITHSCORES);
4918+
args.add(Keyword.WITHATTRIBS);
4919+
addOptionalParams(params, args);
4920+
return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_MAP);
4921+
}
4922+
49054923
public final CommandObject<List<byte[]>> vsim(byte[] key, float[] vector) {
49064924
return vsim(key, vector, null);
49074925
}
@@ -4921,6 +4939,15 @@ public final CommandObject<Map<byte[], Double>> vsimWithScores(byte[] key, float
49214939
return new CommandObject<>(args, BuilderFactory.BINARY_DOUBLE_MAP);
49224940
}
49234941

4942+
public final CommandObject<Map<byte[], VSimScoreAttribs>> vsimWithScoresAndAttribs(byte[] key, float[] vector, VSimParams params) {
4943+
CommandArguments args = commandArguments(Command.VSIM).key(key);
4944+
addVectors(vector, args);
4945+
args.add(Keyword.WITHSCORES);
4946+
args.add(Keyword.WITHATTRIBS);
4947+
addOptionalParams(params, args);
4948+
return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_BINARY_MAP);
4949+
}
4950+
49244951
public final CommandObject<List<byte[]>> vsimByElement(byte[] key, byte[] element) {
49254952
return vsimByElement(key, element, null);
49264953
}
@@ -4940,6 +4967,15 @@ public final CommandObject<Map<byte[], Double>> vsimByElementWithScores(byte[] k
49404967
return new CommandObject<>(args, BuilderFactory.BINARY_DOUBLE_MAP);
49414968
}
49424969

4970+
public final CommandObject<Map<byte[], VSimScoreAttribs>> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element, VSimParams params) {
4971+
CommandArguments args = commandArguments(Command.VSIM).key(key);
4972+
args.add(Keyword.ELE).add(element);
4973+
args.add(Keyword.WITHSCORES);
4974+
args.add(Keyword.WITHATTRIBS);
4975+
addOptionalParams(params, args);
4976+
return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_BINARY_MAP);
4977+
}
4978+
49434979
public final CommandObject<Long> vdim(String key) {
49444980
return new CommandObject<>(commandArguments(Command.VDIM).key(key), BuilderFactory.LONG);
49454981
}

‎src/main/java/redis/clients/jedis/Jedis.java

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10048,6 +10048,12 @@ public Map<String, Double> vsimWithScores(String key, float[] vector, VSimParams
1004810048
return connection.executeCommand(commandObjects.vsimWithScores(key, vector, params));
1004910049
}
1005010050

10051+
@Override
10052+
public Map<String, VSimScoreAttribs> vsimWithScoresAndAttribs(String key, float[] vector, VSimParams params) {
10053+
checkIsInMultiOrPipeline();
10054+
return connection.executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
10055+
}
10056+
1005110057
@Override
1005210058
public List<String> vsimByElement(String key, String element) {
1005310059
checkIsInMultiOrPipeline();
@@ -10066,6 +10072,12 @@ public Map<String, Double> vsimByElementWithScores(String key, String element, V
1006610072
return connection.executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
1006710073
}
1006810074

10075+
@Override
10076+
public Map<String, VSimScoreAttribs> vsimByElementWithScoresAndAttribs(String key, String element, VSimParams params) {
10077+
checkIsInMultiOrPipeline();
10078+
return connection.executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
10079+
}
10080+
1006910081
@Override
1007010082
public long vdim(String key) {
1007110083
checkIsInMultiOrPipeline();
@@ -10193,6 +10205,12 @@ public Map<byte[], Double> vsimWithScores(byte[] key, float[] vector, VSimParams
1019310205
return connection.executeCommand(commandObjects.vsimWithScores(key, vector, params));
1019410206
}
1019510207

10208+
@Override
10209+
public Map<byte[], VSimScoreAttribs> vsimWithScoresAndAttribs(byte[] key, float[] vector, VSimParams params) {
10210+
checkIsInMultiOrPipeline();
10211+
return connection.executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
10212+
}
10213+
1019610214
@Override
1019710215
public List<byte[]> vsimByElement(byte[] key, byte[] element) {
1019810216
checkIsInMultiOrPipeline();
@@ -10211,6 +10229,12 @@ public Map<byte[], Double> vsimByElementWithScores(byte[] key, byte[] element, V
1021110229
return connection.executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
1021210230
}
1021310231

10232+
@Override
10233+
public Map<byte[], VSimScoreAttribs> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element, VSimParams params) {
10234+
checkIsInMultiOrPipeline();
10235+
return connection.executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
10236+
}
10237+
1021410238
@Override
1021510239
public long vdim(byte[] key) {
1021610240
checkIsInMultiOrPipeline();

‎src/main/java/redis/clients/jedis/Protocol.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -342,7 +342,7 @@ public static enum Keyword implements Rawable {
342342
ARGS, RANK, NOW, VERSION, ADDR, SKIPME, USER, LADDR, FIELDS,
343343
CHANNELS, NUMPAT, NUMSUB, SHARDCHANNELS, SHARDNUMSUB, NOVALUES, MAXAGE, FXX, FNX,
344344
// Vector set keywords
345-
REDUCE, CAS, NOQUANT, Q8, BIN, EF, SETATTR, M, VALUES, FP32, ELE, FILTER, FILTER_EF, TRUTH, NOTHREAD, RAW, EPSILON;
345+
REDUCE, CAS, NOQUANT, Q8, BIN, EF, SETATTR, M, VALUES, FP32, ELE, FILTER, FILTER_EF, TRUTH, NOTHREAD, RAW, EPSILON, WITHATTRIBS;
346346

347347
private final byte[] raw;
348348

‎src/main/java/redis/clients/jedis/UnifiedJedis.java

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5223,6 +5223,11 @@ public Map<String, Double> vsimWithScores(String key, float[] vector, VSimParams
52235223
return executeCommand(commandObjects.vsimWithScores(key, vector, params));
52245224
}
52255225

5226+
@Override
5227+
public Map<String, VSimScoreAttribs> vsimWithScoresAndAttribs(String key, float[] vector, VSimParams params) {
5228+
return executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
5229+
}
5230+
52265231
@Override
52275232
public List<String> vsimByElement(String key, String element) {
52285233
return executeCommand(commandObjects.vsimByElement(key, element));
@@ -5238,6 +5243,11 @@ public Map<String, Double> vsimByElementWithScores(String key, String element, V
52385243
return executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
52395244
}
52405245

5246+
@Override
5247+
public Map<String, VSimScoreAttribs> vsimByElementWithScoresAndAttribs(String key, String element, VSimParams params) {
5248+
return executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
5249+
}
5250+
52415251
@Override
52425252
public long vdim(String key) {
52435253
return executeCommand(commandObjects.vdim(key));
@@ -5344,6 +5354,11 @@ public Map<byte[], Double> vsimWithScores(byte[] key, float[] vector, VSimParams
53445354
return executeCommand(commandObjects.vsimWithScores(key, vector, params));
53455355
}
53465356

5357+
@Override
5358+
public Map<byte[], VSimScoreAttribs> vsimWithScoresAndAttribs(byte[] key, float[] vector, VSimParams params) {
5359+
return executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
5360+
}
5361+
53475362
@Override
53485363
public List<byte[]> vsimByElement(byte[] key, byte[] element) {
53495364
return executeCommand(commandObjects.vsimByElement(key, element));
@@ -5359,6 +5374,11 @@ public Map<byte[], Double> vsimByElementWithScores(byte[] key, byte[] element, V
53595374
return executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
53605375
}
53615376

5377+
@Override
5378+
public Map<byte[], VSimScoreAttribs> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element, VSimParams params) {
5379+
return executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
5380+
}
5381+
53625382
@Override
53635383
public long vdim(byte[] key) {
53645384
return executeCommand(commandObjects.vdim(key));

‎src/main/java/redis/clients/jedis/commands/VectorSetBinaryCommands.java

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
package redis.clients.jedis.commands;
22

3-
import java.util.List;
4-
import java.util.Map;
5-
63
import redis.clients.jedis.annots.Experimental;
74
import redis.clients.jedis.params.VAddParams;
85
import redis.clients.jedis.params.VSimParams;
96
import redis.clients.jedis.resps.RawVector;
7+
import redis.clients.jedis.resps.VSimScoreAttribs;
8+
9+
import java.util.List;
10+
import java.util.Map;
1011

1112
/**
1213
* Interface for Redis Vector Set binary commands. Vector sets are a new data type introduced in
@@ -144,6 +145,21 @@ public interface VectorSetBinaryCommands {
144145
@Experimental
145146
Map<byte[], Double> vsimWithScores(byte[] key, float[] vector, VSimParams params);
146147

148+
/**
149+
* <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
150+
* similar to a given vector with their similarity scores and attributes.
151+
* <p>
152+
* Time complexity: O(log(N)) where N is the number of elements in the vector set.
153+
* @param key the name of the key that holds the vector set data
154+
* @param vector the vector to use as similarity reference
155+
* @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
156+
* automatically added)
157+
* @return map of element names to their similarity scores and attributes
158+
*/
159+
@Experimental
160+
Map<byte[], VSimScoreAttribs> vsimWithScoresAndAttribs(byte[] key, float[] vector,
161+
VSimParams params);
162+
147163
/**
148164
* <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
149165
* similar to a given element in the vector set.
@@ -183,6 +199,21 @@ public interface VectorSetBinaryCommands {
183199
@Experimental
184200
Map<byte[], Double> vsimByElementWithScores(byte[] key, byte[] element, VSimParams params);
185201

202+
/**
203+
* <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
204+
* similar to a given element in the vector set with their similarity scores and attributes.
205+
* <p>
206+
* Time complexity: O(log(N)) where N is the number of elements in the vector set.
207+
* @param key the name of the key that holds the vector set data
208+
* @param element the name of the element to use as similarity reference
209+
* @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
210+
* automatically added)
211+
* @return map of element names to their similarity scores and attributes
212+
*/
213+
@Experimental
214+
Map<byte[], VSimScoreAttribs> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element,
215+
VSimParams params);
216+
186217
/**
187218
* <b><a href="https://redis.io/docs/latest/commands/vdim/">VDIM Command</a></b> Return the number
188219
* of dimensions of the vectors in the specified vector set.

‎src/main/java/redis/clients/jedis/commands/VectorSetCommands.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import redis.clients.jedis.params.VAddParams;
88
import redis.clients.jedis.params.VSimParams;
99
import redis.clients.jedis.resps.RawVector;
10+
import redis.clients.jedis.resps.VSimScoreAttribs;
1011
import redis.clients.jedis.resps.VectorInfo;
1112

1213
/**
@@ -145,6 +146,21 @@ public interface VectorSetCommands {
145146
@Experimental
146147
Map<String, Double> vsimWithScores(String key, float[] vector, VSimParams params);
147148

149+
/**
150+
* <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
151+
* similar to a given vector with their similarity scores and attributes.
152+
* <p>
153+
* Time complexity: O(log(N)) where N is the number of elements in the vector set.
154+
* @param key the name of the key that holds the vector set data
155+
* @param vector the vector to use as similarity reference
156+
* @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
157+
* automatically added)
158+
* @return map of element names to their similarity scores and attributes
159+
*/
160+
@Experimental
161+
Map<String, VSimScoreAttribs> vsimWithScoresAndAttribs(String key, float[] vector,
162+
VSimParams params);
163+
148164
/**
149165
* <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
150166
* similar to a given element in the vector set.
@@ -184,6 +200,21 @@ public interface VectorSetCommands {
184200
@Experimental
185201
Map<String, Double> vsimByElementWithScores(String key, String element, VSimParams params);
186202

203+
/**
204+
* <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
205+
* similar to a given element in the vector set with their similarity scores and attributes.
206+
* <p>
207+
* Time complexity: O(log(N)) where N is the number of elements in the vector set.
208+
* @param key the name of the key that holds the vector set data
209+
* @param element the name of the element to use as similarity reference
210+
* @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
211+
* automatically added)
212+
* @return map of element names to their similarity scores and attributes
213+
*/
214+
@Experimental
215+
Map<String, VSimScoreAttribs> vsimByElementWithScoresAndAttribs(String key, String element,
216+
VSimParams params);
217+
187218
/**
188219
* <b><a href="https://redis.io/docs/latest/commands/vdim/">VDIM Command</a></b> Return the number
189220
* of dimensions of the vectors in the specified vector set.

0 commit comments

Comments
 (0)