lucene/spatial/src/java/org/apache/lucene/spatial/search/GeoPointTermQueryConstantScoreWrapper.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.lucene.spatial.search;

import java.io.IOException;

import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.index.Terms;
import org.apache.lucene.search.BulkScorer;
import org.apache.lucene.search.ConstantScoreScorer;
import org.apache.lucene.search.ConstantScoreWeight;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Weight;
import org.apache.lucene.util.AttributeSource;
import org.apache.lucene.util.DocIdSetBuilder;
import org.apache.lucene.spatial.util.GeoUtils;

/**
 * Custom ConstantScoreWrapper for {@code GeoPointTermQuery} that cuts over to DocValues
 * for post filtering boundary ranges. Multi-valued GeoPoint documents are supported.
 *
 * @lucene.experimental
 */
final class GeoPointTermQueryConstantScoreWrapper <Q extends GeoPointTermQuery> extends Query {
  protected final Q query;

  protected GeoPointTermQueryConstantScoreWrapper(Q query) {
    this.query = query;
  }

  @Override
  public String toString(String field) {
    return query.toString();
  }

  @Override
  public final boolean equals(final Object o) {
    if (super.equals(o) == false) {
      return false;
    }
    final GeoPointTermQueryConstantScoreWrapper<?> that = (GeoPointTermQueryConstantScoreWrapper<?>) o;
    return this.query.equals(that.query);
  }

  @Override
  public final int hashCode() {
    return 31 * super.hashCode() + query.hashCode();
  }

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

      private DocIdSet getDocIDs(LeafReaderContext context) throws IOException {
        final Terms terms = context.reader().terms(query.getField());
        if (terms == null) {
          return DocIdSet.EMPTY;
        }

        final GeoPointTermsEnum termsEnum = (GeoPointTermsEnum)(query.getTermsEnum(terms, new AttributeSource()));
        assert termsEnum != null;

        LeafReader reader = context.reader();
        DocIdSetBuilder builder = new DocIdSetBuilder(reader.maxDoc());
        PostingsEnum docs = null;
        SortedNumericDocValues sdv = reader.getSortedNumericDocValues(query.getField());

        while (termsEnum.next() != null) {
          docs = termsEnum.postings(docs, PostingsEnum.NONE);
          // boundary terms need post filtering by
          if (termsEnum.boundaryTerm()) {
            int docId = docs.nextDoc();
            long hash;
            do {
              sdv.setDocument(docId);
              for (int i=0; i<sdv.count(); ++i) {
                hash = sdv.valueAt(i);
                if (termsEnum.postFilter(GeoUtils.mortonUnhashLon(hash), GeoUtils.mortonUnhashLat(hash))) {
                  builder.add(docId);
                  break;
                }
              }
            } while ((docId = docs.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS);
          } else {
            builder.add(docs);
          }
        }

        return builder.build();
      }

      private Scorer scorer(DocIdSet set) throws IOException {
        if (set == null) {
          return null;
        }
        final DocIdSetIterator disi = set.iterator();
        if (disi == null) {
          return null;
        }
        return new ConstantScoreScorer(this, score(), disi);
      }

      @Override
      public BulkScorer bulkScorer(LeafReaderContext context) throws IOException {
        final Scorer scorer = scorer(getDocIDs(context));
        if (scorer == null) {
          return null;
        }
        return new DefaultBulkScorer(scorer);
      }

      @Override
      public Scorer scorer(LeafReaderContext context) throws IOException {
        return scorer(getDocIDs(context));
      }
    };
  }
}