Push limit term into secondary index queries

Change-Id: Id4ec5fe0375229ff8b82ec46fbe300c78098dbbe

gerrit-lucene/src/main/java/com/google/gerrit/lucene/LuceneChangeIndex.java

@@ -231,7 +231,7 @@ public class LuceneChangeIndex implements ChangeIndex {
   }
 
   @Override
-  public ChangeDataSource getSource(Predicate<ChangeData> p)
+  public ChangeDataSource getSource(Predicate<ChangeData> p, int limit)
       throws QueryParseException {
     Set<Change.Status> statuses = IndexRewriteImpl.getPossibleStatus(p);
     List<SubIndex> indexes = Lists.newArrayListWithCapacity(2);
@@ -241,7 +241,7 @@ public class LuceneChangeIndex implements ChangeIndex {
     if (!Sets.intersection(statuses, CLOSED_STATUSES).isEmpty()) {
       indexes.add(closedIndex);
     }
-    return new QuerySource(indexes, QueryBuilder.toQuery(p));
+    return new QuerySource(indexes, QueryBuilder.toQuery(p), limit);
   }
 
   @Override
@@ -256,16 +256,16 @@ public class LuceneChangeIndex implements ChangeIndex {
   }
 
   private static class QuerySource implements ChangeDataSource {
-    // TODO(dborowitz): Push limit down from predicate tree.
-    private static final int LIMIT = 1000;
     private static final ImmutableSet<String> FIELDS = ImmutableSet.of(ID_FIELD);
 
     private final List<SubIndex> indexes;
     private final Query query;
+    private final int limit;
 
-    public QuerySource(List<SubIndex> indexes, Query query) {
+    public QuerySource(List<SubIndex> indexes, Query query, int limit) {
       this.indexes = indexes;
       this.query = query;
+      this.limit = limit;
     }
 
     @Override
@@ -295,9 +295,9 @@ public class LuceneChangeIndex implements ChangeIndex {
         TopDocs[] hits = new TopDocs[indexes.size()];
         for (int i = 0; i < indexes.size(); i++) {
           searchers[i] = indexes.get(i).acquire();
-          hits[i] = searchers[i].search(query, LIMIT, sort);
+          hits[i] = searchers[i].search(query, limit, sort);
         }
-        TopDocs docs = TopDocs.merge(sort, LIMIT, hits);
+        TopDocs docs = TopDocs.merge(sort, limit, hits);
 
         List<ChangeData> result =
             Lists.newArrayListWithCapacity(docs.scoreDocs.length);

gerrit-server/src/main/java/com/google/gerrit/server/index/ChangeIndex.java

@@ -62,7 +62,7 @@ public interface ChangeIndex {
     }
 
     @Override
-    public ChangeDataSource getSource(Predicate<ChangeData> p) {
+    public ChangeDataSource getSource(Predicate<ChangeData> p, int limit) {
       throw new UnsupportedOperationException();
     }
 
@@ -133,12 +133,13 @@ public interface ChangeIndex {
    * @param p the predicate to match. Must be a tree containing only AND, OR,
    *     or NOT predicates as internal nodes, and {@link IndexPredicate}s as
    *     leaves.
+   * @param limit maximum number of results to return.
    * @return a source of documents matching the predicate.
    *
    * @throws QueryParseException if the predicate could not be converted to an
    *     indexed data source.
    */
-  public ChangeDataSource getSource(Predicate<ChangeData> p)
+  public ChangeDataSource getSource(Predicate<ChangeData> p, int limit)
       throws QueryParseException;
 
   /**

gerrit-server/src/main/java/com/google/gerrit/server/index/IndexRewriteImpl.java

@@ -14,6 +14,7 @@
 
 package com.google.gerrit.server.index;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import com.google.gerrit.reviewdb.client.Change;
@@ -28,6 +29,7 @@ import com.google.gerrit.server.query.QueryRewriter;
 import com.google.gerrit.server.query.change.AndSource;
 import com.google.gerrit.server.query.change.BasicChangeRewrites;
 import com.google.gerrit.server.query.change.ChangeData;
+import com.google.gerrit.server.query.change.ChangeQueryBuilder;
 import com.google.gerrit.server.query.change.ChangeQueryRewriter;
 import com.google.gerrit.server.query.change.ChangeStatusPredicate;
 import com.google.gerrit.server.query.change.OrSource;
@@ -62,6 +64,9 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
     CLOSED_STATUSES = Sets.immutableEnumSet(closed);
   }
 
+  @VisibleForTesting
+  static final int MAX_LIMIT = 1000;
+
   /**
    * Get the set of statuses that changes matching the given predicate may have.
    *
@@ -134,11 +139,15 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
   @Override
   public Predicate<ChangeData> rewrite(Predicate<ChangeData> in) {
     in = basicRewrites.rewrite(in);
+    // Add 1 to specified limit to match behavior of QueryProcessor.
+    int limit = ChangeQueryBuilder.hasLimit(in)
+        ? ChangeQueryBuilder.getLimit(in) + 1
+        : MAX_LIMIT;
 
     ChangeIndex index = indexes.getSearchIndex();
-    Predicate<ChangeData> out = rewriteImpl(in, index);
+    Predicate<ChangeData> out = rewriteImpl(in, index, limit);
     if (in == out || out instanceof IndexPredicate) {
-      return query(out, index);
+      return query(out, index, limit);
     } else if (out == null /* cannot rewrite */) {
       return in;
     } else {
@@ -151,6 +160,7 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
    *
    * @param in predicate to rewrite.
    * @param index index whose schema determines which fields are indexed.
+   * @param limit maximum number of results to return.
    * @return {@code null} if no part of this subtree can be queried in the
    *     index directly. {@code in} if this subtree and all its children can be
    *     queried directly in the index. Otherwise, a predicate that is
@@ -158,7 +168,7 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
    *     index directly.
    */
   private Predicate<ChangeData> rewriteImpl(Predicate<ChangeData> in,
-      ChangeIndex index) {
+      ChangeIndex index, int limit) {
     if (isIndexPredicate(in, index)) {
       return in;
     } else if (!isRewritePossible(in)) {
@@ -172,7 +182,7 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
     List<Predicate<ChangeData>> newChildren = Lists.newArrayListWithCapacity(n);
     for (int i = 0; i < n; i++) {
       Predicate<ChangeData> c = in.getChild(i);
-      Predicate<ChangeData> nc = rewriteImpl(c, index);
+      Predicate<ChangeData> nc = rewriteImpl(c, index, limit);
       if (nc == c) {
         isIndexed.set(i);
         newChildren.add(c);
@@ -192,7 +202,7 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
     } else if (rewritten.cardinality() == n) {
       return in.copy(newChildren); // All children were rewritten.
     }
-    return partitionChildren(in, newChildren, isIndexed, index);
+    return partitionChildren(in, newChildren, isIndexed, index, limit);
   }
 
   private boolean isIndexPredicate(Predicate<ChangeData> in, ChangeIndex index) {
@@ -207,10 +217,11 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
       Predicate<ChangeData> in,
       List<Predicate<ChangeData>> newChildren,
       BitSet isIndexed,
-      ChangeIndex index) {
+      ChangeIndex index,
+      int limit) {
     if (isIndexed.cardinality() == 1) {
       int i = isIndexed.nextSetBit(0);
-      newChildren.add(0, query(newChildren.remove(i), index));
+      newChildren.add(0, query(newChildren.remove(i), index, limit));
       return copy(in, newChildren);
     }
 
@@ -230,7 +241,7 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
         all.add(c);
       }
     }
-    all.add(0, query(in.copy(indexed), index));
+    all.add(0, query(in.copy(indexed), index, limit));
     return copy(in, all);
   }
 
@@ -245,9 +256,10 @@ public class IndexRewriteImpl implements ChangeQueryRewriter {
     return in.copy(all);
   }
 
-  private IndexedChangeQuery query(Predicate<ChangeData> p, ChangeIndex index) {
+  private IndexedChangeQuery query(Predicate<ChangeData> p, ChangeIndex index,
+      int limit) {
     try {
-      return new IndexedChangeQuery(index, p);
+      return new IndexedChangeQuery(index, p, limit);
     } catch (QueryParseException e) {
       throw new IllegalStateException(
           "Failed to convert " + p + " to index predicate", e);

gerrit-server/src/main/java/com/google/gerrit/server/index/IndexedChangeQuery.java

@@ -15,6 +15,8 @@
 package com.google.gerrit.server.index;
 
 import com.google.common.base.Function;
+
+import com.google.common.base.Objects;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.gerrit.server.query.Predicate;
@@ -39,12 +41,14 @@ import java.util.List;
 public class IndexedChangeQuery extends Predicate<ChangeData>
     implements ChangeDataSource {
   private final Predicate<ChangeData> pred;
+  private final int limit;
   private final ChangeDataSource source;
 
-  public IndexedChangeQuery(ChangeIndex index, Predicate<ChangeData> pred)
+  public IndexedChangeQuery(ChangeIndex index, Predicate<ChangeData> pred, int limit)
       throws QueryParseException {
     this.pred = pred;
-    this.source = index.getSource(pred);
+    this.limit = limit;
+    this.source = index.getSource(pred, limit);
   }
 
   @Override
@@ -132,13 +136,19 @@ public class IndexedChangeQuery extends Predicate<ChangeData>
 
   @Override
   public boolean equals(Object other) {
-    return other != null
+    if (other == null || getClass() != other.getClass()) {
-        && getClass() == other.getClass()
+      return false;
-        && pred.equals(((IndexedChangeQuery) other).pred);
+    }
+    IndexedChangeQuery o = (IndexedChangeQuery) other;
+    return pred.equals(o.pred)
+        && limit == o.limit;
   }
 
   @Override
   public String toString() {
-    return "index(" + source + ")";
+    return Objects.toStringHelper("index")
+        .add("p", source)
+        .add("limit", limit)
+        .toString();
   }
 }

gerrit-server/src/main/java/com/google/gerrit/server/query/change/CommentPredicate.java

@@ -38,7 +38,7 @@ class CommentPredicate extends IndexPredicate<ChangeData> {
   public boolean match(ChangeData object) throws OrmException {
     try {
       for (ChangeData cData : index.getSource(
-          Predicate.and(new LegacyChangeIdPredicate(db, object.getId()), this))
+          Predicate.and(new LegacyChangeIdPredicate(db, object.getId()), this), 1)
           .read()) {
         if (cData.getId().equals(object.getId())) {
           return true;

gerrit-server/src/main/java/com/google/gerrit/server/query/change/MessagePredicate.java

@@ -42,7 +42,7 @@ class MessagePredicate extends IndexPredicate<ChangeData> {
   public boolean match(ChangeData object) throws OrmException {
     try {
       for (ChangeData cData : index.getSource(
-          Predicate.and(new LegacyChangeIdPredicate(db, object.getId()), this))
+          Predicate.and(new LegacyChangeIdPredicate(db, object.getId()), this), 1)
           .read()) {
         if (cData.getId().equals(object.getId())) {
           return true;

gerrit-server/src/test/java/com/google/gerrit/server/index/IndexRewriteTest.java

@@ -79,7 +79,7 @@ public class IndexRewriteTest extends TestCase {
     }
 
     @Override
-    public ChangeDataSource getSource(Predicate<ChangeData> p)
+    public ChangeDataSource getSource(Predicate<ChangeData> p, int limit)
         throws QueryParseException {
       return new Source(p);
     }
@@ -269,6 +269,16 @@ public class IndexRewriteTest extends TestCase {
         out.getChildren());
   }
 
+  public void testLimit() throws Exception {
+    Predicate<ChangeData> in = parse("file:a limit:3");
+    Predicate<ChangeData> out = rewrite(in);
+    assertSame(AndSource.class, out.getClass());
+    assertEquals(ImmutableList.of(
+          query(in.getChild(0), 4),
+          in.getChild(1)),
+        out.getChildren());
+  }
+
   public void testGetPossibleStatus() throws Exception {
     assertEquals(EnumSet.allOf(Change.Status.class), status("file:a"));
     assertEquals(EnumSet.of(NEW), status("is:new"));
@@ -308,7 +318,12 @@ public class IndexRewriteTest extends TestCase {
 
   private IndexedChangeQuery query(Predicate<ChangeData> p)
       throws QueryParseException {
-    return new IndexedChangeQuery(index, p);
+    return query(p, IndexRewriteImpl.MAX_LIMIT);
+  }
+
+  private IndexedChangeQuery query(Predicate<ChangeData> p, int limit)
+      throws QueryParseException {
+    return new IndexedChangeQuery(index, p, limit);
   }
 
   private Set<Change.Status> status(String query) throws QueryParseException {

gerrit-solr/src/main/java/com/google/gerrit/solr/SolrChangeIndex.java

@@ -194,7 +194,7 @@ class SolrChangeIndex implements ChangeIndex, LifecycleListener {
   }
 
   @Override
-  public ChangeDataSource getSource(Predicate<ChangeData> p)
+  public ChangeDataSource getSource(Predicate<ChangeData> p, int limit)
       throws QueryParseException {
     Set<Change.Status> statuses = IndexRewriteImpl.getPossibleStatus(p);
     List<SolrServer> indexes = Lists.newArrayListWithCapacity(2);
@@ -204,7 +204,7 @@ class SolrChangeIndex implements ChangeIndex, LifecycleListener {
     if (!Sets.intersection(statuses, CLOSED_STATUSES).isEmpty()) {
       indexes.add(closedIndex);
     }
-    return new QuerySource(indexes, QueryBuilder.toQuery(p));
+    return new QuerySource(indexes, QueryBuilder.toQuery(p), limit);
   }
 
   private void commit(SolrServer server) throws IOException {
@@ -219,11 +219,12 @@ class SolrChangeIndex implements ChangeIndex, LifecycleListener {
     private final List<SolrServer> indexes;
     private final SolrQuery query;
 
-    public QuerySource(List<SolrServer> indexes, Query q) {
+    public QuerySource(List<SolrServer> indexes, Query q, int limit) {
       this.indexes = indexes;
 
       query = new SolrQuery(q.toString());
       query.setParam("shards.tolerant", true);
+      query.setParam("rows", Integer.toString(limit));
       query.setFields(ID_FIELD);
       query.setSort(
           ChangeField.UPDATED.getName(),