[vector sets] Support for VSIM WITHATTRIBS (CAE-1421)

pom.xml

@@ -352,6 +352,7 @@
 						<include>**/VectorTestUtils.java</include>
 						<include>**/VAddParams.java</include>
 						<include>**/VSimParams.java</include>
+						<include>**/VSimScoreAttribs.java</include>
 						<include>**/*FunctionCommandsTest*</include>
 					</includes>
 				</configuration>

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

@@ -514,15 +514,14 @@ public Map<byte[], Double> build(Object data) {
       final List<Object> list = (List<Object>) data;
       if (list.isEmpty()) return Collections.emptyMap();
 
+      final JedisByteMap<Double> map = new JedisByteMap<>();
       if (list.get(0) instanceof KeyValue) {
-        final Map<byte[], Double> map = new LinkedHashMap<>(list.size(), 1f);
         for (Object o : list) {
           KeyValue<?, ?> kv = (KeyValue<?, ?>) o;
           map.put(BINARY.build(kv.getKey()), DOUBLE.build(kv.getValue()));
         }
         return map;
       } else {
-        final Map<byte[], Double> map = new LinkedHashMap<>(list.size() / 2, 1f);
         final Iterator iterator = list.iterator();
         while (iterator.hasNext()) {
           map.put(BINARY.build(iterator.next()), DOUBLE.build(iterator.next()));
@@ -2249,7 +2248,73 @@ public String toString() {
   };
 
   // Vector Set builders
+  public static final Builder<Map<String, VSimScoreAttribs>> VSIM_SCORE_ATTRIBS_MAP = new Builder<Map<String, VSimScoreAttribs>>() {
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public Map<String, VSimScoreAttribs> build(Object data) {
+      if (data == null) return null;
+      List<Object> list = (List<Object>) data;
+      if (list.isEmpty()) return Collections.emptyMap();
+
+      if (list.get(0) instanceof KeyValue) {
+        final Map<String, VSimScoreAttribs> result = new LinkedHashMap<>(list.size(), 1f);
+        for (Object o : list) {
+          KeyValue<?, ?> kv = (KeyValue<?, ?>) o;
+          List<Object> scoreAndAttribs = (List<Object>) kv.getValue();
+          result.put(STRING.build(kv.getKey()),
+              new VSimScoreAttribs(DOUBLE.build(scoreAndAttribs.get(0)),
+                  STRING.build(scoreAndAttribs.get(1))));
+        }
+        return result;
+      } else {
+        final Map<String, VSimScoreAttribs> result = new LinkedHashMap<>(list.size() / 3, 1f);
+        for (int i = 0; i < list.size(); i += 3) {
+          result.put(STRING.build(list.get(i)),
+              new VSimScoreAttribs(DOUBLE.build(list.get(i + 1)), STRING.build(list.get(i + 2))));
+        }
+        return result;
+      }
+    }
+
+    @Override
+    public String toString() {
+      return "Map<String, VSimScoreAttribs>";
+    }
+  };
+
+  public static final Builder<Map<byte[], VSimScoreAttribs>> VSIM_SCORE_ATTRIBS_BINARY_MAP = new Builder<Map<byte[], VSimScoreAttribs>>() {
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public Map<byte[], VSimScoreAttribs> build(Object data) {
+      if (data == null) return null;
+      List<Object> list = (List<Object>) data;
+      if (list.isEmpty()) return Collections.emptyMap();
+
+      JedisByteMap<VSimScoreAttribs> result = new JedisByteMap<>();
+      if (list.get(0) instanceof KeyValue) {
+        for (Object o : list) {
+          KeyValue<?, ?> kv = (KeyValue<?, ?>) o;
+          List<Object> scoreAndAttribs = (List<Object>) kv.getValue();
+          result.put(BINARY.build(kv.getKey()),
+              new VSimScoreAttribs(DOUBLE.build(scoreAndAttribs.get(0)),
+                  STRING.build(scoreAndAttribs.get(1))));
+        }
+      } else {
+        for (int i = 0; i < list.size(); i += 3) {
+          result.put(BINARY.build(list.get(i)),
+              new VSimScoreAttribs(DOUBLE.build(list.get(i + 1)), STRING.build(list.get(i + 2))));
+        }
+      }
+      return result;
+    }
+
+    @Override
+    public String toString() {
+      return "Map<byte[], VSimScoreAttribs>";
+    }
+  };
 
   public static final Builder<RawVector> VEMB_RAW_RESULT = new Builder<RawVector>() {
     @Override

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

@@ -4883,6 +4883,15 @@ public final CommandObject<Map<String, Double>> vsimWithScores(String key, float
     return new CommandObject<>(args, BuilderFactory.STRING_DOUBLE_MAP);
   }
 
+  public final CommandObject<Map<String, VSimScoreAttribs>> vsimWithScoresAndAttribs(String key, float[] vector, VSimParams params) {
+    CommandArguments args = commandArguments(Command.VSIM).key(key);
+    addVectors(vector, args);
+    args.add(Keyword.WITHSCORES);
+    args.add(Keyword.WITHATTRIBS);
+    addOptionalParams(params, args);
+    return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_MAP);
+  }
+
   public final CommandObject<List<String>> vsimByElement(String key, String element) {
     return vsimByElement(key, element, null);
   }
@@ -4902,6 +4911,15 @@ public final CommandObject<Map<String, Double>> vsimByElementWithScores(String k
     return new CommandObject<>(args, BuilderFactory.STRING_DOUBLE_MAP);
   }
 
+  public final CommandObject<Map<String, VSimScoreAttribs>> vsimByElementWithScoresAndAttribs(String key, String element, VSimParams params) {
+    CommandArguments args = commandArguments(Command.VSIM).key(key);
+    args.add(Keyword.ELE).add(element);
+    args.add(Keyword.WITHSCORES);
+    args.add(Keyword.WITHATTRIBS);
+    addOptionalParams(params, args);
+    return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_MAP);
+  }
+
   public final CommandObject<List<byte[]>> vsim(byte[] key, float[] vector) {
     return vsim(key, vector, null);
   }
@@ -4921,6 +4939,15 @@ public final CommandObject<Map<byte[], Double>> vsimWithScores(byte[] key, float
     return new CommandObject<>(args, BuilderFactory.BINARY_DOUBLE_MAP);
   }
 
+  public final CommandObject<Map<byte[], VSimScoreAttribs>> vsimWithScoresAndAttribs(byte[] key, float[] vector, VSimParams params) {
+    CommandArguments args = commandArguments(Command.VSIM).key(key);
+    addVectors(vector, args);
+    args.add(Keyword.WITHSCORES);
+    args.add(Keyword.WITHATTRIBS);
+    addOptionalParams(params, args);
+    return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_BINARY_MAP);
+  }
+
   public final CommandObject<List<byte[]>> vsimByElement(byte[] key, byte[] element) {
     return vsimByElement(key, element, null);
   }
@@ -4940,6 +4967,15 @@ public final CommandObject<Map<byte[], Double>> vsimByElementWithScores(byte[] k
     return new CommandObject<>(args, BuilderFactory.BINARY_DOUBLE_MAP);
   }
 
+  public final CommandObject<Map<byte[], VSimScoreAttribs>> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element, VSimParams params) {
+    CommandArguments args = commandArguments(Command.VSIM).key(key);
+    args.add(Keyword.ELE).add(element);
+    args.add(Keyword.WITHSCORES);
+    args.add(Keyword.WITHATTRIBS);
+    addOptionalParams(params, args);
+    return new CommandObject<>(args, BuilderFactory.VSIM_SCORE_ATTRIBS_BINARY_MAP);
+  }
+
   public final CommandObject<Long> vdim(String key) {
     return new CommandObject<>(commandArguments(Command.VDIM).key(key), BuilderFactory.LONG);
   }

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

@@ -10048,6 +10048,12 @@ public Map<String, Double> vsimWithScores(String key, float[] vector, VSimParams
     return connection.executeCommand(commandObjects.vsimWithScores(key, vector, params));
   }
 
+  @Override
+  public Map<String, VSimScoreAttribs> vsimWithScoresAndAttribs(String key, float[] vector, VSimParams params) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
+  }
+
   @Override
   public List<String> vsimByElement(String key, String element) {
     checkIsInMultiOrPipeline();
@@ -10066,6 +10072,12 @@ public Map<String, Double> vsimByElementWithScores(String key, String element, V
     return connection.executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
   }
 
+  @Override
+  public Map<String, VSimScoreAttribs> vsimByElementWithScoresAndAttribs(String key, String element, VSimParams params) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
+  }
+
   @Override
   public long vdim(String key) {
     checkIsInMultiOrPipeline();
@@ -10193,6 +10205,12 @@ public Map<byte[], Double> vsimWithScores(byte[] key, float[] vector, VSimParams
     return connection.executeCommand(commandObjects.vsimWithScores(key, vector, params));
   }
 
+  @Override
+  public Map<byte[], VSimScoreAttribs> vsimWithScoresAndAttribs(byte[] key, float[] vector, VSimParams params) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
+  }
+
   @Override
   public List<byte[]> vsimByElement(byte[] key, byte[] element) {
     checkIsInMultiOrPipeline();
@@ -10211,6 +10229,12 @@ public Map<byte[], Double> vsimByElementWithScores(byte[] key, byte[] element, V
     return connection.executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
   }
 
+  @Override
+  public Map<byte[], VSimScoreAttribs> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element, VSimParams params) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
+  }
+
   @Override
   public long vdim(byte[] key) {
     checkIsInMultiOrPipeline();

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

@@ -342,7 +342,7 @@ public static enum Keyword implements Rawable {
     ARGS, RANK, NOW, VERSION, ADDR, SKIPME, USER, LADDR, FIELDS,
     CHANNELS, NUMPAT, NUMSUB, SHARDCHANNELS, SHARDNUMSUB, NOVALUES, MAXAGE, FXX, FNX,
     // Vector set keywords
-    REDUCE, CAS, NOQUANT, Q8, BIN, EF, SETATTR, M, VALUES, FP32, ELE, FILTER, FILTER_EF, TRUTH, NOTHREAD, RAW, EPSILON;
+    REDUCE, CAS, NOQUANT, Q8, BIN, EF, SETATTR, M, VALUES, FP32, ELE, FILTER, FILTER_EF, TRUTH, NOTHREAD, RAW, EPSILON, WITHATTRIBS;
 
     private final byte[] raw;
 

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

@@ -5223,6 +5223,11 @@ public Map<String, Double> vsimWithScores(String key, float[] vector, VSimParams
     return executeCommand(commandObjects.vsimWithScores(key, vector, params));
   }
 
+  @Override
+  public Map<String, VSimScoreAttribs> vsimWithScoresAndAttribs(String key, float[] vector, VSimParams params) {
+    return executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
+  }
+
   @Override
   public List<String> vsimByElement(String key, String element) {
     return executeCommand(commandObjects.vsimByElement(key, element));
@@ -5238,6 +5243,11 @@ public Map<String, Double> vsimByElementWithScores(String key, String element, V
     return executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
   }
 
+  @Override
+  public Map<String, VSimScoreAttribs> vsimByElementWithScoresAndAttribs(String key, String element, VSimParams params) {
+    return executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
+  }
+
   @Override
   public long vdim(String key) {
     return executeCommand(commandObjects.vdim(key));
@@ -5344,6 +5354,11 @@ public Map<byte[], Double> vsimWithScores(byte[] key, float[] vector, VSimParams
     return executeCommand(commandObjects.vsimWithScores(key, vector, params));
   }
 
+  @Override
+  public Map<byte[], VSimScoreAttribs> vsimWithScoresAndAttribs(byte[] key, float[] vector, VSimParams params) {
+    return executeCommand(commandObjects.vsimWithScoresAndAttribs(key, vector, params));
+  }
+
   @Override
   public List<byte[]> vsimByElement(byte[] key, byte[] element) {
     return executeCommand(commandObjects.vsimByElement(key, element));
@@ -5359,6 +5374,11 @@ public Map<byte[], Double> vsimByElementWithScores(byte[] key, byte[] element, V
     return executeCommand(commandObjects.vsimByElementWithScores(key, element, params));
   }
 
+  @Override
+  public Map<byte[], VSimScoreAttribs> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element, VSimParams params) {
+    return executeCommand(commandObjects.vsimByElementWithScoresAndAttribs(key, element, params));
+  }
+
   @Override
   public long vdim(byte[] key) {
     return executeCommand(commandObjects.vdim(key));

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

@@ -1,12 +1,13 @@
 package redis.clients.jedis.commands;
 
-import java.util.List;
-import java.util.Map;
-
 import redis.clients.jedis.annots.Experimental;
 import redis.clients.jedis.params.VAddParams;
 import redis.clients.jedis.params.VSimParams;
 import redis.clients.jedis.resps.RawVector;
+import redis.clients.jedis.resps.VSimScoreAttribs;
+
+import java.util.List;
+import java.util.Map;
 
 /**
  * Interface for Redis Vector Set binary commands. Vector sets are a new data type introduced in
@@ -144,6 +145,21 @@ public interface VectorSetBinaryCommands {
   @Experimental
   Map<byte[], Double> vsimWithScores(byte[] key, float[] vector, VSimParams params);
 
+  /**
+   * <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
+   * similar to a given vector with their similarity scores and attributes.
+   * <p>
+   * Time complexity: O(log(N)) where N is the number of elements in the vector set.
+   * @param key the name of the key that holds the vector set data
+   * @param vector the vector to use as similarity reference
+   * @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
+   *          automatically added)
+   * @return map of element names to their similarity scores and attributes
+   */
+  @Experimental
+  Map<byte[], VSimScoreAttribs> vsimWithScoresAndAttribs(byte[] key, float[] vector,
+      VSimParams params);
+
   /**
    * <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
    * similar to a given element in the vector set.
@@ -183,6 +199,21 @@ public interface VectorSetBinaryCommands {
   @Experimental
   Map<byte[], Double> vsimByElementWithScores(byte[] key, byte[] element, VSimParams params);
 
+  /**
+   * <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
+   * similar to a given element in the vector set with their similarity scores and attributes.
+   * <p>
+   * Time complexity: O(log(N)) where N is the number of elements in the vector set.
+   * @param key the name of the key that holds the vector set data
+   * @param element the name of the element to use as similarity reference
+   * @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
+   *          automatically added)
+   * @return map of element names to their similarity scores and attributes
+   */
+  @Experimental
+  Map<byte[], VSimScoreAttribs> vsimByElementWithScoresAndAttribs(byte[] key, byte[] element,
+      VSimParams params);
+
   /**
    * <b><a href="https://redis.io/docs/latest/commands/vdim/">VDIM Command</a></b> Return the number
    * of dimensions of the vectors in the specified vector set.

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

@@ -7,6 +7,7 @@
 import redis.clients.jedis.params.VAddParams;
 import redis.clients.jedis.params.VSimParams;
 import redis.clients.jedis.resps.RawVector;
+import redis.clients.jedis.resps.VSimScoreAttribs;
 import redis.clients.jedis.resps.VectorInfo;
 
 /**
@@ -145,6 +146,21 @@ public interface VectorSetCommands {
   @Experimental
   Map<String, Double> vsimWithScores(String key, float[] vector, VSimParams params);
 
+  /**
+   * <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
+   * similar to a given vector with their similarity scores and attributes.
+   * <p>
+   * Time complexity: O(log(N)) where N is the number of elements in the vector set.
+   * @param key the name of the key that holds the vector set data
+   * @param vector the vector to use as similarity reference
+   * @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
+   *          automatically added)
+   * @return map of element names to their similarity scores and attributes
+   */
+  @Experimental
+  Map<String, VSimScoreAttribs> vsimWithScoresAndAttribs(String key, float[] vector,
+      VSimParams params);
+
   /**
    * <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
    * similar to a given element in the vector set.
@@ -184,6 +200,21 @@ public interface VectorSetCommands {
   @Experimental
   Map<String, Double> vsimByElementWithScores(String key, String element, VSimParams params);
 
+  /**
+   * <b><a href="https://redis.io/docs/latest/commands/vsim/">VSIM Command</a></b> Return elements
+   * similar to a given element in the vector set with their similarity scores and attributes.
+   * <p>
+   * Time complexity: O(log(N)) where N is the number of elements in the vector set.
+   * @param key the name of the key that holds the vector set data
+   * @param element the name of the element to use as similarity reference
+   * @param params additional parameters for the VSIM command (WITHSCORES and WITHATTRIBS will be
+   *          automatically added)
+   * @return map of element names to their similarity scores and attributes
+   */
+  @Experimental
+  Map<String, VSimScoreAttribs> vsimByElementWithScoresAndAttribs(String key, String element,
+      VSimParams params);
+
   /**
    * <b><a href="https://redis.io/docs/latest/commands/vdim/">VDIM Command</a></b> Return the number
    * of dimensions of the vectors in the specified vector set.