Refactor Search Aggregation codes (#3451)

* Refactor Search Aggregation codes * use keywords more * breaking doc * separate option for SORTBY MAX
redis · Jun 1, 2023 · ef21a05 · ef21a05
1 parent 0f760bc
commit ef21a05
Show file tree

Hide file tree

Showing 11 changed files with 157 additions and 241 deletions.
diff --git a/docs/jedis5-breaking.md b/docs/jedis5-breaking.md
@@ -45,6 +45,7 @@
   - `KeyedListElement`
   - `TSKeyValue`
   - `TSKeyedElements`
+  - `Limit`
 
 - `STREAM_AUTO_CLAIM_ID_RESPONSE` in BuilderFactory has been renamed to `STREAM_AUTO_CLAIM_JUSTID_RESPONSE`.
 
@@ -67,6 +68,20 @@
 
 - `getParams()` method is removed from `SortingParams` class.
 
+- `addCommandEncodedArguments` and `addCommandBinaryArguments` methods have been removed from `FieldName` class.
+
+- `limit` and `getArgs` methods have been removed from `Group` class.
+
+- `Reducer` abstract class is refactored:
+  - `Reducer(String field)` constructored is removed; `Reducer(String name, String field)` constructor is added.
+  - `Reducer(String name)` constructored is added; it will cause runtime error with older `Reducer(String field)` constructor.
+  - `getName` method is removed.
+  - `getAlias` method is removed.
+  - `setAlias` method is removed; use `as` method.
+  - `setAliasAsField` method is removed.
+  - `getOwnArgs` method is now abstract.
+  - `getArgs` method is removed.
+
 - All variants of `blmpop` and `bzmpop` methods now take `double timeout` parameter instead of `long timeout` parameter.
   This is breaking ONLY IF you are using `Long` for timeout.
 

diff --git a/src/main/java/redis/clients/jedis/search/FTSearchParams.java b/src/main/java/redis/clients/jedis/search/FTSearchParams.java
@@ -9,6 +9,7 @@
 import redis.clients.jedis.args.GeoUnit;
 import redis.clients.jedis.args.SortingOrder;
 import redis.clients.jedis.params.IParams;
+import redis.clients.jedis.util.LazyRawable;
 
 /**
  * Query represents query parameters and filters to load results from the engine

diff --git a/src/main/java/redis/clients/jedis/search/FieldName.java b/src/main/java/redis/clients/jedis/search/FieldName.java
@@ -3,15 +3,10 @@
 import java.util.List;
 import redis.clients.jedis.CommandArguments;
 import redis.clients.jedis.params.IParams;
-
-import redis.clients.jedis.util.SafeEncoder;
+import redis.clients.jedis.search.SearchProtocol.SearchKeyword;
 
 public class FieldName implements IParams {
 
-  private static final String AS_ENCODED = "AS";
-  private static final byte[] AS_BINARY = SafeEncoder.encode(AS_ENCODED);
-  private static final byte[] AS = SafeEncoder.encode("AS");
-
   private final String name;
   private String attribute;
 
@@ -35,36 +30,25 @@ public FieldName as(String attribute) {
     return this;
   }
 
-  public int addCommandEncodedArguments(List<String> args) {
+  public int addCommandArguments(List<Object> args) {
     args.add(name);
     if (attribute == null) {
       return 1;
     }
 
-    args.add(AS_ENCODED);
+    args.add(SearchKeyword.AS);
     args.add(attribute);
     return 3;
   }
 
-  public int addCommandBinaryArguments(List<byte[]> args) {
-    args.add(SafeEncoder.encode(name));
-    if (attribute == null) {
-      return 1;
-    }
-
-    args.add(AS_BINARY);
-    args.add(SafeEncoder.encode(attribute));
-    return 3;
-  }
-
   public int addCommandArguments(CommandArguments args) {
-    args.add(SafeEncoder.encode(name));
+    args.add(name);
     if (attribute == null) {
       return 1;
     }
 
-    args.add(AS);
-    args.add(SafeEncoder.encode(attribute));
+    args.add(SearchKeyword.AS);
+    args.add(attribute);
     return 3;
   }
 

diff --git a/src/main/java/redis/clients/jedis/search/Query.java b/src/main/java/redis/clients/jedis/search/Query.java
@@ -9,6 +9,7 @@
 import redis.clients.jedis.Protocol;
 import redis.clients.jedis.params.IParams;
 import redis.clients.jedis.search.SearchProtocol.SearchKeyword;
+import redis.clients.jedis.util.LazyRawable;
 import redis.clients.jedis.util.SafeEncoder;
 
 /**

diff --git a/src/main/java/redis/clients/jedis/search/SearchProtocol.java b/src/main/java/redis/clients/jedis/search/SearchProtocol.java
@@ -56,7 +56,7 @@ public enum SearchKeyword implements Rawable {
     SCORE_FIELD, SCORER, PARAMS, AS, DIALECT, SLOP, TIMEOUT, INORDER, EXPANDER, MAXTEXTFIELDS,
     SKIPINITIALSCAN, WITHSUFFIXTRIE, NOSTEM, NOINDEX, PHONETIC, WEIGHT, CASESENSITIVE,
     LOAD, APPLY, GROUPBY, MAXIDLE, WITHCURSOR, DISTANCE, TERMS, INCLUDE, EXCLUDE,
-    SEARCH, AGGREGATE, QUERY, LIMITED, COUNT;
+    SEARCH, AGGREGATE, QUERY, LIMITED, COUNT, REDUCE;
 
     private final byte[] raw;
 

diff --git a/src/main/java/redis/clients/jedis/search/aggr/AggregationBuilder.java b/src/main/java/redis/clients/jedis/search/aggr/AggregationBuilder.java
@@ -6,17 +6,19 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.StringJoiner;
 
+import redis.clients.jedis.Protocol;
 import redis.clients.jedis.search.FieldName;
-import redis.clients.jedis.util.SafeEncoder;
+import redis.clients.jedis.search.SearchProtocol.SearchKeyword;
+import redis.clients.jedis.util.LazyRawable;
 
 /**
  * @author Guy Korland
  */
 public class AggregationBuilder {
 
-  private final List<String> args = new ArrayList<>();
+  private final List<Object> args = new ArrayList<>();
+  private Integer dialect;
   private boolean isWithCursor = false;
 
   public AggregationBuilder(String query) {
@@ -32,26 +34,27 @@ public AggregationBuilder load(String... fields) {
   }
 
   public AggregationBuilder load(FieldName... fields) {
-    args.add("LOAD");
-    final int loadCountIndex = args.size();
-    args.add(null);
+    args.add(SearchKeyword.LOAD);
+    LazyRawable rawLoadCount = new LazyRawable();
+    args.add(rawLoadCount);
     int loadCount = 0;
     for (FieldName fn : fields) {
-      loadCount += fn.addCommandEncodedArguments(args);
+      loadCount += fn.addCommandArguments(args);
     }
-    args.set(loadCountIndex, Integer.toString(loadCount));
+    rawLoadCount.setRaw(Protocol.toByteArray(loadCount));
     return this;
   }
 
   public AggregationBuilder loadAll() {
-    args.add("LOAD");
-    args.add("*");
+    args.add(SearchKeyword.LOAD);
+    args.add(Protocol.BYTES_ASTERISK);
     return this;
   }
 
   public AggregationBuilder limit(int offset, int count) {
-    Limit limit = new Limit(offset, count);
-    limit.addArgs(args);
+    args.add(SearchKeyword.LIMIT);
+    args.add(offset);
+    args.add(count);
     return this;
   }
 
@@ -60,22 +63,12 @@ public AggregationBuilder limit(int count) {
   }
 
   public AggregationBuilder sortBy(SortedField... fields) {
-    args.add("SORTBY");
+    args.add(SearchKeyword.SORTBY);
     args.add(Integer.toString(fields.length * 2));
     for (SortedField field : fields) {
       args.add(field.getField());
       args.add(field.getOrder());
     }
-
-    return this;
-  }
-
-  public AggregationBuilder sortBy(int max, SortedField... fields) {
-    sortBy(fields);
-    if (max > 0) {
-      args.add("MAX");
-      args.add(Integer.toString(max));
-    }
     return this;
   }
 
@@ -87,20 +80,51 @@ public AggregationBuilder sortByDesc(String field) {
     return sortBy(SortedField.desc(field));
   }
 
+  /**
+   * {@link AggregationBuilder#sortBy(redis.clients.jedis.search.aggr.SortedField...)}
+   * (or {@link AggregationBuilder#sortByAsc(java.lang.String)}
+   * or {@link AggregationBuilder#sortByDesc(java.lang.String)})
+   * MUST BE called JUST BEFORE this.
+   * @param max limit
+   * @return this
+   */
+  public AggregationBuilder sortByMax(int max) {
+    args.add(SearchKeyword.MAX);
+    args.add(max);
+    return this;
+  }
+
+  /**
+   * Shortcut to {@link AggregationBuilder#sortBy(redis.clients.jedis.search.aggr.SortedField...)}
+   * and {@link AggregationBuilder#sortByMax(int)}.
+   * @param max limit
+   * @param fields sorted fields
+   * @return this
+   */
+  public AggregationBuilder sortBy(int max, SortedField... fields) {
+    sortBy(fields);
+    sortByMax(max);
+    return this;
+  }
+
   public AggregationBuilder apply(String projection, String alias) {
-    args.add("APPLY");
+    args.add(SearchKeyword.APPLY);
     args.add(projection);
-    args.add("AS");
+    args.add(SearchKeyword.AS);
     args.add(alias);
     return this;
   }
 
+  public AggregationBuilder groupBy(Group group) {
+    args.add(SearchKeyword.GROUPBY);
+    group.addArgs(args);
+    return this;
+  }
+
   public AggregationBuilder groupBy(Collection<String> fields, Collection<Reducer> reducers) {
     String[] fieldsArr = new String[fields.size()];
     Group g = new Group(fields.toArray(fieldsArr));
-    for (Reducer r : reducers) {
-      g.reduce(r);
-    }
+    reducers.forEach((r) -> g.reduce(r));
     groupBy(g);
     return this;
   }
@@ -109,65 +133,61 @@ public AggregationBuilder groupBy(String field, Reducer... reducers) {
     return groupBy(Collections.singletonList(field), Arrays.asList(reducers));
   }
 
-  public AggregationBuilder groupBy(Group group) {
-    args.add("GROUPBY");
-    group.addArgs(args);
+  public AggregationBuilder filter(String expression) {
+    args.add(SearchKeyword.FILTER);
+    args.add(expression);
     return this;
   }
 
-  public AggregationBuilder filter(String expression) {
-    args.add("FILTER");
-    args.add(expression);
+  public AggregationBuilder cursor(int count) {
+    isWithCursor = true;
+    args.add(SearchKeyword.WITHCURSOR);
+    args.add(SearchKeyword.COUNT);
+    args.add(count);
     return this;
   }
 
   public AggregationBuilder cursor(int count, long maxIdle) {
     isWithCursor = true;
-    if (count > 0) {
-      args.add("WITHCURSOR");
-      args.add("COUNT");
-      args.add(Integer.toString(count));
-      if (maxIdle < Long.MAX_VALUE && maxIdle >= 0) {
-        args.add("MAXIDLE");
-        args.add(Long.toString(maxIdle));
-      }
-    }
+    args.add(SearchKeyword.WITHCURSOR);
+    args.add(SearchKeyword.COUNT);
+    args.add(count);
+    args.add(SearchKeyword.MAXIDLE);
+    args.add(maxIdle);
     return this;
   }
 
   public AggregationBuilder verbatim() {
-    args.add("VERBATIM");
+    args.add(SearchKeyword.VERBATIM);
     return this;
   }
 
   public AggregationBuilder timeout(long timeout) {
-    if (timeout >= 0) {
-      args.add("TIMEOUT");
-      args.add(Long.toString(timeout));
-    }
+    args.add(SearchKeyword.TIMEOUT);
+    args.add(timeout);
     return this;
   }
 
   public AggregationBuilder params(Map<String, Object> params) {
-    if (params.size() >= 1) {
-      args.add("PARAMS");
-      args.add(Integer.toString(params.size() * 2));
-      for (Map.Entry<String, Object> entry : params.entrySet()) {
-        args.add(entry.getKey());
-        args.add(String.valueOf(entry.getValue()));
-      }
-    }
-
+    args.add(SearchKeyword.PARAMS);
+    args.add(params.size() * 2);
+    params.forEach((k, v) -> {
+      args.add(k);
+      args.add(v);
+    });
     return this;
   }
 
   public AggregationBuilder dialect(int dialect) {
-    args.add("DIALECT");
-    args.add(Integer.toString(dialect));
+    this.dialect = dialect;
     return this;
   }
 
-  public List<String> getArgs() {
+  public List<Object> getArgs() {
+    if (dialect != null) {
+      args.add(SearchKeyword.DIALECT);
+      args.add(dialect);
+    }
     return Collections.unmodifiableList(args);
   }
 

diff --git a/src/main/java/redis/clients/jedis/search/aggr/Group.java b/src/main/java/redis/clients/jedis/search/aggr/Group.java
@@ -9,9 +9,8 @@
  */
 public class Group {
 
-  private final List<Reducer> reducers = new ArrayList<>();
   private final List<String> fields = new ArrayList<>();
-  private Limit limit = new Limit(0, 0);
+  private final List<Reducer> reducers = new ArrayList<>();
 
   public Group(String... fields) {
     this.fields.addAll(Arrays.asList(fields));
@@ -22,30 +21,11 @@ public Group reduce(Reducer r) {
     return this;
   }
 
-  public Group limit(Limit limit) {
-    this.limit = limit;
-    return this;
-  }
+  public void addArgs(List<Object> args) {
 
-  public void addArgs(List<String> args) {
-    args.add(Integer.toString(fields.size()));
+    args.add(fields.size());
     args.addAll(fields);
-    for (Reducer r : reducers) {
-      args.add("REDUCE");
-      args.add(r.getName());
-      r.addArgs(args);
-      String alias = r.getAlias();
-      if (alias != null && !alias.isEmpty()) {
-        args.add("AS");
-        args.add(alias);
-      }
-    }
-    args.addAll(limit.getArgs());
-  }
 
-  public List<String> getArgs() {
-    List<String> args = new ArrayList<>();
-    addArgs(args);
-    return args;
+    reducers.forEach((r) -> r.addArgs(args));
   }
 }