diff --git a/src/main/java/org/apache/lucene/spatial/prefix/tree/PackedQuadPrefixTree.java b/src/main/java/org/apache/lucene/spatial/prefix/tree/PackedQuadPrefixTree.java index d4ffec7772afb..65808c041e37b 100644 --- a/src/main/java/org/apache/lucene/spatial/prefix/tree/PackedQuadPrefixTree.java +++ b/src/main/java/org/apache/lucene/spatial/prefix/tree/PackedQuadPrefixTree.java @@ -44,6 +44,8 @@ public class PackedQuadPrefixTree extends QuadPrefixTree { public static final byte[] QUAD = new byte[] {0x00, 0x01, 0x02, 0x03}; public static final int MAX_LEVELS_POSSIBLE = 29; + private boolean leafyPrune = true; + public static class Factory extends QuadPrefixTree.Factory { @Override protected SpatialPrefixTree newSPT() { @@ -121,6 +123,10 @@ public CellIterator getTreeCellIterator(Shape shape, int detailLevel) { return new PrefixTreeIterator(shape); } + public void setPruneLeafyBranches( boolean pruneLeafyBranches ) { + this.leafyPrune = pruneLeafyBranches; + } + /** * PackedQuadCell Binary Representation is as follows * CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCDDDDDL @@ -129,7 +135,7 @@ public CellIterator getTreeCellIterator(Shape shape, int detailLevel) { * D = Depth bits (5 with max of 29 levels) * L = isLeaf bit */ - public class PackedQuadCell extends QuadPrefixTree.QuadCell { + public class PackedQuadCell extends QuadCell { private long term; PackedQuadCell(long term) { @@ -170,7 +176,7 @@ public boolean isEnd(final int level, final int shift) { * setting this to false will step to the next sibling. * Note: This complies with lexicographical ordering, once you've moved to the next sibling there is no backtracking. */ - public Cell nextCell(boolean descend) { + public PackedQuadCell nextCell(boolean descend) { final int level = getLevel(); final int shift = getShiftForLevel(level); // base case: can't go further @@ -248,7 +254,7 @@ protected Collection getSubCells() { } @Override - protected QuadPrefixTree.QuadCell getSubCell(Point p) { + protected QuadCell getSubCell(Point p) { return (PackedQuadCell) PackedQuadPrefixTree.this.getCell(p, getLevel() + 1);//not performant! } @@ -348,11 +354,18 @@ public String toString() { protected class PrefixTreeIterator extends CellIterator { private Shape shape; + private PackedQuadCell thisCell; + private PackedQuadCell nextCell; + + private short leaves; + private short level; + private final short maxLevels; + private CellIterator pruneIter; PrefixTreeIterator(Shape shape) { - super(); this.shape = shape; this.thisCell = ((PackedQuadCell)(getWorldCell())).nextCell(true); + this.maxLevels = (short)thisCell.getMaxLevels(); this.nextCell = null; } @@ -366,31 +379,41 @@ public boolean hasNext() { while (thisCell != null) { rel = thisCell.getShape().relate(shape); if (rel == SpatialRelation.DISJOINT) { - thisCell = ((PackedQuadCell) thisCell).nextCell(false); - } else if (rel == SpatialRelation.INTERSECTS || rel == SpatialRelation.CONTAINS) { + thisCell = thisCell.nextCell(false); + } else { // within || intersects || contains thisCell.setShapeRel(rel); nextCell = thisCell; - if (thisCell.getLevel() == getMaxLevels()) { + if (rel == SpatialRelation.WITHIN) { thisCell.setLeaf(); - if (shape instanceof Point) { - thisCell.setShapeRel(SpatialRelation.WITHIN); - thisCell = null; + thisCell = thisCell.nextCell(false); + } else { // intersects || contains + level = (short) (thisCell.getLevel()); + if (level == maxLevels || pruned(rel)) { + thisCell.setLeaf(); + if (shape instanceof Point) { + thisCell.setShapeRel(SpatialRelation.WITHIN); + thisCell = null; + } else { + thisCell = thisCell.nextCell(false); + } break; } + thisCell = thisCell.nextCell(true); } - thisCell = ((PackedQuadCell) thisCell).nextCell(true); - break; - } else if (rel == SpatialRelation.WITHIN ) { - thisCell.setLeaf(); - thisCell.setShapeRel(rel); - nextCell = thisCell; - thisCell = ((PackedQuadCell) thisCell).nextCell(false); break; } } return nextCell != null; } + private boolean pruned(SpatialRelation rel) { + if (rel == SpatialRelation.INTERSECTS && leafyPrune && level == maxLevels-1) { + for (leaves=0, pruneIter=thisCell.getNextLevelCells(shape); pruneIter.hasNext(); pruneIter.next(), ++leaves); + return leaves == 4; + } + return false; + } + @Override public Cell next() { if (nextCell == null) { @@ -409,5 +432,4 @@ public void remove() { //no-op } } - }