Index: lucene/join/src/java/org/apache/lucene/search/join/ToParentFieldComparator.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- lucene/join/src/java/org/apache/lucene/search/join/ToParentFieldComparator.java (revision ) +++ lucene/join/src/java/org/apache/lucene/search/join/ToParentFieldComparator.java (revision ) @@ -0,0 +1,308 @@ +package org.apache.lucene.search.join; + +/* + * 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.AtomicReaderContext; +import org.apache.lucene.search.DocIdSet; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.FieldComparator; +import org.apache.lucene.search.Filter; +import org.apache.lucene.util.FixedBitSet; + +import java.io.IOException; + +/** + * A field comparator that allows parent documents to be sorted by fields + * from the child documents. + */ +public abstract class ToParentFieldComparator extends FieldComparator { + + final Filter parentFilter; + final Filter childFilter; + final int spareSlot; + + FieldComparator wrappedComparator; + FixedBitSet parentDocuments; + FixedBitSet childDocuments; + + ToParentFieldComparator(FieldComparator wrappedComparator, Filter parentFilter, Filter childFilter, int spareSlot) { + this.wrappedComparator = wrappedComparator; + this.parentFilter = parentFilter; + this.childFilter = childFilter; + this.spareSlot = spareSlot; + } + + @Override + public int compare(int slot1, int slot2) { + return wrappedComparator.compare(slot1, slot2); + } + + @Override + public void setBottom(int slot) { + wrappedComparator.setBottom(slot); + } + + @Override + public FieldComparator setNextReader(AtomicReaderContext context) throws IOException { + DocIdSet innerDocuments = childFilter.getDocIdSet(context, null); + if (isEmpty(innerDocuments)) { + this.childDocuments = null; + } else if (innerDocuments instanceof FixedBitSet) { + this.childDocuments = (FixedBitSet) innerDocuments; + } else { + this.childDocuments = toFixedBitSet(innerDocuments.iterator(), context.reader().maxDoc()); + } + DocIdSet rootDocuments = parentFilter.getDocIdSet(context, null); + if (isEmpty(rootDocuments)) { + this.parentDocuments = null; + } else if (rootDocuments instanceof FixedBitSet) { + this.parentDocuments = (FixedBitSet) rootDocuments; + } else { + this.parentDocuments = toFixedBitSet(rootDocuments.iterator(), context.reader().maxDoc()); + } + + wrappedComparator = wrappedComparator.setNextReader(context); + return this; + } + + private static boolean isEmpty(DocIdSet set) { + return set == null || set == DocIdSet.EMPTY_DOCIDSET; + } + + private static FixedBitSet toFixedBitSet(DocIdSetIterator iterator, int numBits) throws IOException { + FixedBitSet set = new FixedBitSet(numBits); + int doc; + while ((doc = iterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) { + set.set(doc); + } + return set; + } + + @Override + public Object value(int slot) { + return wrappedComparator.value(slot); + } + + final static class Lowest extends ToParentFieldComparator { + + Lowest(FieldComparator wrappedComparator, Filter parentFilter, Filter childFilter, int spareSlot) { + super(wrappedComparator, parentFilter, childFilter, spareSlot); + } + + @Override + public int compareBottom(int rootDoc) throws IOException { + if (rootDoc == 0 || parentDocuments == null || childDocuments == null) { + return 0; + } + + // We need to copy the lowest value from all nested docs into slot. + int prevRootDoc = parentDocuments.prevSetBit(rootDoc - 1); + int nestedDoc = childDocuments.nextSetBit(prevRootDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return 0; + } + + // We only need to emit a single cmp value for any matching nested doc + int cmp = wrappedComparator.compareBottom(nestedDoc); + if (cmp > 0) { + return cmp; + } + + while (true) { + nestedDoc = childDocuments.nextSetBit(nestedDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return cmp; + } + int cmp1 = wrappedComparator.compareBottom(nestedDoc); + if (cmp1 > 0) { + return cmp1; + } else { + if (cmp1 == 0) { + cmp = 0; + } + } + } + } + + @Override + public void copy(int slot, int rootDoc) throws IOException { + if (rootDoc == 0 || parentDocuments == null || childDocuments == null) { + return; + } + + // We need to copy the lowest value from all nested docs into slot. + int prevRootDoc = parentDocuments.prevSetBit(rootDoc - 1); + int nestedDoc = childDocuments.nextSetBit(prevRootDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return; + } + wrappedComparator.copy(spareSlot, nestedDoc); + wrappedComparator.copy(slot, nestedDoc); + + while (true) { + nestedDoc = childDocuments.nextSetBit(nestedDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return; + } + wrappedComparator.copy(spareSlot, nestedDoc); + if (wrappedComparator.compare(spareSlot, slot) < 0) { + wrappedComparator.copy(slot, nestedDoc); + } + } + } + + @Override + @SuppressWarnings("unchecked") + public int compareDocToValue(int rootDoc, Object value) throws IOException { + if (rootDoc == 0 || parentDocuments == null || childDocuments == null) { + return 0; + } + + // We need to copy the lowest value from all nested docs into slot. + int prevRootDoc = parentDocuments.prevSetBit(rootDoc - 1); + int nestedDoc = childDocuments.nextSetBit(prevRootDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return 0; + } + + // We only need to emit a single cmp value for any matching nested doc + int cmp = wrappedComparator.compareBottom(nestedDoc); + if (cmp > 0) { + return cmp; + } + + while (true) { + nestedDoc = childDocuments.nextSetBit(nestedDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return cmp; + } + int cmp1 = wrappedComparator.compareDocToValue(nestedDoc, value); + if (cmp1 > 0) { + return cmp1; + } else { + if (cmp1 == 0) { + cmp = 0; + } + } + } + } + + } + + final static class Highest extends ToParentFieldComparator { + + Highest(FieldComparator wrappedComparator, Filter parentFilter, Filter childFilter, int spareSlot) { + super(wrappedComparator, parentFilter, childFilter, spareSlot); + } + + @Override + public int compareBottom(int rootDoc) throws IOException { + if (rootDoc == 0 || parentDocuments == null || childDocuments == null) { + return 0; + } + + int prevRootDoc = parentDocuments.prevSetBit(rootDoc - 1); + int nestedDoc = childDocuments.nextSetBit(prevRootDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return 0; + } + + int cmp = wrappedComparator.compareBottom(nestedDoc); + if (cmp < 0) { + return cmp; + } + + while (true) { + nestedDoc = childDocuments.nextSetBit(nestedDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return cmp; + } + int cmp1 = wrappedComparator.compareBottom(nestedDoc); + if (cmp1 < 0) { + return cmp1; + } else { + if (cmp1 == 0) { + cmp = 0; + } + } + } + } + + @Override + public void copy(int slot, int rootDoc) throws IOException { + if (rootDoc == 0 || parentDocuments == null || childDocuments == null) { + return; + } + + int prevRootDoc = parentDocuments.prevSetBit(rootDoc - 1); + int nestedDoc = childDocuments.nextSetBit(prevRootDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return; + } + wrappedComparator.copy(spareSlot, nestedDoc); + wrappedComparator.copy(slot, nestedDoc); + + while (true) { + nestedDoc = childDocuments.nextSetBit(nestedDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return; + } + wrappedComparator.copy(spareSlot, nestedDoc); + if (wrappedComparator.compare(spareSlot, slot) > 0) { + wrappedComparator.copy(slot, nestedDoc); + } + } + } + + @Override + @SuppressWarnings("unchecked") + public int compareDocToValue(int rootDoc, Object value) throws IOException { + if (rootDoc == 0 || parentDocuments == null || childDocuments == null) { + return 0; + } + + int prevRootDoc = parentDocuments.prevSetBit(rootDoc - 1); + int nestedDoc = childDocuments.nextSetBit(prevRootDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return 0; + } + + int cmp = wrappedComparator.compareBottom(nestedDoc); + if (cmp < 0) { + return cmp; + } + + while (true) { + nestedDoc = childDocuments.nextSetBit(nestedDoc + 1); + if (nestedDoc >= rootDoc || nestedDoc == -1) { + return cmp; + } + int cmp1 = wrappedComparator.compareDocToValue(nestedDoc, value); + if (cmp1 < 0) { + return cmp1; + } else { + if (cmp1 == 0) { + cmp = 0; + } + } + } + } + + } + +} Index: lucene/join/src/java/org/apache/lucene/search/join/ToParentFieldComparatorSource.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- lucene/join/src/java/org/apache/lucene/search/join/ToParentFieldComparatorSource.java (revision ) +++ lucene/join/src/java/org/apache/lucene/search/join/ToParentFieldComparatorSource.java (revision ) @@ -0,0 +1,67 @@ +package org.apache.lucene.search.join; + +/* + * 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.search.FieldComparator; +import org.apache.lucene.search.FieldComparatorSource; +import org.apache.lucene.search.Filter; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.search.QueryWrapperFilter; + +import java.io.IOException; + +/** + * Provides a {@link ToParentFieldComparator}, that allows sorting parent docs based on child level fields. + * Based on the sort order it either takes the document with the lowest or highest field value into account. + */ +public class ToParentFieldComparatorSource extends FieldComparatorSource { + + private final FieldComparatorSource childSource; + private final Filter parentFilter; + private final Filter childFilter; + + /** + * @param childSource Provides the field comparator for fields on the child document level. + * @param parentFilter The filter marking the parent documents + * @param childFilter The filter that child documents need to apply with in order to be taken into account during + * sorting. + */ + public ToParentFieldComparatorSource(FieldComparatorSource childSource, Filter parentFilter, Filter childFilter) { + this.childSource = childSource; + this.parentFilter = parentFilter; + this.childFilter = childFilter; + } + + // In the case + public ToParentFieldComparatorSource(FieldComparatorSource childSource, Filter parentFilter) { + this.childSource = childSource; + this.parentFilter = parentFilter; + this.childFilter = new QueryWrapperFilter(new MatchAllDocsQuery()); + } + + @Override + public FieldComparator newComparator(String fieldname, int numHits, int sortPos, boolean reversed) throws IOException { + // +1: have one spare slot for value comparison between inner documents. + FieldComparator wrappedComparator = childSource.newComparator(fieldname, numHits + 1, sortPos, reversed); + if (reversed) { + return new ToParentFieldComparator.Highest(wrappedComparator, parentFilter, childFilter, numHits); + } else { + return new ToParentFieldComparator.Lowest(wrappedComparator, parentFilter, childFilter, numHits); + } + } +}