org.apache.lucene.spatial.prefix.tree.S2PrefixTree Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of lucene-spatial-extras Show documentation
Advanced Spatial Shape Strategies for Apache Lucene
There is a newer version: 10.1.0
Show newest version
/*
 * 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.prefix.tree;

import com.google.common.geometry.S2CellId;
import com.google.common.geometry.S2LatLng;
import com.google.common.geometry.S2Projections;
import java.util.ArrayList;
import java.util.List;
import org.apache.lucene.util.BytesRef;
import org.locationtech.spatial4j.context.SpatialContext;
import org.locationtech.spatial4j.distance.DistanceUtils;
import org.locationtech.spatial4j.shape.Point;
import org.locationtech.spatial4j.shape.Shape;

/**
 * Spatial prefix tree for S2 Geometry. Shape factories for the
 * given {@link SpatialContext} must implement the interface {@link S2ShapeFactory}.
 *
 * The tree can be configured on how it divided itself by providing an arity. The default arity
 * is 1 which divided every sub-cell in 4 (except the first level that is always divided by 6) .
 * Arity 2 divides sub-cells in 16 and arity 3 in 64 sub-cells.
 *
 * @lucene.experimental
 */
public class S2PrefixTree extends SpatialPrefixTree {

  /** Factory for creating {@link S2PrefixTree} instances with useful defaults */
  protected static class Factory extends SpatialPrefixTreeFactory {

    @Override
    protected int getLevelForDistance(double degrees) {
      S2PrefixTree grid = new S2PrefixTree(ctx, S2PrefixTree.getMaxLevels(1));
      return grid.getLevelForDistance(degrees);
    }

    @Override
    protected SpatialPrefixTree newSPT() {
      return new S2PrefixTree(ctx, maxLevels != null ? maxLevels : S2PrefixTree.getMaxLevels(1));
    }
  }

  // factory to generate S2 cell shapes
  protected final S2ShapeFactory s2ShapeFactory;
  protected final int arity;

  /**
   * Creates a S2 spatial tree with arity 1.
   *
   * @param ctx The provided spatial context. The shape factor of the spatial context must implement
   *     {@link S2ShapeFactory}
   * @param maxLevels The provided maximum level for this tree.
   */
  public S2PrefixTree(SpatialContext ctx, int maxLevels) {
    this(ctx, maxLevels, 1);
  }

  /**
   * Creates a S2 spatial tree with provided arity.
   *
   * @param ctx The provided spatial context. The shape factor of the spatial context must implement
   *     {@link S2ShapeFactory}
   * @param maxLevels The provided maximum level for this tree.
   * @param arity The arity of the tree.
   */
  public S2PrefixTree(SpatialContext ctx, int maxLevels, int arity) {
    super(ctx, maxLevels);
    if (!(ctx.getShapeFactory() instanceof S2ShapeFactory)) {
      throw new IllegalArgumentException("Spatial context does not support S2 spatial index.");
    }
    this.s2ShapeFactory = (S2ShapeFactory) ctx.getShapeFactory();
    if (arity < 1 || arity > 3) {
      throw new IllegalArgumentException(
          "Invalid value for S2 tree arity. Possible values are 1, 2 or 3. Provided value is "
              + arity
              + ".");
    }
    this.arity = arity;
  }

  /**
   * Get max levels for this spatial tree.
   *
   * @param arity The arity of the tree.
   * @return The maximum number of levels by the provided arity.
   */
  public static int getMaxLevels(int arity) {
    return S2CellId.MAX_LEVEL / arity + 1;
  }

  @Override
  public int getLevelForDistance(double dist) {
    if (dist == 0) {
      return maxLevels;
    }
    int level = S2Projections.MAX_WIDTH.getMinLevel(dist * DistanceUtils.DEGREES_TO_RADIANS);
    int roundLevel = level % arity != 0 ? 1 : 0;
    level = level / arity + roundLevel;
    return Math.min(maxLevels, level + 1);
  }

  @Override
  public double getDistanceForLevel(int level) {
    if (level == 0) {
      return 180;
    }
    return S2Projections.MAX_WIDTH.getValue(arity * (level - 1)) * DistanceUtils.RADIANS_TO_DEGREES;
  }

  @Override
  public Cell getWorldCell() {
    return new S2PrefixTreeCell(this, null);
  }

  @Override
  public Cell readCell(BytesRef term, Cell scratch) {
    S2PrefixTreeCell cell = (S2PrefixTreeCell) scratch;
    if (cell == null) {
      cell = (S2PrefixTreeCell) getWorldCell();
    }
    cell.readCell(this, term);
    return cell;
  }

  @Override
  public CellIterator getTreeCellIterator(Shape shape, int detailLevel) {
    if (!(shape instanceof Point)) {
      return super.getTreeCellIterator(shape, detailLevel);
    }
    Point p = (Point) shape;
    S2CellId id =
        S2CellId.fromLatLng(S2LatLng.fromDegrees(p.getY(), p.getX()))
            .parent(arity * (detailLevel - 1));
    List cells = new ArrayList<>(detailLevel);
    for (int i = 0; i < detailLevel - 1; i++) {
      cells.add(new S2PrefixTreeCell(this, id.parent(i * arity)));
    }
    cells.add(new S2PrefixTreeCell(this, id));
    return new FilterCellIterator(cells.iterator(), null);
  }
}