--- /dev/null
+package org.apache.lucene.facet.search;
+
+import java.io.IOException;
+import java.util.ArrayList;
+
+import org.apache.lucene.facet.search.params.FacetRequest;
+import org.apache.lucene.facet.search.results.FacetResult;
+import org.apache.lucene.facet.search.results.FacetResultNode;
+import org.apache.lucene.facet.search.results.MutableFacetResultNode;
+import org.apache.lucene.facet.search.results.IntermediateFacetResult;
+import org.apache.lucene.facet.taxonomy.TaxonomyReader;
+import org.apache.lucene.facet.taxonomy.TaxonomyReader.ChildrenArrays;
+import org.apache.lucene.facet.util.ResultSortUtils;
+
+/**
+ * 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.
+ */
+
+/**
+ * Generate Top-K results for a particular FacetRequest.
+ * <p>
+ * K is global (among all results) and is defined by {@link FacetRequest#getNumResults()}.
+ * <p>
+ * Note: Values of 0 (Zero) are ignored by this results handler.
+ *
+ * @lucene.experimental
+ */
+public class TopKFacetResultsHandler extends FacetResultsHandler {
+
+ /**
+ * Construct top-K results handler.
+ * @param taxonomyReader taxonomy reader
+ * @param facetRequest facet request being served
+ */
+ public TopKFacetResultsHandler(TaxonomyReader taxonomyReader,
+ FacetRequest facetRequest) {
+ super(taxonomyReader, facetRequest);
+ }
+
+ // fetch top K for specific partition.
+ @Override
+ public IntermediateFacetResult fetchPartitionResult(FacetArrays facetArrays, int offset)
+ throws IOException {
+ TopKFacetResult res = null;
+ int ordinal = taxonomyReader.getOrdinal(facetRequest.getCategoryPath());
+ if (ordinal != TaxonomyReader.INVALID_ORDINAL) {
+ double value = 0;
+ if (isSelfPartition(ordinal, facetArrays, offset)) {
+ int partitionSize = facetArrays.getArraysLength();
+ value = facetRequest.getValueOf(facetArrays, ordinal % partitionSize);
+ }
+
+ // TODO (Facet): should initial value of "residue" depend on aggregator if not sum?
+ MutableFacetResultNode parentResultNode =
+ new MutableFacetResultNode(ordinal, value);
+
+ Heap<FacetResultNode> heap = ResultSortUtils.createSuitableHeap(facetRequest);
+ int totalFacets = heapDescendants(ordinal, heap, parentResultNode, facetArrays, offset);
+ res = new TopKFacetResult(facetRequest, parentResultNode, totalFacets);
+ res.setHeap(heap);
+ }
+ return res;
+ }
+
+ // merge given top K results into current
+ @Override
+ public IntermediateFacetResult mergeResults(IntermediateFacetResult... tmpResults) throws IOException {
+
+ int ordinal = taxonomyReader.getOrdinal(facetRequest.getCategoryPath());
+ MutableFacetResultNode resNode = new MutableFacetResultNode(ordinal, 0);
+
+ int totalFacets = 0;
+ Heap<FacetResultNode> heap = null;
+
+ // merge other results in queue
+ for (IntermediateFacetResult tmpFres : tmpResults) {
+ // cast should succeed
+ TopKFacetResult fres = (TopKFacetResult) tmpFres;
+ totalFacets += fres.getNumValidDescendants();
+ // set the value for the result node representing the facet request
+ resNode.increaseValue(fres.getFacetResultNode().getValue());
+ Heap<FacetResultNode> tmpHeap = fres.getHeap();
+ if (heap == null) {
+ heap = tmpHeap;
+ continue;
+ }
+ // bring sub results from heap of tmp res into result heap
+ for (int i = tmpHeap.size(); i > 0; i--) {
+
+ FacetResultNode a = heap.insertWithOverflow(tmpHeap.pop());
+ if (a != null) {
+ resNode.increaseResidue(a.getResidue());
+ }
+ }
+ }
+
+ TopKFacetResult res = new TopKFacetResult(facetRequest, resNode, totalFacets);
+ res.setHeap(heap);
+ return res;
+ }
+
+ /**
+ * Finds the top K descendants of ordinal, which are at most facetRequest.getDepth()
+ * deeper than facetRequest.getCategoryPath (whose ordinal is input parameter ordinal).
+ * Candidates are restricted to current "counting list" and current "partition",
+ * they join the overall priority queue pq of size K.
+ * @return total number of descendants considered here by pq, excluding ordinal itself.
+ */
+ private int heapDescendants(int ordinal, Heap<FacetResultNode> pq,
+ MutableFacetResultNode parentResultNode, FacetArrays facetArrays, int offset) {
+ int partitionSize = facetArrays.getArraysLength();
+ int endOffset = offset + partitionSize;
+ ChildrenArrays childrenArray = taxonomyReader.getChildrenArrays();
+ int[] youngestChild = childrenArray.getYoungestChildArray();
+ int[] olderSibling = childrenArray.getOlderSiblingArray();
+ FacetResultNode reusable = null;
+ int localDepth = 0;
+ int depth = facetRequest.getDepth();
+ int[] ordinalStack = new int[2+Math.min(Short.MAX_VALUE, depth)];
+ int childrenCounter = 0;
+
+ int tosOrdinal; // top of stack element
+
+ int yc = youngestChild[ordinal];
+ while (yc >= endOffset) {
+ yc = olderSibling[yc];
+ }
+ // make use of the fact that TaxonomyReader.INVALID_ORDINAL == -1, < endOffset
+ // and it, too, can stop the loop.
+ ordinalStack[++localDepth] = yc;
+
+ /*
+ * stack holds input parameter ordinal in position 0.
+ * Other elements are < endoffset.
+ * Only top of stack can be TaxonomyReader.INVALID_ORDINAL, and this if and only if
+ * the element below it exhausted all its children: has them all processed.
+ *
+ * stack elements are processed (counted and accumulated) only if they
+ * belong to current partition (between offset and endoffset) and first time
+ * they are on top of stack
+ *
+ * loop as long as stack is not empty of elements other than input ordinal, or for a little while -- it sibling
+ */
+ while (localDepth > 0) {
+ tosOrdinal = ordinalStack[localDepth];
+ if (tosOrdinal == TaxonomyReader.INVALID_ORDINAL) {
+ // element below tos has all its children, and itself, all processed
+ // need to proceed to its sibling
+ localDepth--;
+ // change element now on top of stack to its sibling.
+ ordinalStack[localDepth] = olderSibling[ordinalStack[localDepth]];
+ continue;
+ }
+ // top of stack is not invalid, this is the first time we see it on top of stack.
+ // collect it, if belongs to current partition, and then push its kids on itself, if applicable
+ if (tosOrdinal >= offset) { // tosOrdinal resides in current partition
+ int relativeOrdinal = tosOrdinal % partitionSize;
+ double value = facetRequest.getValueOf(facetArrays, relativeOrdinal);
+ if (value != 0 && !Double.isNaN(value)) {
+ // Count current ordinal -- the TOS
+ if (reusable == null) {
+ reusable = new MutableFacetResultNode(tosOrdinal, value);
+ } else {
+ // it is safe to cast since reusable was created here.
+ ((MutableFacetResultNode)reusable).reset(tosOrdinal, value);
+ }
+ ++childrenCounter;
+ reusable = pq.insertWithOverflow(reusable);
+ if (reusable != null) {
+ // TODO (Facet): is other logic (not add) needed, per aggregator?
+ parentResultNode.increaseResidue(reusable.getValue());
+ }
+ }
+ }
+ if (localDepth < depth) {
+ // push kid of current tos
+ yc = youngestChild[tosOrdinal];
+ while (yc >= endOffset) {
+ yc = olderSibling[yc];
+ }
+ ordinalStack[++localDepth] = yc;
+ } else { // localDepth == depth; current tos exhausted its possible children, mark this by pushing INVALID_ORDINAL
+ ordinalStack[++localDepth] = TaxonomyReader.INVALID_ORDINAL;
+ }
+ } // endof while stack is not empty
+
+ return childrenCounter; // we're done
+ }
+
+ @Override
+ public FacetResult renderFacetResult(IntermediateFacetResult tmpResult) {
+ TopKFacetResult res = (TopKFacetResult) tmpResult; // cast is safe by contract of this class
+ if (res != null) {
+ Heap<FacetResultNode> heap = res.getHeap();
+ MutableFacetResultNode resNode = (MutableFacetResultNode)res.getFacetResultNode(); // cast safe too
+ for (int i = heap.size(); i > 0; i--) {
+ resNode.insertSubResult(heap.pop());
+ }
+ }
+ return res;
+ }
+
+ @Override
+ public FacetResult rearrangeFacetResult(FacetResult facetResult) {
+ TopKFacetResult res = (TopKFacetResult) facetResult; // cast is safe by contract of this class
+ Heap<FacetResultNode> heap = res.getHeap();
+ heap.clear(); // just to be safe
+ MutableFacetResultNode topFrn = (MutableFacetResultNode) res.getFacetResultNode(); // safe cast
+ for (FacetResultNode frn : topFrn.getSubResults()) {
+ heap.add(frn);
+ }
+ int size = heap.size();
+ ArrayList<FacetResultNode> subResults = new ArrayList<FacetResultNode>(size);
+ for (int i = heap.size(); i > 0; i--) {
+ subResults.add(0,heap.pop());
+ }
+ topFrn.setSubResults(subResults);
+ return res;
+ }
+
+ @Override
+ // label top K sub results
+ public void labelResult(FacetResult facetResult) throws IOException {
+ if (facetResult != null) { // any result to label?
+ FacetResultNode facetResultNode = facetResult.getFacetResultNode();
+ if (facetResultNode != null) { // any result to label?
+ facetResultNode.getLabel(taxonomyReader);
+ int num2label = facetRequest.getNumLabel();
+ for (FacetResultNode frn : facetResultNode.getSubResults()) {
+ if (--num2label < 0) {
+ break;
+ }
+ frn.getLabel(taxonomyReader);
+ }
+ }
+ }
+ }
+
+ ////////////////////////////////////////////////////////////////////////////////////
+ ////////////////////////////////////////////////////////////////////////////////////
+
+ /**
+ * Private Mutable implementation of result of faceted search.
+ */
+ private static class TopKFacetResult extends FacetResult implements IntermediateFacetResult {
+
+ // TODO (Facet): is it worth to override PriorityQueue.getSentinelObject()
+ // for any of our PQs?
+ private Heap<FacetResultNode> heap;
+
+ /**
+ * Create a Facet Result.
+ * @param facetRequest Request for which this result was obtained.
+ * @param facetResultNode top result node for this facet result.
+ * @param totalFacets - number of children of the targetFacet, up till the requested depth.
+ */
+ TopKFacetResult(FacetRequest facetRequest, MutableFacetResultNode facetResultNode, int totalFacets) {
+ super(facetRequest, facetResultNode, totalFacets);
+ }
+
+ /**
+ * @return the heap
+ */
+ public Heap<FacetResultNode> getHeap() {
+ return heap;
+ }
+
+ /**
+ * Set the heap for this result.
+ * @param heap heap top be set.
+ */
+ public void setHeap(Heap<FacetResultNode> heap) {
+ this.heap = heap;
+ }
+
+ }
+
+ //////////////////////////////////////////////////////
+}