lucene/spatial-extras/src/test/org/apache/lucene/spatial/prefix/TestHeatmapFacetCounter.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 static com.carrotsearch.randomizedtesting.RandomizedTest.atMost;
 import static com.carrotsearch.randomizedtesting.RandomizedTest.randomIntBetween;

 import com.carrotsearch.randomizedtesting.annotations.Repeat;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHitCountCollector;
 import org.apache.lucene.spatial.StrategyTestCase;
 import org.apache.lucene.spatial.prefix.tree.QuadPrefixTree;
 import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
 import org.apache.lucene.util.Bits;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.locationtech.spatial4j.context.SpatialContext;
 import org.locationtech.spatial4j.context.SpatialContextFactory;
 import org.locationtech.spatial4j.distance.DistanceUtils;
 import org.locationtech.spatial4j.shape.Circle;
 import org.locationtech.spatial4j.shape.Point;
 import org.locationtech.spatial4j.shape.Rectangle;
 import org.locationtech.spatial4j.shape.Shape;
 import org.locationtech.spatial4j.shape.ShapeFactory;
 import org.locationtech.spatial4j.shape.SpatialRelation;
 import org.locationtech.spatial4j.shape.impl.RectangleImpl;

 public class TestHeatmapFacetCounter extends StrategyTestCase {

   SpatialPrefixTree grid;
   ShapeFactory shapeFactory;

   int cellsValidated;
   int cellValidatedNonZero;

   @Before
   public void setUp() throws Exception {
     super.setUp();
     cellsValidated = cellValidatedNonZero = 0;
     ctx = SpatialContext.GEO;
     shapeFactory = ctx.getShapeFactory();
     grid = new QuadPrefixTree(ctx, randomIntBetween(1, 8));
     strategy = new RecursivePrefixTreeStrategy(grid, getTestClass().getSimpleName());
     if (rarely()) {
       ((PrefixTreeStrategy) strategy).setPointsOnly(true);
     }
   }

   @After
   public void after() {
     log.info(
         "Validated " + cellsValidated + " cells, " + cellValidatedNonZero + " non-zero"); // nowarn
   }

   @Test
   public void testStatic() throws IOException {
     // Some specific tests (static, not random).
     adoc("0", shapeFactory.rect(179.8, -170, -90, -80)); // barely crosses equator
     adoc("1", shapeFactory.pointXY(-180, -85)); // a pt within the above rect
     adoc("2", shapeFactory.pointXY(172, -85)); // a pt to left of rect
     commit();

     validateHeatmapResultLoop(shapeFactory.rect(+170, +180, -90, -85), 1, 100);
     validateHeatmapResultLoop(shapeFactory.rect(-180, -160, -89, -50), 1, 100);
     validateHeatmapResultLoop(shapeFactory.rect(179, 179, -89, -50), 1, 100); // line
     // We could test anything and everything at this point... I prefer we leave that to random
     // testing and then
     // add specific tests if we find a bug.
   }

   @Test
   public void testLucene7291Dateline() throws IOException {
     grid = new QuadPrefixTree(ctx, 2); // only 2, and we wind up with some big leaf cells
     strategy = new RecursivePrefixTreeStrategy(grid, getTestClass().getSimpleName());
     adoc("0", shapeFactory.rect(-102, -83, 43, 52));
     commit();
     validateHeatmapResultLoop(shapeFactory.rect(179, -179, 62, 63), 2, 100); // HM crosses dateline
   }

   @Test
   public void testQueryCircle() throws IOException {
     // overwrite setUp; non-geo bounds is more straight-forward; otherwise 88,88 would actually be
     // practically north,
     final SpatialContextFactory spatialContextFactory = new SpatialContextFactory();
     spatialContextFactory.geo = false;
     spatialContextFactory.worldBounds = new RectangleImpl(-90, 90, -90, 90, null);
     ctx = spatialContextFactory.newSpatialContext();
     shapeFactory = ctx.getShapeFactory();
     final int LEVEL = 4;
     grid = new QuadPrefixTree(ctx, LEVEL);
     strategy = new RecursivePrefixTreeStrategy(grid, getTestClass().getSimpleName());
     Circle circle = shapeFactory.circle(0, 0, 89);
     adoc("0", shapeFactory.pointXY(88, 88)); // top-right, inside bbox of circle but not the circle
     adoc("1", shapeFactory.pointXY(0, 0)); // clearly inside; dead center in fact
     commit();
     final HeatmapFacetCounter.Heatmap heatmap =
         HeatmapFacetCounter.calcFacets(
             (PrefixTreeStrategy) strategy,
             indexSearcher.getTopReaderContext(),
             null,
             circle,
             LEVEL,
             1000);
     // assert that only one point is found, not 2
     boolean foundOne = false;
     for (int count : heatmap.counts) {
       switch (count) {
         case 0:
           break;
         case 1:
           assertFalse(foundOne); // this is the first
           foundOne = true;
           break;
         default:
           fail("counts should be 0 or 1: " + count);
       }
     }
     assertTrue(foundOne);
   }

   /**
    * Recursively facet & validate at higher resolutions until we've seen enough. We assume there are
    * some non-zero cells.
    */
   private void validateHeatmapResultLoop(
       Rectangle inputRange, int facetLevel, int cellCountRecursThreshold) throws IOException {
     if (facetLevel > grid.getMaxLevels()) {
       return;
     }
     final int maxCells = 10_000;
     final HeatmapFacetCounter.Heatmap heatmap =
         HeatmapFacetCounter.calcFacets(
             (PrefixTreeStrategy) strategy,
             indexSearcher.getTopReaderContext(),
             null,
             inputRange,
             facetLevel,
             maxCells);
     int preNonZero = cellValidatedNonZero;
     validateHeatmapResult(inputRange, facetLevel, heatmap);
     assert cellValidatedNonZero - preNonZero > 0; // we validated more non-zero cells
     if (heatmap.counts.length < cellCountRecursThreshold) {
       validateHeatmapResultLoop(inputRange, facetLevel + 1, cellCountRecursThreshold);
     }
   }

   @Test
   @Repeat(iterations = 20)
   public void testRandom() throws IOException {
     // Tests using random index shapes & query shapes. This has found all sorts of edge case bugs
     // (e.g. dateline,
     // cell border, overflow(?)).

     final int numIndexedShapes = 1 + atMost(9);
     List<Shape> indexedShapes = new ArrayList<>(numIndexedShapes);
     for (int i = 0; i < numIndexedShapes; i++) {
       indexedShapes.add(randomIndexedShape());
     }

     // Main index loop:
     for (int i = 0; i < indexedShapes.size(); i++) {
       Shape shape = indexedShapes.get(i);
       adoc("" + i, shape);

       if (random().nextInt(10) == 0) {
         commit(); // intermediate commit, produces extra segments
       }
     }
     // delete some documents randomly
     for (int id = 0; id < indexedShapes.size(); id++) {
       if (random().nextInt(10) == 0) {
         deleteDoc("" + id);
         indexedShapes.set(id, null);
       }
     }

     commit();

     // once without dateline wrap
     final Rectangle rect = randomRectangle();
     queryHeatmapRecursive(usually() ? ctx.getWorldBounds() : rect, 1);
     // and once with dateline wrap
     if (rect.getWidth() > 0) {
       double shift = random().nextDouble() % rect.getWidth();
       queryHeatmapRecursive(
           shapeFactory.rect(
               DistanceUtils.normLonDEG(rect.getMinX() - shift),
               DistanceUtils.normLonDEG(rect.getMaxX() - shift),
               rect.getMinY(),
               rect.getMaxY()),
           1);
     }
   }

   /** Build heatmap, validate results, then descend recursively to another facet level. */
   private boolean queryHeatmapRecursive(Rectangle inputRange, int facetLevel) throws IOException {
     if (!inputRange.hasArea()) {
       // Don't test line inputs. It's not that we don't support it but it is more challenging to
       // test if per-chance it
       // coincides with a grid line due due to edge overlap issue for some grid implementations (geo
       // & quad).
       return false;
     }
     Bits filter =
         null; // FYI testing filtering of underlying PrefixTreeFacetCounter is done in another test
     // Calculate facets
     final int maxCells = 10_000;
     final HeatmapFacetCounter.Heatmap heatmap =
         HeatmapFacetCounter.calcFacets(
             (PrefixTreeStrategy) strategy,
             indexSearcher.getTopReaderContext(),
             filter,
             inputRange,
             facetLevel,
             maxCells);

     validateHeatmapResult(inputRange, facetLevel, heatmap);

     boolean foundNonZeroCount = false;
     for (int count : heatmap.counts) {
       if (count > 0) {
         foundNonZeroCount = true;
         break;
       }
     }

     // Test again recursively to higher facetLevel (more detailed cells)
     if (foundNonZeroCount
         && cellsValidated <= 500
         && facetLevel != grid.getMaxLevels()
         && inputRange.hasArea()) {
       for (int i = 0; i < 5; i++) { // try multiple times until we find non-zero counts
         if (queryHeatmapRecursive(randomRectangle(inputRange), facetLevel + 1)) {
           break; // we found data here so we needn't try again
         }
       }
     }
     return foundNonZeroCount;
   }

   private void validateHeatmapResult(
       Rectangle inputRange, int facetLevel, HeatmapFacetCounter.Heatmap heatmap)
       throws IOException {
     final Rectangle heatRect = heatmap.region;
     assertTrue(
         heatRect.relate(inputRange) == SpatialRelation.CONTAINS || heatRect.equals(inputRange));
     final double cellWidth = heatRect.getWidth() / heatmap.columns;
     final double cellHeight = heatRect.getHeight() / heatmap.rows;
     for (int c = 0; c < heatmap.columns; c++) {
       for (int r = 0; r < heatmap.rows; r++) {
         final int facetCount = heatmap.getCount(c, r);
         double x = DistanceUtils.normLonDEG(heatRect.getMinX() + c * cellWidth + cellWidth / 2);
         double y = DistanceUtils.normLatDEG(heatRect.getMinY() + r * cellHeight + cellHeight / 2);
         Point pt = shapeFactory.pointXY(x, y);
         assertEquals(countMatchingDocsAtLevel(pt, facetLevel), facetCount);
       }
     }
   }

   private int countMatchingDocsAtLevel(Point pt, int facetLevel) throws IOException {
     // we use IntersectsPrefixTreeFilter directly so that we can specify the level to go to exactly.
     RecursivePrefixTreeStrategy strategy = (RecursivePrefixTreeStrategy) this.strategy;
     Query filter =
         new IntersectsPrefixTreeQuery(
             pt, strategy.getFieldName(), grid, facetLevel, grid.getMaxLevels());
     final TotalHitCountCollector collector = new TotalHitCountCollector();
     indexSearcher.search(filter, collector);
     cellsValidated++;
     if (collector.getTotalHits() > 0) {
       cellValidatedNonZero++;
     }
     return collector.getTotalHits();
   }

   private Shape randomIndexedShape() {
     if (((PrefixTreeStrategy) strategy).isPointsOnly() || random().nextBoolean()) {
       return randomPoint();
     } else {
       return randomRectangle();
     }
   }
 }
	/*
	* 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 static com.carrotsearch.randomizedtesting.RandomizedTest.atMost;
	import static com.carrotsearch.randomizedtesting.RandomizedTest.randomIntBetween;

	import com.carrotsearch.randomizedtesting.annotations.Repeat;
	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.List;
	import org.apache.lucene.search.Query;
	import org.apache.lucene.search.TotalHitCountCollector;
	import org.apache.lucene.spatial.StrategyTestCase;
	import org.apache.lucene.spatial.prefix.tree.QuadPrefixTree;
	import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
	import org.apache.lucene.util.Bits;
	import org.junit.After;
	import org.junit.Before;
	import org.junit.Test;
	import org.locationtech.spatial4j.context.SpatialContext;
	import org.locationtech.spatial4j.context.SpatialContextFactory;
	import org.locationtech.spatial4j.distance.DistanceUtils;
	import org.locationtech.spatial4j.shape.Circle;
	import org.locationtech.spatial4j.shape.Point;
	import org.locationtech.spatial4j.shape.Rectangle;
	import org.locationtech.spatial4j.shape.Shape;
	import org.locationtech.spatial4j.shape.ShapeFactory;
	import org.locationtech.spatial4j.shape.SpatialRelation;
	import org.locationtech.spatial4j.shape.impl.RectangleImpl;

	public class TestHeatmapFacetCounter extends StrategyTestCase {

	SpatialPrefixTree grid;
	ShapeFactory shapeFactory;

	int cellsValidated;
	int cellValidatedNonZero;

	@Before
	public void setUp() throws Exception {
	super.setUp();
	cellsValidated = cellValidatedNonZero = 0;
	ctx = SpatialContext.GEO;
	shapeFactory = ctx.getShapeFactory();
	grid = new QuadPrefixTree(ctx, randomIntBetween(1, 8));
	strategy = new RecursivePrefixTreeStrategy(grid, getTestClass().getSimpleName());
	if (rarely()) {
	((PrefixTreeStrategy) strategy).setPointsOnly(true);
	}
	}

	@After
	public void after() {
	log.info(
	"Validated " + cellsValidated + " cells, " + cellValidatedNonZero + " non-zero"); // nowarn
	}

	@Test
	public void testStatic() throws IOException {
	// Some specific tests (static, not random).
	adoc("0", shapeFactory.rect(179.8, -170, -90, -80)); // barely crosses equator
	adoc("1", shapeFactory.pointXY(-180, -85)); // a pt within the above rect
	adoc("2", shapeFactory.pointXY(172, -85)); // a pt to left of rect
	commit();

	validateHeatmapResultLoop(shapeFactory.rect(+170, +180, -90, -85), 1, 100);
	validateHeatmapResultLoop(shapeFactory.rect(-180, -160, -89, -50), 1, 100);
	validateHeatmapResultLoop(shapeFactory.rect(179, 179, -89, -50), 1, 100); // line
	// We could test anything and everything at this point... I prefer we leave that to random
	// testing and then
	// add specific tests if we find a bug.
	}

	@Test
	public void testLucene7291Dateline() throws IOException {
	grid = new QuadPrefixTree(ctx, 2); // only 2, and we wind up with some big leaf cells
	strategy = new RecursivePrefixTreeStrategy(grid, getTestClass().getSimpleName());
	adoc("0", shapeFactory.rect(-102, -83, 43, 52));
	commit();
	validateHeatmapResultLoop(shapeFactory.rect(179, -179, 62, 63), 2, 100); // HM crosses dateline
	}

	@Test
	public void testQueryCircle() throws IOException {
	// overwrite setUp; non-geo bounds is more straight-forward; otherwise 88,88 would actually be
	// practically north,
	final SpatialContextFactory spatialContextFactory = new SpatialContextFactory();
	spatialContextFactory.geo = false;
	spatialContextFactory.worldBounds = new RectangleImpl(-90, 90, -90, 90, null);
	ctx = spatialContextFactory.newSpatialContext();
	shapeFactory = ctx.getShapeFactory();
	final int LEVEL = 4;
	grid = new QuadPrefixTree(ctx, LEVEL);
	strategy = new RecursivePrefixTreeStrategy(grid, getTestClass().getSimpleName());
	Circle circle = shapeFactory.circle(0, 0, 89);
	adoc("0", shapeFactory.pointXY(88, 88)); // top-right, inside bbox of circle but not the circle
	adoc("1", shapeFactory.pointXY(0, 0)); // clearly inside; dead center in fact
	commit();
	final HeatmapFacetCounter.Heatmap heatmap =
	HeatmapFacetCounter.calcFacets(
	(PrefixTreeStrategy) strategy,
	indexSearcher.getTopReaderContext(),
	null,
	circle,
	LEVEL,
	1000);
	// assert that only one point is found, not 2
	boolean foundOne = false;
	for (int count : heatmap.counts) {
	switch (count) {
	case 0:
	break;
	case 1:
	assertFalse(foundOne); // this is the first
	foundOne = true;
	break;
	default:
	fail("counts should be 0 or 1: " + count);
	}
	}
	assertTrue(foundOne);
	}

	/**
	* Recursively facet & validate at higher resolutions until we've seen enough. We assume there are
	* some non-zero cells.
	*/
	private void validateHeatmapResultLoop(
	Rectangle inputRange, int facetLevel, int cellCountRecursThreshold) throws IOException {
	if (facetLevel > grid.getMaxLevels()) {
	return;
	}
	final int maxCells = 10_000;
	final HeatmapFacetCounter.Heatmap heatmap =
	HeatmapFacetCounter.calcFacets(
	(PrefixTreeStrategy) strategy,
	indexSearcher.getTopReaderContext(),
	null,
	inputRange,
	facetLevel,
	maxCells);
	int preNonZero = cellValidatedNonZero;
	validateHeatmapResult(inputRange, facetLevel, heatmap);
	assert cellValidatedNonZero - preNonZero > 0; // we validated more non-zero cells
	if (heatmap.counts.length < cellCountRecursThreshold) {
	validateHeatmapResultLoop(inputRange, facetLevel + 1, cellCountRecursThreshold);
	}
	}

	@Test
	@Repeat(iterations = 20)
	public void testRandom() throws IOException {
	// Tests using random index shapes & query shapes. This has found all sorts of edge case bugs
	// (e.g. dateline,
	// cell border, overflow(?)).

	final int numIndexedShapes = 1 + atMost(9);
	List<Shape> indexedShapes = new ArrayList<>(numIndexedShapes);
	for (int i = 0; i < numIndexedShapes; i++) {
	indexedShapes.add(randomIndexedShape());
	}

	// Main index loop:
	for (int i = 0; i < indexedShapes.size(); i++) {
	Shape shape = indexedShapes.get(i);
	adoc("" + i, shape);

	if (random().nextInt(10) == 0) {
	commit(); // intermediate commit, produces extra segments
	}
	}
	// delete some documents randomly
	for (int id = 0; id < indexedShapes.size(); id++) {
	if (random().nextInt(10) == 0) {
	deleteDoc("" + id);
	indexedShapes.set(id, null);
	}
	}

	commit();

	// once without dateline wrap
	final Rectangle rect = randomRectangle();
	queryHeatmapRecursive(usually() ? ctx.getWorldBounds() : rect, 1);
	// and once with dateline wrap
	if (rect.getWidth() > 0) {
	double shift = random().nextDouble() % rect.getWidth();
	queryHeatmapRecursive(
	shapeFactory.rect(
	DistanceUtils.normLonDEG(rect.getMinX() - shift),
	DistanceUtils.normLonDEG(rect.getMaxX() - shift),
	rect.getMinY(),
	rect.getMaxY()),
	1);
	}
	}

	/** Build heatmap, validate results, then descend recursively to another facet level. */
	private boolean queryHeatmapRecursive(Rectangle inputRange, int facetLevel) throws IOException {
	if (!inputRange.hasArea()) {
	// Don't test line inputs. It's not that we don't support it but it is more challenging to
	// test if per-chance it
	// coincides with a grid line due due to edge overlap issue for some grid implementations (geo
	// & quad).
	return false;
	}
	Bits filter =
	null; // FYI testing filtering of underlying PrefixTreeFacetCounter is done in another test
	// Calculate facets
	final int maxCells = 10_000;
	final HeatmapFacetCounter.Heatmap heatmap =
	HeatmapFacetCounter.calcFacets(
	(PrefixTreeStrategy) strategy,
	indexSearcher.getTopReaderContext(),
	filter,
	inputRange,
	facetLevel,
	maxCells);

	validateHeatmapResult(inputRange, facetLevel, heatmap);

	boolean foundNonZeroCount = false;
	for (int count : heatmap.counts) {
	if (count > 0) {
	foundNonZeroCount = true;
	break;
	}
	}

	// Test again recursively to higher facetLevel (more detailed cells)
	if (foundNonZeroCount
	&& cellsValidated <= 500
	&& facetLevel != grid.getMaxLevels()
	&& inputRange.hasArea()) {
	for (int i = 0; i < 5; i++) { // try multiple times until we find non-zero counts
	if (queryHeatmapRecursive(randomRectangle(inputRange), facetLevel + 1)) {
	break; // we found data here so we needn't try again
	}
	}
	}
	return foundNonZeroCount;
	}

	private void validateHeatmapResult(
	Rectangle inputRange, int facetLevel, HeatmapFacetCounter.Heatmap heatmap)
	throws IOException {
	final Rectangle heatRect = heatmap.region;
	assertTrue(
	heatRect.relate(inputRange) == SpatialRelation.CONTAINS \|\| heatRect.equals(inputRange));
	final double cellWidth = heatRect.getWidth() / heatmap.columns;
	final double cellHeight = heatRect.getHeight() / heatmap.rows;
	for (int c = 0; c < heatmap.columns; c++) {
	for (int r = 0; r < heatmap.rows; r++) {
	final int facetCount = heatmap.getCount(c, r);
	double x = DistanceUtils.normLonDEG(heatRect.getMinX() + c * cellWidth + cellWidth / 2);
	double y = DistanceUtils.normLatDEG(heatRect.getMinY() + r * cellHeight + cellHeight / 2);
	Point pt = shapeFactory.pointXY(x, y);
	assertEquals(countMatchingDocsAtLevel(pt, facetLevel), facetCount);
	}
	}
	}

	private int countMatchingDocsAtLevel(Point pt, int facetLevel) throws IOException {
	// we use IntersectsPrefixTreeFilter directly so that we can specify the level to go to exactly.
	RecursivePrefixTreeStrategy strategy = (RecursivePrefixTreeStrategy) this.strategy;
	Query filter =
	new IntersectsPrefixTreeQuery(
	pt, strategy.getFieldName(), grid, facetLevel, grid.getMaxLevels());
	final TotalHitCountCollector collector = new TotalHitCountCollector();
	indexSearcher.search(filter, collector);
	cellsValidated++;
	if (collector.getTotalHits() > 0) {
	cellValidatedNonZero++;
	}
	return collector.getTotalHits();
	}

	private Shape randomIndexedShape() {
	if (((PrefixTreeStrategy) strategy).isPointsOnly() \|\| random().nextBoolean()) {
	return randomPoint();
	} else {
	return randomRectangle();
	}
	}
	}