org.hibernate.search.spatial.impl.DistanceCollector Maven / Gradle / Ivy

Go to download
/*
 * Hibernate Search, full-text search for your domain model
 *
 * License: GNU Lesser General Public License (LGPL), version 2.1 or later
 * See the lgpl.txt file in the root directory or .
 */
package org.hibernate.search.spatial.impl;

import static org.hibernate.search.spatial.impl.CoordinateHelper.coordinate;

import java.io.IOException;
import java.util.ArrayList;

import org.apache.lucene.index.DocValues;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.util.Bits;
import org.hibernate.search.exception.SearchException;
import org.hibernate.search.spatial.Coordinates;

/**
 * A Lucene distance {@code Collector} for spatial searches.
 *
 * @author Sanne Grinovero
 * @author Nicolas Helleringer
 */
public class DistanceCollector implements Collector {

	private final Point center;
	private final SpatialResultsCollector distances;
	private final String latitudeField;
	private final String longitudeField;

	public DistanceCollector(Coordinates centerCoordinates, int hitsCount, String fieldname) {
		this.center = Point.fromCoordinates( centerCoordinates );
		this.distances = new SpatialResultsCollector( hitsCount );
		this.latitudeField = SpatialHelper.formatLatitude( fieldname );
		this.longitudeField = SpatialHelper.formatLongitude( fieldname );
	}

	public Double getDistance(final int index) {
		return distances.get( index, center );
	}

	@Override
	public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
		return new DistanceLeafCollector( context );
	}

	@Override
	public boolean needsScores() {
		return false;
	}

	/**
	 * We'll store matching hits in HitEntry instances so to allow retrieving results
	 * in a second phase after the Collector has run.
	 * Also take the opportunity to lazily calculate the actual distance: only store
	 * latitude and longitude coordinates.
	 */
	private abstract static class HitEntry {
		int documentId;

		private HitEntry(int documentId) {
			this.documentId = documentId;
		}

		abstract Double getDistance(Point center);
	}

	private static final class CompleteHitEntry extends HitEntry {
		double latitude;
		double longitude;

		private CompleteHitEntry(int documentId, double latitude, double longitude) {
			super( documentId );
			this.latitude = latitude;
			this.longitude = longitude;
		}

		@Override
		Double getDistance(final Point center) {
			return center.getDistanceTo( latitude, longitude );
		}
	}

	private static final class IncompleteHitEntry extends HitEntry {
		private IncompleteHitEntry(int documentId) {
			super( documentId );
		}

		@Override
		Double getDistance(final Point center) {
			return null;
		}
	}

	/**
	 * A custom structure to store all HitEntry instances.
	 * Based on an array, as in most cases we'll append sequentially and iterate the
	 * results in the same order too. The size is well known in most situations so we can
	 * also guess an appropriate allocation size.
	 *
	 * Iteration of the results will in practice follow a monotonic index, unless a non natural
	 * Sort is specified. So by keeping track of the position in the array of the last request,
	 * and look from that pointer first, the cost of get operations is O(1) in most common use
	 * cases.
	 */
	private static class SpatialResultsCollector {
		final ArrayList orderedEntries;
		int currentIterator = 0;

		private SpatialResultsCollector(int size) {
			orderedEntries = new ArrayList( size );
		}

		public Double get(int index, Point center) {
			//Optimize for an iteration having a monotonic index:
			int startingPoint = currentIterator;
			for ( ; currentIterator < orderedEntries.size(); currentIterator++ ) {
				HitEntry currentEntry = orderedEntries.get( currentIterator );
				if ( currentEntry == null ) {
					break;
				}
				if ( currentEntry.documentId == index ) {
					return currentEntry.getDistance( center );
				}
			}

			//No match yet! scan the remaining section from the beginning:
			for ( currentIterator = 0; currentIterator < startingPoint; currentIterator++ ) {
				HitEntry currentEntry = orderedEntries.get( currentIterator );
				if ( currentEntry == null ) {
					break;
				}
				if ( currentEntry.documentId == index ) {
					return currentEntry.getDistance( center );
				}
			}
			throw new SearchException( "Unexpected index: this documentId was not collected" );
		}

		void put(int documentId, double latitude, double longitude) {
			orderedEntries.add( new CompleteHitEntry( documentId, latitude, longitude ) );
		}

		void putIncomplete(int documentId) {
			orderedEntries.add( new IncompleteHitEntry( documentId ) );
		}
	}

	private class DistanceLeafCollector implements LeafCollector {

		private final int docBase;
		private final Bits docsWithLatitude;
		private final Bits docsWithLongitude;
		private final NumericDocValues latitudeValues;
		private final NumericDocValues longitudeValues;

		DistanceLeafCollector(LeafReaderContext context) throws IOException {
			final LeafReader atomicReader = context.reader();
			this.docsWithLatitude = DocValues.getDocsWithField( atomicReader, latitudeField );
			this.docsWithLongitude = DocValues.getDocsWithField( atomicReader, longitudeField );
			this.latitudeValues = DocValues.getNumeric( atomicReader, latitudeField );
			this.longitudeValues = DocValues.getNumeric( atomicReader, longitudeField );
			this.docBase = context.docBase;
		}

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

		@Override
		public void collect(int doc) throws IOException {
			final int absolute = docBase + doc;
			if ( docsWithLatitude.get( doc ) && docsWithLongitude.get( doc ) ) {
				double lat = coordinate( latitudeValues, doc );
				double lon = coordinate( longitudeValues, doc );
				distances.put( absolute, lat, lon );
			}
			else {
				distances.putIncomplete( absolute );
			}
		}
	}
}