Skip to content

Commit cf937c6

Browse files
dweissdweiss
committed
Revert "Update the IOContext rather than the ReadAdvice on IndexInput (#14702)"
This reverts commit 9936b20.
1 parent 823718f commit cf937c6

File tree

10 files changed

+78
-59
lines changed

10 files changed

+78
-59
lines changed

‎lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsReader.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ public abstract void search(
124124
*
125125
* <p>The default implementation returns {@code this}
126126
*/
127-
public KnnVectorsReader getMergeInstance() throws IOException {
127+
public KnnVectorsReader getMergeInstance() {
128128
return this;
129129
}
130130

‎lucene/core/src/java/org/apache/lucene/codecs/hnsw/FlatVectorsReader.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ public abstract RandomVectorScorer getRandomVectorScorer(String field, byte[] ta
9696
* <p>The default implementation returns {@code this}
9797
*/
9898
@Override
99-
public FlatVectorsReader getMergeInstance() throws IOException {
99+
public FlatVectorsReader getMergeInstance() {
100100
return this;
101101
}
102102
}

‎lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsReader.java‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
2222

2323
import java.io.IOException;
24+
import java.io.UncheckedIOException;
2425
import java.util.Map;
2526
import org.apache.lucene.codecs.CodecUtil;
2627
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
@@ -44,6 +45,7 @@
4445
import org.apache.lucene.store.FileTypeHint;
4546
import org.apache.lucene.store.IOContext;
4647
import org.apache.lucene.store.IndexInput;
48+
import org.apache.lucene.store.ReadAdvice;
4749
import org.apache.lucene.util.IOUtils;
4850
import org.apache.lucene.util.RamUsageEstimator;
4951
import org.apache.lucene.util.hnsw.RandomVectorScorer;
@@ -61,25 +63,23 @@ public final class Lucene99FlatVectorsReader extends FlatVectorsReader {
6163
private final IntObjectHashMap<FieldEntry> fields = new IntObjectHashMap<>();
6264
private final IndexInput vectorData;
6365
private final FieldInfos fieldInfos;
64-
private final IOContext dataContext;
6566

6667
public Lucene99FlatVectorsReader(SegmentReadState state, FlatVectorsScorer scorer)
6768
throws IOException {
6869
super(scorer);
6970
int versionMeta = readMetadata(state);
7071
this.fieldInfos = state.fieldInfos;
71-
// Flat formats are used to randomly access vectors from their node ID that is stored
72-
// in the HNSW graph.
73-
dataContext =
74-
state.context.withHints(FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM);
7572
try {
7673
vectorData =
7774
openDataInput(
7875
state,
7976
versionMeta,
8077
Lucene99FlatVectorsFormat.VECTOR_DATA_EXTENSION,
8178
Lucene99FlatVectorsFormat.VECTOR_DATA_CODEC_NAME,
82-
dataContext);
79+
// Flat formats are used to randomly access vectors from their node ID that is stored
80+
// in the HNSW graph.
81+
state.context.withHints(
82+
FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM));
8383
} catch (Throwable t) {
8484
IOUtils.closeWhileSuppressingExceptions(t, this);
8585
throw t;
@@ -177,10 +177,14 @@ public void checkIntegrity() throws IOException {
177177
}
178178

179179
@Override
180-
public FlatVectorsReader getMergeInstance() throws IOException {
181-
// Update the read advice since vectors are guaranteed to be accessed sequentially for merge
182-
vectorData.updateIOContext(dataContext.withHints(DataAccessHint.SEQUENTIAL));
183-
return this;
180+
public FlatVectorsReader getMergeInstance() {
181+
try {
182+
// Update the read advice since vectors are guaranteed to be accessed sequentially for merge
183+
this.vectorData.updateReadAdvice(ReadAdvice.SEQUENTIAL);
184+
return this;
185+
} catch (IOException exception) {
186+
throw new UncheckedIOException(exception);
187+
}
184188
}
185189

186190
private FieldEntry getFieldEntryOrThrow(String field) {
@@ -272,7 +276,7 @@ public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) thr
272276
public void finishMerge() throws IOException {
273277
// This makes sure that the access pattern hint is reverted back since HNSW implementation
274278
// needs it
275-
vectorData.updateIOContext(dataContext);
279+
this.vectorData.updateReadAdvice(ReadAdvice.RANDOM);
276280
}
277281

278282
@Override

‎lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ private Lucene99HnswVectorsReader(
131131
}
132132

133133
@Override
134-
public KnnVectorsReader getMergeInstance() throws IOException {
134+
public KnnVectorsReader getMergeInstance() {
135135
return new Lucene99HnswVectorsReader(this, this.flatVectorsReader.getMergeInstance());
136136
}
137137

‎lucene/core/src/java/org/apache/lucene/codecs/perfield/PerFieldKnnVectorsFormat.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ public FieldsReader(final SegmentReadState readState) throws IOException {
238238
}
239239
}
240240

241-
private FieldsReader(final FieldsReader fieldsReader) throws IOException {
241+
private FieldsReader(final FieldsReader fieldsReader) {
242242
this.fieldInfos = fieldsReader.fieldInfos;
243243
for (FieldInfo fi : this.fieldInfos) {
244244
if (fi.hasVectorValues() && fieldsReader.fields.containsKey(fi.number)) {
@@ -248,7 +248,7 @@ private FieldsReader(final FieldsReader fieldsReader) throws IOException {
248248
}
249249

250250
@Override
251-
public KnnVectorsReader getMergeInstance() throws IOException {
251+
public KnnVectorsReader getMergeInstance() {
252252
return new FieldsReader(this);
253253
}
254254

‎lucene/core/src/java/org/apache/lucene/store/IndexInput.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -227,12 +227,12 @@ public String toString() {
227227
public void prefetch(long offset, long length) throws IOException {}
228228

229229
/**
230-
* Optional method: Updates the {@code IOContext} to specify a new read access pattern. IndexInput
230+
* Optional method: Give a hint to this input about the change in read access pattern. IndexInput
231231
* implementations may take advantage of this hint to optimize reads from storage.
232232
*
233233
* <p>The default implementation is a no-op.
234234
*/
235-
public void updateIOContext(IOContext context) throws IOException {}
235+
public void updateReadAdvice(ReadAdvice readAdvice) throws IOException {}
236236

237237
/**
238238
* Returns a hint whether all the contents of this input are resident in physical memory. It's a

‎lucene/core/src/java/org/apache/lucene/store/MemorySegmentIndexInput.java‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -357,11 +357,7 @@ public void prefetch(long offset, long length) throws IOException {
357357
}
358358

359359
@Override
360-
public void updateIOContext(IOContext context) throws IOException {
361-
updateReadAdvice(toReadAdvice.apply(context));
362-
}
363-
364-
private void updateReadAdvice(ReadAdvice readAdvice) throws IOException {
360+
public void updateReadAdvice(ReadAdvice readAdvice) throws IOException {
365361
if (NATIVE_ACCESS.isEmpty()) {
366362
return;
367363
}

‎lucene/test-framework/src/java/org/apache/lucene/tests/codecs/asserting/AssertingKnnVectorsFormat.java‎

Lines changed: 19 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -113,17 +113,16 @@ public long ramBytesUsed() {
113113
public static class AssertingKnnVectorsReader extends KnnVectorsReader
114114
implements HnswGraphProvider {
115115
public final KnnVectorsReader delegate;
116-
private final FieldInfos fis;
117-
private final boolean mergeInstance;
118-
private final AtomicInteger mergeInstanceCount = new AtomicInteger();
119-
private final AtomicInteger finishMergeCount = new AtomicInteger();
116+
final FieldInfos fis;
117+
final boolean mergeInstance;
118+
AtomicInteger mergeInstanceCount = new AtomicInteger();
119+
AtomicInteger finishMergeCount = new AtomicInteger();
120120

121-
private AssertingKnnVectorsReader(KnnVectorsReader delegate, FieldInfos fis) {
121+
AssertingKnnVectorsReader(KnnVectorsReader delegate, FieldInfos fis) {
122122
this(delegate, fis, false);
123123
}
124124

125-
private AssertingKnnVectorsReader(
126-
KnnVectorsReader delegate, FieldInfos fis, boolean mergeInstance) {
125+
AssertingKnnVectorsReader(KnnVectorsReader delegate, FieldInfos fis, boolean mergeInstance) {
127126
assert delegate != null;
128127
this.delegate = delegate;
129128
this.fis = fis;
@@ -137,8 +136,6 @@ public void checkIntegrity() throws IOException {
137136

138137
@Override
139138
public FloatVectorValues getFloatVectorValues(String field) throws IOException {
140-
assert mergeInstanceCount.get() == finishMergeCount.get() || mergeInstance
141-
: "Called on the wrong instance";
142139
FieldInfo fi = fis.fieldInfo(field);
143140
assert fi != null
144141
&& fi.getVectorDimension() > 0
@@ -153,8 +150,6 @@ public FloatVectorValues getFloatVectorValues(String field) throws IOException {
153150

154151
@Override
155152
public ByteVectorValues getByteVectorValues(String field) throws IOException {
156-
assert mergeInstanceCount.get() == finishMergeCount.get() || mergeInstance
157-
: "Called on the wrong instance";
158153
FieldInfo fi = fis.fieldInfo(field);
159154
assert fi != null
160155
&& fi.getVectorDimension() > 0
@@ -170,7 +165,7 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
170165
@Override
171166
public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs)
172167
throws IOException {
173-
assert mergeInstanceCount.get() == finishMergeCount.get() : "There is an open merge instance";
168+
assert !mergeInstance;
174169
FieldInfo fi = fis.fieldInfo(field);
175170
assert fi != null
176171
&& fi.getVectorDimension() > 0
@@ -181,7 +176,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
181176
@Override
182177
public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs)
183178
throws IOException {
184-
assert mergeInstanceCount.get() == finishMergeCount.get() : "There is an open merge instance";
179+
assert !mergeInstance;
185180
FieldInfo fi = fis.fieldInfo(field);
186181
assert fi != null
187182
&& fi.getVectorDimension() > 0
@@ -190,28 +185,15 @@ public void search(String field, byte[] target, KnnCollector knnCollector, Bits
190185
}
191186

192187
@Override
193-
public KnnVectorsReader getMergeInstance() throws IOException {
188+
public KnnVectorsReader getMergeInstance() {
189+
assert !mergeInstance;
194190
var mergeVectorsReader = delegate.getMergeInstance();
195191
assert mergeVectorsReader != null;
196192
mergeInstanceCount.incrementAndGet();
197-
AtomicInteger parentMergeFinishCount = this.finishMergeCount;
198193

194+
final var parent = this;
199195
return new AssertingKnnVectorsReader(
200196
mergeVectorsReader, AssertingKnnVectorsReader.this.fis, true) {
201-
private boolean finished;
202-
203-
@Override
204-
public void search(
205-
String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) {
206-
assert false : "This instance should only be used for merging";
207-
}
208-
209-
@Override
210-
public void search(
211-
String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) {
212-
assert false : "This instance should only be used for merging";
213-
}
214-
215197
@Override
216198
public KnnVectorsReader getMergeInstance() {
217199
assert false; // merging from a merge instance it not allowed
@@ -220,10 +202,9 @@ public KnnVectorsReader getMergeInstance() {
220202

221203
@Override
222204
public void finishMerge() throws IOException {
223-
assert !finished : "Merging already finished";
224-
finished = true;
205+
assert mergeInstance;
225206
delegate.finishMerge();
226-
parentMergeFinishCount.incrementAndGet();
207+
parent.finishMergeCount.incrementAndGet();
227208
}
228209

229210
@Override
@@ -235,7 +216,9 @@ public void close() {
235216

236217
@Override
237218
public void finishMerge() throws IOException {
238-
assert false; // can only finish merge on the merge instance
219+
assert mergeInstance;
220+
delegate.finishMerge();
221+
finishMergeCount.incrementAndGet();
239222
}
240223

241224
@Override
@@ -245,8 +228,10 @@ public Map<String, Long> getOffHeapByteSize(FieldInfo fieldInfo) {
245228

246229
@Override
247230
public void close() throws IOException {
248-
assert mergeInstanceCount.get() == finishMergeCount.get();
231+
assert !mergeInstance;
232+
delegate.close();
249233
delegate.close();
234+
assert finishMergeCount.get() <= 0 || mergeInstanceCount.get() == finishMergeCount.get();
250235
}
251236

252237
@Override

‎lucene/test-framework/src/java/org/apache/lucene/tests/store/BaseDirectoryTestCase.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import org.apache.lucene.store.IndexOutput;
6262
import org.apache.lucene.store.MMapDirectory;
6363
import org.apache.lucene.store.RandomAccessInput;
64+
import org.apache.lucene.store.ReadAdvice;
6465
import org.apache.lucene.tests.mockfile.ExtrasFS;
6566
import org.apache.lucene.tests.util.LuceneTestCase;
6667
import org.apache.lucene.tests.util.TestUtil;
@@ -1570,6 +1571,38 @@ public void testPrefetchOnSlice() throws IOException {
15701571
doTestPrefetch(TestUtil.nextInt(random(), 1, 1024));
15711572
}
15721573

1574+
public void testUpdateReadAdvice() throws IOException {
1575+
try (Directory dir = getDirectory(createTempDir("testUpdateReadAdvice"))) {
1576+
final int totalLength = TestUtil.nextInt(random(), 16384, 65536);
1577+
byte[] arr = new byte[totalLength];
1578+
random().nextBytes(arr);
1579+
try (IndexOutput out = dir.createOutput("temp.bin", IOContext.DEFAULT)) {
1580+
out.writeBytes(arr, arr.length);
1581+
}
1582+
1583+
try (IndexInput orig = dir.openInput("temp.bin", IOContext.DEFAULT)) {
1584+
IndexInput in = random().nextBoolean() ? orig.clone() : orig;
1585+
// Read advice updated at start
1586+
in.updateReadAdvice(randomFrom(random(), ReadAdvice.values()));
1587+
for (int i = 0; i < totalLength; i++) {
1588+
int offset = TestUtil.nextInt(random(), 0, (int) in.length() - 1);
1589+
in.seek(offset);
1590+
assertEquals(arr[offset], in.readByte());
1591+
}
1592+
1593+
// Updating readAdvice in the middle
1594+
for (int i = 0; i < 10_000; ++i) {
1595+
int offset = TestUtil.nextInt(random(), 0, (int) in.length() - 1);
1596+
in.seek(offset);
1597+
assertEquals(arr[offset], in.readByte());
1598+
if (random().nextBoolean()) {
1599+
in.updateReadAdvice(randomFrom(random(), ReadAdvice.values()));
1600+
}
1601+
}
1602+
}
1603+
}
1604+
}
1605+
15731606
private void doTestPrefetch(int startOffset) throws IOException {
15741607
try (Directory dir = getDirectory(createTempDir())) {
15751608
final int totalLength = startOffset + TestUtil.nextInt(random(), 16384, 65536);

‎lucene/test-framework/src/java/org/apache/lucene/tests/store/MockIndexInputWrapper.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.apache.lucene.store.FilterIndexInput;
2626
import org.apache.lucene.store.IOContext;
2727
import org.apache.lucene.store.IndexInput;
28+
import org.apache.lucene.store.ReadAdvice;
2829

2930
/**
3031
* Used by MockDirectoryWrapper to create an input stream that keeps track of when it's been closed.
@@ -185,10 +186,10 @@ public Optional<Boolean> isLoaded() {
185186
}
186187

187188
@Override
188-
public void updateIOContext(IOContext context) throws IOException {
189+
public void updateReadAdvice(ReadAdvice readAdvice) throws IOException {
189190
ensureOpen();
190191
ensureAccessible();
191-
in.updateIOContext(context);
192+
in.updateReadAdvice(readAdvice);
192193
}
193194

194195
@Override

0 commit comments

Comments
 (0)