Skip to content

Commit

Permalink
Adding KNN codec that is based on Lucene92 codec (#444)
Browse files Browse the repository at this point in the history
* Adding KNN codec that is based on Lucene92 codec

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
  • Loading branch information
martin-gaievski authored Jul 13, 2022
1 parent fe9f293 commit 7499375
Show file tree
Hide file tree
Showing 9 changed files with 126 additions and 94 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/
package org.opensearch.knn.index.codec.KNN920Codec;

import lombok.Builder;
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.opensearch.knn.index.codec.KNNFormatFacade;
import org.opensearch.knn.index.codec.KNNFormatFactory;

import static org.opensearch.knn.index.codec.KNNCodecFactory.CodecDelegateFactory.createKNN92DefaultDelegate;

/**
* KNN codec that is based on Lucene92 codec
*/
public final class KNN920Codec extends FilterCodec {

private static final String KNN920 = "KNN920Codec";
private final KNNFormatFacade knnFormatFacade;

/**
* No arg constructor that uses Lucene91 as the delegate
*/
public KNN920Codec() {
this(createKNN92DefaultDelegate());
}

/**
* 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
*/
@Builder
public KNN920Codec(Codec delegate) {
super(KNN920, delegate);
knnFormatFacade = KNNFormatFactory.createKNN920Format(delegate);
}

@Override
public DocValuesFormat docValuesFormat() {
return knnFormatFacade.docValuesFormat();
}

@Override
public CompoundFormat compoundFormat() {
return knnFormatFacade.compoundFormat();
}
}
40 changes: 9 additions & 31 deletions src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java
Original file line number Diff line number Diff line change
Expand Up @@ -4,42 +4,23 @@
*/
package org.opensearch.knn.index.codec;

import com.google.common.collect.ImmutableMap;
import lombok.AllArgsConstructor;
import org.apache.lucene.codecs.Codec;
import org.apache.lucene.backward_codecs.lucene91.Lucene91Codec;
import org.apache.lucene.codecs.lucene92.Lucene92Codec;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.knn.index.codec.util.CodecBuilder;

import java.util.Map;
import org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec;

/**
* Factory abstraction for KNN codec
*/
@AllArgsConstructor
public class KNNCodecFactory {

private final Map<KNNCodecVersion, CodecBuilder> codecByVersion;

private static final KNNCodecVersion LATEST_KNN_CODEC_VERSION = KNNCodecVersion.KNN910;

public KNNCodecFactory(MapperService mapperService) {
codecByVersion = ImmutableMap.of(KNNCodecVersion.KNN910, new CodecBuilder.KNN91CodecBuilder(mapperService));
}
private final MapperService mapperService;

public Codec createKNNCodec(final Codec userCodec) {
return getCodec(LATEST_KNN_CODEC_VERSION, userCodec);
}

public Codec createKNNCodec(final KNNCodecVersion knnCodecVersion, final Codec userCodec) {
return getCodec(knnCodecVersion, userCodec);
}

private Codec getCodec(final KNNCodecVersion knnCodecVersion, final Codec userCodec) {
try {
final CodecBuilder codecBuilder = codecByVersion.getOrDefault(knnCodecVersion, codecByVersion.get(LATEST_KNN_CODEC_VERSION));
return codecBuilder.userCodec(userCodec).build();
} catch (Exception ex) {
throw new RuntimeException("Cannot create instance of KNN codec", ex);
}
return KNN920Codec.builder().delegate(userCodec).build();
}

/**
Expand All @@ -50,12 +31,9 @@ public static class CodecDelegateFactory {
public static Codec createKNN91DefaultDelegate() {
return new Lucene91Codec();
}
}

/**
* Collection of supported coded versions
*/
enum KNNCodecVersion {
KNN910
public static Codec createKNN92DefaultDelegate() {
return new Lucene92Codec();
}
}
}
18 changes: 18 additions & 0 deletions src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,29 @@
*/
public class KNNFormatFactory {

/**
* Return facade class that abstracts format specific to KNN910 codec
* @param delegate delegate codec that is wrapped by KNN codec
* @return
*/
public static KNNFormatFacade createKNN910Format(final Codec delegate) {
final KNNFormatFacade knnFormatFacade = new KNNFormatFacade(
new KNN80DocValuesFormat(delegate.docValuesFormat()),
new KNN80CompoundFormat(delegate.compoundFormat())
);
return knnFormatFacade;
}

/**
* Return facade class that abstracts format specific to KNN920 codec
* @param delegate delegate codec that is wrapped by KNN codec
* @return
*/
public static KNNFormatFacade createKNN920Format(final Codec delegate) {
final KNNFormatFacade knnFormatFacade = new KNNFormatFacade(
new KNN80DocValuesFormat(delegate.docValuesFormat()),
new KNN80CompoundFormat(delegate.compoundFormat())
);
return knnFormatFacade;
}
}

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,5 @@ org.opensearch.knn.index.codec.KNN80Codec.KNN80Codec
org.opensearch.knn.index.codec.KNN84Codec.KNN84Codec
org.opensearch.knn.index.codec.KNN86Codec.KNN86Codec
org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec
org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec
org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec
org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.knn.index.codec.KNN920Codec;

import org.opensearch.knn.index.codec.KNNCodecTestCase;

import java.io.IOException;
import java.util.concurrent.ExecutionException;

public class KNN920CodecTests extends KNNCodecTestCase {

public void testMultiFieldsKnnIndex() throws Exception {
testMultiFieldsKnnIndex(new KNN920Codec());
}

public void testBuildFromModelTemplate() throws InterruptedException, ExecutionException, IOException {
testBuildFromModelTemplate(new KNN920Codec());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,10 @@

import org.apache.lucene.codecs.Codec;
import org.apache.lucene.backward_codecs.lucene91.Lucene91Codec;
import org.apache.lucene.codecs.lucene92.Lucene92Codec;
import org.opensearch.index.mapper.MapperService;
import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec;
import org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec;

import static org.mockito.Mockito.mock;

Expand All @@ -21,21 +22,18 @@ public void testKNN91DefaultDelegate() {
assertTrue(knn91DefaultDelegate instanceof Lucene91Codec);
}

public void testKNN91DefaultCodec() {
Lucene91Codec lucene91CodecDelegate = new Lucene91Codec();
MapperService mapperService = mock(MapperService.class);
KNNCodecFactory knnCodecFactory = new KNNCodecFactory(mapperService);
Codec knnCodec = knnCodecFactory.createKNNCodec(lucene91CodecDelegate);
assertNotNull(knnCodec);
assertTrue(knnCodec instanceof KNN910Codec);
public void testKNN92DefaultDelegate() {
Codec knn92DefaultDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN92DefaultDelegate();
assertNotNull(knn92DefaultDelegate);
assertTrue(knn92DefaultDelegate instanceof Lucene92Codec);
}

public void testKNN91CodecByVersion() {
public void testKNNDefaultCodec() {
Lucene91Codec lucene91CodecDelegate = new Lucene91Codec();
MapperService mapperService = mock(MapperService.class);
KNNCodecFactory knnCodecFactory = new KNNCodecFactory(mapperService);
Codec knnCodec = knnCodecFactory.createKNNCodec(KNNCodecFactory.KNNCodecVersion.KNN910, lucene91CodecDelegate);
Codec knnCodec = knnCodecFactory.createKNNCodec(lucene91CodecDelegate);
assertNotNull(knnCodec);
assertTrue(knnCodec instanceof KNN910Codec);
assertTrue(knnCodec instanceof KNN920Codec);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec;
import org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec;
import org.opensearch.knn.jni.JNIService;
import org.opensearch.knn.index.KNNQuery;
import org.opensearch.knn.index.KNNSettings;
Expand Down Expand Up @@ -61,7 +61,7 @@
*/
public class KNNCodecTestCase extends KNNTestCase {

private static final KNN910Codec ACTUAL_CODEC = new KNN910Codec();
private static final KNN920Codec ACTUAL_CODEC = new KNN920Codec();
private static FieldType sampleFieldType;
static {
sampleFieldType = new FieldType(KNNVectorFieldMapper.Defaults.FIELD_TYPE);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,4 +24,16 @@ public void testKNN91Format() {
assertNotNull(knnFormatFacade.compoundFormat());
assertNotNull(knnFormatFacade.docValuesFormat());
}

public void testKNN92Format() {
final Codec lucene92CodecDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN92DefaultDelegate();
MapperService mapperService = mock(MapperService.class);
KNNCodecFactory knnCodecFactory = new KNNCodecFactory(mapperService);
final Codec knnCodec = knnCodecFactory.createKNNCodec(lucene92CodecDelegate);
KNNFormatFacade knnFormatFacade = KNNFormatFactory.createKNN920Format(knnCodec);

assertNotNull(knnFormatFacade);
assertNotNull(knnFormatFacade.compoundFormat());
assertNotNull(knnFormatFacade.docValuesFormat());
}
}

0 comments on commit 7499375

Please sign in to comment.