LUCENE-8190: Specialized cell interface to allow any spatial prefix tree to benefit from the setting setPruneLeafyBranches on RecursivePrefixTreeStrategy

lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/RecursivePrefixTreeStrategy.java

@@ -24,8 +24,8 @@ 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.LegacyCell;
 import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
 import org.apache.lucene.spatial.query.SpatialArgs;
 import org.apache.lucene.spatial.query.SpatialOperation;
@@ -53,7 +53,7 @@ public class RecursivePrefixTreeStrategy extends PrefixTreeStrategy {
   protected int prefixGridScanLevel;
 
   //Formerly known as simplifyIndexedCells. Eventually will be removed. Only compatible with RPT
-  // and a LegacyPrefixTree.
+  // and cells implementing CellCanPrune, otherwise ignored.
   protected boolean pruneLeafyBranches = true;
 
   protected boolean multiOverlappingIndexedShapes = true;
@@ -93,8 +93,10 @@ public class RecursivePrefixTreeStrategy extends PrefixTreeStrategy {
   }
 
   /**
-   * An optional hint affecting non-point shapes: it will
-   * prune away a complete set sibling leaves to their parent (recursively), resulting in ~20-50%
+   * 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
@@ -132,10 +134,6 @@ public class RecursivePrefixTreeStrategy extends PrefixTreeStrategy {
 
   /** 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) {
-    // Important: this logic assumes Cells don't share anything with other cells when
-    // calling cell.getNextLevelCells(). This is only true for LegacyCell.
-    if (!(cell instanceof LegacyCell))
-      throw new IllegalStateException("pruneLeafyBranches must be disabled for use with grid "+grid);
 
     if (cell.getLevel() == detailLevel) {
       cell.setLeaf();//FYI might already be a leaf
@@ -154,8 +152,14 @@ public class RecursivePrefixTreeStrategy extends PrefixTreeStrategy {
       if (recursiveTraverseAndPrune(subCell, shape, detailLevel, result))
         leaves++;
     }
+
+    if (!(cell instanceof CellCanPrune)) {
+      //Cannot prune so return false
+      return false;
+    }
+
     //can we prune?
-    if (leaves == ((LegacyCell)cell).getSubCellsSize() && cell.getLevel() != 0) {
+    if (leaves == ((CellCanPrune)cell).getSubCellsSize() && cell.getLevel() != 0) {
       //Optimization: substitute the parent as a leaf instead of adding all
       // children as leaves
 

lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/tree/CellCanPrune.java

@@ -0,0 +1,37 @@
+/*
+ * 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 org.locationtech.spatial4j.shape.Shape;
+
+/**
+ *  Grid cells that share nothing with other cells when calling {@link #getNextLevelCells(Shape)}
+ *  might implement this interface. Children cells for this cell will be eligible for pruning via
+ *  {@link org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy#setPruneLeafyBranches(boolean)}.
+ *
+ * @lucene.experimental
+ */
+public interface CellCanPrune extends Cell{
+
+  /**
+   * Returns the number of children for this cell.
+   *
+   * @return the number of children.
+   */
+  int getSubCellsSize();
+}

lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/tree/LegacyCell.java

@@ -27,7 +27,7 @@ import org.apache.lucene.util.StringHelper;
 /** The base for the original two SPT's: Geohash and Quad. Don't subclass this for new SPTs.
  * @lucene.internal */
 //public for RPT pruneLeafyBranches code
-public abstract class LegacyCell implements Cell {
+public abstract class LegacyCell implements CellCanPrune {
 
   // Important: A LegacyCell doesn't share state for getNextLevelCells(), and
   //  LegacySpatialPrefixTree assumes this in its simplify tree logic.
@@ -159,11 +159,6 @@ public abstract class LegacyCell implements Cell {
    */
   protected abstract Collection<Cell> getSubCells();
 
-  /**
-   * {@link #getSubCells()}.size() -- usually a constant. Should be &gt;=2
-   */
-  public abstract int getSubCellsSize();
-
   @Override
   public boolean isPrefixOf(Cell c) {
     //Note: this only works when each level uses a whole number of bytes.

lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/tree/S2PrefixTreeCell.java

@@ -33,7 +33,7 @@ import org.locationtech.spatial4j.shape.SpatialRelation;
  *
  * @lucene.internal
  */
-class S2PrefixTreeCell implements Cell {
+class S2PrefixTreeCell implements CellCanPrune {
 
     //Faces of S2 Geometry
     private static S2CellId[] FACES = new S2CellId[6];
@@ -265,6 +265,14 @@ class S2PrefixTreeCell implements Cell {
         bref.offset = 0;
     }
 
+    @Override
+    public int getSubCellsSize() {
+        if (cellId == null) { //root node
+            return 6;
+        }
+        return  (int) Math.pow(4, tree.arity);
+    }
+
     @Override
     public int hashCode() {
         if (cellId == null) {

lucene/spatial-extras/src/test/org/apache/lucene/spatial/spatial4j/Geo3dRptTest.java

@@ -56,17 +56,14 @@ public class Geo3dRptTest extends RandomSpatialOpStrategyTestCase {
     int type = random().nextInt(4);
     if (type == 0) {
       this.grid = new GeohashPrefixTree(ctx, 2);
-      this.rptStrategy = newRPT();
     } else if (type == 1) {
       this.grid = new QuadPrefixTree(ctx, 5);
-      this.rptStrategy = newRPT();
-    }
-    else {
+    } else {
       int arity = random().nextInt(3) + 1;
       this.grid = new S2PrefixTree(ctx, 5 - arity, arity);
-      this.rptStrategy = newRPT();
-      this.rptStrategy.setPruneLeafyBranches(false);
     }
+    this.rptStrategy = newRPT();
+    this.rptStrategy.setPruneLeafyBranches(random().nextBoolean());
   }
 
   protected RecursivePrefixTreeStrategy newRPT() {