-
Notifications
You must be signed in to change notification settings - Fork 128
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Initial integration with OS 2.0 Alpha1
Signed-off-by: Martin Gaievski <[email protected]>
- Loading branch information
1 parent
a45dc81
commit 12f9bec
Showing
31 changed files
with
1,226 additions
and
36 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
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
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
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
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
64 changes: 64 additions & 0 deletions
64
src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91BinaryDocValues.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,64 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.knn.index.codec.KNN91Codec; | ||
|
||
import org.apache.lucene.index.BinaryDocValues; | ||
import org.apache.lucene.index.DocIDMerger; | ||
import org.apache.lucene.util.BytesRef; | ||
import org.opensearch.knn.index.codec.util.BinaryDocValuesSub; | ||
|
||
import java.io.IOException; | ||
|
||
/** | ||
* A per-document kNN numeric value. | ||
*/ | ||
class KNN91BinaryDocValues extends BinaryDocValues { | ||
|
||
private DocIDMerger<BinaryDocValuesSub> docIDMerger; | ||
|
||
KNN91BinaryDocValues(DocIDMerger<BinaryDocValuesSub> docIdMerger) { | ||
this.docIDMerger = docIdMerger; | ||
} | ||
|
||
private BinaryDocValuesSub current; | ||
private int docID = -1; | ||
|
||
@Override | ||
public int docID() { | ||
return docID; | ||
} | ||
|
||
@Override | ||
public int nextDoc() throws IOException { | ||
current = docIDMerger.next(); | ||
if (current == null) { | ||
docID = NO_MORE_DOCS; | ||
} else { | ||
docID = current.mappedDocID; | ||
} | ||
return docID; | ||
} | ||
|
||
@Override | ||
public int advance(int target) throws IOException { | ||
throw new UnsupportedOperationException(); | ||
} | ||
|
||
@Override | ||
public boolean advanceExact(int target) throws IOException { | ||
throw new UnsupportedOperationException(); | ||
} | ||
|
||
@Override | ||
public long cost() { | ||
throw new UnsupportedOperationException(); | ||
} | ||
|
||
@Override | ||
public BytesRef binaryValue() throws IOException { | ||
return current.getValues().binaryValue(); | ||
} | ||
}; |
58 changes: 58 additions & 0 deletions
58
src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91Codec.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,58 @@ | ||
/* | ||
* 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.KNN91Codec; | ||
|
||
import org.apache.lucene.codecs.Codec; | ||
import org.apache.lucene.codecs.CompoundFormat; | ||
import org.apache.lucene.codecs.DocValuesFormat; | ||
import org.apache.lucene.codecs.FilterCodec; | ||
import org.apache.lucene.codecs.lucene91.Lucene91Codec; | ||
|
||
/** | ||
* Extends the Codec to support a new file format for KNN index | ||
* based on the mappings. | ||
* | ||
*/ | ||
public final class KNN91Codec extends FilterCodec { | ||
|
||
private final DocValuesFormat docValuesFormat; | ||
private final CompoundFormat compoundFormat; | ||
|
||
public static final String KNN_91 = "KNN91Codec"; | ||
|
||
/** | ||
* No arg constructor that uses Lucene91 as the delegate | ||
*/ | ||
public KNN91Codec() { | ||
this(new Lucene91Codec()); | ||
} | ||
/** | ||
* Constructor that takes a Codec delegate to delegate all methods this code does not implement to. | ||
* | ||
* @param delegate codec that will perform all operations this codec does not override | ||
*/ | ||
public KNN91Codec(Codec delegate) { | ||
super(KNN_91, delegate); | ||
this.docValuesFormat = new KNN91DocValuesFormat(delegate.docValuesFormat()); | ||
this.compoundFormat = new KNN91CompoundFormat(delegate.compoundFormat()); | ||
} | ||
|
||
@Override | ||
public DocValuesFormat docValuesFormat() { | ||
return this.docValuesFormat; | ||
} | ||
|
||
@Override | ||
public CompoundFormat compoundFormat() { | ||
return this.compoundFormat; | ||
} | ||
} |
68 changes: 68 additions & 0 deletions
68
src/main/java/org/opensearch/knn/index/codec/KNN91Codec/KNN91CompoundFormat.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,68 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.knn.index.codec.KNN91Codec; | ||
|
||
import org.apache.lucene.codecs.CompoundDirectory; | ||
import org.apache.lucene.codecs.CompoundFormat; | ||
import org.apache.lucene.index.SegmentInfo; | ||
import org.apache.lucene.store.Directory; | ||
import org.apache.lucene.store.IOContext; | ||
import org.opensearch.knn.common.KNNConstants; | ||
import org.opensearch.knn.index.util.KNNEngine; | ||
|
||
import java.io.IOException; | ||
import java.util.HashSet; | ||
import java.util.Set; | ||
import java.util.stream.Collectors; | ||
|
||
/** | ||
* Class to encode/decode compound file | ||
*/ | ||
public class KNN91CompoundFormat extends CompoundFormat { | ||
|
||
private final CompoundFormat delegate; | ||
|
||
/** | ||
* Constructor that takes a delegate to handle non-overridden methods | ||
* | ||
* @param delegate CompoundFormat that will handle non-overridden methods | ||
*/ | ||
public KNN91CompoundFormat(CompoundFormat delegate) { | ||
this.delegate = delegate; | ||
} | ||
|
||
@Override | ||
public CompoundDirectory getCompoundReader(Directory dir, SegmentInfo si, IOContext context) throws IOException { | ||
return delegate.getCompoundReader(dir, si, context); | ||
} | ||
|
||
@Override | ||
public void write(Directory dir, SegmentInfo si, IOContext context) throws IOException { | ||
for (KNNEngine knnEngine : KNNEngine.values()) { | ||
writeEngineFiles(dir, si, context, knnEngine.getExtension()); | ||
} | ||
delegate.write(dir, si, context); | ||
} | ||
|
||
private void writeEngineFiles(Directory dir, SegmentInfo si, IOContext context, String engineExtension) throws IOException { | ||
/* | ||
* If engine file present, remove it from the compounding file list to avoid header/footer checks | ||
* and create a new compounding file format with extension engine + c. | ||
*/ | ||
Set<String> engineFiles = si.files().stream().filter(file -> file.endsWith(engineExtension)).collect(Collectors.toSet()); | ||
|
||
Set<String> segmentFiles = new HashSet<>(si.files()); | ||
|
||
if (!engineFiles.isEmpty()) { | ||
for (String engineFile : engineFiles) { | ||
String engineCompoundFile = engineFile + KNNConstants.COMPOUND_EXTENSION; | ||
dir.copyFrom(dir, engineFile, engineCompoundFile, context); | ||
} | ||
segmentFiles.removeAll(engineFiles); | ||
si.setFiles(segmentFiles); | ||
} | ||
} | ||
} |
Oops, something went wrong.