Support FT.PROFILE command (#3188)

* Support FT.PROFILE command

* edit

* Support recursive child iterators

* Test single child iterator

* Test profile expansion

* Test vector similarity

* Support LIMITED option

* Support 'Child iterator' in reply

irrespective of RediSearch/RediSearch#3205 getting addressed.

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

@@ -1667,6 +1667,103 @@ public AggregationResult build(Object data) {
     }
   };
 
+  public static final Builder<Map<String, Object>> SEARCH_PROFILE_PROFILE = new Builder<Map<String, Object>>() {
+
+    private final String ITERATORS_PROFILE_STR = "Iterators profile";
+    private final String CHILD_ITERATORS_STR = "Child iterators";
+    private final String RESULT_PROCESSORS_PROFILE_STR = "Result processors profile";
+
+    @Override
+    public Map<String, Object> build(Object data) {
+      List<Object> list = (List<Object>) SafeEncoder.encodeObject(data);
+      Map<String, Object> profileMap = new HashMap<>(list.size(), 1f);
+
+      for (Object listObject : list) {
+        List<Object> attributeList = (List<Object>) listObject;
+        String attributeName = (String) attributeList.get(0);
+        Object attributeValue;
+
+        if (attributeList.size() == 2) {
+
+          Object value = attributeList.get(1);
+          if (attributeName.equals(ITERATORS_PROFILE_STR)) {
+            attributeValue = parseIterators(value);
+          } else if (attributeName.endsWith(" time")) {
+            attributeValue = DoublePrecision.parseFloatingPointNumber((String) value);
+          } else {
+            attributeValue = value;
+          }
+
+        } else if (attributeList.size() > 2) {
+
+          if (attributeName.equals(RESULT_PROCESSORS_PROFILE_STR)) {
+            List<Map<String, Object>> resultProcessorsProfileList = new ArrayList<>(attributeList.size() - 1);
+            for (int i = 1; i < attributeList.size(); i++) {
+              resultProcessorsProfileList.add(parseResultProcessors(attributeList.get(i)));
+            }
+            attributeValue = resultProcessorsProfileList;
+          } else {
+            attributeValue = attributeList.subList(1, attributeList.size());
+          }
+
+        } else {
+          attributeValue = null;
+        }
+
+        profileMap.put(attributeName, attributeValue);
+      }
+      return profileMap;
+    }
+
+    private Map<String, Object> parseResultProcessors(Object data) {
+      List<Object> list = (List<Object>) data;
+      Map<String, Object> map = new HashMap<>(list.size() / 2, 1f);
+      for (int i = 0; i < list.size(); i += 2) {
+        String key = (String) list.get(i);
+        Object value = list.get(i + 1);
+        if (key.equals("Time")) {
+          value = DoublePrecision.parseFloatingPointNumber((String) value);
+        }
+        map.put(key, value);
+      }
+      return map;
+    }
+
+    private Object parseIterators(Object data) {
+      if (!(data instanceof List)) return data;
+      List iteratorsAttributeList = (List) data;
+      int childIteratorsIndex = iteratorsAttributeList.indexOf(CHILD_ITERATORS_STR);
+      // https://github.com/RediSearch/RediSearch/issues/3205 patch. TODO: Undo if resolved in RediSearch.
+      if (childIteratorsIndex < 0) childIteratorsIndex = iteratorsAttributeList.indexOf("Child iterator");
+
+      Map<String, Object> iteratorsProfile;
+      if (childIteratorsIndex < 0) {
+        childIteratorsIndex = iteratorsAttributeList.size();
+        iteratorsProfile = new HashMap<>(childIteratorsIndex / 2, 1f);
+      } else {
+        iteratorsProfile = new HashMap<>(1 + childIteratorsIndex / 2, 1f);
+      }
+
+      for (int i = 0; i < childIteratorsIndex; i += 2) {
+        String key = (String) iteratorsAttributeList.get(i);
+        Object value = iteratorsAttributeList.get(i + 1);
+        if (key.equals("Time")) {
+          value = DoublePrecision.parseFloatingPointNumber((String) value);
+        }
+        iteratorsProfile.put(key, value);
+      }
+
+      if (childIteratorsIndex + 1 < iteratorsAttributeList.size()) {
+        List childIteratorsList = new ArrayList(iteratorsAttributeList.size() - childIteratorsIndex - 1);
+        for (int i = childIteratorsIndex + 1; i < iteratorsAttributeList.size(); i++) {
+          childIteratorsList.add(parseIterators(iteratorsAttributeList.get(i)));
+        }
+        iteratorsProfile.put(CHILD_ITERATORS_STR, childIteratorsList);
+      }
+      return iteratorsProfile;
+    }
+  };
+
   public static final Builder<Map<String, List<String>>> SEARCH_SYNONYM_GROUPS = new Builder<Map<String, List<String>>>() {
     @Override
     public Map<String, List<String>> build(Object data) {

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

@@ -3132,6 +3132,31 @@ public CommandObject<String> ftCursorDel(String indexName, long cursorId) {
         .add(indexName).add(cursorId), BuilderFactory.STRING);
   }
 
+  public CommandObject<Map.Entry<AggregationResult, Map<String, Object>>> ftProfileAggregate(
+      String indexName, FTProfileParams profileParams, AggregationBuilder aggr) {
+    return new CommandObject<>(commandArguments(SearchCommand.PROFILE).add(indexName)
+        .add(SearchKeyword.AGGREGATE).addParams(profileParams).add(SearchKeyword.QUERY)
+        .addObjects(aggr.getArgs()), new SearchProfileResponseBuilder<>(!aggr.isWithCursor()
+            ? BuilderFactory.SEARCH_AGGREGATION_RESULT
+            : BuilderFactory.SEARCH_AGGREGATION_RESULT_WITH_CURSOR));
+  }
+
+  public CommandObject<Map.Entry<SearchResult, Map<String, Object>>> ftProfileSearch(
+      String indexName, FTProfileParams profileParams, Query query) {
+    return new CommandObject<>(commandArguments(SearchCommand.PROFILE).add(indexName)
+        .add(SearchKeyword.SEARCH).addParams(profileParams).add(SearchKeyword.QUERY)
+        .addParams(query), new SearchProfileResponseBuilder<>(new SearchResultBuilder(
+            !query.getNoContent(), query.getWithScores(), query.getWithPayloads(), true)));
+  }
+
+  public CommandObject<Map.Entry<SearchResult, Map<String, Object>>> ftProfileSearch(
+      String indexName, FTProfileParams profileParams, String query, FTSearchParams searchParams) {
+    return new CommandObject<>(commandArguments(SearchCommand.PROFILE).add(indexName)
+        .add(SearchKeyword.SEARCH).addParams(profileParams).add(SearchKeyword.QUERY).add(query)
+        .addParams(searchParams), new SearchProfileResponseBuilder<>(new SearchResultBuilder(
+            !searchParams.getNoContent(), searchParams.getWithScores(), false, true)));
+  }
+
   public CommandObject<String> ftDropIndex(String indexName) {
     return new CommandObject<>(commandArguments(SearchCommand.DROPINDEX).add(indexName), BuilderFactory.STRING);
   }
@@ -3992,6 +4017,22 @@ public final CommandObject<Map<String, Object>> graphConfigGet(String configName
   }
   // RedisGraph commands
 
+  private class SearchProfileResponseBuilder<T> extends Builder<Map.Entry<T, Map<String, Object>>> {
+
+    private final Builder<T> replyBuilder;
+
+    public SearchProfileResponseBuilder(Builder<T> replyBuilder) {
+      this.replyBuilder = replyBuilder;
+    }
+
+    @Override
+    public Map.Entry<T, Map<String, Object>> build(Object data) {
+      List<Object> list = (List<Object>) data;
+      return KeyValue.of(replyBuilder.build(list.get(0)),
+          BuilderFactory.SEARCH_PROFILE_PROFILE.build(list.get(1)));
+    }
+  }
+
   private static final Gson GSON = new Gson();
 
   private class GsonObjectBuilder<T> extends Builder<T> {

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

@@ -3511,6 +3511,24 @@ public String ftCursorDel(String indexName, long cursorId) {
     return executeCommand(commandObjects.ftCursorDel(indexName, cursorId));
   }
 
+  @Override
+  public Map.Entry<AggregationResult, Map<String, Object>> ftProfileAggregate(String indexName,
+      FTProfileParams profileParams, AggregationBuilder aggr) {
+    return executeCommand(commandObjects.ftProfileAggregate(indexName, profileParams, aggr));
+  }
+
+  @Override
+  public Map.Entry<SearchResult, Map<String, Object>> ftProfileSearch(String indexName,
+      FTProfileParams profileParams, Query query) {
+    return executeCommand(commandObjects.ftProfileSearch(indexName, profileParams, query));
+  }
+
+  @Override
+  public Map.Entry<SearchResult, Map<String, Object>> ftProfileSearch(String indexName,
+      FTProfileParams profileParams, String query, FTSearchParams searchParams) {
+    return executeCommand(commandObjects.ftProfileSearch(indexName, profileParams, query, searchParams));
+  }
+
   @Override
   public String ftDropIndex(String indexName) {
     return executeCommand(commandObjects.ftDropIndex(indexName));

src/main/java/redis/clients/jedis/search/FTProfileParams.java

@@ -0,0 +1,34 @@
+package redis.clients.jedis.search;
+
+import static redis.clients.jedis.search.SearchProtocol.SearchKeyword.LIMITED;
+
+import redis.clients.jedis.CommandArguments;
+import redis.clients.jedis.params.IParams;
+
+public class FTProfileParams implements IParams {
+
+  private boolean limited;
+
+  public FTProfileParams() {
+  }
+
+  public static FTProfileParams profileParams() {
+    return new FTProfileParams();
+  }
+
+  /**
+   * Removes details of {@code reader} iterator.
+   */
+  public FTProfileParams limited() {
+    this.limited = true;
+    return this;
+  }
+
+  @Override
+  public void addParams(CommandArguments args) {
+
+    if (limited) {
+      args.add(LIMITED);
+    }
+  }
+}

src/main/java/redis/clients/jedis/search/RediSearchCommands.java

@@ -62,6 +62,15 @@ default SearchResult ftSearch(String indexName) {
 
   String ftCursorDel(String indexName, long cursorId);
 
+  Map.Entry<AggregationResult, Map<String, Object>> ftProfileAggregate(String indexName,
+      FTProfileParams profileParams, AggregationBuilder aggr);
+
+  Map.Entry<SearchResult, Map<String, Object>> ftProfileSearch(String indexName,
+      FTProfileParams profileParams, Query query);
+
+  Map.Entry<SearchResult, Map<String, Object>> ftProfileSearch(String indexName,
+      FTProfileParams profileParams, String query, FTSearchParams searchParams);
+
   String ftDropIndex(String indexName);
 
   String ftDropIndexDD(String indexName);
@@ -84,7 +93,8 @@ default SearchResult ftSearch(String indexName) {
 
   Map<String, Map<String, Double>> ftSpellCheck(String index, String query);
 
-  Map<String, Map<String, Double>> ftSpellCheck(String index, String query, FTSpellCheckParams spellCheckParams);
+  Map<String, Map<String, Double>> ftSpellCheck(String index, String query,
+      FTSpellCheckParams spellCheckParams);
 
   Map<String, Object> ftInfo(String indexName);
 

src/main/java/redis/clients/jedis/search/RediSearchPipelineCommands.java

@@ -85,7 +85,8 @@ default Response<SearchResult> ftSearch(String indexName) {
 
   Response<Map<String, Map<String, Double>>> ftSpellCheck(String index, String query);
 
-  Response<Map<String, Map<String, Double>>> ftSpellCheck(String index, String query, FTSpellCheckParams spellCheckParams);
+  Response<Map<String, Map<String, Double>>> ftSpellCheck(String index, String query,
+      FTSpellCheckParams spellCheckParams);
 
   Response<Map<String, Object>> ftInfo(String indexName);
 

src/main/java/redis/clients/jedis/search/SearchProtocol.java

@@ -31,7 +31,8 @@ public enum SearchCommand implements ProtocolCommand {
     DICTDEL("FT.DICTDEL"),
     DICTDUMP("FT.DICTDUMP"),
     SPELLCHECK("FT.SPELLCHECK"),
-    TAGVALS("FT.TAGVALS");
+    TAGVALS("FT.TAGVALS"),
+    PROFILE("FT.PROFILE");
 
     private final byte[] raw;
 
@@ -54,6 +55,7 @@ public enum SearchKeyword implements Rawable {
     SCORE_FIELD, SCORER, PARAMS, AS, DIALECT, SLOP, TIMEOUT, INORDER, EXPANDER, MAXTEXTFIELDS,
     SKIPINITIALSCAN, WITHSUFFIXTRIE, NOSTEM, NOINDEX, PHONETIC, WEIGHT, CASESENSITIVE, /*EXPLAINSCORE,*/
     LOAD, APPLY, GROUPBY, MAXIDLE, WITHCURSOR, DISTANCE, TERMS, INCLUDE, EXCLUDE,
+    SEARCH, AGGREGATE, QUERY, LIMITED,
     @Deprecated ASYNC, @Deprecated PAYLOAD_FIELD, @Deprecated WITHPAYLOADS, @Deprecated PAYLOAD,
     @Deprecated COUNT;
 

src/test/java/redis/clients/jedis/modules/search/SearchTest.java

@@ -993,10 +993,10 @@ public void timeout() {
     Schema sc = new Schema().addTextField("field1", 1.0).addTextField("field2", 1.0);
     assertEquals("OK", client.ftCreate(index, IndexOptions.defaultOptions(), sc));
 
-    Map<String, Object> doc = new HashMap<>();
-    doc.put("field1", "value");
-    doc.put("field2", "not");
-    addDocument("doc1", doc);
+    Map<String, String> map = new HashMap<>();
+    map.put("field1", "value");
+    map.put("field2", "not");
+    client.hset("doc1", map);
 
     SearchResult res = client.ftSearch(index, new Query("value").timeout(1000));
     assertEquals(1, res.getTotalResults());
@@ -1010,11 +1010,11 @@ public void inOrder() {
     Schema sc = new Schema().addTextField("field1", 1.0).addTextField("field2", 1.0);
     assertEquals("OK", client.ftCreate(index, IndexOptions.defaultOptions(), sc));
 
-    Map<String, Object> doc = new HashMap<>();
-    doc.put("field1", "value");
-    doc.put("field2", "not");
-    addDocument("doc2", doc);
-    addDocument("doc1", doc);
+    Map<String, String> map = new HashMap<>();
+    map.put("field1", "value");
+    map.put("field2", "not");
+    client.hset("doc2", map);
+    client.hset("doc1", map);
 
     SearchResult res = client.ftSearch(index, new Query("value").setInOrder());
     assertEquals(2, res.getTotalResults());
@@ -1144,6 +1144,15 @@ public void testHNSWVVectorSimilarity() {
     Document doc1 = client.ftSearch(index, query).getDocuments().get(0);
     assertEquals("a", doc1.getId());
     assertEquals("0", doc1.get("__v_score"));
+
+    // profile
+    Map.Entry<SearchResult, Map<String, Object>> profile
+        = client.ftProfileSearch(index, FTProfileParams.profileParams(), query);
+    doc1 = profile.getKey().getDocuments().get(0);
+    assertEquals("a", doc1.getId());
+    assertEquals("0", doc1.get("__v_score"));
+    Map<String, Object> iteratorsProfile = (Map<String, Object>) profile.getValue().get("Iterators profile");
+    assertEquals("VECTOR", iteratorsProfile.get("Type"));
   }
 
   @Test
@@ -1168,5 +1177,35 @@ public void testFlatVectorSimilarity() {
     Document doc1 = client.ftSearch(index, query).getDocuments().get(0);
     assertEquals("a", doc1.getId());
     assertEquals("0", doc1.get("__v_score"));
+
+    // profile
+    Map.Entry<SearchResult, Map<String, Object>> profile
+        = client.ftProfileSearch(index, FTProfileParams.profileParams(), query);
+    doc1 = profile.getKey().getDocuments().get(0);
+    assertEquals("a", doc1.getId());
+    assertEquals("0", doc1.get("__v_score"));
+    Map<String, Object> iteratorsProfile = (Map<String, Object>) profile.getValue().get("Iterators profile");
+    assertEquals("VECTOR", iteratorsProfile.get("Type"));
+  }
+
+  @Test
+  public void searchProfile() {
+    Schema sc = new Schema().addTextField("t1", 1.0).addTextField("t2", 1.0);
+    assertEquals("OK", client.ftCreate(index, IndexOptions.defaultOptions(), sc));
+
+    Map<String, String> map = new HashMap<>();
+    map.put("t1", "foo");
+    map.put("t2", "bar");
+    client.hset("doc1", map);
+
+    Map.Entry<SearchResult, Map<String, Object>> profile = client.ftProfileSearch(index,
+        FTProfileParams.profileParams(), new Query("foo"));
+    // Iterators profile={Type=TEXT, Time=0.0, Term=foo, Counter=1, Size=1}
+    Map<String, Object> iteratorsProfile = (Map<String, Object>) profile.getValue().get("Iterators profile");
+    assertEquals("TEXT", iteratorsProfile.get("Type"));
+    assertEquals("foo", iteratorsProfile.get("Term"));
+    assertEquals(1L, iteratorsProfile.get("Counter"));
+    assertEquals(1L, iteratorsProfile.get("Size"));
+    assertSame(Double.class, iteratorsProfile.get("Time").getClass());
   }
 }