Skip to content

Commit

Permalink
Specialize DisiPriorityQueue for the 2-clauses case. (#14070)
Browse files Browse the repository at this point in the history 
Disjunctions with 2 clauses are rather common. Specializing this case enables
some shortcuts.
  • Loading branch information
@jpountz
jpountz authored Jan 28, 2025
1 parent cab88ca commit 71256cc
Show file tree
Hide file tree
Showing 9 changed files with 448 additions and 194 deletions.
216 changes: 30 additions & 186 deletions lucene/core/src/java/org/apache/lucene/search/DisiPriorityQueue.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@
*/
package org.apache.lucene.search;

import java.util.Arrays;
import java.util.Iterator;
import org.apache.lucene.util.PriorityQueue;

/**
Expand All @@ -27,205 +25,51 @@
*
* @lucene.internal
*/
public final class DisiPriorityQueue implements Iterable<DisiWrapper> {

static int leftNode(int node) {
return ((node + 1) << 1) - 1;
}

static int rightNode(int leftNode) {
return leftNode + 1;
}

static int parentNode(int node) {
return ((node + 1) >>> 1) - 1;
public abstract sealed class DisiPriorityQueue implements Iterable<DisiWrapper>
permits DisiPriorityQueue2, DisiPriorityQueueN {

/** Create a {@link DisiPriorityQueue} of the given maximum size. */
public static DisiPriorityQueue ofMaxSize(int maxSize) {
if (maxSize <= 2) {
return new DisiPriorityQueue2();
} else {
return new DisiPriorityQueueN(maxSize);
}
}

private final DisiWrapper[] heap;
private int size;
/** Return the number of entries in this heap. */
public abstract int size();

public DisiPriorityQueue(int maxSize) {
heap = new DisiWrapper[maxSize];
size = 0;
}

public int size() {
return size;
}

public DisiWrapper top() {
return heap[0];
}
/** Return top value in this heap, or null if the heap is empty. */
public abstract DisiWrapper top();

/** Return the 2nd least value in this heap, or null if the heap contains less than 2 values. */
public DisiWrapper top2() {
switch (size()) {
case 0:
case 1:
return null;
case 2:
return heap[1];
default:
if (heap[1].doc <= heap[2].doc) {
return heap[1];
} else {
return heap[2];
}
}
}
public abstract DisiWrapper top2();

/** Get the list of scorers which are on the current doc. */
public DisiWrapper topList() {
final DisiWrapper[] heap = this.heap;
final int size = this.size;
DisiWrapper list = heap[0];
list.next = null;
if (size >= 3) {
list = topList(list, heap, size, 1);
list = topList(list, heap, size, 2);
} else if (size == 2 && heap[1].doc == list.doc) {
list = prepend(heap[1], list);
}
return list;
}

// prepend w1 (iterator) to w2 (list)
private DisiWrapper prepend(DisiWrapper w1, DisiWrapper w2) {
w1.next = w2;
return w1;
}

private DisiWrapper topList(DisiWrapper list, DisiWrapper[] heap, int size, int i) {
final DisiWrapper w = heap[i];
if (w.doc == list.doc) {
list = prepend(w, list);
final int left = leftNode(i);
final int right = left + 1;
if (right < size) {
list = topList(list, heap, size, left);
list = topList(list, heap, size, right);
} else if (left < size && heap[left].doc == list.doc) {
list = prepend(heap[left], list);
}
}
return list;
}
public abstract DisiWrapper topList();

public DisiWrapper add(DisiWrapper entry) {
final DisiWrapper[] heap = this.heap;
final int size = this.size;
heap[size] = entry;
upHeap(size);
this.size = size + 1;
return heap[0];
}
/** Add a {@link DisiWrapper} to this queue and return the top entry. */
public abstract DisiWrapper add(DisiWrapper entry);

/** Bulk add. */
public void addAll(DisiWrapper[] entries, int offset, int len) {
// Nothing to do if empty:
if (len == 0) {
return;
}

// Fail early if we're going to over-fill:
if (size + len > heap.length) {
throw new IndexOutOfBoundsException(
"Cannot add "
+ len
+ " elements to a queue with remaining capacity "
+ (heap.length - size));
}

// Copy the entries over to our heap array:
System.arraycopy(entries, offset, heap, size, len);
size += len;

// Heapify in bulk:
final int firstLeafIndex = size >>> 1;
for (int rootIndex = firstLeafIndex - 1; rootIndex >= 0; rootIndex--) {
int parentIndex = rootIndex;
DisiWrapper parent = heap[parentIndex];
while (parentIndex < firstLeafIndex) {
int childIndex = leftNode(parentIndex);
int rightChildIndex = rightNode(childIndex);
DisiWrapper child = heap[childIndex];
if (rightChildIndex < size && heap[rightChildIndex].doc < child.doc) {
child = heap[rightChildIndex];
childIndex = rightChildIndex;
}
if (child.doc >= parent.doc) {
break;
}
heap[parentIndex] = child;
parentIndex = childIndex;
}
heap[parentIndex] = parent;
for (int i = 0; i < len; ++i) {
add(entries[offset + i]);
}
}

public DisiWrapper pop() {
final DisiWrapper[] heap = this.heap;
final DisiWrapper result = heap[0];
final int i = --size;
heap[0] = heap[i];
heap[i] = null;
downHeap(i);
return result;
}
/** Remove the top entry and return it. */
public abstract DisiWrapper pop();

public DisiWrapper updateTop() {
downHeap(size);
return heap[0];
}
/** Rebalance this heap and return the top entry. */
public abstract DisiWrapper updateTop();

DisiWrapper updateTop(DisiWrapper topReplacement) {
heap[0] = topReplacement;
return updateTop();
}
/**
* Replace the top entry with the given entry, rebalance the heap, and return the new top entry.
*/
abstract DisiWrapper updateTop(DisiWrapper topReplacement);

/** Clear the heap. */
public void clear() {
Arrays.fill(heap, null);
size = 0;
}

void upHeap(int i) {
final DisiWrapper node = heap[i];
final int nodeDoc = node.doc;
int j = parentNode(i);
while (j >= 0 && nodeDoc < heap[j].doc) {
heap[i] = heap[j];
i = j;
j = parentNode(j);
}
heap[i] = node;
}

void downHeap(int size) {
int i = 0;
final DisiWrapper node = heap[0];
int j = leftNode(i);
if (j < size) {
int k = rightNode(j);
if (k < size && heap[k].doc < heap[j].doc) {
j = k;
}
if (heap[j].doc < node.doc) {
do {
heap[i] = heap[j];
i = j;
j = leftNode(i);
k = rightNode(j);
if (k < size && heap[k].doc < heap[j].doc) {
j = k;
}
} while (j < size && heap[j].doc < node.doc);
heap[i] = node;
}
}
}

@Override
public Iterator<DisiWrapper> iterator() {
return Arrays.asList(heap).subList(0, size).iterator();
}
public abstract void clear();
}
110 changes: 110 additions & 0 deletions lucene/core/src/java/org/apache/lucene/search/DisiPriorityQueue2.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
/*
* 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.search;

import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;

/** {@link DisiPriorityQueue} of two entries or less. */
final class DisiPriorityQueue2 extends DisiPriorityQueue {

private DisiWrapper top, top2;

@Override
public Iterator<DisiWrapper> iterator() {
if (top2 != null) {
return Arrays.asList(top, top2).iterator();
} else if (top != null) {
return Collections.singleton(top).iterator();
} else {
return Collections.emptyIterator();
}
}

@Override
public int size() {
return top2 == null ? (top == null ? 0 : 1) : 2;
}

@Override
public DisiWrapper top() {
return top;
}

@Override
public DisiWrapper top2() {
return top2;
}

@Override
public DisiWrapper topList() {
DisiWrapper topList = null;
if (top != null) {
top.next = null;
topList = top;
if (top2 != null && top.doc == top2.doc) {
top2.next = topList;
topList = top2;
}
}
return topList;
}

@Override
public DisiWrapper add(DisiWrapper entry) {
if (top == null) {
return top = entry;
} else if (top2 == null) {
top2 = entry;
return updateTop();
} else {
throw new IllegalStateException(
"Trying to add a 3rd element to a DisiPriorityQueue configured with a max size of 2");
}
}

@Override
public DisiWrapper pop() {
DisiWrapper ret = top;
top = top2;
top2 = null;
return ret;
}

@Override
public DisiWrapper updateTop() {
if (top2 != null && top2.doc < top.doc) {
DisiWrapper tmp = top;
top = top2;
top2 = tmp;
}
return top;
}

@Override
DisiWrapper updateTop(DisiWrapper topReplacement) {
top = topReplacement;
return updateTop();
}

@Override
public void clear() {
top = null;
top2 = null;
}
}
Loading

0 comments on commit 71256cc

Please sign in to comment.