Support iteration for SCAN, FT.SEARCH, FT.AGGREGATE commands (#3378)

* Support iteration for SCAN, FT.SEARCH, FT.AGGREGATE commands

* format

* Support fluent collect()

Author: sazzad16

src/main/java/redis/clients/jedis/ScanRoundRobin.java renamed to src/main/java/redis/clients/jedis/ScanIteration.java

@@ -1,33 +1,35 @@
 package redis.clients.jedis;
 
+import java.util.Collection;
 import java.util.function.Function;
+
 import redis.clients.jedis.params.ScanParams;
 import redis.clients.jedis.Protocol.Keyword;
 import redis.clients.jedis.providers.ConnectionProvider;
 import redis.clients.jedis.resps.ScanResult;
-import redis.clients.jedis.util.JedisRoundRobinBase;
+import redis.clients.jedis.util.JedisCommandIterationBase;
 
-public class ScanRoundRobin extends JedisRoundRobinBase<ScanResult<String>> {
+public class ScanIteration extends JedisCommandIterationBase<ScanResult<String>, String> {
 
   private final int count;
   private final Function<String, CommandArguments> args;
 
-  public ScanRoundRobin(ConnectionProvider connectionProvider, int batchCount, String match) {
+  public ScanIteration(ConnectionProvider connectionProvider, int batchCount, String match) {
     super(connectionProvider, BuilderFactory.SCAN_RESPONSE);
     this.count = batchCount;
     this.args = (cursor) -> new CommandArguments(Protocol.Command.SCAN).add(cursor)
         .add(Keyword.MATCH).add(match).add(Keyword.COUNT).add(count);
   }
 
-  public ScanRoundRobin(ConnectionProvider connectionProvider, int batchCount, String match, String type) {
+  public ScanIteration(ConnectionProvider connectionProvider, int batchCount, String match, String type) {
     super(connectionProvider, BuilderFactory.SCAN_RESPONSE);
     this.count = batchCount;
     this.args = (cursor) -> new CommandArguments(Protocol.Command.SCAN).add(cursor)
         .add(Keyword.MATCH).add(match).add(Keyword.COUNT).add(count).add(Keyword.TYPE).add(type);
   }
 
   @Override
-  protected boolean isIterationCompleted(ScanResult<String> reply) {
+  protected boolean isNodeCompleted(ScanResult<String> reply) {
     return reply.isCompleteIteration();
   }
 
@@ -40,4 +42,9 @@ protected CommandArguments initCommandArguments() {
   protected CommandArguments nextCommandArguments(ScanResult<String> lastReply) {
     return args.apply(lastReply.getCursor());
   }
+
+  @Override
+  protected Collection<String> convertBatchToData(ScanResult<String> batch) {
+    return batch.getResult();
+  }
 }

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

@@ -30,7 +30,7 @@
 import redis.clients.jedis.search.*;
 import redis.clients.jedis.search.aggr.AggregationBuilder;
 import redis.clients.jedis.search.aggr.AggregationResult;
-import redis.clients.jedis.search.aggr.FtAggregateRoundRobin;
+import redis.clients.jedis.search.aggr.FtAggregateIteration;
 import redis.clients.jedis.search.schemafields.SchemaField;
 import redis.clients.jedis.timeseries.*;
 import redis.clients.jedis.util.IOUtils;
@@ -595,12 +595,23 @@ public ScanResult<String> scan(String cursor, ScanParams params, String type) {
     return executeCommand(commandObjects.scan(cursor, params, type));
   }
 
-  public ScanRoundRobin scan(int batchCount, String match) {
-    return new ScanRoundRobin(provider, batchCount, match);
+  /**
+   * @param batchCount COUNT for each batch execution
+   * @param match pattern
+   * @return scan iteration
+   */
+  public ScanIteration scanIteration(int batchCount, String match) {
+    return new ScanIteration(provider, batchCount, match);
   }
 
-  public ScanRoundRobin scan(int batchCount, String match, String type) {
-    return new ScanRoundRobin(provider, batchCount, match, type);
+  /**
+   * @param batchCount COUNT for each batch execution
+   * @param match pattern
+   * @param type key type
+   * @return scan iteration
+   */
+  public ScanIteration scanIteration(int batchCount, String match, String type) {
+    return new ScanIteration(provider, batchCount, match, type);
   }
 
   @Override
@@ -3559,14 +3570,14 @@ public SearchResult ftSearch(String indexName, String query, FTSearchParams para
 
   /**
    * {@link FTSearchParams#limit(int, int)} will be ignored.
-   * @param batchSize
-   * @param indexName
-   * @param query
+   * @param batchSize batch size
+   * @param indexName index name
+   * @param query query
    * @param params limit will be ignored
-   * @return search
+   * @return search iteration
    */
-  public FtSearchRoundRobin ftSearch(int batchSize, String indexName, String query, FTSearchParams params) {
-    return new FtSearchRoundRobin(provider, batchSize, indexName, query, params);
+  public FtSearchIteration ftSearchIteration(int batchSize, String indexName, String query, FTSearchParams params) {
+    return new FtSearchIteration(provider, batchSize, indexName, query, params);
   }
 
   @Override
@@ -3576,13 +3587,13 @@ public SearchResult ftSearch(String indexName, Query query) {
 
   /**
    * {@link Query#limit(java.lang.Integer, java.lang.Integer)} will be ignored.
-   * @param batchSize
-   * @param indexName
+   * @param batchSize batch size
+   * @param indexName index name
    * @param query limit will be ignored
-   * @return search
+   * @return search iteration
    */
-  public FtSearchRoundRobin ftSearch(int batchSize, String indexName, Query query) {
-    return new FtSearchRoundRobin(provider, batchSize, indexName, query);
+  public FtSearchIteration ftSearchIteration(int batchSize, String indexName, Query query) {
+    return new FtSearchIteration(provider, batchSize, indexName, query);
   }
 
   @Override
@@ -3615,8 +3626,14 @@ public String ftCursorDel(String indexName, long cursorId) {
     return executeCommand(commandObjects.ftCursorDel(indexName, cursorId));
   }
 
-  public FtAggregateRoundRobin ftAggregateRoundRobin(String indexName, AggregationBuilder aggr) {
-    return new FtAggregateRoundRobin(provider, indexName, aggr);
+  /**
+   * {@link AggregationBuilder#cursor(int, long) CURSOR} must be set.
+   * @param indexName index name
+   * @param aggr cursor must be set
+   * @return aggregate iteration
+   */
+  public FtAggregateIteration ftAggregateIteration(String indexName, AggregationBuilder aggr) {
+    return new FtAggregateIteration(provider, indexName, aggr);
   }
 
   @Override

src/main/java/redis/clients/jedis/search/FtSearchRoundRobin.java renamed to src/main/java/redis/clients/jedis/search/FtSearchIteration.java

@@ -1,13 +1,14 @@
 package redis.clients.jedis.search;
 
+import java.util.Collection;
 import java.util.function.IntFunction;
 
 import redis.clients.jedis.CommandArguments;
 import redis.clients.jedis.providers.ConnectionProvider;
 import redis.clients.jedis.search.SearchResult.SearchResultBuilder;
-import redis.clients.jedis.util.JedisRoundRobinBase;
+import redis.clients.jedis.util.JedisCommandIterationBase;
 
-public class FtSearchRoundRobin extends JedisRoundRobinBase<SearchResult> {
+public class FtSearchIteration extends JedisCommandIterationBase<SearchResult, Document> {
 
   private int batchStart;
   private final int batchSize;
@@ -16,7 +17,7 @@ public class FtSearchRoundRobin extends JedisRoundRobinBase<SearchResult> {
   /**
    * {@link FTSearchParams#limit(int, int)} will be ignored.
    */
-  public FtSearchRoundRobin(ConnectionProvider connectionProvider, int batchSize, String indexName, String query, FTSearchParams params) {
+  public FtSearchIteration(ConnectionProvider connectionProvider, int batchSize, String indexName, String query, FTSearchParams params) {
     super(connectionProvider, new SearchResultBuilder(!params.getNoContent(), params.getWithScores(), false, true));
     this.batchSize = batchSize;
     this.args = (limitFirst) -> new CommandArguments(SearchProtocol.SearchCommand.SEARCH)
@@ -26,15 +27,15 @@ public FtSearchRoundRobin(ConnectionProvider connectionProvider, int batchSize,
   /**
    * {@link Query#limit(java.lang.Integer, java.lang.Integer)} will be ignored.
    */
-  public FtSearchRoundRobin(ConnectionProvider connectionProvider, int batchSize, String indexName, Query query) {
+  public FtSearchIteration(ConnectionProvider connectionProvider, int batchSize, String indexName, Query query) {
     super(connectionProvider, new SearchResultBuilder(!query.getNoContent(), query.getWithScores(), query.getWithPayloads(), true));
     this.batchSize = batchSize;
     this.args = (limitFirst) -> new CommandArguments(SearchProtocol.SearchCommand.SEARCH)
         .add(indexName).addParams(query.limit(limitFirst, this.batchSize));
   }
 
   @Override
-  protected boolean isIterationCompleted(SearchResult reply) {
+  protected boolean isNodeCompleted(SearchResult reply) {
     return batchStart >= reply.getTotalResults() - batchSize;
   }
 
@@ -49,4 +50,9 @@ protected CommandArguments nextCommandArguments(SearchResult lastReply) {
     batchStart += batchSize;
     return args.apply(batchStart);
   }
+
+  @Override
+  protected Collection<Document> convertBatchToData(SearchResult batch) {
+    return batch.getDocuments();
+  }
 }

src/main/java/redis/clients/jedis/search/aggr/FtAggregateRoundRobin.java renamed to src/main/java/redis/clients/jedis/search/aggr/FtAggregateIteration.java

@@ -1,25 +1,33 @@
 package redis.clients.jedis.search.aggr;
 
+import java.util.Collection;
+
 import redis.clients.jedis.CommandArguments;
 import redis.clients.jedis.providers.ConnectionProvider;
 import redis.clients.jedis.search.SearchBuilderFactory;
 import redis.clients.jedis.search.SearchProtocol;
-import redis.clients.jedis.util.JedisRoundRobinBase;
+import redis.clients.jedis.util.JedisCommandIterationBase;
 
-public class FtAggregateRoundRobin extends JedisRoundRobinBase<AggregationResult> {
+public class FtAggregateIteration extends JedisCommandIterationBase<AggregationResult, Row> {
 
   private final String indexName;
   private final CommandArguments args;
 
-  public FtAggregateRoundRobin(ConnectionProvider connectionProvider, String indexName, AggregationBuilder aggr) {
+  /**
+   * {@link AggregationBuilder#cursor(int, long) CURSOR} must be set.
+   * @param connectionProvider connection provider
+   * @param indexName index name
+   * @param aggr cursor must be set
+   */
+  public FtAggregateIteration(ConnectionProvider connectionProvider, String indexName, AggregationBuilder aggr) {
     super(connectionProvider, SearchBuilderFactory.SEARCH_AGGREGATION_RESULT_WITH_CURSOR);
     if (!aggr.isWithCursor()) throw new IllegalArgumentException("cursor must be set");
     this.indexName = indexName;
     this.args = new CommandArguments(SearchProtocol.SearchCommand.AGGREGATE).add(this.indexName).addObjects(aggr.getArgs());
   }
 
   @Override
-  protected boolean isIterationCompleted(AggregationResult reply) {
+  protected boolean isNodeCompleted(AggregationResult reply) {
     return reply.getCursorId() == 0L;
   }
 
@@ -33,4 +41,9 @@ protected CommandArguments nextCommandArguments(AggregationResult lastReply) {
     return new CommandArguments(SearchProtocol.SearchCommand.CURSOR).add(SearchProtocol.SearchKeyword.READ)
         .add(indexName).add(lastReply.getCursorId());
   }
+
+  @Override
+  protected Collection<Row> convertBatchToData(AggregationResult batch) {
+    return batch.getRows();
+  }
 }

src/main/java/redis/clients/jedis/util/JedisRoundRobinBase.java renamed to src/main/java/redis/clients/jedis/util/JedisCommandIterationBase.java

@@ -1,31 +1,37 @@
 package redis.clients.jedis.util;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Queue;
+import java.util.function.Supplier;
 
 import redis.clients.jedis.Builder;
 import redis.clients.jedis.CommandArguments;
 import redis.clients.jedis.Connection;
 import redis.clients.jedis.providers.ConnectionProvider;
 
-public abstract class JedisRoundRobinBase<R> {
+/**
+ * @param <B> Type of each batch reply
+ * @param <D> Type of each data
+ */
+public abstract class JedisCommandIterationBase<B, D> {
 
-  private final Builder<R> builder;
+  private final Builder<B> builder;
 
   private final Queue<Map.Entry> connections;
 
   private Map.Entry connection;
 
-  private R lastReply;
+  private B lastReply;
 
   private boolean roundRobinCompleted;
   private boolean iterationCompleted;
 
-  protected JedisRoundRobinBase(ConnectionProvider connectionProvider, Builder<R> responseBuilder) {
+  protected JedisCommandIterationBase(ConnectionProvider connectionProvider, Builder<B> responseBuilder) {
     Map connectionMap = connectionProvider.getConnectionMap();
     ArrayList<Map.Entry> connectionList = new ArrayList<>(connectionMap.entrySet());
     Collections.shuffle(connectionList);
@@ -35,17 +41,17 @@ protected JedisRoundRobinBase(ConnectionProvider connectionProvider, Builder<R>
     this.roundRobinCompleted = this.connections.isEmpty();
   }
 
-  public final boolean isRoundRobinCompleted() {
+  public final boolean isIterationCompleted() {
     return roundRobinCompleted;
   }
 
-  protected abstract boolean isIterationCompleted(R reply);
+  protected abstract boolean isNodeCompleted(B reply);
 
   protected abstract CommandArguments initCommandArguments();
 
-  protected abstract CommandArguments nextCommandArguments(R lastReply);
+  protected abstract CommandArguments nextCommandArguments(B lastReply);
 
-  public final R get() {
+  public final B nextBatch() {
     if (roundRobinCompleted) {
       throw new NoSuchElementException();
     }
@@ -70,12 +76,25 @@ public final R get() {
     }
 
     lastReply = builder.build(rawReply);
-    iterationCompleted = isIterationCompleted(lastReply);
+    iterationCompleted = isNodeCompleted(lastReply);
     if (iterationCompleted) {
       if (connections.isEmpty()) {
         roundRobinCompleted = true;
       }
     }
     return lastReply;
   }
+
+  protected abstract Collection<D> convertBatchToData(B batch);
+
+  public final Collection<D> nextBatchList() {
+    return convertBatchToData(nextBatch());
+  }
+
+  public final Collection<D> collect(Collection<D> c) {
+    while (!isIterationCompleted()) {
+      c.addAll(nextBatchList());
+    }
+    return c;
+  }
 }

src/test/java/redis/clients/jedis/commands/jedis/ClusterValuesCommandsTest.java

@@ -18,7 +18,7 @@
 import redis.clients.jedis.GeoCoordinate;
 import redis.clients.jedis.JedisPubSub;
 import redis.clients.jedis.Protocol;
-import redis.clients.jedis.ScanRoundRobin;
+import redis.clients.jedis.ScanIteration;
 import redis.clients.jedis.args.GeoUnit;
 import redis.clients.jedis.params.GeoRadiusParam;
 import redis.clients.jedis.params.GeoRadiusStoreParam;
@@ -138,7 +138,7 @@ public void flushAllBroadcast() {
   }
 
   @Test
-  public void scanRoundRobin() {
+  public void scanIteration() {
     Set<String> allIn = new HashSet<>(26 * 26);
     char[] arr = new char[2];
     for (int i = 0; i < 26; i++) {
@@ -152,11 +152,36 @@ public void scanRoundRobin() {
     }
 
     Set<String> allScan = new HashSet<>();
-    ScanRoundRobin scan = cluster.scan(10, "*");
-    while (!scan.isRoundRobinCompleted()) {
-      ScanResult<String> batch = scan.get();
+    ScanIteration scan = cluster.scanIteration(10, "*");
+    while (!scan.isIterationCompleted()) {
+      ScanResult<String> batch = scan.nextBatch();
       allScan.addAll(batch.getResult());
     }
     assertEquals(allIn, allScan);
+
+    Set<String> allTypeScan = new HashSet<>();
+    ScanIteration typeScan = cluster.scanIteration(10, "*", "string");
+    while (!typeScan.isIterationCompleted()) {
+      ScanResult<String> batch = typeScan.nextBatch();
+      allTypeScan.addAll(batch.getResult());
+    }
+    assertEquals(allIn, allTypeScan);
+  }
+
+  @Test
+  public void scanIterationCollect() {
+    Set<String> allIn = new HashSet<>(26 * 26);
+    char[] arr = new char[2];
+    for (int i = 0; i < 26; i++) {
+      arr[0] = (char) ('a' + i);
+      for (int j = 0; j < 26; j++) {
+        arr[1] = (char) ('a' + j);
+        String str = new String(arr);
+        cluster.incr(str);
+        allIn.add(str);
+      }
+    }
+
+    assertEquals(allIn, cluster.scanIteration(100, "*").collect(new HashSet<>(26 * 26)));
   }
 }