Examples with UnitNRShape org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape used on opensource projects

Example 1 with UnitNRShape

use of org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape in project lucene-solr by apache.

the class NumberRangeFacetsTest method test.

@Repeat(iterations = 20)
@Test
public void test() throws IOException {
    //generate test data
    List<Shape> indexedShapes = new ArrayList<>();
    final int numIndexedShapes = random().nextInt(15);
    for (int i = 0; i < numIndexedShapes; i++) {
        indexedShapes.add(randomShape());
    }
    //Main index loop:
    for (int i = 0; i < indexedShapes.size(); i++) {
        Shape shape = indexedShapes.get(i);
        adoc("" + i, shape);
        if (random().nextInt(10) == 0)
            //intermediate commit, produces extra segments
            commit();
    }
    //delete some documents randomly
    for (int id = 0; id < indexedShapes.size(); id++) {
        if (random().nextInt(10) == 0) {
            deleteDoc("" + id);
            indexedShapes.set(id, null);
        }
    }
    commit();
    //Main query loop:
    for (int queryIdx = 0; queryIdx < 10; queryIdx++) {
        preQueryHavoc();
        // We need to have a facet range window to do the facets between (a start time & end time). We randomly
        // pick a date, decide the level we want to facet on, and then pick a right end time that is up to 2 thousand
        // values later.
        int calFieldFacet = randomCalWindowField - 1;
        if (calFieldFacet > 1 && rarely()) {
            calFieldFacet--;
        }
        final Calendar leftCal = randomCalendar();
        leftCal.add(calFieldFacet, -1 * randomInt(1000));
        Calendar rightCal = (Calendar) leftCal.clone();
        rightCal.add(calFieldFacet, randomInt(2000));
        // Pick facet detail level based on cal field.
        int detailLevel = tree.getTreeLevelForCalendarField(calFieldFacet);
        if (detailLevel < 0) {
            //no exact match
            detailLevel = -1 * detailLevel;
        }
        //Randomly pick a filter/acceptDocs
        Bits topAcceptDocs = null;
        List<Integer> acceptFieldIds = new ArrayList<>();
        if (usually()) {
            // replace the list.
            for (int i = 0; i < indexedShapes.size(); i++) {
                if (indexedShapes.get(i) == null) {
                    // we deleted this one
                    continue;
                }
                acceptFieldIds.add(i);
            }
            Collections.shuffle(acceptFieldIds, random());
            acceptFieldIds = acceptFieldIds.subList(0, randomInt(acceptFieldIds.size()));
            if (!acceptFieldIds.isEmpty()) {
                List<BytesRef> terms = new ArrayList<>();
                for (Integer acceptDocId : acceptFieldIds) {
                    terms.add(new BytesRef(acceptDocId.toString()));
                }
                topAcceptDocs = searchForDocBits(new TermInSetQuery("id", terms));
            }
        }
        //Lets do it!
        NumberRangePrefixTree.NRShape facetRange = tree.toRangeShape(tree.toShape(leftCal), tree.toShape(rightCal));
        Facets facets = ((NumberRangePrefixTreeStrategy) strategy).calcFacets(indexSearcher.getTopReaderContext(), topAcceptDocs, facetRange, detailLevel);
        //System.out.println("Q: " + queryIdx + " " + facets);
        //Verify results. We do it by looping over indexed shapes and reducing the facet counts.
        Shape facetShapeRounded = facetRange.roundToLevel(detailLevel);
        for (int indexedShapeId = 0; indexedShapeId < indexedShapes.size(); indexedShapeId++) {
            if (topAcceptDocs != null && !acceptFieldIds.contains(indexedShapeId)) {
                // this doc was filtered out via acceptDocs
                continue;
            }
            Shape indexedShape = indexedShapes.get(indexedShapeId);
            if (indexedShape == null) {
                //was deleted
                continue;
            }
            Shape indexedShapeRounded = ((NumberRangePrefixTree.NRShape) indexedShape).roundToLevel(detailLevel);
            if (!indexedShapeRounded.relate(facetShapeRounded).intersects()) {
                // no intersection at all
                continue;
            }
            // walk the cells
            final CellIterator cellIterator = tree.getTreeCellIterator(indexedShape, detailLevel);
            while (cellIterator.hasNext()) {
                Cell cell = cellIterator.next();
                if (!cell.getShape().relate(facetShapeRounded).intersects()) {
                    //no intersection; prune
                    cellIterator.remove();
                    continue;
                }
                assert cell.getLevel() <= detailLevel;
                if (cell.getLevel() == detailLevel) {
                    //count it
                    UnitNRShape shape = (UnitNRShape) cell.getShape();
                    //get parent
                    final UnitNRShape parentShape = shape.getShapeAtLevel(detailLevel - 1);
                    final Facets.FacetParentVal facetParentVal = facets.parents.get(parentShape);
                    assertNotNull(facetParentVal);
                    int index = shape.getValAtLevel(shape.getLevel());
                    assertNotNull(facetParentVal.childCounts);
                    assert facetParentVal.childCounts[index] > 0;
                    facetParentVal.childCounts[index]--;
                } else if (cell.isLeaf()) {
                    //count it, and remove/prune.
                    if (cell.getLevel() < detailLevel - 1) {
                        assert facets.topLeaves > 0;
                        facets.topLeaves--;
                    } else {
                        UnitNRShape shape = (UnitNRShape) cell.getShape();
                        //get parent
                        final UnitNRShape parentShape = shape.getShapeAtLevel(detailLevel - 1);
                        final Facets.FacetParentVal facetParentVal = facets.parents.get(parentShape);
                        assertNotNull(facetParentVal);
                        assert facetParentVal.parentLeaves > 0;
                        facetParentVal.parentLeaves--;
                    }
                    cellIterator.remove();
                }
            }
        }
        // At this point; all counts should be down to zero.
        assertTrue(facets.topLeaves == 0);
        for (Facets.FacetParentVal facetParentVal : facets.parents.values()) {
            assertTrue(facetParentVal.parentLeaves == 0);
            if (facetParentVal.childCounts != null) {
                for (int childCount : facetParentVal.childCounts) {
                    assertTrue(childCount == 0);
                }
            }
        }
    }
}

Example 2 with UnitNRShape

use of org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape in project lucene-solr by apache.

the class NumberRangeFacetsTest method randomShape.

protected Shape randomShape() {
    Calendar cal1 = randomCalendar();
    UnitNRShape s1 = tree.toShape(cal1);
    if (rarely()) {
        return s1;
    }
    try {
        Calendar cal2 = randomCalendar();
        UnitNRShape s2 = tree.toShape(cal2);
        if (cal1.compareTo(cal2) < 0) {
            return tree.toRangeShape(s1, s2);
        } else {
            return tree.toRangeShape(s2, s1);
        }
    } catch (IllegalArgumentException e) {
        assert e.getMessage().startsWith("Differing precision");
        return s1;
    }
}

Example 3 with UnitNRShape

use of org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape in project lucene-solr by apache.

the class DateNRStrategyTest method randomShape.

private Shape randomShape() {
    Calendar cal1 = randomCalendar();
    UnitNRShape s1 = tree.toShape(cal1);
    if (rarely()) {
        return s1;
    }
    try {
        Calendar cal2 = randomCalendar();
        UnitNRShape s2 = tree.toShape(cal2);
        if (cal1.compareTo(cal2) < 0) {
            return tree.toRangeShape(s1, s2);
        } else {
            return tree.toRangeShape(s2, s1);
        }
    } catch (IllegalArgumentException e) {
        assert e.getMessage().startsWith("Differing precision");
        return s1;
    }
}

Example 4 with UnitNRShape

use of org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape in project lucene-solr by apache.

the class NumberRangePrefixTreeStrategy method calcFacets.

/**
   * Calculates facets (aggregated counts) given a range shape (start-end span) and a level, which specifies the detail.
   * To get the level of an existing shape, say a Calendar, call
   * {@link org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree#toUnitShape(Object)} then call
   * {@link org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape#getLevel()}.
   * Facet computation is implemented by navigating the underlying indexed terms efficiently.
   */
public Facets calcFacets(IndexReaderContext context, Bits topAcceptDocs, Shape facetRange, final int level) throws IOException {
    final Facets facets = new Facets(level);
    PrefixTreeFacetCounter.compute(this, context, topAcceptDocs, facetRange, level, new PrefixTreeFacetCounter.FacetVisitor() {

        Facets.FacetParentVal parentFacet;

        UnitNRShape parentShape;

        @Override
        public void visit(Cell cell, int count) {
            if (cell.getLevel() < level - 1) {
                //some ancestor of parent facet level, direct or distant
                //reset
                parentFacet = null;
                //reset
                parentShape = null;
                facets.topLeaves += count;
            } else if (cell.getLevel() == level - 1) {
                //parent
                //set up FacetParentVal
                setupParent((UnitNRShape) cell.getShape());
                parentFacet.parentLeaves += count;
            } else {
                //at facet level
                UnitNRShape unitShape = (UnitNRShape) cell.getShape();
                UnitNRShape unitShapeParent = unitShape.getShapeAtLevel(unitShape.getLevel() - 1);
                if (parentFacet == null || !parentShape.equals(unitShapeParent)) {
                    setupParent(unitShapeParent);
                }
                //lazy init childCounts
                if (parentFacet.childCounts == null) {
                    parentFacet.childCounts = new int[parentFacet.childCountsLen];
                }
                parentFacet.childCounts[unitShape.getValAtLevel(cell.getLevel())] += count;
            }
        }

        private void setupParent(UnitNRShape unitShape) {
            parentShape = unitShape.clone();
            //Look for existing parentFacet (from previous segment), or create anew if needed
            parentFacet = facets.parents.get(parentShape);
            if (parentFacet == null) {
                //didn't find one; make a new one
                parentFacet = new Facets.FacetParentVal();
                parentFacet.childCountsLen = getGrid().getNumSubCells(parentShape);
                facets.parents.put(parentShape, parentFacet);
            }
        }
    });
    return facets;
}

Example 5 with UnitNRShape

use of org.apache.lucene.spatial.prefix.tree.NumberRangePrefixTree.UnitNRShape in project lucene-solr by apache.

the class DateRangePrefixTreeTest method roundTrip.

private void roundTrip(Calendar calOrig) throws ParseException {
    Calendar cal = (Calendar) calOrig.clone();
    String lastString = null;
    while (true) {
        String calString;
        {
            Calendar preToStringCalClone = (Calendar) cal.clone();
            calString = tree.toString(cal);
            //ensure toString doesn't modify cal state
            assertEquals(preToStringCalClone, cal);
        }
        //test parseCalendar
        assertEquals(cal, tree.parseCalendar(calString));
        //to Shape and back to Cal
        UnitNRShape shape = tree.toShape(cal);
        Calendar cal2 = tree.toCalendar(shape);
        assertEquals(calString, tree.toString(cal2));
        if (!calString.equals("*")) {
            //not world cell
            //to Term and back to Cell
            Cell cell = (Cell) shape;
            BytesRef term = cell.getTokenBytesNoLeaf(null);
            Cell cell2 = tree.readCell(BytesRef.deepCopyOf(term), null);
            assertEquals(calString, cell, cell2);
            Calendar cal3 = tree.toCalendar((UnitNRShape) cell2.getShape());
            assertEquals(calString, tree.toString(cal3));
            // setLeaf comparison
            cell2.setLeaf();
            BytesRef termLeaf = cell2.getTokenBytesWithLeaf(null);
            assertTrue(term.compareTo(termLeaf) < 0);
            assertEquals(termLeaf.length, term.length + 1);
            assertEquals(0, termLeaf.bytes[termLeaf.offset + termLeaf.length - 1]);
            assertTrue(cell.isPrefixOf(cell2));
        }
        //end of loop; decide if should loop again with lower precision
        final int calPrecField = tree.getCalPrecisionField(cal);
        if (calPrecField == -1)
            break;
        int fieldIdx = Arrays.binarySearch(CAL_FIELDS, calPrecField);
        assert fieldIdx >= 0;
        int prevPrecField = (fieldIdx == 0 ? -1 : CAL_FIELDS[--fieldIdx]);
        try {
            tree.clearFieldsAfter(cal, prevPrecField);
        } catch (AssertionError e) {
            if (e.getMessage().equals("Calendar underflow"))
                return;
            throw e;
        }
        lastString = calString;
    }
}