Index: CHANGES.txt
===================================================================
--- CHANGES.txt	(revision 891367)
+++ CHANGES.txt	(working copy)
@@ -27,6 +27,13 @@
  * LUCENE-2158: At high indexing rates, NRT reader could temporarily
    lose deletions.  (Mike McCandless)
   
+API Changes
+
+  * LUCENE-2190: Added setNextReader method to CustomScoreQuery, which
+    is necessary with per-segment searching to notify the subclass
+    which reader the int doc, passed to customScore, refers to.  (Paul
+    chez Jamespot via Mike McCandless)
+
 Optimizations
 
  * LUCENE-2086: When resolving deleted terms, do so in term sort order
Index: src/test/org/apache/lucene/search/function/TestOrdValues.java
===================================================================
--- src/test/org/apache/lucene/search/function/TestOrdValues.java	(revision 891366)
+++ src/test/org/apache/lucene/search/function/TestOrdValues.java	(working copy)
@@ -40,7 +40,7 @@
 
   /* @override constructor */
   public TestOrdValues(String name) {
-    super(name);
+    super(name, false);
   }
 
   /** Test OrdFieldSource */
Index: src/test/org/apache/lucene/search/function/TestFieldScoreQuery.java
===================================================================
--- src/test/org/apache/lucene/search/function/TestFieldScoreQuery.java	(revision 891366)
+++ src/test/org/apache/lucene/search/function/TestFieldScoreQuery.java	(working copy)
@@ -42,7 +42,7 @@
 
   /* @override constructor */
   public TestFieldScoreQuery(String name) {
-    super(name);
+    super(name, true);
   }
 
   /** Test that FieldScoreQuery of Type.BYTE returns docs in expected order. */
@@ -164,7 +164,7 @@
     expectedArrayTypes.put(FieldScoreQuery.Type.FLOAT, new float[0]);
     
     IndexSearcher s = new IndexSearcher(dir);
-    Object innerArray = null;
+    Object[] innerArray = new Object[s.getIndexReader().getSequentialSubReaders().length];
 
     boolean warned = false; // print warning once.
     for (int i=0; i<10; i++) {
@@ -176,16 +176,16 @@
         IndexReader reader = readers[j];
         try {
           if (i == 0) {
-            innerArray = q.valSrc.getValues(reader).getInnerArray();
-            log(i + ".  compare: " + innerArray.getClass() + " to "
+            innerArray[j] = q.valSrc.getValues(reader).getInnerArray();
+            log(i + ".  compare: " + innerArray[j].getClass() + " to "
                 + expectedArrayTypes.get(tp).getClass());
             assertEquals(
                 "field values should be cached in the correct array type!",
-                innerArray.getClass(), expectedArrayTypes.get(tp).getClass());
+                innerArray[j].getClass(), expectedArrayTypes.get(tp).getClass());
           } else {
-            log(i + ".  compare: " + innerArray + " to "
+            log(i + ".  compare: " + innerArray[j] + " to "
                 + q.valSrc.getValues(reader).getInnerArray());
-            assertSame("field values should be cached and reused!", innerArray,
+            assertSame("field values should be cached and reused!", innerArray[j],
                 q.valSrc.getValues(reader).getInnerArray());
           }
         } catch (UnsupportedOperationException e) {
Index: src/test/org/apache/lucene/search/function/FunctionTestSetup.java
===================================================================
--- src/test/org/apache/lucene/search/function/FunctionTestSetup.java	(revision 891366)
+++ src/test/org/apache/lucene/search/function/FunctionTestSetup.java	(working copy)
@@ -70,7 +70,14 @@
   
   /* @override constructor */
   public FunctionTestSetup(String name) {
+    this(name, false);
+  }
+
+  private final boolean doMultiSegment;
+
+  public FunctionTestSetup(String name, boolean doMultiSegment) {
     super(name);
+    this.doMultiSegment = doMultiSegment;
   }
 
   /* @override */
@@ -100,6 +107,9 @@
       addDoc(iw,i);
       done[i] = true;
       i = (i+4)%N_DOCS;
+      if (doMultiSegment && remaining % 3 == 0) {
+        iw.commit();
+      }
       remaining --;
     }
     iw.close();
Index: src/test/org/apache/lucene/search/function/TestCustomScoreQuery.java
===================================================================
--- src/test/org/apache/lucene/search/function/TestCustomScoreQuery.java	(revision 891366)
+++ src/test/org/apache/lucene/search/function/TestCustomScoreQuery.java	(working copy)
@@ -28,6 +28,8 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryUtils;
 import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.FieldCache;
+import org.apache.lucene.index.IndexReader;
 
 /**
  * Test CustomScoreQuery search.
@@ -36,7 +38,7 @@
 
   /* @override constructor */
   public TestCustomScoreQuery(String name) {
-    super(name);
+    super(name, true);
   }
 
   /** Test that CustomScoreQuery of Type.BYTE returns the expected scores. */
@@ -64,7 +66,7 @@
     // INT field can be parsed as float
     doTestCustomScore(INT_FIELD,FieldScoreQuery.Type.FLOAT,1.0);
     doTestCustomScore(INT_FIELD,FieldScoreQuery.Type.FLOAT,5.0);
-    // same values, but in flot format
+    // same values, but in float format
     doTestCustomScore(FLOAT_FIELD,FieldScoreQuery.Type.FLOAT,1.0);
     doTestCustomScore(FLOAT_FIELD,FieldScoreQuery.Type.FLOAT,6.0);
   }
@@ -112,6 +114,8 @@
       }
       if (valSrcScores.length == 1) {
         return subQueryScore + valSrcScores[0];
+        // confirm that skipping beyond the last doc, on the
+        // previous reader, hits NO_MORE_DOCS
       }
       return (subQueryScore + valSrcScores[0]) * valSrcScores[1]; // we know there are two
     } 
@@ -133,6 +137,44 @@
       return exp2;      
     } 
   }
+
+  private final class CustomExternalQuery extends CustomScoreQuery {
+    private IndexReader reader;
+    private int[] values;
+
+    public float customScore(int doc, float subScore, float valSrcScore) {
+      assertTrue(doc <= reader.maxDoc());
+      return (float) values[doc];
+    }
+
+    public void setNextReader(IndexReader r) throws IOException {
+      reader = r;
+      values = FieldCache.DEFAULT.getInts(r, INT_FIELD);
+    }
+
+    public CustomExternalQuery(Query q) {
+      super(q);
+    }
+  }
+
+  public void testCustomExternalQuery() throws Exception {
+    QueryParser qp = new QueryParser(TEXT_FIELD,anlzr); 
+    String qtxt = "first aid text"; // from the doc texts in FunctionQuerySetup.
+    Query q1 = qp.parse(qtxt); 
+    
+    final Query q = new CustomExternalQuery(q1);
+    log(q);
+
+    IndexSearcher s = new IndexSearcher(dir);
+    TopDocs hits = s.search(q, 1000);
+    assertEquals(N_DOCS, hits.totalHits);
+    for(int i=0;i<N_DOCS;i++) {
+      final int doc = hits.scoreDocs[i].doc;
+      final float score = hits.scoreDocs[i].score;
+      assertEquals("doc=" + doc, (float) 1+(4*doc) % N_DOCS, score, 0.0001);
+    }
+    s.close();
+  }
   
   // Test that FieldScoreQuery returns docs with expected score.
   private void doTestCustomScore (String field, FieldScoreQuery.Type tp, double dboost) throws CorruptIndexException, Exception {
Index: src/test/org/apache/lucene/search/QueryUtils.java
===================================================================
--- src/test/org/apache/lucene/search/QueryUtils.java	(revision 891366)
+++ src/test/org/apache/lucene/search/QueryUtils.java	(working copy)
@@ -5,8 +5,6 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
-import java.util.ArrayList;
-import java.util.List;
 
 import junit.framework.Assert;
 
@@ -17,7 +15,6 @@
 import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.IndexWriter.MaxFieldLength;
 import org.apache.lucene.store.RAMDirectory;
-import org.apache.lucene.util.ReaderUtil;
 
 /**
  * Copyright 2005 Apache Software Foundation
@@ -357,6 +354,8 @@
     //System.out.println("checkFirstSkipTo: "+q);
     final float maxDiff = 1e-5f;
     final int lastDoc[] = {-1};
+    final IndexReader lastReader[] = {null};
+
     s.search(q,new Collector() {
       private Scorer scorer;
       private IndexReader reader;
@@ -382,30 +381,37 @@
           throw new RuntimeException(e);
         }
       }
-      public void setNextReader(IndexReader reader, int docBase) {
-        this.reader = reader;
+      public void setNextReader(IndexReader reader, int docBase) throws IOException {
+        // confirm that skipping beyond the last doc, on the
+        // previous reader, hits NO_MORE_DOCS
+        if (lastReader[0] != null) {
+          final IndexReader previousReader = lastReader[0];
+          Weight w = q.weight(new IndexSearcher(previousReader));
+          Scorer scorer = w.scorer(previousReader, true, false);
+          if (scorer != null) {
+            boolean more = scorer.advance(lastDoc[0] + 1) != DocIdSetIterator.NO_MORE_DOCS;
+            Assert.assertFalse("query's last doc was "+ lastDoc[0] +" but skipTo("+(lastDoc[0]+1)+") got to "+scorer.docID(),more);
+          }
+        }
+
+        this.reader = lastReader[0] = reader;
         lastDoc[0] = -1;
       }
       public boolean acceptsDocsOutOfOrder() {
         return false;
       }
     });
-    
-    List readerList = new ArrayList();
-    ReaderUtil.gatherSubReaders(readerList, s.getIndexReader());
-    IndexReader[] readers = (IndexReader[]) readerList.toArray(new IndexReader[0]);
-    for(int i = 0; i < readers.length; i++) {
-      IndexReader reader = readers[i];
-      Weight w = q.weight(s);
-      Scorer scorer = w.scorer(reader, true, false);
-      
+
+    if (lastReader[0] != null) {
+      // confirm that skipping beyond the last doc, on the
+      // previous reader, hits NO_MORE_DOCS
+      final IndexReader previousReader = lastReader[0];
+      Weight w = q.weight(new IndexSearcher(previousReader));
+      Scorer scorer = w.scorer(previousReader, true, false);
       if (scorer != null) {
         boolean more = scorer.advance(lastDoc[0] + 1) != DocIdSetIterator.NO_MORE_DOCS;
-  
-        if (more && lastDoc[0] != -1) 
-          Assert.assertFalse("query's last doc was "+ lastDoc[0] +" but skipTo("+(lastDoc[0]+1)+") got to "+scorer.docID(),more);
+        Assert.assertFalse("query's last doc was "+ lastDoc[0] +" but skipTo("+(lastDoc[0]+1)+") got to "+scorer.docID(),more);
       }
     }
-
   }
 }
Index: src/java/org/apache/lucene/search/function/CustomScoreQuery.java
===================================================================
--- src/java/org/apache/lucene/search/function/CustomScoreQuery.java	(revision 891366)
+++ src/java/org/apache/lucene/search/function/CustomScoreQuery.java	(working copy)
@@ -207,6 +207,9 @@
    * <pre>
    *     ModifiedScore = subQueryScore * valSrcScore
    * </pre>
+   *
+   * <p><b>NOTE</b>: The doc is relative to the current
+   * reader, last passed to {@link #setNextReader}.
    * 
    * @param doc id of scored doc. 
    * @param subQueryScore score of that doc by the subQuery.
@@ -218,6 +221,15 @@
   }
 
   /**
+   * Called when the scoring switches to another reader.
+   * 
+   * @param reader
+   *          next IndexReader
+   */
+  public void setNextReader(IndexReader reader) throws IOException {
+  }
+
+  /**
    * Explain the custom score.
    * Whenever overriding {@link #customScore(int, float, float[])}, 
    * this method should also be overridden to provide the correct explanation
@@ -399,6 +411,7 @@
       this.valSrcScorers = valSrcScorers;
       this.reader = reader;
       this.vScores = new float[valSrcScorers.length];
+      setNextReader(reader);
     }
 
     /** @deprecated use {@link #nextDoc()} instead. */