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

Example 16 with BitSetIterator

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

the class PointRangeQuery method createWeight.

@Override
public final Weight createWeight(IndexSearcher searcher, boolean needsScores, float boost) throws IOException {
    return new ConstantScoreWeight(this, boost) {

        private IntersectVisitor getIntersectVisitor(DocIdSetBuilder result) {
            return 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) {
                    for (int dim = 0; dim < numDims; dim++) {
                        int offset = dim * bytesPerDim;
                        if (StringHelper.compare(bytesPerDim, packedValue, offset, lowerPoint, offset) < 0) {
                            // Doc's value is too low, in this dimension
                            return;
                        }
                        if (StringHelper.compare(bytesPerDim, packedValue, offset, upperPoint, offset) > 0) {
                            // Doc's value is too high, in this dimension
                            return;
                        }
                    }
                    // Doc is in-bounds
                    adder.add(docID);
                }

                @Override
                public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
                    boolean crosses = false;
                    for (int dim = 0; dim < numDims; dim++) {
                        int offset = dim * bytesPerDim;
                        if (StringHelper.compare(bytesPerDim, minPackedValue, offset, upperPoint, offset) > 0 || StringHelper.compare(bytesPerDim, maxPackedValue, offset, lowerPoint, offset) < 0) {
                            return Relation.CELL_OUTSIDE_QUERY;
                        }
                        crosses |= StringHelper.compare(bytesPerDim, minPackedValue, offset, lowerPoint, offset) < 0 || StringHelper.compare(bytesPerDim, maxPackedValue, offset, upperPoint, offset) > 0;
                    }
                    if (crosses) {
                        return Relation.CELL_CROSSES_QUERY;
                    } else {
                        return Relation.CELL_INSIDE_QUERY;
                    }
                }
            };
        }

        /**
       * Create a visitor that clears documents that do NOT match the range.
       */
        private IntersectVisitor getInverseIntersectVisitor(FixedBitSet result, int[] cost) {
            return new IntersectVisitor() {

                @Override
                public void visit(int docID) {
                    result.clear(docID);
                    cost[0]--;
                }

                @Override
                public void visit(int docID, byte[] packedValue) {
                    for (int dim = 0; dim < numDims; dim++) {
                        int offset = dim * bytesPerDim;
                        if (StringHelper.compare(bytesPerDim, packedValue, offset, lowerPoint, offset) < 0) {
                            // Doc's value is too low, in this dimension
                            result.clear(docID);
                            cost[0]--;
                            return;
                        }
                        if (StringHelper.compare(bytesPerDim, packedValue, offset, upperPoint, offset) > 0) {
                            // Doc's value is too high, in this dimension
                            result.clear(docID);
                            cost[0]--;
                            return;
                        }
                    }
                }

                @Override
                public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
                    boolean crosses = false;
                    for (int dim = 0; dim < numDims; dim++) {
                        int offset = dim * bytesPerDim;
                        if (StringHelper.compare(bytesPerDim, minPackedValue, offset, upperPoint, offset) > 0 || StringHelper.compare(bytesPerDim, maxPackedValue, offset, lowerPoint, offset) < 0) {
                            // This dim is not in the range
                            return Relation.CELL_INSIDE_QUERY;
                        }
                        crosses |= StringHelper.compare(bytesPerDim, minPackedValue, offset, lowerPoint, offset) < 0 || StringHelper.compare(bytesPerDim, maxPackedValue, offset, upperPoint, offset) > 0;
                    }
                    if (crosses) {
                        return Relation.CELL_CROSSES_QUERY;
                    } else {
                        return Relation.CELL_OUTSIDE_QUERY;
                    }
                }
            };
        }

        @Override
        public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
            LeafReader reader = context.reader();
            PointValues values = reader.getPointValues(field);
            if (values == null) {
                // No docs in this segment/field indexed any points
                return null;
            }
            if (values.getNumDimensions() != numDims) {
                throw new IllegalArgumentException("field=\"" + field + "\" was indexed with numDims=" + values.getNumDimensions() + " but this query has numDims=" + numDims);
            }
            if (bytesPerDim != values.getBytesPerDimension()) {
                throw new IllegalArgumentException("field=\"" + field + "\" was indexed with bytesPerDim=" + values.getBytesPerDimension() + " but this query has bytesPerDim=" + bytesPerDim);
            }
            boolean allDocsMatch;
            if (values.getDocCount() == reader.maxDoc()) {
                final byte[] fieldPackedLower = values.getMinPackedValue();
                final byte[] fieldPackedUpper = values.getMaxPackedValue();
                allDocsMatch = true;
                for (int i = 0; i < numDims; ++i) {
                    int offset = i * bytesPerDim;
                    if (StringHelper.compare(bytesPerDim, lowerPoint, offset, fieldPackedLower, offset) > 0 || StringHelper.compare(bytesPerDim, upperPoint, offset, fieldPackedUpper, offset) < 0) {
                        allDocsMatch = false;
                        break;
                    }
                }
            } else {
                allDocsMatch = false;
            }
            final Weight weight = this;
            if (allDocsMatch) {
                // all docs have a value and all points are within bounds, so everything matches
                return new ScorerSupplier() {

                    @Override
                    public Scorer get(boolean randomAccess) {
                        return new ConstantScoreScorer(weight, score(), DocIdSetIterator.all(reader.maxDoc()));
                    }

                    @Override
                    public long cost() {
                        return reader.maxDoc();
                    }
                };
            } else {
                return new ScorerSupplier() {

                    final DocIdSetBuilder result = new DocIdSetBuilder(reader.maxDoc(), values, field);

                    final IntersectVisitor visitor = getIntersectVisitor(result);

                    long cost = -1;

                    @Override
                    public Scorer get(boolean randomAccess) throws IOException {
                        if (values.getDocCount() == reader.maxDoc() && values.getDocCount() == values.size() && cost() > reader.maxDoc() / 2) {
                            // If all docs have exactly one value and the cost is greater
                            // than half the leaf size then maybe we can make things faster
                            // by computing the set of documents that do NOT match the range
                            final FixedBitSet result = new FixedBitSet(reader.maxDoc());
                            result.set(0, reader.maxDoc());
                            int[] cost = new int[] { reader.maxDoc() };
                            values.intersect(getInverseIntersectVisitor(result, cost));
                            final DocIdSetIterator iterator = new BitSetIterator(result, cost[0]);
                            return new ConstantScoreScorer(weight, score(), iterator);
                        }
                        values.intersect(visitor);
                        DocIdSetIterator iterator = result.build().iterator();
                        return new ConstantScoreScorer(weight, score(), iterator);
                    }

                    @Override
                    public long cost() {
                        if (cost == -1) {
                            // Computing the cost may be expensive, so only do it if necessary
                            cost = values.estimatePointCount(visitor);
                            assert cost >= 0;
                        }
                        return cost;
                    }
                };
            }
        }

        @Override
        public Scorer scorer(LeafReaderContext context) throws IOException {
            ScorerSupplier scorerSupplier = scorerSupplier(context);
            if (scorerSupplier == null) {
                return null;
            }
            return scorerSupplier.get(false);
        }
    };
}

Example 17 with BitSetIterator

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

the class TestIndexedDISI method doTest.

private void doTest(FixedBitSet set, Directory dir) throws IOException {
    final int cardinality = set.cardinality();
    long length;
    try (IndexOutput out = dir.createOutput("foo", IOContext.DEFAULT)) {
        IndexedDISI.writeBitSet(new BitSetIterator(set, cardinality), out);
        length = out.getFilePointer();
    }
    try (IndexInput in = dir.openInput("foo", IOContext.DEFAULT)) {
        IndexedDISI disi = new IndexedDISI(in, 0L, length, cardinality);
        BitSetIterator disi2 = new BitSetIterator(set, cardinality);
        int i = 0;
        for (int doc = disi2.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = disi2.nextDoc()) {
            assertEquals(doc, disi.nextDoc());
            assertEquals(i++, disi.index());
        }
        assertEquals(DocIdSetIterator.NO_MORE_DOCS, disi.nextDoc());
    }
    for (int step : new int[] { 1, 10, 100, 1000, 10000, 100000 }) {
        try (IndexInput in = dir.openInput("foo", IOContext.DEFAULT)) {
            IndexedDISI disi = new IndexedDISI(in, 0L, length, cardinality);
            BitSetIterator disi2 = new BitSetIterator(set, cardinality);
            int index = -1;
            while (true) {
                int target = disi2.docID() + step;
                int doc;
                do {
                    doc = disi2.nextDoc();
                    index++;
                } while (doc < target);
                assertEquals(doc, disi.advance(target));
                if (doc == DocIdSetIterator.NO_MORE_DOCS) {
                    break;
                }
                assertEquals(index, disi.index());
            }
        }
    }
    for (int step : new int[] { 10, 100, 1000, 10000, 100000 }) {
        try (IndexInput in = dir.openInput("foo", IOContext.DEFAULT)) {
            IndexedDISI disi = new IndexedDISI(in, 0L, length, cardinality);
            BitSetIterator disi2 = new BitSetIterator(set, cardinality);
            int index = -1;
            for (int target = 0; target < set.length(); ) {
                target += TestUtil.nextInt(random(), 0, step);
                int doc = disi2.docID();
                while (doc < target) {
                    doc = disi2.nextDoc();
                    index++;
                }
                boolean exists = disi.advanceExact(target);
                assertEquals(doc == target, exists);
                if (exists) {
                    assertEquals(index, disi.index());
                } else if (random().nextBoolean()) {
                    assertEquals(doc, disi.nextDoc());
                    assertEquals(index, disi.index());
                    target = doc;
                }
            }
        }
    }
    dir.deleteFile("foo");
}