Examples with DocIdSetBuilder org.apache.lucene.util.DocIdSetBuilder used on opensource projects

Example 11 with DocIdSetBuilder

use of org.apache.lucene.util.DocIdSetBuilder in project lucene-solr by apache.

the class LatLonPointInPolygonQuery method createWeight.

@Override
public Weight createWeight(IndexSearcher searcher, boolean needsScores, float boost) throws IOException {
    // I don't use RandomAccessWeight here: it's no good to approximate with "match all docs"; this is an inverted structure and should be
    // used in the first pass:
    // bounding box over all polygons, this can speed up tree intersection/cheaply improve approximation for complex multi-polygons
    // these are pre-encoded with LatLonPoint's encoding
    final Rectangle box = Rectangle.fromPolygon(polygons);
    final byte[] minLat = new byte[Integer.BYTES];
    final byte[] maxLat = new byte[Integer.BYTES];
    final byte[] minLon = new byte[Integer.BYTES];
    final byte[] maxLon = new byte[Integer.BYTES];
    NumericUtils.intToSortableBytes(encodeLatitude(box.minLat), minLat, 0);
    NumericUtils.intToSortableBytes(encodeLatitude(box.maxLat), maxLat, 0);
    NumericUtils.intToSortableBytes(encodeLongitude(box.minLon), minLon, 0);
    NumericUtils.intToSortableBytes(encodeLongitude(box.maxLon), maxLon, 0);
    final Polygon2D tree = Polygon2D.create(polygons);
    final GeoEncodingUtils.PolygonPredicate polygonPredicate = GeoEncodingUtils.createPolygonPredicate(polygons, tree);
    return new ConstantScoreWeight(this, boost) {

        @Override
        public Scorer scorer(LeafReaderContext context) throws IOException {
            LeafReader reader = context.reader();
            PointValues values = reader.getPointValues(field);
            if (values == null) {
                // No docs in this segment had any points fields
                return null;
            }
            FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(field);
            if (fieldInfo == null) {
                // No docs in this segment indexed this field at all
                return null;
            }
            LatLonPoint.checkCompatible(fieldInfo);
            // matching docids
            DocIdSetBuilder result = new DocIdSetBuilder(reader.maxDoc(), values, field);
            values.intersect(new IntersectVisitor() {

                DocIdSetBuilder.BulkAdder adder;

                @Override
                public void grow(int count) {
                    adder = result.grow(count);
                }

                @Override
                public void visit(int docID) {
                    adder.add(docID);
                }

                @Override
                public void visit(int docID, byte[] packedValue) {
                    if (polygonPredicate.test(NumericUtils.sortableBytesToInt(packedValue, 0), NumericUtils.sortableBytesToInt(packedValue, Integer.BYTES))) {
                        adder.add(docID);
                    }
                }

                @Override
                public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
                    if (StringHelper.compare(Integer.BYTES, minPackedValue, 0, maxLat, 0) > 0 || StringHelper.compare(Integer.BYTES, maxPackedValue, 0, minLat, 0) < 0 || StringHelper.compare(Integer.BYTES, minPackedValue, Integer.BYTES, maxLon, 0) > 0 || StringHelper.compare(Integer.BYTES, maxPackedValue, Integer.BYTES, minLon, 0) < 0) {
                        // outside of global bounding box range
                        return Relation.CELL_OUTSIDE_QUERY;
                    }
                    double cellMinLat = decodeLatitude(minPackedValue, 0);
                    double cellMinLon = decodeLongitude(minPackedValue, Integer.BYTES);
                    double cellMaxLat = decodeLatitude(maxPackedValue, 0);
                    double cellMaxLon = decodeLongitude(maxPackedValue, Integer.BYTES);
                    return tree.relate(cellMinLat, cellMaxLat, cellMinLon, cellMaxLon);
                }
            });
            return new ConstantScoreScorer(this, score(), result.build().iterator());
        }
    };
}

Example 12 with DocIdSetBuilder

use of org.apache.lucene.util.DocIdSetBuilder in project lucene-solr by apache.

the class TestReqExclBulkScorer method doTestRandom.

public void doTestRandom() throws IOException {
    final int maxDoc = TestUtil.nextInt(random(), 1, 1000);
    DocIdSetBuilder reqBuilder = new DocIdSetBuilder(maxDoc);
    DocIdSetBuilder exclBuilder = new DocIdSetBuilder(maxDoc);
    final int numIncludedDocs = TestUtil.nextInt(random(), 1, maxDoc);
    final int numExcludedDocs = TestUtil.nextInt(random(), 1, maxDoc);
    DocIdSetBuilder.BulkAdder reqAdder = reqBuilder.grow(numIncludedDocs);
    for (int i = 0; i < numIncludedDocs; ++i) {
        reqAdder.add(random().nextInt(maxDoc));
    }
    DocIdSetBuilder.BulkAdder exclAdder = exclBuilder.grow(numExcludedDocs);
    for (int i = 0; i < numExcludedDocs; ++i) {
        exclAdder.add(random().nextInt(maxDoc));
    }
    final DocIdSet req = reqBuilder.build();
    final DocIdSet excl = exclBuilder.build();
    final BulkScorer reqBulkScorer = new BulkScorer() {

        final DocIdSetIterator iterator = req.iterator();

        @Override
        public int score(LeafCollector collector, Bits acceptDocs, int min, int max) throws IOException {
            int doc = iterator.docID();
            if (iterator.docID() < min) {
                doc = iterator.advance(min);
            }
            while (doc < max) {
                if (acceptDocs == null || acceptDocs.get(doc)) {
                    collector.collect(doc);
                }
                doc = iterator.nextDoc();
            }
            return doc;
        }

        @Override
        public long cost() {
            return iterator.cost();
        }
    };
    ReqExclBulkScorer reqExcl = new ReqExclBulkScorer(reqBulkScorer, excl.iterator());
    final FixedBitSet actualMatches = new FixedBitSet(maxDoc);
    if (random().nextBoolean()) {
        reqExcl.score(new LeafCollector() {

            @Override
            public void setScorer(Scorer scorer) throws IOException {
            }

            @Override
            public void collect(int doc) throws IOException {
                actualMatches.set(doc);
            }
        }, null);
    } else {
        int next = 0;
        while (next < maxDoc) {
            final int min = next;
            final int max = min + random().nextInt(10);
            next = reqExcl.score(new LeafCollector() {

                @Override
                public void setScorer(Scorer scorer) throws IOException {
                }

                @Override
                public void collect(int doc) throws IOException {
                    actualMatches.set(doc);
                }
            }, null, min, max);
            assertTrue(next >= max);
        }
    }
    final FixedBitSet expectedMatches = new FixedBitSet(maxDoc);
    expectedMatches.or(req.iterator());
    FixedBitSet excludedSet = new FixedBitSet(maxDoc);
    excludedSet.or(excl.iterator());
    expectedMatches.andNot(excludedSet);
    assertArrayEquals(expectedMatches.getBits(), actualMatches.getBits());
}