lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/tree/QuadPrefixTree.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.tree;

 import java.io.PrintStream;
 import java.text.NumberFormat;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Locale;

 import org.apache.lucene.util.Version;
 import org.locationtech.spatial4j.context.SpatialContext;
 import org.locationtech.spatial4j.shape.Point;
 import org.locationtech.spatial4j.shape.Rectangle;
 import org.locationtech.spatial4j.shape.Shape;
 import org.locationtech.spatial4j.shape.SpatialRelation;
 import org.apache.lucene.util.BytesRef;

 /**
  * A {@link SpatialPrefixTree} which uses a
  * <a href="http://en.wikipedia.org/wiki/Quadtree">quad tree</a> in which an
  * indexed term will be generated for each cell, 'A', 'B', 'C', 'D'.
  *
  * @lucene.experimental
  */
 public class QuadPrefixTree extends LegacyPrefixTree {

   /**
    * Factory for creating {@link QuadPrefixTree} instances with useful defaults
    */
   public static class Factory extends SpatialPrefixTreeFactory {
     @Override
     protected int getLevelForDistance(double degrees) {
       return newSPT().getLevelForDistance(degrees);
     }

     @Override
     protected SpatialPrefixTree newSPT() {
       QuadPrefixTree tree = new QuadPrefixTree(ctx,
           maxLevels != null ? maxLevels : MAX_LEVELS_POSSIBLE);
       tree.robust = getVersion().onOrAfter(Version.LUCENE_8_3_0);
       return tree;
     }
   }

   public static final int MAX_LEVELS_POSSIBLE = 50;//not really sure how big this should be

   public static final int DEFAULT_MAX_LEVELS = 12;
   protected final double xmin;
   protected final double xmax;
   protected final double ymin;
   protected final double ymax;
   protected final double xmid;
   protected final double ymid;

   protected final double gridW;
   public final double gridH;

   final double[] levelW;
   final double[] levelH;

   protected boolean robust = true; // for backward compatibility, use the old method if user specified old version.

   public QuadPrefixTree(
       SpatialContext ctx, Rectangle bounds, int maxLevels) {
     super(ctx, maxLevels);
     this.xmin = bounds.getMinX();
     this.xmax = bounds.getMaxX();
     this.ymin = bounds.getMinY();
     this.ymax = bounds.getMaxY();

     levelW = new double[maxLevels + 1];
     levelH = new double[maxLevels + 1];

     gridW = xmax - xmin;
     gridH = ymax - ymin;
     this.xmid = xmin + gridW/2.0;
     this.ymid = ymin + gridH/2.0;
     levelW[0] = gridW/2.0;
     levelH[0] = gridH/2.0;

     for (int i = 1; i < levelW.length; i++) {
       levelW[i] = levelW[i - 1] / 2.0;
       levelH[i] = levelH[i - 1] / 2.0;
     }
   }

   public QuadPrefixTree(SpatialContext ctx) {
     this(ctx, DEFAULT_MAX_LEVELS);
   }

   public QuadPrefixTree(
       SpatialContext ctx, int maxLevels) {
     this(ctx, ctx.getWorldBounds(), maxLevels);
   }

   @Override
   public Cell getWorldCell() {
     return new QuadCell(BytesRef.EMPTY_BYTES, 0, 0);
   }

   public void printInfo(PrintStream out) {
     NumberFormat nf = NumberFormat.getNumberInstance(Locale.ROOT);
     nf.setMaximumFractionDigits(5);
     nf.setMinimumFractionDigits(5);
     nf.setMinimumIntegerDigits(3);

     for (int i = 0; i < maxLevels; i++) {
       out.println(i + "]\t" + nf.format(levelW[i]) + "\t" + nf.format(levelH[i]));
     }
   }

   @Override
   public int getLevelForDistance(double dist) {
     if (dist == 0)//short circuit
       return maxLevels;
     for (int i = 0; i < maxLevels-1; i++) {
       //note: level[i] is actually a lookup for level i+1
       if(dist > levelW[i] && dist > levelH[i]) {
         return i+1;
       }
     }
     return maxLevels;
   }

   @Override
   public Cell getCell(Point p, int level) {
     if (!robust) { // old method
       List<Cell> cells = new ArrayList<>(1);
       buildNotRobustly(xmid, ymid, 0, cells, new BytesRef(maxLevels+1), ctx.makePoint(p.getX(),p.getY()), level);
       if (!cells.isEmpty()) {
         return cells.get(0);//note cells could be longer if p on edge
       }
     }

     double currentXmid = xmid;
     double currentYmid = ymid;
     double xp = p.getX();
     double yp = p.getY();
     BytesRef str = new BytesRef(maxLevels+1);
     int levelLimit = level > maxLevels ? maxLevels : level;
     SpatialRelation rel = SpatialRelation.CONTAINS;
     for (int lvl = 0; lvl < levelLimit; lvl++){
       int c = battenberg(currentXmid, currentYmid, xp, yp);
       double halfWidth = levelW[lvl + 1];
       double halfHeight = levelH[lvl + 1];
       switch(c){
         case 0:
           currentXmid -= halfWidth;
           currentYmid += halfHeight;
           break;
         case 1:
           currentXmid += halfWidth;
           currentYmid += halfHeight;
           break;
         case 2:
           currentXmid -= halfWidth;
           currentYmid -= halfHeight;
           break;
         case 3:
           currentXmid += halfWidth;
           currentYmid -= halfHeight;
           break;
         default:
       }
       str.bytes[str.length++] = (byte)('A' + c);
     }
     return new QuadCell(str, rel);
   }

   private void buildNotRobustly(
       double x,
       double y,
       int level,
       List<Cell> matches,
       BytesRef str,
       Shape shape,
       int maxLevel) {
     assert str.length == level;
     double w = levelW[level] / 2;
     double h = levelH[level] / 2;

     // Z-Order
     // http://en.wikipedia.org/wiki/Z-order_%28curve%29
     checkBattenbergNotRobustly('A', x - w, y + h, level, matches, str, shape, maxLevel);
     checkBattenbergNotRobustly('B', x + w, y + h, level, matches, str, shape, maxLevel);
     checkBattenbergNotRobustly('C', x - w, y - h, level, matches, str, shape, maxLevel);
     checkBattenbergNotRobustly('D', x + w, y - h, level, matches, str, shape, maxLevel);

     // possibly consider hilbert curve
     // http://en.wikipedia.org/wiki/Hilbert_curve
     // http://blog.notdot.net/2009/11/Damn-Cool-Algorithms-Spatial-indexing-with-Quadtrees-and-Hilbert-Curves
     // if we actually use the range property in the query, this could be useful
   }

   protected void checkBattenbergNotRobustly(
       char c,
       double cx,
       double cy,
       int level,
       List<Cell> matches,
       BytesRef str,
       Shape shape,
       int maxLevel) {
     assert str.length == level;
     assert str.offset == 0;
     double w = levelW[level] / 2;
     double h = levelH[level] / 2;

     int strlen = str.length;
     Rectangle rectangle = ctx.makeRectangle(cx - w, cx + w, cy - h, cy + h);
     SpatialRelation v = shape.relate(rectangle);
     if (SpatialRelation.CONTAINS == v) {
       str.bytes[str.length++] = (byte)c;//append
       //str.append(SpatialPrefixGrid.COVER);
       matches.add(new QuadCell(BytesRef.deepCopyOf(str), v.transpose()));
     } else if (SpatialRelation.DISJOINT == v) {
       // nothing
     } else { // SpatialRelation.WITHIN, SpatialRelation.INTERSECTS
       str.bytes[str.length++] = (byte)c;//append

       int nextLevel = level+1;
       if (nextLevel >= maxLevel) {
         //str.append(SpatialPrefixGrid.INTERSECTS);
         matches.add(new QuadCell(BytesRef.deepCopyOf(str), v.transpose()));
       } else {
         buildNotRobustly(cx, cy, nextLevel, matches, str, shape, maxLevel);
       }
     }
     str.length = strlen;
   }

   /** Returns a Z-Order quadrant [0-3]. */
   protected int battenberg(double xmid, double ymid, double xp, double yp){
     // http://en.wikipedia.org/wiki/Z-order_%28curve%29
     if (ymid <= yp) {
       if (xmid >= xp) {
         return 0;
       }
       return 1;
     } else {
       if (xmid >= xp) {
         return 2;
       }
       return 3;
     }
     // possibly consider hilbert curve
     // http://en.wikipedia.org/wiki/Hilbert_curve
     // http://blog.notdot.net/2009/11/Damn-Cool-Algorithms-Spatial-indexing-with-Quadtrees-and-Hilbert-Curves
     // if we actually use the range property in the query, this could be useful
   }

   protected class QuadCell extends LegacyCell {

     QuadCell(byte[] bytes, int off, int len) {
       super(bytes, off, len);
     }

     QuadCell(BytesRef str, SpatialRelation shapeRel) {
       this(str.bytes, str.offset, str.length);
       this.shapeRel = shapeRel;
     }

     @Override
     protected QuadPrefixTree getGrid() { return QuadPrefixTree.this; }

     @Override
     protected int getMaxLevels() { return maxLevels; }

     @Override
     protected Collection<Cell> getSubCells() {
       BytesRef source = getTokenBytesNoLeaf(null);

       List<Cell> cells = new ArrayList<>(4);
       cells.add(new QuadCell(concat(source, (byte)'A'), null));
       cells.add(new QuadCell(concat(source, (byte)'B'), null));
       cells.add(new QuadCell(concat(source, (byte)'C'), null));
       cells.add(new QuadCell(concat(source, (byte)'D'), null));
       return cells;
     }

     protected BytesRef concat(BytesRef source, byte b) {
       //+2 for new char + potential leaf
       final byte[] buffer = new byte[source.length + 2];
       System.arraycopy(source.bytes, source.offset, buffer, 0, source.length);
       BytesRef target = new BytesRef(buffer);
       target.length = source.length;
       target.bytes[target.length++] = b;
       return target;
     }

     @Override
     public int getSubCellsSize() {
       return 4;
     }

     @Override
     protected QuadCell getSubCell(Point p) {
       return (QuadCell) QuadPrefixTree.this.getCell(p, getLevel() + 1);//not performant!
     }

     @Override
     public Shape getShape() {
       if (shape == null)
         shape = makeShape();
       return shape;
     }

     protected Rectangle makeShape() {
       BytesRef token = getTokenBytesNoLeaf(null);
       double xmin = QuadPrefixTree.this.xmin;
       double ymin = QuadPrefixTree.this.ymin;

       for (int i = 0; i < token.length; i++) {
         byte c = token.bytes[token.offset + i];
         switch (c) {
           case 'A':
             ymin += levelH[i];
             break;
           case 'B':
             xmin += levelW[i];
             ymin += levelH[i];
             break;
           case 'C':
             break;//nothing really
           case 'D':
             xmin += levelW[i];
             break;
           default:
             throw new RuntimeException("unexpected char: " + c);
         }
       }
       int len = token.length;
       double width, height;
       if (len > 0) {
         width = levelW[len-1];
         height = levelH[len-1];
       } else {
         width = gridW;
         height = gridH;
       }
       return ctx.makeRectangle(xmin, xmin + width, ymin, ymin + height);
     }
   }//QuadCell
 }
	/*
	* 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 java.io.PrintStream;
	import java.text.NumberFormat;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.List;
	import java.util.Locale;

	import org.apache.lucene.util.Version;
	import org.locationtech.spatial4j.context.SpatialContext;
	import org.locationtech.spatial4j.shape.Point;
	import org.locationtech.spatial4j.shape.Rectangle;
	import org.locationtech.spatial4j.shape.Shape;
	import org.locationtech.spatial4j.shape.SpatialRelation;
	import org.apache.lucene.util.BytesRef;

	/**
	* A {@link SpatialPrefixTree} which uses a
	* <a href="http://en.wikipedia.org/wiki/Quadtree">quad tree</a> in which an
	* indexed term will be generated for each cell, 'A', 'B', 'C', 'D'.
	*
	* @lucene.experimental
	*/
	public class QuadPrefixTree extends LegacyPrefixTree {

	/**
	* Factory for creating {@link QuadPrefixTree} instances with useful defaults
	*/
	public static class Factory extends SpatialPrefixTreeFactory {
	@Override
	protected int getLevelForDistance(double degrees) {
	return newSPT().getLevelForDistance(degrees);
	}

	@Override
	protected SpatialPrefixTree newSPT() {
	QuadPrefixTree tree = new QuadPrefixTree(ctx,
	maxLevels != null ? maxLevels : MAX_LEVELS_POSSIBLE);
	tree.robust = getVersion().onOrAfter(Version.LUCENE_8_3_0);
	return tree;
	}
	}

	public static final int MAX_LEVELS_POSSIBLE = 50;//not really sure how big this should be

	public static final int DEFAULT_MAX_LEVELS = 12;
	protected final double xmin;
	protected final double xmax;
	protected final double ymin;
	protected final double ymax;
	protected final double xmid;
	protected final double ymid;

	protected final double gridW;
	public final double gridH;

	final double[] levelW;
	final double[] levelH;

	protected boolean robust = true; // for backward compatibility, use the old method if user specified old version.

	public QuadPrefixTree(
	SpatialContext ctx, Rectangle bounds, int maxLevels) {
	super(ctx, maxLevels);
	this.xmin = bounds.getMinX();
	this.xmax = bounds.getMaxX();
	this.ymin = bounds.getMinY();
	this.ymax = bounds.getMaxY();

	levelW = new double[maxLevels + 1];
	levelH = new double[maxLevels + 1];

	gridW = xmax - xmin;
	gridH = ymax - ymin;
	this.xmid = xmin + gridW/2.0;
	this.ymid = ymin + gridH/2.0;
	levelW[0] = gridW/2.0;
	levelH[0] = gridH/2.0;

	for (int i = 1; i < levelW.length; i++) {
	levelW[i] = levelW[i - 1] / 2.0;
	levelH[i] = levelH[i - 1] / 2.0;
	}
	}

	public QuadPrefixTree(SpatialContext ctx) {
	this(ctx, DEFAULT_MAX_LEVELS);
	}

	public QuadPrefixTree(
	SpatialContext ctx, int maxLevels) {
	this(ctx, ctx.getWorldBounds(), maxLevels);
	}

	@Override
	public Cell getWorldCell() {
	return new QuadCell(BytesRef.EMPTY_BYTES, 0, 0);
	}

	public void printInfo(PrintStream out) {
	NumberFormat nf = NumberFormat.getNumberInstance(Locale.ROOT);
	nf.setMaximumFractionDigits(5);
	nf.setMinimumFractionDigits(5);
	nf.setMinimumIntegerDigits(3);

	for (int i = 0; i < maxLevels; i++) {
	out.println(i + "]\t" + nf.format(levelW[i]) + "\t" + nf.format(levelH[i]));
	}
	}

	@Override
	public int getLevelForDistance(double dist) {
	if (dist == 0)//short circuit
	return maxLevels;
	for (int i = 0; i < maxLevels-1; i++) {
	//note: level[i] is actually a lookup for level i+1
	if(dist > levelW[i] && dist > levelH[i]) {
	return i+1;
	}
	}
	return maxLevels;
	}

	@Override
	public Cell getCell(Point p, int level) {
	if (!robust) { // old method
	List<Cell> cells = new ArrayList<>(1);
	buildNotRobustly(xmid, ymid, 0, cells, new BytesRef(maxLevels+1), ctx.makePoint(p.getX(),p.getY()), level);
	if (!cells.isEmpty()) {
	return cells.get(0);//note cells could be longer if p on edge
	}
	}

	double currentXmid = xmid;
	double currentYmid = ymid;
	double xp = p.getX();
	double yp = p.getY();
	BytesRef str = new BytesRef(maxLevels+1);
	int levelLimit = level > maxLevels ? maxLevels : level;
	SpatialRelation rel = SpatialRelation.CONTAINS;
	for (int lvl = 0; lvl < levelLimit; lvl++){
	int c = battenberg(currentXmid, currentYmid, xp, yp);
	double halfWidth = levelW[lvl + 1];
	double halfHeight = levelH[lvl + 1];
	switch(c){
	case 0:
	currentXmid -= halfWidth;
	currentYmid += halfHeight;
	break;
	case 1:
	currentXmid += halfWidth;
	currentYmid += halfHeight;
	break;
	case 2:
	currentXmid -= halfWidth;
	currentYmid -= halfHeight;
	break;
	case 3:
	currentXmid += halfWidth;
	currentYmid -= halfHeight;
	break;
	default:
	}
	str.bytes[str.length++] = (byte)('A' + c);
	}
	return new QuadCell(str, rel);
	}

	private void buildNotRobustly(
	double x,
	double y,
	int level,
	List<Cell> matches,
	BytesRef str,
	Shape shape,
	int maxLevel) {
	assert str.length == level;
	double w = levelW[level] / 2;
	double h = levelH[level] / 2;

	// Z-Order
	// http://en.wikipedia.org/wiki/Z-order_%28curve%29
	checkBattenbergNotRobustly('A', x - w, y + h, level, matches, str, shape, maxLevel);
	checkBattenbergNotRobustly('B', x + w, y + h, level, matches, str, shape, maxLevel);
	checkBattenbergNotRobustly('C', x - w, y - h, level, matches, str, shape, maxLevel);
	checkBattenbergNotRobustly('D', x + w, y - h, level, matches, str, shape, maxLevel);

	// possibly consider hilbert curve
	// http://en.wikipedia.org/wiki/Hilbert_curve
	// http://blog.notdot.net/2009/11/Damn-Cool-Algorithms-Spatial-indexing-with-Quadtrees-and-Hilbert-Curves
	// if we actually use the range property in the query, this could be useful
	}

	protected void checkBattenbergNotRobustly(
	char c,
	double cx,
	double cy,
	int level,
	List<Cell> matches,
	BytesRef str,
	Shape shape,
	int maxLevel) {
	assert str.length == level;
	assert str.offset == 0;
	double w = levelW[level] / 2;
	double h = levelH[level] / 2;

	int strlen = str.length;
	Rectangle rectangle = ctx.makeRectangle(cx - w, cx + w, cy - h, cy + h);
	SpatialRelation v = shape.relate(rectangle);
	if (SpatialRelation.CONTAINS == v) {
	str.bytes[str.length++] = (byte)c;//append
	//str.append(SpatialPrefixGrid.COVER);
	matches.add(new QuadCell(BytesRef.deepCopyOf(str), v.transpose()));
	} else if (SpatialRelation.DISJOINT == v) {
	// nothing
	} else { // SpatialRelation.WITHIN, SpatialRelation.INTERSECTS
	str.bytes[str.length++] = (byte)c;//append

	int nextLevel = level+1;
	if (nextLevel >= maxLevel) {
	//str.append(SpatialPrefixGrid.INTERSECTS);
	matches.add(new QuadCell(BytesRef.deepCopyOf(str), v.transpose()));
	} else {
	buildNotRobustly(cx, cy, nextLevel, matches, str, shape, maxLevel);
	}
	}
	str.length = strlen;
	}

	/** Returns a Z-Order quadrant [0-3]. */
	protected int battenberg(double xmid, double ymid, double xp, double yp){
	// http://en.wikipedia.org/wiki/Z-order_%28curve%29
	if (ymid <= yp) {
	if (xmid >= xp) {
	return 0;
	}
	return 1;
	} else {
	if (xmid >= xp) {
	return 2;
	}
	return 3;
	}
	// possibly consider hilbert curve
	// http://en.wikipedia.org/wiki/Hilbert_curve
	// http://blog.notdot.net/2009/11/Damn-Cool-Algorithms-Spatial-indexing-with-Quadtrees-and-Hilbert-Curves
	// if we actually use the range property in the query, this could be useful
	}

	protected class QuadCell extends LegacyCell {

	QuadCell(byte[] bytes, int off, int len) {
	super(bytes, off, len);
	}

	QuadCell(BytesRef str, SpatialRelation shapeRel) {
	this(str.bytes, str.offset, str.length);
	this.shapeRel = shapeRel;
	}

	@Override
	protected QuadPrefixTree getGrid() { return QuadPrefixTree.this; }

	@Override
	protected int getMaxLevels() { return maxLevels; }

	@Override
	protected Collection<Cell> getSubCells() {
	BytesRef source = getTokenBytesNoLeaf(null);

	List<Cell> cells = new ArrayList<>(4);
	cells.add(new QuadCell(concat(source, (byte)'A'), null));
	cells.add(new QuadCell(concat(source, (byte)'B'), null));
	cells.add(new QuadCell(concat(source, (byte)'C'), null));
	cells.add(new QuadCell(concat(source, (byte)'D'), null));
	return cells;
	}

	protected BytesRef concat(BytesRef source, byte b) {
	//+2 for new char + potential leaf
	final byte[] buffer = new byte[source.length + 2];
	System.arraycopy(source.bytes, source.offset, buffer, 0, source.length);
	BytesRef target = new BytesRef(buffer);
	target.length = source.length;
	target.bytes[target.length++] = b;
	return target;
	}

	@Override
	public int getSubCellsSize() {
	return 4;
	}

	@Override
	protected QuadCell getSubCell(Point p) {
	return (QuadCell) QuadPrefixTree.this.getCell(p, getLevel() + 1);//not performant!
	}

	@Override
	public Shape getShape() {
	if (shape == null)
	shape = makeShape();
	return shape;
	}

	protected Rectangle makeShape() {
	BytesRef token = getTokenBytesNoLeaf(null);
	double xmin = QuadPrefixTree.this.xmin;
	double ymin = QuadPrefixTree.this.ymin;

	for (int i = 0; i < token.length; i++) {
	byte c = token.bytes[token.offset + i];
	switch (c) {
	case 'A':
	ymin += levelH[i];
	break;
	case 'B':
	xmin += levelW[i];
	ymin += levelH[i];
	break;
	case 'C':
	break;//nothing really
	case 'D':
	xmin += levelW[i];
	break;
	default:
	throw new RuntimeException("unexpected char: " + c);
	}
	}
	int len = token.length;
	double width, height;
	if (len > 0) {
	width = levelW[len-1];
	height = levelH[len-1];
	} else {
	width = gridW;
	height = gridH;
	}
	return ctx.makeRectangle(xmin, xmin + width, ymin, ymin + height);
	}
	}//QuadCell
	}