lucene-field-sort

import static org.junit.jupiter.api.Assertions.*;
import java.io.IOException;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.SortedDocValuesField;
import org.apache.lucene.document.StoredField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.FilterCollector;
import org.apache.lucene.search.FilterLeafCollector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.SortField.Type;
import org.apache.lucene.search.TopFieldCollector;
import org.apache.lucene.search.TotalHits.Relation;
import org.apache.lucene.store.ByteBuffersDirectory;
import org.apache.lucene.util.BytesRef;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

class TestFieldCollector {

  static final int NUM_DOCS = 10000;
  static ByteBuffersDirectory directory;
  static final long seed = 123;

  @BeforeAll
  static void setUpBeforeClass() throws Exception {
    directory = new ByteBuffersDirectory();
    Random rand = new Random(seed);
    try (IndexWriter indexWriter =
        new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer()))) {
      for (int i = 0; i < NUM_DOCS; i++) {
        Document doc = new Document();
        doc.add(new StoredField("id", String.valueOf(i)));
        doc.add(new SortedDocValuesField("sortedDocValues", new BytesRef("sdv" + rand.nextInt())));
        indexWriter.addDocument(doc);
        if (i % 999 == 0) {
          indexWriter.commit();
        }
      }
    }
  }

  @AfterAll
  static void tearDownAfterClass() throws Exception {}

  @BeforeEach
  void setUp() throws Exception {}

  @AfterEach
  void tearDown() throws Exception {}

  @Test
  void testFieldSort() throws IOException {
    SortField [] sf = { new SortField("sortedDocValues", Type.STRING) };
    Sort sort = new Sort(sf);
    AtomicInteger totalCollected = new AtomicInteger();
    TopDocsCollector topDocsCollector = TopFieldCollector.create(sort, 100, 100);
    FilterCollector collector = new FilterCollector(topDocsCollector) {

      @Override
      public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
        return new FilterLeafCollector(super.getLeafCollector(context)) {
          @Override
          public void collect(int doc) throws IOException {
            totalCollected.incrementAndGet();
            super.collect(doc);
          }

          @Override
          public DocIdSetIterator competitiveIterator() throws IOException {
            return in.competitiveIterator();
          }
        };
      }
    };
    IndexSearcher searcher = new IndexSearcher(DirectoryReader.open(directory));
    searcher.search(new MatchAllDocsQuery(), collector);
    TopDocs docs = topDocsCollector.topDocs();
    assertTrue(totalCollected.get() < NUM_DOCS, String.valueOf(totalCollected.get()));
    assertEquals(100, docs.scoreDocs.length);
    assertEquals(Relation.GREATER_THAN_OR_EQUAL_TO, docs.totalHits.relation);
    assertTrue(docs.totalHits.value < NUM_DOCS, String.valueOf(docs.totalHits.value));
  }
}