From 9979ddd1e0ffd569e3d533998e478166718ded24 Mon Sep 17 00:00:00 2001
From: Yonik Seeley <yonik@apache.org>
Date: Thu, 14 May 2009 18:37:25 +0000
Subject: [PATCH] SOLR-1111: make distributed search use leaf readers to access
 sort field values

git-svn-id: https://svn.apache.org/repos/asf/lucene/solr/trunk@774873 13f79535-47bb-0310-9956-ffa450edef68
---
 .../handler/component/QueryComponent.java     | 85 ++++++++++++------
 .../apache/solr/search/SolrIndexReader.java   | 87 ++++++++++++++++++-
 2 files changed, 143 insertions(+), 29 deletions(-)

diff --git a/src/java/org/apache/solr/handler/component/QueryComponent.java b/src/java/org/apache/solr/handler/component/QueryComponent.java
index 9d160ad48ca..738f1dd5af8 100644
--- a/src/java/org/apache/solr/handler/component/QueryComponent.java
+++ b/src/java/org/apache/solr/handler/component/QueryComponent.java
@@ -179,6 +179,10 @@ public class QueryComponent extends SearchComponent
     rsp.add("response",rb.getResults().docList);
     rsp.getToLog().add("hits", rb.getResults().docList.matches());
 
+
+    // The query cache doesn't currently store sort field values, and SolrIndexSearcher doesn't
+    // currently have an option to return sort field values.  Because of this, we
+    // take the documents given and re-derive the sort values.
     boolean fsv = req.getParams().getBool(ResponseBuilder.FIELD_SORT_VALUES,false);
     if(fsv){
       Sort sort = rb.getSortSpec().getSort();
@@ -187,46 +191,46 @@ public class QueryComponent extends SearchComponent
       NamedList sortVals = new NamedList(); // order is important for the sort fields
       Field field = new Field("dummy", "", Field.Store.YES, Field.Index.NO); // a dummy Field
 
+      SolrIndexReader reader = searcher.getReader();
+      SolrIndexReader[] readers = reader.getLeafReaders();
+      if (readers.length==1) readers=null;
+      int[] offsets = reader.getLeafOffsets();
+
       for (SortField sortField: sortFields) {
         int type = sortField.getType();
         if (type==SortField.SCORE || type==SortField.DOC) continue;
 
         ScoreDocComparator comparator = null;
-        IndexReader reader = searcher.getReader();
+        ScoreDocComparator comparators[] = (readers==null) ? null : new ScoreDocComparator[readers.length];
+
         String fieldname = sortField.getField();
         FieldType ft = fieldname==null ? null : req.getSchema().getFieldTypeNoEx(fieldname);
 
-
-        switch (type) {
-          case SortField.INT:
-            comparator = comparatorInt (reader, fieldname, sortField.getParser());
-            break;
-          case SortField.FLOAT:
-            comparator = comparatorFloat (reader, fieldname, sortField.getParser());
-            break;
-          case SortField.LONG:
-            comparator = comparatorLong(reader, fieldname, sortField.getParser());
-            break;
-          case SortField.DOUBLE:
-            comparator = comparatorDouble(reader, fieldname, sortField.getParser());
-            break;
-          case SortField.STRING:
-            if (sortField.getLocale() != null) comparator = comparatorStringLocale (reader, fieldname, sortField.getLocale());
-            else comparator = comparatorString (reader, fieldname);
-            break;
-          case SortField.CUSTOM:
-            comparator = sortField.getFactory().newComparator (reader, fieldname);
-            break;
-          default:
-            throw new RuntimeException ("unknown field type: "+type);
-        }
-
         DocList docList = rb.getResults().docList;
         ArrayList<Object> vals = new ArrayList<Object>(docList.size());
         DocIterator it = rb.getResults().docList.iterator();
+        SolrIndexReader subReader = reader;
+        int offset = 0;
+        int idx = 0;
+
         while(it.hasNext()) {
           sd.doc = it.nextDoc();
+          if (readers != null) {
+            idx = SolrIndexReader.readerIndex(sd.doc, offsets);
+            subReader = readers[idx];
+            offset = offsets[idx];
+            comparator = comparators[idx];
+          }
+
+          if (comparator == null) {
+            comparator = getComparator(subReader, sortField);
+            if (comparators != null)
+              comparators[idx] = comparator;
+          }
+
+          sd.doc -= offset;  // adjust for what segment this is in
           Object val = comparator.sortValue(sd);
+          
           // Sortable float, double, int, long types all just use a string
           // comparator. For these, we need to put the type into a readable
           // format.  One reason for this is that XML can't represent all
@@ -253,6 +257,35 @@ public class QueryComponent extends SearchComponent
     }
   }
 
+  private ScoreDocComparator getComparator(SolrIndexReader reader, SortField sortField) throws IOException {
+    ScoreDocComparator comparator = null;
+    String fieldname = sortField.getField();
+    switch (sortField.getType()) {
+      case SortField.INT:
+        comparator = comparatorInt (reader, fieldname, sortField.getParser());
+        break;
+      case SortField.FLOAT:
+        comparator = comparatorFloat (reader, fieldname, sortField.getParser());
+        break;
+      case SortField.LONG:
+        comparator = comparatorLong(reader, fieldname, sortField.getParser());
+        break;
+      case SortField.DOUBLE:
+        comparator = comparatorDouble(reader, fieldname, sortField.getParser());
+        break;
+      case SortField.STRING:
+        if (sortField.getLocale() != null) comparator = comparatorStringLocale (reader, fieldname, sortField.getLocale());
+        else comparator = comparatorString (reader, fieldname);
+        break;
+      case SortField.CUSTOM:
+        comparator = sortField.getFactory().newComparator (reader, fieldname);
+        break;
+      default:
+        throw new RuntimeException ("unknown field type: "+sortField.getType());
+    }
+    return comparator;
+  }
+
   @Override  
   public int distributedProcess(ResponseBuilder rb) throws IOException {
     if (rb.stage < ResponseBuilder.STAGE_PARSE_QUERY)
diff --git a/src/java/org/apache/solr/search/SolrIndexReader.java b/src/java/org/apache/solr/search/SolrIndexReader.java
index 6134ef2b6ec..5ec64ee16e6 100755
--- a/src/java/org/apache/solr/search/SolrIndexReader.java
+++ b/src/java/org/apache/solr/search/SolrIndexReader.java
@@ -33,9 +33,13 @@ import java.util.HashMap;
  */
 public class SolrIndexReader extends FilterIndexReader {
   private final SolrIndexReader[] subReaders;
+  private final SolrIndexReader[] leafReaders;
+  private int[] leafOffsets;
   private final SolrIndexReader parent;
   private final int base; // docid offset of this reader within parent
 
+  private static int[] zeroIntArray = new int[]{0};
+
   // top level searcher for this reader tree
   // a bit if a hack currently... searcher needs to set
   SolrIndexSearcher searcher;
@@ -54,14 +58,91 @@ public class SolrIndexReader extends FilterIndexReader {
     this.parent = parent;
     this.base = base;
     IndexReader subs[] = in.getSequentialSubReaders();
-    subReaders = subs == null ? null : new SolrIndexReader[subs.length];
     if (subs != null) {
+      subReaders = new SolrIndexReader[subs.length]; 
+      int numLeaves = subs.length;
+      leafOffsets = new int[numLeaves];
       int b=0;
       for (int i=0; i<subReaders.length; i++) {
-        subReaders[i] = new SolrIndexReader(subs[i], this, b);
-        b += subReaders[i].maxDoc();
+        SolrIndexReader sir = subReaders[i] = new SolrIndexReader(subs[i], this, b);
+        leafOffsets[i] = b;
+        b += sir.maxDoc();
+        IndexReader subLeaves[] = sir.leafReaders;
+        numLeaves += subLeaves.length - 1;  // subtract 1 for the parent
+      }
+      leafReaders = getLeaves(numLeaves);
+    } else {
+      subReaders = null;
+      leafReaders = new SolrIndexReader[]{this};
+      leafOffsets = zeroIntArray;
+    }
+
+  }
+
+  private SolrIndexReader[] getLeaves(int numLeaves) {
+    // fast path for a normal multiReader
+    if (subReaders==null || numLeaves == subReaders.length) return subReaders;
+
+    SolrIndexReader[] leaves = new SolrIndexReader[numLeaves];
+    leafOffsets = new int[numLeaves];
+    
+    int i=0;
+    int b = 0;
+    for (SolrIndexReader sir : subReaders) {
+      SolrIndexReader subLeaves[] = sir.leafReaders;
+      if (subLeaves == null) {
+        leafOffsets[i] = b;
+        b += sir.maxDoc();
+        leaves[i++] = sir;
+      } else {
+        for (SolrIndexReader subLeaf : subLeaves) {
+          leafOffsets[i] = b;
+          b += subLeaf.maxDoc();
+          leaves[i++] = subLeaf;
+        }
       }
     }
+    assert(i == numLeaves && b == maxDoc());
+    return leaves;
+  }
+
+  /** return the leaf readers in this reader tree, or an array of size 1 containing "this" if "this" is a leaf */
+  public SolrIndexReader[] getLeafReaders() {
+    return leafReaders;
+  }
+
+  /** Return the doc id offsets for each leaf reader.  This will be different than getBase() for
+   * any leaf reader who is not a direct descendant of "this".
+   */
+  public int[] getLeafOffsets() {
+    return leafOffsets;
+  }
+
+  /** Given an array of IndexReader offsets, find which contains the given doc */
+  public static int readerIndex(int doc, int[] offsets) {    // find reader for doc doc:
+    int high = offsets.length - 1;
+
+    // fast-path for a big optimized index and a bunch of smaller ones.
+    if (high <= 0 || doc < offsets[1]) return 0;
+
+    int low = 1;
+
+    while (high >= low) {
+      int mid = (low + high) >>> 1;
+      int offset = offsets[mid];
+      // check low first since first segments are normally bigger.
+      if (doc < offset)
+        high = mid - 1;
+      else if (doc > offset) {
+        low = mid + 1;
+      }
+      else {
+        // exact match on the offset.
+        return mid;
+      }
+    }
+    // low is the insertion point, high should be just below that (and the segment we want).
+    return high;
   }
 
   static String shortName(Object o) {