lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/RecursivePrefixTreeStrategy.java - lucene-solr - Git at Google

 /*
  * 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;

 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;

 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.spatial.prefix.tree.Cell;
 import org.apache.lucene.spatial.prefix.tree.CellCanPrune;
 import org.apache.lucene.spatial.prefix.tree.CellIterator;
 import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
 import org.apache.lucene.spatial.query.SpatialArgs;
 import org.apache.lucene.spatial.query.SpatialOperation;
 import org.apache.lucene.spatial.query.UnsupportedSpatialOperation;
 import org.locationtech.spatial4j.shape.Shape;

 /**
  * A {@link PrefixTreeStrategy} which uses {@link AbstractVisitingPrefixTreeQuery}.
  * This strategy has support for searching non-point shapes (note: not tested).
  * Even a query shape with distErrPct=0 (fully precise to the grid) should have
  * good performance for typical data, unless there is a lot of indexed data
  * coincident with the shape's edge.
  *
  * @lucene.experimental
  */
 public class RecursivePrefixTreeStrategy extends PrefixTreeStrategy {
   /* Future potential optimizations:

     Each shape.relate(otherShape) result could be cached since much of the same relations will be invoked when
     multiple segments are involved. Do this for "complex" shapes, not cheap ones, and don't cache when disjoint to
     bbox because it's a cheap calc. This is one advantage TermQueryPrefixTreeStrategy has over RPT.

    */

   protected int prefixGridScanLevel;

   //Formerly known as simplifyIndexedCells. Eventually will be removed. Only compatible with RPT
   // and cells implementing CellCanPrune, otherwise ignored.
   protected boolean pruneLeafyBranches = true;

   protected boolean multiOverlappingIndexedShapes = true;

   public RecursivePrefixTreeStrategy(SpatialPrefixTree grid, String fieldName) {
     super(grid, fieldName);
     prefixGridScanLevel = grid.getMaxLevels() - 4;//TODO this default constant is dependent on the prefix grid size
   }

   public int getPrefixGridScanLevel() {
     return prefixGridScanLevel;
   }

   /**
    * Sets the grid level [1-maxLevels] at which indexed terms are scanned brute-force
    * instead of by grid decomposition.  By default this is maxLevels - 4.  The
    * final level, maxLevels, is always scanned.
    *
    * @param prefixGridScanLevel 1 to maxLevels
    */
   public void setPrefixGridScanLevel(int prefixGridScanLevel) {
     //TODO if negative then subtract from maxlevels
     this.prefixGridScanLevel = prefixGridScanLevel;
   }

   public boolean isMultiOverlappingIndexedShapes() {
     return multiOverlappingIndexedShapes;
   }

   /** See {@link ContainsPrefixTreeQuery#multiOverlappingIndexedShapes}. */
   public void setMultiOverlappingIndexedShapes(boolean multiOverlappingIndexedShapes) {
     this.multiOverlappingIndexedShapes = multiOverlappingIndexedShapes;
   }

   public boolean isPruneLeafyBranches() {
     return pruneLeafyBranches;
   }

   /**
    * An optional hint affecting non-point shapes and tree cells implementing {@link CellCanPrune}, otherwise
    * ignored.
    * <p>
    * It will prune away a complete set sibling leaves to their parent (recursively), resulting in ~20-50%
    * fewer indexed cells, and consequently that much less disk and that much faster indexing.
    * So if it's a quad tree and all 4 sub-cells are there marked as a leaf, then they will be
    * removed (pruned) and the parent is marked as a leaf instead.  This occurs recursively on up.  Unfortunately, the
    * current implementation will buffer all cells to do this, so consider disabling for high precision (low distErrPct)
    * shapes. (default=true)
    */
   public void setPruneLeafyBranches(boolean pruneLeafyBranches) {
     this.pruneLeafyBranches = pruneLeafyBranches;
   }

   @Override
   public String toString() {
     StringBuilder str = new StringBuilder(getClass().getSimpleName()).append('(');
     str.append("SPG:(").append(grid.toString()).append(')');
     if (pointsOnly)
       str.append(",pointsOnly");
     if (pruneLeafyBranches)
       str.append(",pruneLeafyBranches");
     if (prefixGridScanLevel != grid.getMaxLevels() - 4)
       str.append(",prefixGridScanLevel:").append("").append(prefixGridScanLevel);
     if (!multiOverlappingIndexedShapes)
       str.append(",!multiOverlappingIndexedShapes");
     return str.append(')').toString();
   }

   @Override
   protected Iterator<Cell> createCellIteratorToIndex(Shape shape, int detailLevel, Iterator<Cell> reuse) {
     if (!pruneLeafyBranches || isGridAlignedShape(shape))
       return super.createCellIteratorToIndex(shape, detailLevel, reuse);

     List<Cell> cells = new ArrayList<>(4096);
     recursiveTraverseAndPrune(grid.getWorldCell(), shape, detailLevel, cells);
     return cells.iterator();
   }

   /** Returns true if cell was added as a leaf. If it wasn't it recursively descends. */
   private boolean recursiveTraverseAndPrune(Cell cell, Shape shape, int detailLevel, List<Cell> result) {

     if (cell.getLevel() == detailLevel) {
       cell.setLeaf();//FYI might already be a leaf
     }
     if (cell.isLeaf()) {
       result.add(cell);
       return true;
     }
     if (cell.getLevel() != 0)
       result.add(cell);

     int leaves = 0;
     CellIterator subCells = cell.getNextLevelCells(shape);
     while (subCells.hasNext()) {
       Cell subCell = subCells.next();
       if (recursiveTraverseAndPrune(subCell, shape, detailLevel, result))
         leaves++;
     }

     if (!(cell instanceof CellCanPrune)) {
       //Cannot prune so return false
       return false;
     }

     //can we prune?
     if (leaves == ((CellCanPrune)cell).getSubCellsSize() && cell.getLevel() != 0) {
       //Optimization: substitute the parent as a leaf instead of adding all
       // children as leaves

       //remove the leaves
       do {
         result.remove(result.size() - 1);//remove last
       } while (--leaves > 0);
       //add cell as the leaf
       cell.setLeaf();
       return true;
     }
     return false;
   }

   @Override
   public Query makeQuery(SpatialArgs args) {
     final SpatialOperation op = args.getOperation();

     Shape shape = args.getShape();
     int detailLevel = grid.getLevelForDistance(args.resolveDistErr(ctx, distErrPct));

     if (op == SpatialOperation.Intersects) {
       if (isGridAlignedShape(args.getShape())) {
         return makeGridShapeIntersectsQuery(args.getShape());
       }
       return new IntersectsPrefixTreeQuery(
           shape, getFieldName(), grid, detailLevel, prefixGridScanLevel);
     } else if (op == SpatialOperation.IsWithin) {
       return new WithinPrefixTreeQuery(
           shape, getFieldName(), grid, detailLevel, prefixGridScanLevel,
           -1);//-1 flag is slower but ensures correct results
     } else if (op == SpatialOperation.Contains) {
       return new ContainsPrefixTreeQuery(shape, getFieldName(), grid, detailLevel,
           multiOverlappingIndexedShapes);
     }
     throw new UnsupportedSpatialOperation(op);
   }

   /**
    * A quick check of the shape to see if it is perfectly aligned to a grid.
    * Points always are as they are indivisible.  It's okay to return false
    * if the shape actually is aligned; this is an optimization hint.
    */
   protected boolean isGridAlignedShape(Shape shape) {
     return isPointShape(shape);
   }

   /** {@link #makeQuery(SpatialArgs)} specialized for the query being a grid square. */
   protected Query makeGridShapeIntersectsQuery(Shape gridShape) {
     assert isGridAlignedShape(gridShape);
     if (isPointsOnly()) {
       // Awesome; this will be equivalent to a TermQuery.
       Iterator<Cell> cellIterator = grid.getTreeCellIterator(gridShape, grid.getMaxLevels());
       // get last cell
       Cell cell = cellIterator.next();
       while (cellIterator.hasNext()) {
         int prevLevel = cell.getLevel();
         cell = cellIterator.next();
         assert prevLevel < cell.getLevel();
       }
       assert cell.isLeaf();
       return new TermQuery(new Term(getFieldName(), cell.getTokenBytesWithLeaf(null)));
     } else {
       // Well there could be parent cells. But we can reduce the "scan level" which will be slower for a point query.
       // TODO: AVPTQ will still scan the bottom nonetheless; file an issue to eliminate that
       return new IntersectsPrefixTreeQuery(
           gridShape, getFieldName(), grid, getGrid().getMaxLevels(), getGrid().getMaxLevels() + 1);
     }
   }
 }
	/*
	* 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;

	import java.util.ArrayList;
	import java.util.Iterator;
	import java.util.List;

	import org.apache.lucene.index.Term;
	import org.apache.lucene.search.Query;
	import org.apache.lucene.search.TermQuery;
	import org.apache.lucene.spatial.prefix.tree.Cell;
	import org.apache.lucene.spatial.prefix.tree.CellCanPrune;
	import org.apache.lucene.spatial.prefix.tree.CellIterator;
	import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
	import org.apache.lucene.spatial.query.SpatialArgs;
	import org.apache.lucene.spatial.query.SpatialOperation;
	import org.apache.lucene.spatial.query.UnsupportedSpatialOperation;
	import org.locationtech.spatial4j.shape.Shape;

	/**
	* A {@link PrefixTreeStrategy} which uses {@link AbstractVisitingPrefixTreeQuery}.
	* This strategy has support for searching non-point shapes (note: not tested).
	* Even a query shape with distErrPct=0 (fully precise to the grid) should have
	* good performance for typical data, unless there is a lot of indexed data
	* coincident with the shape's edge.
	*
	* @lucene.experimental
	*/
	public class RecursivePrefixTreeStrategy extends PrefixTreeStrategy {
	/* Future potential optimizations:

	Each shape.relate(otherShape) result could be cached since much of the same relations will be invoked when
	multiple segments are involved. Do this for "complex" shapes, not cheap ones, and don't cache when disjoint to
	bbox because it's a cheap calc. This is one advantage TermQueryPrefixTreeStrategy has over RPT.

	*/

	protected int prefixGridScanLevel;

	//Formerly known as simplifyIndexedCells. Eventually will be removed. Only compatible with RPT
	// and cells implementing CellCanPrune, otherwise ignored.
	protected boolean pruneLeafyBranches = true;

	protected boolean multiOverlappingIndexedShapes = true;

	public RecursivePrefixTreeStrategy(SpatialPrefixTree grid, String fieldName) {
	super(grid, fieldName);
	prefixGridScanLevel = grid.getMaxLevels() - 4;//TODO this default constant is dependent on the prefix grid size
	}

	public int getPrefixGridScanLevel() {
	return prefixGridScanLevel;
	}

	/**
	* Sets the grid level [1-maxLevels] at which indexed terms are scanned brute-force
	* instead of by grid decomposition. By default this is maxLevels - 4. The
	* final level, maxLevels, is always scanned.
	*
	* @param prefixGridScanLevel 1 to maxLevels
	*/
	public void setPrefixGridScanLevel(int prefixGridScanLevel) {
	//TODO if negative then subtract from maxlevels
	this.prefixGridScanLevel = prefixGridScanLevel;
	}

	public boolean isMultiOverlappingIndexedShapes() {
	return multiOverlappingIndexedShapes;
	}

	/** See {@link ContainsPrefixTreeQuery#multiOverlappingIndexedShapes}. */
	public void setMultiOverlappingIndexedShapes(boolean multiOverlappingIndexedShapes) {
	this.multiOverlappingIndexedShapes = multiOverlappingIndexedShapes;
	}

	public boolean isPruneLeafyBranches() {
	return pruneLeafyBranches;
	}

	/**
	* An optional hint affecting non-point shapes and tree cells implementing {@link CellCanPrune}, otherwise
	* ignored.
	* <p>
	* It will prune away a complete set sibling leaves to their parent (recursively), resulting in ~20-50%
	* fewer indexed cells, and consequently that much less disk and that much faster indexing.
	* So if it's a quad tree and all 4 sub-cells are there marked as a leaf, then they will be
	* removed (pruned) and the parent is marked as a leaf instead. This occurs recursively on up. Unfortunately, the
	* current implementation will buffer all cells to do this, so consider disabling for high precision (low distErrPct)
	* shapes. (default=true)
	*/
	public void setPruneLeafyBranches(boolean pruneLeafyBranches) {
	this.pruneLeafyBranches = pruneLeafyBranches;
	}

	@Override
	public String toString() {
	StringBuilder str = new StringBuilder(getClass().getSimpleName()).append('(');
	str.append("SPG:(").append(grid.toString()).append(')');
	if (pointsOnly)
	str.append(",pointsOnly");
	if (pruneLeafyBranches)
	str.append(",pruneLeafyBranches");
	if (prefixGridScanLevel != grid.getMaxLevels() - 4)
	str.append(",prefixGridScanLevel:").append("").append(prefixGridScanLevel);
	if (!multiOverlappingIndexedShapes)
	str.append(",!multiOverlappingIndexedShapes");
	return str.append(')').toString();
	}

	@Override
	protected Iterator<Cell> createCellIteratorToIndex(Shape shape, int detailLevel, Iterator<Cell> reuse) {
	if (!pruneLeafyBranches \|\| isGridAlignedShape(shape))
	return super.createCellIteratorToIndex(shape, detailLevel, reuse);

	List<Cell> cells = new ArrayList<>(4096);
	recursiveTraverseAndPrune(grid.getWorldCell(), shape, detailLevel, cells);
	return cells.iterator();
	}

	/** Returns true if cell was added as a leaf. If it wasn't it recursively descends. */
	private boolean recursiveTraverseAndPrune(Cell cell, Shape shape, int detailLevel, List<Cell> result) {

	if (cell.getLevel() == detailLevel) {
	cell.setLeaf();//FYI might already be a leaf
	}
	if (cell.isLeaf()) {
	result.add(cell);
	return true;
	}
	if (cell.getLevel() != 0)
	result.add(cell);

	int leaves = 0;
	CellIterator subCells = cell.getNextLevelCells(shape);
	while (subCells.hasNext()) {
	Cell subCell = subCells.next();
	if (recursiveTraverseAndPrune(subCell, shape, detailLevel, result))
	leaves++;
	}

	if (!(cell instanceof CellCanPrune)) {
	//Cannot prune so return false
	return false;
	}

	//can we prune?
	if (leaves == ((CellCanPrune)cell).getSubCellsSize() && cell.getLevel() != 0) {
	//Optimization: substitute the parent as a leaf instead of adding all
	// children as leaves

	//remove the leaves
	do {
	result.remove(result.size() - 1);//remove last
	} while (--leaves > 0);
	//add cell as the leaf
	cell.setLeaf();
	return true;
	}
	return false;
	}

	@Override
	public Query makeQuery(SpatialArgs args) {
	final SpatialOperation op = args.getOperation();

	Shape shape = args.getShape();
	int detailLevel = grid.getLevelForDistance(args.resolveDistErr(ctx, distErrPct));

	if (op == SpatialOperation.Intersects) {
	if (isGridAlignedShape(args.getShape())) {
	return makeGridShapeIntersectsQuery(args.getShape());
	}
	return new IntersectsPrefixTreeQuery(
	shape, getFieldName(), grid, detailLevel, prefixGridScanLevel);
	} else if (op == SpatialOperation.IsWithin) {
	return new WithinPrefixTreeQuery(
	shape, getFieldName(), grid, detailLevel, prefixGridScanLevel,
	-1);//-1 flag is slower but ensures correct results
	} else if (op == SpatialOperation.Contains) {
	return new ContainsPrefixTreeQuery(shape, getFieldName(), grid, detailLevel,
	multiOverlappingIndexedShapes);
	}
	throw new UnsupportedSpatialOperation(op);
	}

	/**
	* A quick check of the shape to see if it is perfectly aligned to a grid.
	* Points always are as they are indivisible. It's okay to return false
	* if the shape actually is aligned; this is an optimization hint.
	*/
	protected boolean isGridAlignedShape(Shape shape) {
	return isPointShape(shape);
	}

	/** {@link #makeQuery(SpatialArgs)} specialized for the query being a grid square. */
	protected Query makeGridShapeIntersectsQuery(Shape gridShape) {
	assert isGridAlignedShape(gridShape);
	if (isPointsOnly()) {
	// Awesome; this will be equivalent to a TermQuery.
	Iterator<Cell> cellIterator = grid.getTreeCellIterator(gridShape, grid.getMaxLevels());
	// get last cell
	Cell cell = cellIterator.next();
	while (cellIterator.hasNext()) {
	int prevLevel = cell.getLevel();
	cell = cellIterator.next();
	assert prevLevel < cell.getLevel();
	}
	assert cell.isLeaf();
	return new TermQuery(new Term(getFieldName(), cell.getTokenBytesWithLeaf(null)));
	} else {
	// Well there could be parent cells. But we can reduce the "scan level" which will be slower for a point query.
	// TODO: AVPTQ will still scan the bottom nonetheless; file an issue to eliminate that
	return new IntersectsPrefixTreeQuery(
	gridShape, getFieldName(), grid, getGrid().getMaxLevels(), getGrid().getMaxLevels() + 1);
	}
	}
	}