ToParentBlockJoinQuery_bug

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.core.WhitespaceTokenizer;
import org.apache.lucene.document.BinaryDocValuesField;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.BinaryDocValues;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.KeepOnlyLastCommitDeletionPolicy;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SnapshotDeletionPolicy;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorable;
import org.apache.lucene.search.SearcherManager;
import org.apache.lucene.search.TermInSetQuery;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.WildcardQuery;
import org.apache.lucene.search.join.BitSetProducer;
import org.apache.lucene.search.join.QueryBitSetProducer;
import org.apache.lucene.search.join.ScoreMode;
import org.apache.lucene.search.join.ToParentBlockJoinQuery;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.BytesRef;
import org.junit.Rule;
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;
import org.junit.rules.TemporaryFolder;


class ParentChildTest {

    @Rule
    public TemporaryFolder dir = new TemporaryFolder();

    SearcherManager manager = null;
    IndexSearcher searcher = null;


    @BeforeAll
    static void setUpBeforeClass() throws Exception {
    }

    @AfterAll
    static void tearDownAfterClass() throws Exception {
    }

    @BeforeEach
    void setUp() throws Exception {
        if (manager == null) {
            IndexWriterConfig config = new IndexWriterConfig(analyzer);
            config.setIndexDeletionPolicy(new SnapshotDeletionPolicy(new KeepOnlyLastCommitDeletionPolicy()));
            config.setCommitOnClose(true);
            dir.create();
            Path path = dir.newFolder().toPath();
            try (IndexWriter indexWriter = new IndexWriter(FSDirectory.open(path), config)) {
                for (int i = 0; i < 6; i++) {
                    List<org.apache.lucene.document.Document> docs = new ArrayList<>();
                    org.apache.lucene.document.Document parent = new org.apache.lucene.document.Document();
                    String id = String.format("id%05d", i);
                    parent.add(new StringField("id", id, Field.Store.YES));
                    parent.add(new BinaryDocValuesField("idBinary", new BytesRef(id)));
                    switch (i % 3) {
                    case 2: {
                        org.apache.lucene.document.Document doc = new org.apache.lucene.document.Document();
                        doc.add(new StringField("program", "P2", Field.Store.YES));
                        docs.add(doc);
                    }
                        //$FALL-THROUGH$
                    case 1: {
                        org.apache.lucene.document.Document doc = new org.apache.lucene.document.Document();
                        doc.add(new StringField("program", "P1", Field.Store.YES));
                        docs.add(doc);
                    }
                        //$FALL-THROUGH$
                    default:
                        break;
                    }
                    docs.add(parent);
                    indexWriter.updateDocuments(new Term("id", id), docs);
                }
                indexWriter.commit();
            }
            manager = new SearcherManager(FSDirectory.open(path), null);
        }
        searcher = manager.acquire();
    }

    @AfterEach
    void tearDown() throws Exception {
        if (searcher != null) {
            manager.release(searcher);
        }
    }

    Set<BytesRef> toSet(String ... val) {
        Set<BytesRef> filter = new TreeSet<BytesRef>();
        for (String v : val) {
            filter.add(new BytesRef(v));
        }
        return filter;
    }

    @Test
    void testParentChild() throws IOException {
        SimpleCollector collector = new SimpleCollector();
        BitSetProducer parentSet = new QueryBitSetProducer(new TermInSetQuery("id", toSet("id00000", "id00001", "id00002", "id00003", "id00004", "id00005")));
        Query q = new ToParentBlockJoinQuery(new TermInSetQuery("program", toSet("P1", "P2")), parentSet, ScoreMode.None);
        searcher.search(q, collector);
        assertEquals(4, collector.results.size());
        assertFalse(collector.results.contains("id00000"), collector.results.toString());
        assertFalse(collector.results.contains("id00003"), collector.results.toString());
    }

    @Test
    void testParentChildWildcard() throws IOException {
        SimpleCollector collector = new SimpleCollector();
        BitSetProducer parentSet = new QueryBitSetProducer(new WildcardQuery(new Term("id", "*")));
        Query q = new ToParentBlockJoinQuery(new WildcardQuery(new Term("program", "*")), parentSet, ScoreMode.None);
        searcher.search(q, collector);
        assertFalse(collector.results.contains("id00000"), collector.results.toString());
        assertFalse(collector.results.contains("id00003"), collector.results.toString());
        assertEquals(4, collector.results.size());
    }

    @Test
    void testChildWildcardMatch() throws IOException {
        SimpleCollector collector = new SimpleCollector();
        BitSetProducer parentSet = new QueryBitSetProducer(new TermInSetQuery("id", toSet("id00000", "id00001", "id00002", "id00003", "id00004", "id00005")));
        Query q = new ToParentBlockJoinQuery(new WildcardQuery(new Term("program", "*")), parentSet, ScoreMode.None);
        searcher.search(q, collector);
        assertEquals(4, collector.results.size());
        assertFalse(collector.results.contains("id00000"), collector.results.toString());
        assertFalse(collector.results.contains("id00003"), collector.results.toString());
    }


    @Test
    void testChildWildcardNoMatch() throws IOException {
        SimpleCollector collector = new SimpleCollector();
        BitSetProducer parentSet = new QueryBitSetProducer(new TermInSetQuery("id", toSet("id00000", "id00001", "id00003")));
        Query q = new ToParentBlockJoinQuery(new WildcardQuery(new Term("program", "*")), parentSet, org.apache.lucene.search.join.ScoreMode.None);
        searcher.search(q, collector);
        //assertEquals(1, collector.results.size());
        assertFalse(collector.results.contains("id00000"), collector.results.toString());
        assertTrue(collector.results.contains("id00001"), collector.results.toString());
        assertFalse(collector.results.contains("id00003"), collector.results.toString());
    }

    @Test
    void testChildWildcardNoMatch2() throws IOException {
        SimpleCollector collector = new SimpleCollector();
        BitSetProducer parentSet = new QueryBitSetProducer(new TermQuery(new Term("id", "id00003")));
        Query q = new ToParentBlockJoinQuery(new WildcardQuery(new Term("program", "*")), parentSet, org.apache.lucene.search.join.ScoreMode.None);
        searcher.search(q, collector);
        //assertEquals(1, collector.results.size());
        assertFalse(collector.results.contains("id00003"), collector.results.toString());
    }

    @Test
    void testChildWildcardMatch2() throws IOException {
        SimpleCollector collector = new SimpleCollector();
        BitSetProducer parentSet = new QueryBitSetProducer(new TermInSetQuery("id", toSet( "id00003", "id00002")));
        Query q = new ToParentBlockJoinQuery(new WildcardQuery(new Term("program", "*")), parentSet, org.apache.lucene.search.join.ScoreMode.None);
        searcher.search(q, collector);
        assertEquals(1, collector.results.size());
        assertFalse(collector.results.contains("id00003"), collector.results.toString());
        assertTrue(collector.results.contains("id00002"), collector.results.toString());
    }

    @Test
    void testChildQueryOnly() throws IOException {
        //assertEquals(4, searcher.count(new WildcardQuery(new Term("program", "*"))));
        //assertEquals(4, searcher.count(new TermRangeQuery("program", new BytesRef("*"), new BytesRef("*"), true, true)));
        assertEquals(4, searcher.count(new TermQuery(new Term("program", "P1"))));
        assertEquals(2, searcher.count(new TermQuery(new Term("program", "P2"))));
    }

    static class SimpleCollector implements Collector, LeafCollector {
        BinaryDocValues ids = null;
        Set<String> results = new TreeSet<>();


        @Override
        public void collect(int arg0) throws IOException {
            this.ids.advanceExact(arg0);
            results.add(this.ids.binaryValue().utf8ToString());
        }

        @Override
        public void setScorer(Scorable arg0) throws IOException {
        }

        @Override
        public LeafCollector getLeafCollector(LeafReaderContext leafContext) throws IOException {
            this.ids = leafContext.reader().getBinaryDocValues("idBinary");
            return this;
        }

        @Override
        public org.apache.lucene.search.ScoreMode scoreMode() {
            return org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES;
        }

    }

    private static Analyzer analyzer = new Analyzer() {

        class StandardFilter extends TokenFilter {

            /** Sole constructor */
            public StandardFilter(TokenStream in) {
                super(in);
            }

            @Override
            public final boolean incrementToken() throws IOException {
                return input.incrementToken(); // TODO: add some niceties for the new grammar
            }
        }

        @Override
        protected TokenStreamComponents createComponents(final String fieldName) {
            final WhitespaceTokenizer src = new WhitespaceTokenizer();
            return new TokenStreamComponents(src, new StandardFilter(src));
        }
    };
}