Skip to content

Commit

Permalink
tidy, lint, restore test setup
Browse files Browse the repository at this point in the history
  • Loading branch information
Michael Sokolov committed Jan 2, 2025
1 parent 7314314 commit 69e6dbf
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 53 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,27 @@ public abstract class AbstractBPReorderer implements IndexReorderer {
* this number of documents: 32.
*/
public static final int DEFAULT_MIN_PARTITION_SIZE = 32;

/**
* Default maximum number of iterations per recursion level: 20. Higher numbers of iterations
* typically don't help significantly.
*/
public static final int DEFAULT_MAX_ITERS = 20;
protected int minPartitionSize;
protected int maxIters;

protected int minPartitionSize = DEFAULT_MIN_PARTITION_SIZE;
protected int maxIters = DEFAULT_MAX_ITERS;
protected double ramBudgetMB;

/**
* Set the minimum partition size, when the algorithm stops recursing, 32 by default.
*/
public AbstractBPReorderer() {
// 10% of the available heap size by default
setRAMBudgetMB(Runtime.getRuntime().totalMemory() / 1024d / 1024d / 10d);
}

/** Set the minimum partition size, when the algorithm stops recursing, 32 by default. */
public void setMinPartitionSize(int minPartitionSize) {
if (minPartitionSize < 1) {
throw new IllegalArgumentException(
"minPartitionSize must be at least 1, got " + minPartitionSize);
"minPartitionSize must be at least 1, got " + minPartitionSize);
}
this.minPartitionSize = minPartitionSize;
}
Expand All @@ -47,9 +52,7 @@ public void setRAMBudgetMB(double ramBudgetMB) {
this.ramBudgetMB = ramBudgetMB;
}

/**
* Exception that is thrown when not enough RAM is available.
*/
/** Exception that is thrown when not enough RAM is available. */
public static class NotEnoughRAMException extends RuntimeException {
NotEnoughRAMException(String message) {
super(message);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -109,10 +109,6 @@ public final class BPIndexReorderer extends AbstractBPReorderer {
public BPIndexReorderer() {
setMinDocFreq(DEFAULT_MIN_DOC_FREQ);
setMaxDocFreq(1f);
setMinPartitionSize(DEFAULT_MIN_PARTITION_SIZE);
setMaxIters(DEFAULT_MAX_ITERS);
// 10% of the available heap size by default
setRAMBudgetMB(Runtime.getRuntime().totalMemory() / 1024d / 1024d / 10d);
setFields(null);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -99,36 +99,6 @@ public BpVectorReorderer(String partitionField) {
this.partitionField = partitionField;
}

/** Set the minimum partition size, when the algorithm stops recursing, 32 by default. */
public void setMinPartitionSize(int minPartitionSize) {
if (minPartitionSize < 1) {
throw new IllegalArgumentException(
"minPartitionSize must be at least 1, got " + minPartitionSize);
}
this.minPartitionSize = minPartitionSize;
}

/**
* Set the maximum number of iterations on each recursion level, 20 by default. Experiments
* suggests that values above 20 do not help much. However, values below 20 can be used to trade
* effectiveness for faster reordering.
*/
public void setMaxIters(int maxIters) {
if (maxIters < 1) {
throw new IllegalArgumentException("maxIters must be at least 1, got " + maxIters);
}
this.maxIters = maxIters;
}

/**
* Set the amount of RAM that graph partitioning is allowed to use. More RAM allows running
* faster. If not enough RAM is provided, an exception (TODO: NotEnoughRAMException) will be
* thrown. This is 10% of the total heap size by default.
*/
public void setRAMBudgetMB(double ramBudgetMB) {
this.ramBudgetMB = ramBudgetMB;
}

private static class PerThreadState {

final FloatVectorValues vectors;
Expand Down Expand Up @@ -710,7 +680,6 @@ public static void main(String... args) throws IOException {
}

void reorderIndexDirectory(Directory directory, Executor executor) throws IOException {
VectorSimilarityFunction similarity = null;
try (IndexReader reader = DirectoryReader.open(directory)) {
IndexWriterConfig iwc = new IndexWriterConfig();
iwc.setOpenMode(IndexWriterConfig.OpenMode.CREATE);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.index.SerialMergeScheduler;
import org.apache.lucene.index.SlowCodecReaderWrapper;
import org.apache.lucene.index.StoredFields;
import org.apache.lucene.index.Term;
Expand All @@ -43,10 +42,11 @@ public class TestBPReorderingMergePolicy extends LuceneTestCase {

AbstractBPReorderer reorderer;

@Override
@Before
public void setUp() throws Exception {
super.setUp();
if (random().nextBoolean() && false) {
if (random().nextBoolean()) {
BPIndexReorderer bpIndexReorderer = new BPIndexReorderer();
bpIndexReorderer.setMinDocFreq(2);
reorderer = bpIndexReorderer;
Expand All @@ -61,7 +61,7 @@ public void testReorderOnMerge() throws IOException {
Directory dir1 = newDirectory();
Directory dir2 = newDirectory();
IndexWriter w1 =
new IndexWriter(dir1, newIndexWriterConfig().setMergePolicy(newLogMergePolicy()).setMergeScheduler(new SerialMergeScheduler()));
new IndexWriter(dir1, newIndexWriterConfig().setMergePolicy(newLogMergePolicy()));
BPReorderingMergePolicy mp = new BPReorderingMergePolicy(newLogMergePolicy(), reorderer);
mp.setMinNaturalMergeNumDocs(2);
IndexWriter w2 = new IndexWriter(dir2, newIndexWriterConfig().setMergePolicy(mp));
Expand All @@ -70,14 +70,14 @@ public void testReorderOnMerge() throws IOException {
doc.add(idField);
StringField bodyField = new StringField("body", "", Store.YES);
doc.add(bodyField);
KnnFloatVectorField vectorField = new KnnFloatVectorField("vector", new float[]{0});
KnnFloatVectorField vectorField = new KnnFloatVectorField("vector", new float[] {0});
doc.add(vectorField);

for (int i = 0; i < 10000; ++i) {
idField.setStringValue(Integer.toString(i));
int intValue = i % 2 == 0 ? 0 : i % 10;
bodyField.setStringValue(Integer.toString(intValue));
vectorField.setVectorValue(new float[]{intValue});
vectorField.setVectorValue(new float[] {intValue});
w1.addDocument(doc);
w2.addDocument(doc);

Expand Down Expand Up @@ -151,14 +151,14 @@ public void testReorderOnAddIndexes() throws IOException {
doc.add(idField);
StringField bodyField = new StringField("body", "", Store.YES);
doc.add(bodyField);
KnnFloatVectorField vectorField = new KnnFloatVectorField("vector", new float[]{0});
KnnFloatVectorField vectorField = new KnnFloatVectorField("vector", new float[] {0});
doc.add(vectorField);

for (int i = 0; i < 10000; ++i) {
idField.setStringValue(Integer.toString(i));
int intValue = i % 2 == 0 ? 0 : i % 10;
bodyField.setStringValue(Integer.toString(intValue));
vectorField.setVectorValue(new float[]{intValue});
vectorField.setVectorValue(new float[] {intValue});

w1.addDocument(doc);

Expand Down Expand Up @@ -253,14 +253,14 @@ public void testReorderDoesntHaveEnoughRAM() throws IOException {
doc.add(idField);
StringField bodyField = new StringField("body", "", Store.YES);
doc.add(bodyField);
KnnFloatVectorField vectorField = new KnnFloatVectorField("vector", new float[]{0});
KnnFloatVectorField vectorField = new KnnFloatVectorField("vector", new float[] {0});
doc.add(vectorField);

for (int i = 0; i < 10; ++i) {
idField.setStringValue(Integer.toString(i));
int intValue = i % 2 == 0 ? 0 : i % 10;
bodyField.setStringValue(Integer.toString(intValue));
vectorField.setVectorValue(new float[]{intValue});
vectorField.setVectorValue(new float[] {intValue});
w.addDocument(doc);
DirectoryReader.open(w).close();
}
Expand Down

0 comments on commit 69e6dbf

Please sign in to comment.