+++ /dev/null
-package org.apache.lucene.search.grouping;
-
-/*
- * 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.
- */
-
-import org.apache.lucene.index.IndexReader;
-import org.apache.lucene.search.*;
-
-import java.io.IOException;
-import java.util.*;
-
-/**
- * A base implementation of {@link AbstractAllGroupHeadsCollector} for retrieving the most relevant groups when grouping
- * on a string based group field. More specifically this all concrete implementations of this base implementation
- * use {@link org.apache.lucene.search.FieldCache.StringIndex}.
- *
- * @lucene.experimental
- */
-public abstract class TermAllGroupHeadsCollector<GH extends AbstractAllGroupHeadsCollector.GroupHead> extends AbstractAllGroupHeadsCollector<GH> {
-
- private static final int DEFAULT_INITIAL_SIZE = 128;
-
- final String groupField;
- FieldCache.StringIndex groupIndex;
- IndexReader indexReader;
- int docBase;
-
- protected TermAllGroupHeadsCollector(String groupField, int numberOfSorts) {
- super(numberOfSorts);
- this.groupField = groupField;
- }
-
- /**
- * Creates an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments.
- * This factory method decides with implementation is best suited.
- *
- * Delegates to {@link #create(String, org.apache.lucene.search.Sort, int)} with an initialSize of 128.
- *
- * @param groupField The field to group by
- * @param sortWithinGroup The sort within each group
- * @return an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments
- * @throws IOException If I/O related errors occur
- */
- public static AbstractAllGroupHeadsCollector create(String groupField, Sort sortWithinGroup) throws IOException {
- return create(groupField, sortWithinGroup, DEFAULT_INITIAL_SIZE);
- }
-
- /**
- * Creates an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments.
- * This factory method decides with implementation is best suited.
- *
- * @param groupField The field to group by
- * @param sortWithinGroup The sort within each group
- * @param initialSize The initial allocation size of the internal int set and group list which should roughly match
- * the total number of expected unique groups. Be aware that the heap usage is
- * 4 bytes * initialSize.
- * @return an <code>AbstractAllGroupHeadsCollector</code> instance based on the supplied arguments
- * @throws IOException If I/O related errors occur
- */
- public static AbstractAllGroupHeadsCollector create(String groupField, Sort sortWithinGroup, int initialSize) throws IOException {
- boolean sortAllScore = true;
- boolean sortAllFieldValue = true;
-
- for (SortField sortField : sortWithinGroup.getSort()) {
- if (sortField.getType() == SortField.SCORE) {
- sortAllFieldValue = false;
- } else if (needGeneralImpl(sortField)) {
- return new GeneralAllGroupHeadsCollector(groupField, sortWithinGroup);
- } else {
- sortAllScore = false;
- }
- }
-
- if (sortAllScore) {
- return new ScoreAllGroupHeadsCollector(groupField, sortWithinGroup, initialSize);
- } else if (sortAllFieldValue) {
- return new OrdAllGroupHeadsCollector(groupField, sortWithinGroup, initialSize);
- } else {
- return new OrdScoreAllGroupHeadsCollector(groupField, sortWithinGroup, initialSize);
- }
- }
-
- // Returns when a sort field needs the general impl.
- private static boolean needGeneralImpl(SortField sortField) {
- int sortType = sortField.getType();
- // Note (MvG): We can also make an optimized impl when sorting is SortField.DOC
- return sortType != SortField.STRING_VAL && sortType != SortField.STRING && sortType != SortField.SCORE;
- }
-
- // A general impl that works for any group sort.
- static class GeneralAllGroupHeadsCollector extends TermAllGroupHeadsCollector<GeneralAllGroupHeadsCollector.GroupHead> {
-
- private final Sort sortWithinGroup;
- private final Map<String, GroupHead> groups;
-
- private Scorer scorer;
-
- GeneralAllGroupHeadsCollector(String groupField, Sort sortWithinGroup) throws IOException {
- super(groupField, sortWithinGroup.getSort().length);
- this.sortWithinGroup = sortWithinGroup;
- groups = new HashMap<String, GroupHead>();
-
- final SortField[] sortFields = sortWithinGroup.getSort();
- for (int i = 0; i < sortFields.length; i++) {
- reversed[i] = sortFields[i].getReverse() ? -1 : 1;
- }
- }
-
- protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
- final int ord = groupIndex.order[doc];
- final String groupValue = ord == 0 ? null : groupIndex.lookup[ord];
- GroupHead groupHead = groups.get(groupValue);
- if (groupHead == null) {
- groupHead = new GroupHead(groupValue, sortWithinGroup, doc);
- groups.put(groupValue == null ? null : groupValue, groupHead);
- temporalResult.stop = true;
- } else {
- temporalResult.stop = false;
- }
- temporalResult.groupHead = groupHead;
- }
-
- protected Collection<GroupHead> getCollectedGroupHeads() {
- return groups.values();
- }
-
- public void setNextReader(IndexReader reader, int docBase) throws IOException {
- this.indexReader = reader;
- this.docBase = docBase;
- groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
-
- for (GroupHead groupHead : groups.values()) {
- for (int i = 0; i < groupHead.comparators.length; i++) {
- groupHead.comparators[i].setNextReader(reader, docBase);
- }
- }
- }
-
- public void setScorer(Scorer scorer) throws IOException {
- this.scorer = scorer;
- for (GroupHead groupHead : groups.values()) {
- for (FieldComparator comparator : groupHead.comparators) {
- comparator.setScorer(scorer);
- }
- }
- }
-
- class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
-
- final FieldComparator[] comparators;
-
- private GroupHead(String groupValue, Sort sort, int doc) throws IOException {
- super(groupValue, doc + docBase);
- final SortField[] sortFields = sort.getSort();
- comparators = new FieldComparator[sortFields.length];
- for (int i = 0; i < sortFields.length; i++) {
- comparators[i] = sortFields[i].getComparator(1, i);
- comparators[i].setNextReader(indexReader, docBase);
- comparators[i].setScorer(scorer);
- comparators[i].copy(0, doc);
- comparators[i].setBottom(0);
- }
- }
-
- public int compare(int compIDX, int doc) throws IOException {
- return comparators[compIDX].compareBottom(doc);
- }
-
- public void updateDocHead(int doc) throws IOException {
- for (FieldComparator comparator : comparators) {
- comparator.copy(0, doc);
- comparator.setBottom(0);
- }
- this.doc = doc + docBase;
- }
- }
- }
-
-
- // AbstractAllGroupHeadsCollector optimized for ord fields and scores.
- static class OrdScoreAllGroupHeadsCollector extends TermAllGroupHeadsCollector<OrdScoreAllGroupHeadsCollector.GroupHead> {
-
- private final SentinelIntSet ordSet;
- private final List<GroupHead> collectedGroups;
- private final SortField[] fields;
-
- private FieldCache.StringIndex[] sortsIndex;
- private Scorer scorer;
- private GroupHead[] segmentGroupHeads;
-
- OrdScoreAllGroupHeadsCollector(String groupField, Sort sortWithinGroup, int initialSize) {
- super(groupField, sortWithinGroup.getSort().length);
- ordSet = new SentinelIntSet(initialSize, -1);
- collectedGroups = new ArrayList<GroupHead>(initialSize);
-
- final SortField[] sortFields = sortWithinGroup.getSort();
- fields = new SortField[sortFields.length];
- sortsIndex = new FieldCache.StringIndex[sortFields.length];
- for (int i = 0; i < sortFields.length; i++) {
- reversed[i] = sortFields[i].getReverse() ? -1 : 1;
- fields[i] = sortFields[i];
- }
- }
-
- protected Collection<GroupHead> getCollectedGroupHeads() {
- return collectedGroups;
- }
-
- public void setScorer(Scorer scorer) throws IOException {
- this.scorer = scorer;
- }
-
- protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
- int key = groupIndex.order[doc];
- GroupHead groupHead;
- if (!ordSet.exists(key)) {
- ordSet.put(key);
- String term = key == 0 ? null : groupIndex.lookup[key];
- groupHead = new GroupHead(doc, term);
- collectedGroups.add(groupHead);
- segmentGroupHeads[key] = groupHead;
- temporalResult.stop = true;
- } else {
- temporalResult.stop = false;
- groupHead = segmentGroupHeads[key];
- }
- temporalResult.groupHead = groupHead;
- }
-
- public void setNextReader(IndexReader reader, int docBase) throws IOException {
- this.indexReader = reader;
- this.docBase = docBase;
- groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
- for (int i = 0; i < fields.length; i++) {
- if (fields[i].getType() == SortField.SCORE) {
- continue;
- }
-
- sortsIndex[i] = FieldCache.DEFAULT.getStringIndex(reader, fields[i].getField());
- }
-
- // Clear ordSet and fill it with previous encountered groups that can occur in the current segment.
- ordSet.clear();
- segmentGroupHeads = new GroupHead[groupIndex.lookup.length];
- for (GroupHead collectedGroup : collectedGroups) {
- int ord = groupIndex.binarySearchLookup(collectedGroup.groupValue);
- if (ord >= 0) {
- ordSet.put(ord);
- segmentGroupHeads[ord] = collectedGroup;
-
- for (int i = 0; i < sortsIndex.length; i++) {
- if (fields[i].getType() == SortField.SCORE) {
- continue;
- }
-
- collectedGroup.sortOrds[i] = sortsIndex[i].binarySearchLookup(collectedGroup.sortValues[i]);
- }
- }
- }
- }
-
- class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
-
- String[] sortValues;
- int[] sortOrds;
- float[] scores;
-
- private GroupHead(int doc, String groupValue) throws IOException {
- super(groupValue, doc + docBase);
- sortValues = new String[sortsIndex.length];
- sortOrds = new int[sortsIndex.length];
- scores = new float[sortsIndex.length];
- for (int i = 0; i < sortsIndex.length; i++) {
- if (fields[i].getType() == SortField.SCORE) {
- scores[i] = scorer.score();
- } else {
- sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
- sortOrds[i] = sortsIndex[i].order[doc];
- }
- }
-
- }
-
- public int compare(int compIDX, int doc) throws IOException {
- if (fields[compIDX].getType() == SortField.SCORE) {
- float score = scorer.score();
- if (scores[compIDX] < score) {
- return 1;
- } else if (scores[compIDX] > score) {
- return -1;
- }
- return 0;
- } else {
- if (sortOrds[compIDX] < 0) {
- // The current segment doesn't contain the sort value we encountered before. Therefore the ord is negative.
- final String val1 = sortValues[compIDX];
- final String val2 = sortsIndex[compIDX].lookup[sortsIndex[compIDX].order[doc]];
- if (val1 == null) {
- if (val2 == null) {
- return 0;
- }
- return -1;
- } else if (val2 == null) {
- return 1;
- }
- return val1.compareTo(val2);
- } else {
- return sortOrds[compIDX] - sortsIndex[compIDX].order[doc];
- }
- }
- }
-
- public void updateDocHead(int doc) throws IOException {
- for (int i = 0; i < sortsIndex.length; i++) {
- if (fields[i].getType() == SortField.SCORE) {
- scores[i] = scorer.score();
- } else {
- sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
- sortOrds[i] = sortsIndex[i].order[doc];
- }
- }
- this.doc = doc + docBase;
- }
-
- }
-
- }
-
-
- // AbstractAllGroupHeadsCollector optimized for ord fields.
- static class OrdAllGroupHeadsCollector extends TermAllGroupHeadsCollector<OrdAllGroupHeadsCollector.GroupHead> {
-
- private final SentinelIntSet ordSet;
- private final List<GroupHead> collectedGroups;
- private final SortField[] fields;
-
- private FieldCache.StringIndex[] sortsIndex;
- private GroupHead[] segmentGroupHeads;
-
- OrdAllGroupHeadsCollector(String groupField, Sort sortWithinGroup, int initialSize) {
- super(groupField, sortWithinGroup.getSort().length);
- ordSet = new SentinelIntSet(initialSize, -1);
- collectedGroups = new ArrayList<GroupHead>(initialSize);
-
- final SortField[] sortFields = sortWithinGroup.getSort();
- fields = new SortField[sortFields.length];
- sortsIndex = new FieldCache.StringIndex[sortFields.length];
- for (int i = 0; i < sortFields.length; i++) {
- reversed[i] = sortFields[i].getReverse() ? -1 : 1;
- fields[i] = sortFields[i];
- }
- }
-
- protected Collection<GroupHead> getCollectedGroupHeads() {
- return collectedGroups;
- }
-
- public void setScorer(Scorer scorer) throws IOException {
- }
-
- protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
- int key = groupIndex.order[doc];
- GroupHead groupHead;
- if (!ordSet.exists(key)) {
- ordSet.put(key);
- String term = key == 0 ? null : groupIndex.lookup[key];
- groupHead = new GroupHead(doc, term);
- collectedGroups.add(groupHead);
- segmentGroupHeads[key] = groupHead;
- temporalResult.stop = true;
- } else {
- temporalResult.stop = false;
- groupHead = segmentGroupHeads[key];
- }
- temporalResult.groupHead = groupHead;
- }
-
- public void setNextReader(IndexReader reader, int docBase) throws IOException {
- this.indexReader = reader;
- this.docBase = docBase;
- groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
- for (int i = 0; i < fields.length; i++) {
- sortsIndex[i] = FieldCache.DEFAULT.getStringIndex(reader, fields[i].getField());
- }
-
- // Clear ordSet and fill it with previous encountered groups that can occur in the current segment.
- ordSet.clear();
- segmentGroupHeads = new GroupHead[groupIndex.lookup.length];
- for (GroupHead collectedGroup : collectedGroups) {
- int groupOrd = groupIndex.binarySearchLookup(collectedGroup.groupValue);
- if (groupOrd >= 0) {
- ordSet.put(groupOrd);
- segmentGroupHeads[groupOrd] = collectedGroup;
-
- for (int i = 0; i < sortsIndex.length; i++) {
- collectedGroup.sortOrds[i] = sortsIndex[i].binarySearchLookup(collectedGroup.sortValues[i]);
- }
- }
- }
- }
-
- class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
-
- String[] sortValues;
- int[] sortOrds;
-
- private GroupHead(int doc, String groupValue) throws IOException {
- super(groupValue, doc + docBase);
- sortValues = new String[sortsIndex.length];
- sortOrds = new int[sortsIndex.length];
- for (int i = 0; i < sortsIndex.length; i++) {
- sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
- sortOrds[i] = sortsIndex[i].order[doc];
- }
- }
-
- public int compare(int compIDX, int doc) throws IOException {
- if (sortOrds[compIDX] < 0) {
- // The current segment doesn't contain the sort value we encountered before. Therefore the ord is negative.
- final String val1 = sortValues[compIDX];
- final String val2 = sortsIndex[compIDX].lookup[sortsIndex[compIDX].order[doc]];
- if (val1 == null) {
- if (val2 == null) {
- return 0;
- }
- return -1;
- } else if (val2 == null) {
- return 1;
- }
- return val1.compareTo(val2);
- } else {
- return sortOrds[compIDX] - sortsIndex[compIDX].order[doc];
- }
- }
-
- public void updateDocHead(int doc) throws IOException {
- for (int i = 0; i < sortsIndex.length; i++) {
- sortValues[i] = sortsIndex[i].lookup[sortsIndex[i].order[doc]];
- sortOrds[i] = sortsIndex[i].order[doc];
- }
- this.doc = doc + docBase;
- }
-
- }
-
- }
-
-
- // AbstractAllGroupHeadsCollector optimized for scores.
- static class ScoreAllGroupHeadsCollector extends TermAllGroupHeadsCollector<ScoreAllGroupHeadsCollector.GroupHead> {
-
- private final SentinelIntSet ordSet;
- private final List<GroupHead> collectedGroups;
- private final SortField[] fields;
-
- private Scorer scorer;
- private GroupHead[] segmentGroupHeads;
-
- ScoreAllGroupHeadsCollector(String groupField, Sort sortWithinGroup, int initialSize) {
- super(groupField, sortWithinGroup.getSort().length);
- ordSet = new SentinelIntSet(initialSize, -1);
- collectedGroups = new ArrayList<GroupHead>(initialSize);
-
- final SortField[] sortFields = sortWithinGroup.getSort();
- fields = new SortField[sortFields.length];
- for (int i = 0; i < sortFields.length; i++) {
- reversed[i] = sortFields[i].getReverse() ? -1 : 1;
- fields[i] = sortFields[i];
- }
- }
-
- protected Collection<GroupHead> getCollectedGroupHeads() {
- return collectedGroups;
- }
-
- public void setScorer(Scorer scorer) throws IOException {
- this.scorer = scorer;
- }
-
- protected void retrieveGroupHeadAndAddIfNotExist(int doc) throws IOException {
- int key = groupIndex.order[doc];
- GroupHead groupHead;
- if (!ordSet.exists(key)) {
- ordSet.put(key);
- String term = key == 0 ? null : groupIndex.lookup[key];
- groupHead = new GroupHead(doc, term);
- collectedGroups.add(groupHead);
- segmentGroupHeads[key] = groupHead;
- temporalResult.stop = true;
- } else {
- temporalResult.stop = false;
- groupHead = segmentGroupHeads[key];
- }
- temporalResult.groupHead = groupHead;
- }
-
- public void setNextReader(IndexReader reader, int docBase) throws IOException {
- this.indexReader = reader;
- this.docBase = docBase;
- groupIndex = FieldCache.DEFAULT.getStringIndex(reader, groupField);
-
- // Clear ordSet and fill it with previous encountered groups that can occur in the current segment.
- ordSet.clear();
- segmentGroupHeads = new GroupHead[groupIndex.lookup.length];
- for (GroupHead collectedGroup : collectedGroups) {
- int ord = groupIndex.binarySearchLookup(collectedGroup.groupValue);
- if (ord >= 0) {
- ordSet.put(ord);
- segmentGroupHeads[ord] = collectedGroup;
- }
- }
- }
-
- class GroupHead extends AbstractAllGroupHeadsCollector.GroupHead<String> {
-
- float[] scores;
-
- private GroupHead(int doc, String groupValue) throws IOException {
- super(groupValue, doc + docBase);
- scores = new float[fields.length];
- float score = scorer.score();
- for (int i = 0; i < scores.length; i++) {
- scores[i] = score;
- }
- }
-
- public int compare(int compIDX, int doc) throws IOException {
- float score = scorer.score();
- if (scores[compIDX] < score) {
- return 1;
- } else if (scores[compIDX] > score) {
- return -1;
- }
- return 0;
- }
-
- public void updateDocHead(int doc) throws IOException {
- float score = scorer.score();
- for (int i = 0; i < scores.length; i++) {
- scores[i] = score;
- }
- this.doc = doc + docBase;
- }
-
- }
-
- }
-
-}
fnp / pylucene.git / blobdiff