Java tutorial
/* * 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.io.IOException; import java.util.HashMap; import java.util.Map; import org.apache.lucene.index.IndexReaderContext; import org.apache.lucene.spatial.prefix.tree.Cell; import org.apache.lucene.spatial.prefix.tree.CellIterator; import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree; import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.Bits; 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; /** * Computes spatial facets in two dimensions as a grid of numbers. The data is often visualized as a so-called * "heatmap", hence the name. * * @lucene.experimental */ public class HeatmapFacetCounter { //TODO where should this code live? It could go to PrefixTreeFacetCounter, or maybe here in its own class is fine. /** Maximum number of supported rows (or columns). */ public static final int MAX_ROWS_OR_COLUMNS = (int) Math.sqrt(ArrayUtil.MAX_ARRAY_LENGTH); static { Math.multiplyExact(MAX_ROWS_OR_COLUMNS, MAX_ROWS_OR_COLUMNS);//will throw if doesn't stay within integer } /** Response structure */ public static class Heatmap { public final int columns; public final int rows; public final int[] counts;//in order of 1st column (all rows) then 2nd column (all rows) etc. public final Rectangle region; public Heatmap(int columns, int rows, Rectangle region) { this.columns = columns; this.rows = rows; this.counts = new int[columns * rows]; this.region = region; } public int getCount(int x, int y) { return counts[x * rows + y]; } @Override public String toString() { return "Heatmap{" + columns + "x" + rows + " " + region + '}'; } } /** * Calculates spatial 2D facets (aggregated counts) in a grid, sometimes called a heatmap. * Facet computation is implemented by navigating the underlying indexed terms efficiently. If you don't know exactly * what facetLevel to go to for a given input box but you have some sense of how many cells there should be relative * to the size of the shape, then consider using the logic that {@link org.apache.lucene.spatial.prefix.PrefixTreeStrategy} * uses when approximating what level to go to when indexing a shape given a distErrPct. * * @param context the IndexReader's context * @param topAcceptDocs a Bits to limit counted docs. If null, live docs are counted. * @param inputShape the shape to gather grid squares for; typically a {@link Rectangle}. * The <em>actual</em> heatmap area will usually be larger since the cells on the edge that overlap * are returned. We always return a rectangle of integers even if the inputShape isn't a rectangle * -- the non-intersecting cells will all be 0. * If null is given, the entire world is assumed. * @param facetLevel the target depth (detail) of cells. * @param maxCells the maximum number of cells to return. If the cells exceed this count, an */ public static Heatmap calcFacets(PrefixTreeStrategy strategy, IndexReaderContext context, Bits topAcceptDocs, Shape inputShape, final int facetLevel, int maxCells) throws IOException { if (maxCells > (MAX_ROWS_OR_COLUMNS * MAX_ROWS_OR_COLUMNS)) { throw new IllegalArgumentException("maxCells (" + maxCells + ") should be <= " + MAX_ROWS_OR_COLUMNS); } if (inputShape == null) { inputShape = strategy.getSpatialContext().getWorldBounds(); } final Rectangle inputRect = inputShape.getBoundingBox(); //First get the rect of the cell at the bottom-left at depth facetLevel final SpatialPrefixTree grid = strategy.getGrid(); final SpatialContext ctx = grid.getSpatialContext(); final Point cornerPt = ctx.makePoint(inputRect.getMinX(), inputRect.getMinY()); final CellIterator cellIterator = grid.getTreeCellIterator(cornerPt, facetLevel); Cell cornerCell = null; while (cellIterator.hasNext()) { cornerCell = cellIterator.next(); } assert cornerCell != null && cornerCell.getLevel() == facetLevel : "Cell not at target level: " + cornerCell; final Rectangle cornerRect = (Rectangle) cornerCell.getShape(); assert cornerRect.hasArea(); //Now calculate the number of columns and rows necessary to cover the inputRect double heatMinX = cornerRect.getMinX();//note: we might change this below... final double cellWidth = cornerRect.getWidth(); final Rectangle worldRect = ctx.getWorldBounds(); final int columns = calcRowsOrCols(cellWidth, heatMinX, inputRect.getWidth(), inputRect.getMinX(), worldRect.getWidth()); final double heatMinY = cornerRect.getMinY(); final double cellHeight = cornerRect.getHeight(); final int rows = calcRowsOrCols(cellHeight, heatMinY, inputRect.getHeight(), inputRect.getMinY(), worldRect.getHeight()); assert rows > 0 && columns > 0; if (columns > MAX_ROWS_OR_COLUMNS || rows > MAX_ROWS_OR_COLUMNS || columns * rows > maxCells) { throw new IllegalArgumentException("Too many cells (" + columns + " x " + rows + ") for level " + facetLevel + " shape " + inputRect); } //Create resulting heatmap bounding rectangle & Heatmap object. final double halfCellWidth = cellWidth / 2.0; // if X world-wraps, use world bounds' range if (columns * cellWidth + halfCellWidth > worldRect.getWidth()) { heatMinX = worldRect.getMinX(); } double heatMaxX = heatMinX + columns * cellWidth; if (Math.abs(heatMaxX - worldRect.getMaxX()) < halfCellWidth) {//numeric conditioning issue heatMaxX = worldRect.getMaxX(); } else if (heatMaxX > worldRect.getMaxX()) {//wraps dateline (won't happen if !geo) heatMaxX = heatMaxX - worldRect.getMaxX() + worldRect.getMinX(); } final double halfCellHeight = cellHeight / 2.0; double heatMaxY = heatMinY + rows * cellHeight; if (Math.abs(heatMaxY - worldRect.getMaxY()) < halfCellHeight) {//numeric conditioning issue heatMaxY = worldRect.getMaxY(); } final Heatmap heatmap = new Heatmap(columns, rows, ctx.makeRectangle(heatMinX, heatMaxX, heatMinY, heatMaxY)); if (topAcceptDocs instanceof Bits.MatchNoBits) { return heatmap; // short-circuit } //All ancestor cell counts (of facetLevel) will be captured during facet visiting and applied later. If the data is // just points then there won't be any ancestors. //Facet count of ancestors covering all of the heatmap: int[] allCellsAncestorCount = new int[1]; // single-element array so it can be accumulated in the inner class //All other ancestors: Map<Rectangle, Integer> ancestors = new HashMap<>(); //Now lets count some facets! PrefixTreeFacetCounter.compute(strategy, context, topAcceptDocs, inputShape, facetLevel, new PrefixTreeFacetCounter.FacetVisitor() { @Override public void visit(Cell cell, int count) { final double heatMinX = heatmap.region.getMinX(); final Rectangle rect = (Rectangle) cell.getShape(); if (cell.getLevel() == facetLevel) {//heatmap level; count it directly //convert to col & row int column; if (rect.getMinX() >= heatMinX) { column = (int) Math.round((rect.getMinX() - heatMinX) / cellWidth); } else { // due to dateline wrap column = (int) Math.round((rect.getMinX() + 360 - heatMinX) / cellWidth); } int row = (int) Math.round((rect.getMinY() - heatMinY) / cellHeight); //note: unfortunately, it's possible for us to visit adjacent cells to the heatmap (if the SpatialPrefixTree // allows adjacent cells to overlap on the seam), so we need to skip them if (column < 0 || column >= heatmap.columns || row < 0 || row >= heatmap.rows) { return; } // increment heatmap.counts[column * heatmap.rows + row] += count; } else if (rect.relate(heatmap.region) == SpatialRelation.CONTAINS) {//containing ancestor allCellsAncestorCount[0] += count; } else { // ancestor // note: not particularly efficient (possible put twice, and Integer wrapper); oh well Integer existingCount = ancestors.put(rect, count); if (existingCount != null) { ancestors.put(rect, count + existingCount); } } } }); //Update the heatmap counts with ancestor counts // Apply allCellsAncestorCount if (allCellsAncestorCount[0] > 0) { for (int i = 0; i < heatmap.counts.length; i++) { heatmap.counts[i] += allCellsAncestorCount[0]; } } // Apply ancestors // note: This approach isn't optimized for a ton of ancestor cells. We'll potentially increment the same cells // multiple times in separate passes if any ancestors overlap. IF this poses a problem, we could optimize it // with additional complication by keeping track of intervals in a sorted tree structure (possible TreeMap/Set) // and iterate them cleverly such that we just make one pass at this stage. int[] pair = new int[2];//output of intersectInterval for (Map.Entry<Rectangle, Integer> entry : ancestors.entrySet()) { Rectangle rect = entry.getKey(); // from a cell (thus doesn't cross DL) final int count = entry.getValue(); //note: we approach this in a way that eliminates int overflow/underflow (think huge cell, tiny heatmap) intersectInterval(heatMinY, heatMaxY, cellHeight, rows, rect.getMinY(), rect.getMaxY(), pair); final int startRow = pair[0]; final int endRow = pair[1]; if (!heatmap.region.getCrossesDateLine()) { intersectInterval(heatMinX, heatMaxX, cellWidth, columns, rect.getMinX(), rect.getMaxX(), pair); final int startCol = pair[0]; final int endCol = pair[1]; incrementRange(heatmap, startCol, endCol, startRow, endRow, count); } else { // note: the cell rect might intersect 2 disjoint parts of the heatmap, so we do the left & right separately final int leftColumns = (int) Math.round((180 - heatMinX) / cellWidth); final int rightColumns = heatmap.columns - leftColumns; //left half of dateline: if (rect.getMaxX() > heatMinX) { intersectInterval(heatMinX, 180, cellWidth, leftColumns, rect.getMinX(), rect.getMaxX(), pair); final int startCol = pair[0]; final int endCol = pair[1]; incrementRange(heatmap, startCol, endCol, startRow, endRow, count); } //right half of dateline if (rect.getMinX() < heatMaxX) { intersectInterval(-180, heatMaxX, cellWidth, rightColumns, rect.getMinX(), rect.getMaxX(), pair); final int startCol = pair[0] + leftColumns; final int endCol = pair[1] + leftColumns; incrementRange(heatmap, startCol, endCol, startRow, endRow, count); } } } return heatmap; } private static void intersectInterval(double heatMin, double heatMax, double heatCellLen, int numCells, double cellMin, double cellMax, int[] out) { assert heatMin < heatMax && cellMin < cellMax; //precondition: we know there's an intersection if (heatMin >= cellMin) { out[0] = 0; } else { out[0] = (int) Math.round((cellMin - heatMin) / heatCellLen); } if (heatMax <= cellMax) { out[1] = numCells - 1; } else { out[1] = (int) Math.round((cellMax - heatMin) / heatCellLen) - 1; } } private static void incrementRange(Heatmap heatmap, int startColumn, int endColumn, int startRow, int endRow, int count) { //startColumn & startRow are not necessarily within the heatmap range; likewise numRows/columns may overlap. if (startColumn < 0) { endColumn += startColumn; startColumn = 0; } endColumn = Math.min(heatmap.columns - 1, endColumn); if (startRow < 0) { endRow += startRow; startRow = 0; } endRow = Math.min(heatmap.rows - 1, endRow); if (startRow > endRow) { return;//short-circuit } for (int c = startColumn; c <= endColumn; c++) { int cBase = c * heatmap.rows; for (int r = startRow; r <= endRow; r++) { heatmap.counts[cBase + r] += count; } } } /** Computes the number of intervals (rows or columns) to cover a range given the sizes. */ private static int calcRowsOrCols(double cellRange, double cellMin, double requestRange, double requestMin, double worldRange) { assert requestMin >= cellMin; //Idealistically this wouldn't be so complicated but we concern ourselves with overflow and edge cases double range = (requestRange + (requestMin - cellMin)); if (range == 0) { return 1; } final double intervals = Math.ceil(range / cellRange); if (intervals > Integer.MAX_VALUE) { return Integer.MAX_VALUE;//should result in an error soon (exceed thresholds) } // ensures we don't have more intervals than world bounds (possibly due to rounding/edge issue) final long intervalsMax = Math.round(worldRange / cellRange); if (intervalsMax > Integer.MAX_VALUE) { //just return intervals return (int) intervals; } return Math.min((int) intervalsMax, (int) intervals); } private HeatmapFacetCounter() { } }