-
Notifications
You must be signed in to change notification settings - Fork 113
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add DocValuesProducers for releasing memory when close index (#1946)
Add DocValuesProducers for releasing memory when close index #1946
- Loading branch information
1 parent
8277bf0
commit 004fcc0
Showing
14 changed files
with
569 additions
and
29 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
63 changes: 63 additions & 0 deletions
63
src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80CompoundDirectory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.KNN80Codec; | ||
|
|
||
| import lombok.Getter; | ||
| import org.apache.lucene.codecs.CompoundDirectory; | ||
| import org.apache.lucene.store.Directory; | ||
| import org.apache.lucene.store.IOContext; | ||
| import org.apache.lucene.store.IndexInput; | ||
| import org.opensearch.knn.index.engine.KNNEngine; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Set; | ||
|
|
||
| public class KNN80CompoundDirectory extends CompoundDirectory { | ||
|
|
||
| @Getter | ||
| private CompoundDirectory delegate; | ||
| @Getter | ||
| private Directory dir; | ||
|
|
||
| public KNN80CompoundDirectory(CompoundDirectory delegate, Directory dir) { | ||
| this.delegate = delegate; | ||
| this.dir = dir; | ||
| } | ||
|
|
||
| @Override | ||
| public void checkIntegrity() throws IOException { | ||
| delegate.checkIntegrity(); | ||
| } | ||
|
|
||
| @Override | ||
| public String[] listAll() throws IOException { | ||
| return delegate.listAll(); | ||
| } | ||
|
|
||
| @Override | ||
| public long fileLength(String name) throws IOException { | ||
| return delegate.fileLength(name); | ||
| } | ||
|
|
||
| @Override | ||
| public IndexInput openInput(String name, IOContext context) throws IOException { | ||
| if (KNNEngine.getEnginesThatCreateCustomSegmentFiles().stream().anyMatch(engine -> name.endsWith(engine.getCompoundExtension()))) { | ||
| return dir.openInput(name, context); | ||
| } | ||
| return delegate.openInput(name, context); | ||
| } | ||
|
|
||
| @Override | ||
| public void close() throws IOException { | ||
| delegate.close(); | ||
| } | ||
|
|
||
| @Override | ||
| public Set<String> getPendingDeletions() throws IOException { | ||
| return delegate.getPendingDeletions(); | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
143 changes: 143 additions & 0 deletions
143
src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,143 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| * | ||
| * Modifications Copyright OpenSearch Contributors. See | ||
| * GitHub history for details. | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.KNN80Codec; | ||
|
|
||
| import lombok.NonNull; | ||
| import lombok.extern.log4j.Log4j2; | ||
| import org.apache.lucene.codecs.DocValuesProducer; | ||
| import org.apache.lucene.index.BinaryDocValues; | ||
| import org.apache.lucene.index.FieldInfo; | ||
| import org.apache.lucene.index.NumericDocValues; | ||
| import org.apache.lucene.index.SegmentReadState; | ||
| import org.apache.lucene.index.SortedDocValues; | ||
| import org.apache.lucene.index.SortedNumericDocValues; | ||
| import org.apache.lucene.index.SortedSetDocValues; | ||
| import org.apache.lucene.store.Directory; | ||
| import org.apache.lucene.store.FSDirectory; | ||
| import org.apache.lucene.store.FilterDirectory; | ||
| import org.opensearch.common.io.PathUtils; | ||
| import org.opensearch.knn.common.FieldInfoExtractor; | ||
| import org.opensearch.knn.index.codec.util.KNNCodecUtil; | ||
| import org.opensearch.knn.index.engine.KNNEngine; | ||
| import org.opensearch.knn.index.memory.NativeMemoryCacheManager; | ||
|
|
||
| import java.io.IOException; | ||
| import java.nio.file.Path; | ||
| import java.util.ArrayList; | ||
| import java.util.HashMap; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
|
|
||
| import static org.opensearch.knn.common.KNNConstants.MODEL_ID; | ||
| import static org.opensearch.knn.index.mapper.KNNVectorFieldMapper.KNN_FIELD; | ||
|
|
||
| @Log4j2 | ||
| public class KNN80DocValuesProducer extends DocValuesProducer { | ||
|
|
||
| private final SegmentReadState state; | ||
| private final DocValuesProducer delegate; | ||
| private final NativeMemoryCacheManager nativeMemoryCacheManager; | ||
| private final Map<String, String> indexPathMap = new HashMap(); | ||
|
|
||
| public KNN80DocValuesProducer(DocValuesProducer delegate, SegmentReadState state) { | ||
| this.delegate = delegate; | ||
| this.state = state; | ||
| this.nativeMemoryCacheManager = NativeMemoryCacheManager.getInstance(); | ||
|
|
||
| Directory directory = state.directory; | ||
| // directory would be CompoundDirectory, we need get directory firstly and then unwrap | ||
| if (state.directory instanceof KNN80CompoundDirectory) { | ||
| directory = ((KNN80CompoundDirectory) state.directory).getDir(); | ||
| } | ||
|
|
||
| Directory dir = FilterDirectory.unwrap(directory); | ||
| if (!(dir instanceof FSDirectory)) { | ||
| log.warn("{} can not casting to FSDirectory", directory); | ||
| return; | ||
| } | ||
| String directoryPath = ((FSDirectory) dir).getDirectory().toString(); | ||
| for (FieldInfo field : state.fieldInfos) { | ||
| if (!field.attributes().containsKey(KNN_FIELD)) { | ||
| continue; | ||
| } | ||
| // Only Native Engine put into indexPathMap | ||
| KNNEngine knnEngine = getNativeKNNEngine(field); | ||
| if (knnEngine == null) { | ||
| continue; | ||
| } | ||
| List<String> engineFiles = KNNCodecUtil.getEngineFiles(knnEngine.getExtension(), field.name, state.segmentInfo); | ||
| Path indexPath = PathUtils.get(directoryPath, engineFiles.get(0)); | ||
| indexPathMap.putIfAbsent(field.getName(), indexPath.toString()); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public BinaryDocValues getBinary(FieldInfo field) throws IOException { | ||
| return delegate.getBinary(field); | ||
| } | ||
|
|
||
| @Override | ||
| public NumericDocValues getNumeric(FieldInfo field) throws IOException { | ||
| return delegate.getNumeric(field); | ||
| } | ||
|
|
||
| @Override | ||
| public SortedDocValues getSorted(FieldInfo field) throws IOException { | ||
| return delegate.getSorted(field); | ||
| } | ||
|
|
||
| @Override | ||
| public SortedNumericDocValues getSortedNumeric(FieldInfo field) throws IOException { | ||
| return delegate.getSortedNumeric(field); | ||
| } | ||
|
|
||
| @Override | ||
| public SortedSetDocValues getSortedSet(FieldInfo field) throws IOException { | ||
| return delegate.getSortedSet(field); | ||
| } | ||
|
|
||
| @Override | ||
| public void checkIntegrity() throws IOException { | ||
| delegate.checkIntegrity(); | ||
| } | ||
|
|
||
| @Override | ||
| public void close() throws IOException { | ||
| for (String path : indexPathMap.values()) { | ||
| nativeMemoryCacheManager.invalidate(path); | ||
| } | ||
| delegate.close(); | ||
| } | ||
|
|
||
| public final List<String> getOpenedIndexPath() { | ||
| return new ArrayList<>(indexPathMap.values()); | ||
| } | ||
|
|
||
| /** | ||
| * Get KNNEngine From FieldInfo | ||
| * | ||
| * @param field which field we need produce from engine | ||
| * @return if and only if Native Engine we return specific engine, else return null | ||
| */ | ||
| private KNNEngine getNativeKNNEngine(@NonNull FieldInfo field) { | ||
|
|
||
| final String modelId = field.attributes().get(MODEL_ID); | ||
| if (modelId != null) { | ||
| return null; | ||
| } | ||
| KNNEngine engine = FieldInfoExtractor.extractKNNEngine(field); | ||
| if (KNNEngine.getEnginesThatCreateCustomSegmentFiles().contains(engine)) { | ||
| return engine; | ||
| } | ||
| return null; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.