Skip to content

Commit 38cb684

Browse files
jed326Jay Deng
jed326
authored and
Jay Deng
committed
Introduce RemoteIndexBuilder, refactor NativeIndexWriter into a separate interface
Signed-off-by: Jay Deng <jayd0104@gmail.com>
1 parent 349a715 commit 38cb684

16 files changed

+732
-341
lines changed

‎src/main/java/org/opensearch/knn/common/featureflags/KNNFeatureFlags.java

+16-1
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ public class KNNFeatureFlags {
2626

2727
// Feature flags
2828
private static final String KNN_FORCE_EVICT_CACHE_ENABLED = "knn.feature.cache.force_evict.enabled";
29+
private static final String KNN_REMOTE_VECTOR_BUILD = "knn.feature.remote_index_build.enabled";
2930

3031
@VisibleForTesting
3132
public static final Setting<Boolean> KNN_FORCE_EVICT_CACHE_ENABLED_SETTING = Setting.boolSetting(
@@ -35,8 +36,15 @@ public class KNNFeatureFlags {
3536
Dynamic
3637
);
3738

39+
public static final Setting<Boolean> KNN_REMOTE_VECTOR_BUILD_SETTING = Setting.boolSetting(
40+
KNN_REMOTE_VECTOR_BUILD,
41+
false,
42+
NodeScope,
43+
Dynamic
44+
);
45+
3846
public static List<Setting<?>> getFeatureFlags() {
39-
return ImmutableList.of(KNN_FORCE_EVICT_CACHE_ENABLED_SETTING);
47+
return ImmutableList.of(KNN_FORCE_EVICT_CACHE_ENABLED_SETTING, KNN_REMOTE_VECTOR_BUILD_SETTING);
4048
}
4149

4250
/**
@@ -46,4 +54,11 @@ public static List<Setting<?>> getFeatureFlags() {
4654
public static boolean isForceEvictCacheEnabled() {
4755
return Booleans.parseBoolean(KNNSettings.state().getSettingValue(KNN_FORCE_EVICT_CACHE_ENABLED).toString(), false);
4856
}
57+
58+
/**
59+
* @return true if remote vector index build feature flag is enabled
60+
*/
61+
public static boolean isKNNRemoteVectorBuildEnabled() {
62+
return Booleans.parseBooleanStrict(KNNSettings.state().getSettingValue(KNN_REMOTE_VECTOR_BUILD).toString(), false);
63+
}
4964
}

‎src/main/java/org/opensearch/knn/index/KNNSettings.java

+23-2
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,9 @@
4141

4242
import static java.util.stream.Collectors.toUnmodifiableMap;
4343
import static org.opensearch.common.settings.Setting.Property.Dynamic;
44+
import static org.opensearch.common.settings.Setting.Property.Final;
4445
import static org.opensearch.common.settings.Setting.Property.IndexScope;
4546
import static org.opensearch.common.settings.Setting.Property.NodeScope;
46-
import static org.opensearch.common.settings.Setting.Property.Final;
4747
import static org.opensearch.common.settings.Setting.Property.UnmodifiableOnRestore;
4848
import static org.opensearch.common.unit.MemorySizeValue.parseBytesSizeValueOrHeapRatio;
4949
import static org.opensearch.core.common.unit.ByteSizeValue.parseBytesSizeValue;
@@ -94,6 +94,8 @@ public class KNNSettings {
9494
public static final String KNN_FAISS_AVX512_SPR_DISABLED = "knn.faiss.avx512_spr.disabled";
9595
public static final String KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED = "index.knn.disk.vector.shard_level_rescoring_disabled";
9696
public static final String KNN_DERIVED_SOURCE_ENABLED = "index.knn.derived_source.enabled";
97+
public static final String KNN_INDEX_REMOTE_VECTOR_BUILD = "index.knn.remote_index_build.enabled";
98+
public static final String KNN_REMOTE_VECTOR_REPO = "knn.remote_index_build.vector_repo";
9799

98100
/**
99101
* Default setting values
@@ -371,6 +373,15 @@ public class KNNSettings {
371373
NodeScope
372374
);
373375

376+
public static final Setting<Boolean> KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING = Setting.boolSetting(
377+
KNN_INDEX_REMOTE_VECTOR_BUILD,
378+
false,
379+
Dynamic,
380+
IndexScope
381+
);
382+
383+
public static final Setting<String> KNN_REMOTE_VECTOR_REPO_SETTING = Setting.simpleString(KNN_REMOTE_VECTOR_REPO, Dynamic, NodeScope);
384+
374385
/**
375386
* Dynamic settings
376387
*/
@@ -525,6 +536,14 @@ private Setting<?> getSetting(String key) {
525536
return KNN_DERIVED_SOURCE_ENABLED_SETTING;
526537
}
527538

539+
if (KNN_INDEX_REMOTE_VECTOR_BUILD.equals(key)) {
540+
return KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING;
541+
}
542+
543+
if (KNN_REMOTE_VECTOR_REPO.equals(key)) {
544+
return KNN_REMOTE_VECTOR_REPO_SETTING;
545+
}
546+
528547
throw new IllegalArgumentException("Cannot find setting by key [" + key + "]");
529548
}
530549

@@ -550,7 +569,9 @@ public List<Setting<?>> getSettings() {
550569
QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING,
551570
QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING,
552571
KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED_SETTING,
553-
KNN_DERIVED_SOURCE_ENABLED_SETTING
572+
KNN_DERIVED_SOURCE_ENABLED_SETTING,
573+
KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING,
574+
KNN_REMOTE_VECTOR_REPO_SETTING
554575
);
555576
return Stream.concat(settings.stream(), Stream.concat(getFeatureFlags().stream(), dynamicCacheSettings.values().stream()))
556577
.collect(Collectors.toList());

‎src/main/java/org/opensearch/knn/index/codec/BasePerFieldKnnVectorsFormat.java

+24-6
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.opensearch.knn.index.engine.KNNMethodContext;
2222
import org.opensearch.knn.index.mapper.KNNMappingConfig;
2323
import org.opensearch.knn.index.mapper.KNNVectorFieldType;
24+
import org.opensearch.knn.index.remote.RemoteIndexBuilder;
2425

2526
import java.util.Map;
2627
import java.util.Optional;
@@ -44,6 +45,7 @@ public abstract class BasePerFieldKnnVectorsFormat extends PerFieldKnnVectorsFor
4445
private final Supplier<KnnVectorsFormat> defaultFormatSupplier;
4546
private final Function<KNNVectorsFormatParams, KnnVectorsFormat> vectorsFormatSupplier;
4647
private Function<KNNScalarQuantizedVectorsFormatParams, KnnVectorsFormat> scalarQuantizedVectorsFormatSupplier;
48+
private final RemoteIndexBuilder remoteIndexBuilder;
4749
private static final String MAX_CONNECTIONS = "max_connections";
4850
private static final String BEAM_WIDTH = "beam_width";
4951

@@ -54,11 +56,26 @@ public BasePerFieldKnnVectorsFormat(
5456
Supplier<KnnVectorsFormat> defaultFormatSupplier,
5557
Function<KNNVectorsFormatParams, KnnVectorsFormat> vectorsFormatSupplier
5658
) {
57-
this.mapperService = mapperService;
58-
this.defaultMaxConnections = defaultMaxConnections;
59-
this.defaultBeamWidth = defaultBeamWidth;
60-
this.defaultFormatSupplier = defaultFormatSupplier;
61-
this.vectorsFormatSupplier = vectorsFormatSupplier;
59+
this(mapperService, defaultMaxConnections, defaultBeamWidth, defaultFormatSupplier, vectorsFormatSupplier, null, null);
60+
}
61+
62+
public BasePerFieldKnnVectorsFormat(
63+
Optional<MapperService> mapperService,
64+
int defaultMaxConnections,
65+
int defaultBeamWidth,
66+
Supplier<KnnVectorsFormat> defaultFormatSupplier,
67+
Function<KNNVectorsFormatParams, KnnVectorsFormat> vectorsFormatSupplier,
68+
Function<KNNScalarQuantizedVectorsFormatParams, KnnVectorsFormat> scalarQuantizedVectorsFormatSupplier
69+
) {
70+
this(
71+
mapperService,
72+
defaultMaxConnections,
73+
defaultBeamWidth,
74+
defaultFormatSupplier,
75+
vectorsFormatSupplier,
76+
scalarQuantizedVectorsFormatSupplier,
77+
null
78+
);
6279
}
6380

6481
@Override
@@ -141,7 +158,8 @@ private NativeEngines990KnnVectorsFormat nativeEngineVectorsFormat() {
141158
int approximateThreshold = getApproximateThresholdValue();
142159
return new NativeEngines990KnnVectorsFormat(
143160
new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()),
144-
approximateThreshold
161+
approximateThreshold,
162+
remoteIndexBuilder
145163
);
146164
}
147165

‎src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/KNN9120PerFieldKnnVectorsFormat.java

+7-1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.opensearch.knn.index.SpaceType;
1414
import org.opensearch.knn.index.codec.BasePerFieldKnnVectorsFormat;
1515
import org.opensearch.knn.index.engine.KNNEngine;
16+
import org.opensearch.knn.index.remote.RemoteIndexBuilder;
1617

1718
import java.util.Optional;
1819
import java.util.concurrent.ExecutorService;
@@ -25,6 +26,10 @@ public class KNN9120PerFieldKnnVectorsFormat extends BasePerFieldKnnVectorsForma
2526
private static final Tuple<Integer, ExecutorService> DEFAULT_MERGE_THREAD_COUNT_AND_EXECUTOR_SERVICE = Tuple.tuple(1, null);
2627

2728
public KNN9120PerFieldKnnVectorsFormat(final Optional<MapperService> mapperService) {
29+
this(mapperService, null);
30+
}
31+
32+
public KNN9120PerFieldKnnVectorsFormat(final Optional<MapperService> mapperService, RemoteIndexBuilder remoteIndexBuilder) {
2833
super(
2934
mapperService,
3035
Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN,
@@ -67,7 +72,8 @@ public KNN9120PerFieldKnnVectorsFormat(final Optional<MapperService> mapperServi
6772
// Executor service
6873
mergeThreadCountAndExecutorService.v2()
6974
);
70-
}
75+
},
76+
remoteIndexBuilder
7177
);
7278
}
7379

‎src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java

+12-1
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.apache.lucene.index.SegmentWriteState;
2222
import org.opensearch.knn.index.KNNSettings;
2323
import org.opensearch.knn.index.engine.KNNEngine;
24+
import org.opensearch.knn.index.remote.RemoteIndexBuilder;
2425

2526
import java.io.IOException;
2627

@@ -33,6 +34,7 @@ public class NativeEngines990KnnVectorsFormat extends KnnVectorsFormat {
3334
private static FlatVectorsFormat flatVectorsFormat;
3435
private static final String FORMAT_NAME = "NativeEngines990KnnVectorsFormat";
3536
private static int approximateThreshold;
37+
private final RemoteIndexBuilder remoteIndexBuilder;
3638

3739
public NativeEngines990KnnVectorsFormat() {
3840
this(new Lucene99FlatVectorsFormat(new DefaultFlatVectorScorer()));
@@ -47,9 +49,18 @@ public NativeEngines990KnnVectorsFormat(final FlatVectorsFormat flatVectorsForma
4749
}
4850

4951
public NativeEngines990KnnVectorsFormat(final FlatVectorsFormat flatVectorsFormat, int approximateThreshold) {
52+
this(flatVectorsFormat, approximateThreshold, null);
53+
}
54+
55+
public NativeEngines990KnnVectorsFormat(
56+
final FlatVectorsFormat flatVectorsFormat,
57+
int approximateThreshold,
58+
RemoteIndexBuilder remoteIndexBuilder
59+
) {
5060
super(FORMAT_NAME);
5161
NativeEngines990KnnVectorsFormat.flatVectorsFormat = flatVectorsFormat;
5262
NativeEngines990KnnVectorsFormat.approximateThreshold = approximateThreshold;
63+
this.remoteIndexBuilder = remoteIndexBuilder;
5364
}
5465

5566
/**
@@ -59,7 +70,7 @@ public NativeEngines990KnnVectorsFormat(final FlatVectorsFormat flatVectorsForma
5970
*/
6071
@Override
6172
public KnnVectorsWriter fieldsWriter(final SegmentWriteState state) throws IOException {
62-
return new NativeEngines990KnnVectorsWriter(state, flatVectorsFormat.fieldsWriter(state), approximateThreshold);
73+
return new NativeEngines990KnnVectorsWriter(state, flatVectorsFormat.fieldsWriter(state), approximateThreshold, remoteIndexBuilder);
6374
}
6475

6576
/**

‎src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java

+27-2
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.opensearch.knn.index.VectorDataType;
2929
import org.opensearch.knn.index.codec.nativeindex.NativeIndexWriter;
3030
import org.opensearch.knn.index.quantizationservice.QuantizationService;
31+
import org.opensearch.knn.index.remote.RemoteIndexBuilder;
3132
import org.opensearch.knn.index.vectorvalues.KNNVectorValues;
3233
import org.opensearch.knn.plugin.stats.KNNGraphValue;
3334
import org.opensearch.knn.quantization.models.quantizationParams.QuantizationParams;
@@ -54,15 +55,26 @@ public class NativeEngines990KnnVectorsWriter extends KnnVectorsWriter {
5455
private final List<NativeEngineFieldVectorsWriter<?>> fields = new ArrayList<>();
5556
private boolean finished;
5657
private final Integer approximateThreshold;
58+
private final RemoteIndexBuilder remoteIndexBuilder;
5759

5860
public NativeEngines990KnnVectorsWriter(
5961
SegmentWriteState segmentWriteState,
6062
FlatVectorsWriter flatVectorsWriter,
6163
Integer approximateThreshold
64+
) {
65+
this(segmentWriteState, flatVectorsWriter, approximateThreshold, null);
66+
}
67+
68+
public NativeEngines990KnnVectorsWriter(
69+
SegmentWriteState segmentWriteState,
70+
FlatVectorsWriter flatVectorsWriter,
71+
Integer approximateThreshold,
72+
RemoteIndexBuilder remoteIndexBuilder
6273
) {
6374
this.segmentWriteState = segmentWriteState;
6475
this.flatVectorsWriter = flatVectorsWriter;
6576
this.approximateThreshold = approximateThreshold;
77+
this.remoteIndexBuilder = remoteIndexBuilder;
6678
}
6779

6880
/**
@@ -114,7 +126,14 @@ public void flush(int maxDoc, final Sorter.DocMap sortMap) throws IOException {
114126
);
115127
continue;
116128
}
117-
final NativeIndexWriter writer = NativeIndexWriter.getWriter(fieldInfo, segmentWriteState, quantizationState);
129+
log.info("GOT HERE - FLUSH");
130+
final NativeIndexWriter writer = NativeIndexWriter.getWriter(
131+
fieldInfo,
132+
segmentWriteState,
133+
quantizationState,
134+
remoteIndexBuilder,
135+
knnVectorValuesSupplier
136+
);
118137
final KNNVectorValues<?> knnVectorValues = knnVectorValuesSupplier.get();
119138

120139
StopWatch stopWatch = new StopWatch().start();
@@ -153,7 +172,13 @@ public void mergeOneField(final FieldInfo fieldInfo, final MergeState mergeState
153172
);
154173
return;
155174
}
156-
final NativeIndexWriter writer = NativeIndexWriter.getWriter(fieldInfo, segmentWriteState, quantizationState);
175+
final NativeIndexWriter writer = NativeIndexWriter.getWriter(
176+
fieldInfo,
177+
segmentWriteState,
178+
quantizationState,
179+
remoteIndexBuilder,
180+
knnVectorValuesSupplier
181+
);
157182
final KNNVectorValues<?> knnVectorValues = knnVectorValuesSupplier.get();
158183

159184
StopWatch stopWatch = new StopWatch().start();

‎src/main/java/org/opensearch/knn/index/codec/KNNCodecService.java

+6-3
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,24 @@
55

66
package org.opensearch.knn.index.codec;
77

8-
import org.opensearch.index.codec.CodecServiceConfig;
98
import org.apache.lucene.codecs.Codec;
109
import org.opensearch.index.codec.CodecService;
10+
import org.opensearch.index.codec.CodecServiceConfig;
1111
import org.opensearch.index.mapper.MapperService;
12+
import org.opensearch.knn.index.remote.RemoteIndexBuilder;
1213

1314
/**
1415
* KNNCodecService to inject the right KNNCodec version
1516
*/
1617
public class KNNCodecService extends CodecService {
1718

1819
private final MapperService mapperService;
20+
private final RemoteIndexBuilder remoteIndexBuilder;
1921

20-
public KNNCodecService(CodecServiceConfig codecServiceConfig) {
22+
public KNNCodecService(CodecServiceConfig codecServiceConfig, RemoteIndexBuilder remoteIndexBuilder) {
2123
super(codecServiceConfig.getMapperService(), codecServiceConfig.getIndexSettings(), codecServiceConfig.getLogger());
2224
mapperService = codecServiceConfig.getMapperService();
25+
this.remoteIndexBuilder = remoteIndexBuilder;
2326
}
2427

2528
/**
@@ -30,6 +33,6 @@ public KNNCodecService(CodecServiceConfig codecServiceConfig) {
3033
*/
3134
@Override
3235
public Codec codec(String name) {
33-
return KNNCodecVersion.current().getKnnCodecSupplier().apply(super.codec(name), mapperService);
36+
return KNNCodecVersion.current().getKnnCodecSupplier().apply(super.codec(name), mapperService, remoteIndexBuilder);
3437
}
3538
}

0 commit comments

Comments
 (0)