Skip to content

Commit 278f1d0

Browse files
benwtrentbenwtrent
authored
Cleaning up exitable vector value impls (elastic#140190)
Need to run CI, but this cleans up our exitable directory impls by extracting out the vector values and providing `FilterFloat/ByteVectorValues`. I think these could be use elsewhere in the code base, but just doing this initial refactor for now. Also, I noticed that the copy wasn't actually copying the query cancellation logic here 🤦 . Which is likely a latent bug in cancellation we haven't really found yet.
1 parent f8ddd6d commit 278f1d0

File tree

3 files changed

+150
-77
lines changed

3 files changed

+150
-77
lines changed

‎server/src/main/java/org/elasticsearch/search/internal/ExitableDirectoryReader.java‎

Lines changed: 10 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import org.elasticsearch.index.codec.vectors.BulkScorableFloatVectorValues;
3535

3636
import java.io.IOException;
37-
import java.util.Objects;
3837

3938
/**
4039
* Wraps an {@link IndexReader} with a {@link QueryCancellation}
@@ -138,7 +137,7 @@ public ByteVectorValues getByteVectorValues(String field) throws IOException {
138137
if (vectorValues == null) {
139138
return null;
140139
}
141-
return queryCancellation.isEnabled() ? new ExitableByteVectorValues(queryCancellation, vectorValues) : vectorValues;
140+
return queryCancellation.isEnabled() ? new ExitableByteVectorValues(vectorValues, queryCancellation) : vectorValues;
142141
}
143142

144143
@Override
@@ -455,33 +454,12 @@ public void grow(int count) {
455454
}
456455
}
457456

458-
private static class ExitableByteVectorValues extends ByteVectorValues {
457+
private static class ExitableByteVectorValues extends FilterByteVectorValues {
459458
private final QueryCancellation queryCancellation;
460-
private final ByteVectorValues in;
461459

462-
private ExitableByteVectorValues(QueryCancellation queryCancellation, ByteVectorValues in) {
460+
private ExitableByteVectorValues(ByteVectorValues in, QueryCancellation queryCancellation) {
461+
super(in);
463462
this.queryCancellation = queryCancellation;
464-
this.in = in;
465-
}
466-
467-
@Override
468-
public int dimension() {
469-
return in.dimension();
470-
}
471-
472-
@Override
473-
public int size() {
474-
return in.size();
475-
}
476-
477-
@Override
478-
public byte[] vectorValue(int ord) throws IOException {
479-
return in.vectorValue(ord);
480-
}
481-
482-
@Override
483-
public int ordToDoc(int ord) {
484-
return in.ordToDoc(ord);
485463
}
486464

487465
@Override
@@ -513,7 +491,7 @@ public DocIndexIterator iterator() {
513491

514492
@Override
515493
public ByteVectorValues copy() throws IOException {
516-
return in.copy();
494+
return new ExitableByteVectorValues(in.copy(), queryCancellation);
517495
}
518496
}
519497

@@ -544,16 +522,6 @@ private static class ExitableBulkScorableFloatVectorValues extends FilterFloatVe
544522
this.bsfvv = bsfvv;
545523
}
546524

547-
@Override
548-
public float[] vectorValue(int ord) throws IOException {
549-
return in.vectorValue(ord);
550-
}
551-
552-
@Override
553-
public int ordToDoc(int ord) {
554-
return in.ordToDoc(ord);
555-
}
556-
557525
@Override
558526
public VectorScorer scorer(float[] target) throws IOException {
559527
VectorScorer scorer = in.scorer(target);
@@ -583,7 +551,10 @@ public DocIndexIterator iterator() {
583551

584552
@Override
585553
public FloatVectorValues copy() throws IOException {
586-
return in.copy();
554+
assert in instanceof BulkScorableFloatVectorValues;
555+
FloatVectorValues copy = this.in.copy();
556+
BulkScorableFloatVectorValues bulkScorableCopy = (BulkScorableFloatVectorValues) copy;
557+
return new ExitableBulkScorableFloatVectorValues(copy, bulkScorableCopy, queryCancellation);
587558
}
588559

589560
@Override
@@ -645,7 +616,7 @@ public DocIndexIterator iterator() {
645616

646617
@Override
647618
public FloatVectorValues copy() throws IOException {
648-
return in.copy();
619+
return new ExitableFloatVectorValues(in.copy(), queryCancellation);
649620
}
650621
}
651622

@@ -745,42 +716,4 @@ private void checkAndThrowWithSampling() {
745716
}
746717
}
747718

748-
/** Delegates all methods to a wrapped {@link FloatVectorValues}. */
749-
private abstract static class FilterFloatVectorValues extends FloatVectorValues {
750-
751-
/** Wrapped values */
752-
protected final FloatVectorValues in;
753-
754-
/** Sole constructor */
755-
protected FilterFloatVectorValues(FloatVectorValues in) {
756-
Objects.requireNonNull(in);
757-
this.in = in;
758-
}
759-
760-
@Override
761-
public DocIndexIterator iterator() {
762-
return in.iterator();
763-
}
764-
765-
@Override
766-
public float[] vectorValue(int ord) throws IOException {
767-
return in.vectorValue(ord);
768-
}
769-
770-
@Override
771-
public FloatVectorValues copy() throws IOException {
772-
return in.copy();
773-
}
774-
775-
@Override
776-
public int dimension() {
777-
return in.dimension();
778-
}
779-
780-
@Override
781-
public int size() {
782-
return in.size();
783-
}
784-
785-
}
786719
}

‎server/src/main/java/org/elasticsearch/search/internal/FilterByteVectorValues.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.search.internal;
11+
12+
import org.apache.lucene.index.ByteVectorValues;
13+
import org.apache.lucene.index.KnnVectorValues;
14+
import org.apache.lucene.search.VectorScorer;
15+
import org.apache.lucene.util.Bits;
16+
17+
import java.io.IOException;
18+
import java.util.Objects;
19+
20+
/**
21+
* A wrapper on {@link ByteVectorValues}.
22+
*/
23+
public abstract class FilterByteVectorValues extends ByteVectorValues {
24+
25+
/** Wrapped values */
26+
protected final ByteVectorValues in;
27+
28+
/** Sole constructor */
29+
protected FilterByteVectorValues(ByteVectorValues in) {
30+
this.in = Objects.requireNonNull(in);
31+
}
32+
33+
@Override
34+
public KnnVectorValues.DocIndexIterator iterator() {
35+
return in.iterator();
36+
}
37+
38+
@Override
39+
public byte[] vectorValue(int ord) throws IOException {
40+
return in.vectorValue(ord);
41+
}
42+
43+
@Override
44+
public abstract ByteVectorValues copy() throws IOException;
45+
46+
@Override
47+
public int dimension() {
48+
return in.dimension();
49+
}
50+
51+
@Override
52+
public int size() {
53+
return in.size();
54+
}
55+
56+
@Override
57+
public VectorScorer scorer(byte[] target) throws IOException {
58+
return in.scorer(target);
59+
}
60+
61+
@Override
62+
public int ordToDoc(int ord) {
63+
return in.ordToDoc(ord);
64+
}
65+
66+
@Override
67+
public Bits getAcceptOrds(Bits acceptDocs) {
68+
return in.getAcceptOrds(acceptDocs);
69+
}
70+
}

‎server/src/main/java/org/elasticsearch/search/internal/FilterFloatVectorValues.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.search.internal;
11+
12+
import org.apache.lucene.index.FloatVectorValues;
13+
import org.apache.lucene.index.KnnVectorValues;
14+
import org.apache.lucene.search.VectorScorer;
15+
import org.apache.lucene.util.Bits;
16+
17+
import java.io.IOException;
18+
import java.util.Objects;
19+
20+
/**
21+
* A wrapper on {@link FloatVectorValues}.
22+
*/
23+
public abstract class FilterFloatVectorValues extends FloatVectorValues {
24+
25+
/** Wrapped values */
26+
protected final FloatVectorValues in;
27+
28+
/** Sole constructor */
29+
protected FilterFloatVectorValues(FloatVectorValues in) {
30+
this.in = Objects.requireNonNull(in);
31+
}
32+
33+
@Override
34+
public KnnVectorValues.DocIndexIterator iterator() {
35+
return in.iterator();
36+
}
37+
38+
@Override
39+
public float[] vectorValue(int ord) throws IOException {
40+
return in.vectorValue(ord);
41+
}
42+
43+
@Override
44+
public abstract FloatVectorValues copy() throws IOException;
45+
46+
@Override
47+
public int dimension() {
48+
return in.dimension();
49+
}
50+
51+
@Override
52+
public int size() {
53+
return in.size();
54+
}
55+
56+
@Override
57+
public VectorScorer scorer(float[] target) throws IOException {
58+
return in.scorer(target);
59+
}
60+
61+
@Override
62+
public int ordToDoc(int ord) {
63+
return in.ordToDoc(ord);
64+
}
65+
66+
@Override
67+
public Bits getAcceptOrds(Bits acceptDocs) {
68+
return in.getAcceptOrds(acceptDocs);
69+
}
70+
}

0 commit comments

Comments
 (0)