Examples with IndexTree org.apache.lucene.util.bkd.BKDReader.IndexTree used on opensource projects

Example 1 with IndexTree

use of org.apache.lucene.util.bkd.BKDReader.IndexTree in project lucene-solr by apache.

the class NearestNeighbor method nearest.

// TODO: can we somehow share more with, or simply directly use, the LatLonPointDistanceComparator?  It's really doing the same thing as
// our hitQueue...
public static NearestHit[] nearest(double pointLat, double pointLon, List<BKDReader> readers, List<Bits> liveDocs, List<Integer> docBases, final int n) throws IOException {
    //System.out.println("NEAREST: readers=" + readers + " liveDocs=" + liveDocs + " pointLat=" + pointLat + " pointLon=" + pointLon);
    // Holds closest collected points seen so far:
    // TODO: if we used lucene's PQ we could just updateTop instead of poll/offer:
    final PriorityQueue<NearestHit> hitQueue = new PriorityQueue<>(n, new Comparator<NearestHit>() {

        @Override
        public int compare(NearestHit a, NearestHit b) {
            // sort by opposite distanceMeters natural order
            int cmp = Double.compare(a.distanceMeters, b.distanceMeters);
            if (cmp != 0) {
                return -cmp;
            }
            // tie-break by higher docID:
            return b.docID - a.docID;
        }
    });
    // Holds all cells, sorted by closest to the point:
    PriorityQueue<Cell> cellQueue = new PriorityQueue<>();
    NearestVisitor visitor = new NearestVisitor(hitQueue, n, pointLat, pointLon);
    List<BKDReader.IntersectState> states = new ArrayList<>();
    // Add root cell for each reader into the queue:
    int bytesPerDim = -1;
    for (int i = 0; i < readers.size(); i++) {
        BKDReader reader = readers.get(i);
        if (bytesPerDim == -1) {
            bytesPerDim = reader.getBytesPerDimension();
        } else if (bytesPerDim != reader.getBytesPerDimension()) {
            throw new IllegalStateException("bytesPerDim changed from " + bytesPerDim + " to " + reader.getBytesPerDimension() + " across readers");
        }
        byte[] minPackedValue = reader.getMinPackedValue();
        byte[] maxPackedValue = reader.getMaxPackedValue();
        IntersectState state = reader.getIntersectState(visitor);
        states.add(state);
        cellQueue.offer(new Cell(state.index, i, reader.getMinPackedValue(), reader.getMaxPackedValue(), approxBestDistance(minPackedValue, maxPackedValue, pointLat, pointLon)));
    }
    while (cellQueue.size() > 0) {
        Cell cell = cellQueue.poll();
        //System.out.println("  visit " + cell);
        // TODO: if we replace approxBestDistance with actualBestDistance, we can put an opto here to break once this "best" cell is fully outside of the hitQueue bottom's radius:
        BKDReader reader = readers.get(cell.readerIndex);
        if (cell.index.isLeafNode()) {
            //System.out.println("    leaf");
            // Leaf block: visit all points and possibly collect them:
            visitor.curDocBase = docBases.get(cell.readerIndex);
            visitor.curLiveDocs = liveDocs.get(cell.readerIndex);
            reader.visitLeafBlockValues(cell.index, states.get(cell.readerIndex));
        //System.out.println("    now " + hitQueue.size() + " hits");
        } else {
            //System.out.println("    non-leaf");
            // Non-leaf block: split into two cells and put them back into the queue:
            double cellMinLat = decodeLatitude(cell.minPacked, 0);
            double cellMinLon = decodeLongitude(cell.minPacked, Integer.BYTES);
            double cellMaxLat = decodeLatitude(cell.maxPacked, 0);
            double cellMaxLon = decodeLongitude(cell.maxPacked, Integer.BYTES);
            if (cellMaxLat < visitor.minLat || visitor.maxLat < cellMinLat || ((cellMaxLon < visitor.minLon || visitor.maxLon < cellMinLon) && cellMaxLon < visitor.minLon2)) {
                // this cell is outside our search bbox; don't bother exploring any more
                continue;
            }
            BytesRef splitValue = BytesRef.deepCopyOf(cell.index.getSplitDimValue());
            int splitDim = cell.index.getSplitDim();
            // we must clone the index so that we we can recurse left and right "concurrently":
            IndexTree newIndex = cell.index.clone();
            byte[] splitPackedValue = cell.maxPacked.clone();
            System.arraycopy(splitValue.bytes, splitValue.offset, splitPackedValue, splitDim * bytesPerDim, bytesPerDim);
            cell.index.pushLeft();
            cellQueue.offer(new Cell(cell.index, cell.readerIndex, cell.minPacked, splitPackedValue, approxBestDistance(cell.minPacked, splitPackedValue, pointLat, pointLon)));
            splitPackedValue = cell.minPacked.clone();
            System.arraycopy(splitValue.bytes, splitValue.offset, splitPackedValue, splitDim * bytesPerDim, bytesPerDim);
            newIndex.pushRight();
            cellQueue.offer(new Cell(newIndex, cell.readerIndex, splitPackedValue, cell.maxPacked, approxBestDistance(splitPackedValue, cell.maxPacked, pointLat, pointLon)));
        }
    }
    NearestHit[] hits = new NearestHit[hitQueue.size()];
    int downTo = hitQueue.size() - 1;
    while (hitQueue.size() != 0) {
        hits[downTo] = hitQueue.poll();
        downTo--;
    }
    return hits;
}