Refactor kNN codec related classes #582
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.extern.log4j.Log4j2; | ||
| import org.apache.lucene.codecs.KnnVectorsFormat; | ||
| import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; | ||
| import org.opensearch.index.mapper.MapperService; | ||
| import org.opensearch.knn.common.KNNConstants; | ||
| import org.opensearch.knn.index.mapper.KNNVectorFieldMapper; | ||
|
|
||
| import java.util.Map; | ||
| import java.util.Optional; | ||
| import java.util.function.BiFunction; | ||
| import java.util.function.Supplier; | ||
|
|
||
| /** | ||
| * Base class for PerFieldKnnVectorsFormat, builds KnnVectorsFormat based on specific Lucene version | ||
| */ | ||
| @AllArgsConstructor | ||
| @Log4j2 | ||
| public abstract class BasePerFieldKnnVectorsFormat extends PerFieldKnnVectorsFormat { | ||
|
|
||
| private final Optional<MapperService> mapperService; | ||
| private final int defaultMaxConnections; | ||
| private final int defaultBeamWidth; | ||
| private final Supplier<KnnVectorsFormat> defaultFormatSupplier; | ||
| private final BiFunction<Integer, Integer, KnnVectorsFormat> formatSupplier; | ||
|
|
||
| @Override | ||
| public KnnVectorsFormat getKnnVectorsFormatForField(final String field) { | ||
| if (isKnnVectorFieldType(field) == false) { | ||
| log.debug( | ||
| String.format( | ||
| "Initialize KNN vector format for field [%s] with default params [max_connections] = \"%d\" and [beam_width] = \"%d\"", | ||
| field, | ||
| defaultMaxConnections, | ||
| defaultBeamWidth | ||
| ) | ||
| ); | ||
| return defaultFormatSupplier.get(); | ||
| } | ||
| var type = (KNNVectorFieldMapper.KNNVectorFieldType) mapperService.orElseThrow( | ||
| () -> new IllegalStateException( | ||
| String.format("Cannot read field type for field [%s] because mapper service is not available", field) | ||
| ) | ||
| ).fieldType(field); | ||
| var params = type.getKnnMethodContext().getMethodComponent().getParameters(); | ||
| int maxConnections = getMaxConnections(params); | ||
| int beamWidth = getBeamWidth(params); | ||
| log.debug( | ||
| String.format( | ||
|
There was a problem hiding this comment. No need to use String.format |
||
| "Initialize KNN vector format for field [%s] with params [max_connections] = \"%d\" and [beam_width] = \"%d\"", | ||
| field, | ||
| maxConnections, | ||
| beamWidth | ||
| ) | ||
| ); | ||
| return formatSupplier.apply(maxConnections, beamWidth); | ||
| } | ||
|
|
||
| private boolean isKnnVectorFieldType(final String field) { | ||
| return mapperService.isPresent() && mapperService.get().fieldType(field) instanceof KNNVectorFieldMapper.KNNVectorFieldType; | ||
| } | ||
|
|
||
| private int getMaxConnections(final Map<String, Object> params) { | ||
| if (params != null && params.containsKey(KNNConstants.METHOD_PARAMETER_M)) { | ||
| return (int) params.get(KNNConstants.METHOD_PARAMETER_M); | ||
| } | ||
| return defaultMaxConnections; | ||
| } | ||
|
|
||
| private int getBeamWidth(final Map<String, Object> params) { | ||
| if (params != null && params.containsKey(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION)) { | ||
| return (int) params.get(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION); | ||
| } | ||
| return defaultBeamWidth; | ||
| } | ||
| } | ||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,74 +5,24 @@ | |
|
|
||
| package org.opensearch.knn.index.codec.KNN920Codec; | ||
|
|
||
| import org.apache.lucene.backward_codecs.lucene92.Lucene92HnswVectorsFormat; | ||
| import org.opensearch.index.mapper.MapperService; | ||
| import org.opensearch.knn.index.codec.BasePerFieldKnnVectorsFormat; | ||
|
|
||
| import java.util.Optional; | ||
|
|
||
| /** | ||
| * Class provides per field format implementation for Lucene Knn vector type | ||
| */ | ||
| public class KNN920PerFieldKnnVectorsFormat extends BasePerFieldKnnVectorsFormat { | ||
|
|
||
| public KNN920PerFieldKnnVectorsFormat(final Optional<MapperService> mapperService) { | ||
| super( | ||
| mapperService, | ||
| Lucene92HnswVectorsFormat.DEFAULT_MAX_CONN, | ||
| Lucene92HnswVectorsFormat.DEFAULT_BEAM_WIDTH, | ||
| () -> new Lucene92HnswVectorsFormat(), | ||
| (maxConnm, beamWidth) -> new Lucene92HnswVectorsFormat(maxConnm, beamWidth) | ||
| ); | ||
| } | ||
| } | ||
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 | ||||
|---|---|---|---|---|---|---|
|
|
@@ -12,23 +12,19 @@ | |||||
| import org.apache.lucene.codecs.FilterCodec; | ||||||
| import org.apache.lucene.codecs.KnnVectorsFormat; | ||||||
| import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; | ||||||
| import org.opensearch.knn.index.codec.KNNCodecVersion; | ||||||
| import org.opensearch.knn.index.codec.KNNFormatFacade; | ||||||
|
|
||||||
| public class KNN940Codec extends FilterCodec { | ||||||
| private static final KNNCodecVersion version = KNNCodecVersion.V_9_4_0; | ||||||
|
There was a problem hiding this comment.
Suggested change
|
||||||
| private final KNNFormatFacade knnFormatFacade; | ||||||
| private final PerFieldKnnVectorsFormat perFieldKnnVectorsFormat; | ||||||
|
|
||||||
| /** | ||||||
| * No arg constructor that uses Lucene94 as the delegate | ||||||
| */ | ||||||
| public KNN940Codec() { | ||||||
|
There was a problem hiding this comment. FAR: why not pass KNNCodecVersion as parameter? There was a problem hiding this comment. We have to have default constructor as per SPI contract, all codec classes are defined in corresponding registry file. In such case approach with parameter will not work |
||||||
| this(version.getDefaultCodecDelegate(), version.getPerFieldKnnVectorsFormat()); | ||||||
| } | ||||||
|
|
||||||
| /** | ||||||
|
|
@@ -40,8 +36,8 @@ public KNN940Codec() { | |||||
| */ | ||||||
| @Builder | ||||||
| protected KNN940Codec(Codec delegate, PerFieldKnnVectorsFormat knnVectorsFormat) { | ||||||
| super(version.getCodecName(), delegate); | ||||||
| knnFormatFacade = version.getKnnFormatFacadeSupplier().apply(delegate); | ||||||
| perFieldKnnVectorsFormat = knnVectorsFormat; | ||||||
| } | ||||||
|
|
||||||
|
|
||||||
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 |
|---|---|---|
|
|
@@ -5,70 +5,24 @@ | |
|
|
||
| package org.opensearch.knn.index.codec.KNN940Codec; | ||
|
|
||
| import org.apache.lucene.codecs.lucene94.Lucene94HnswVectorsFormat; | ||
| import org.opensearch.index.mapper.MapperService; | ||
| import org.opensearch.knn.index.codec.BasePerFieldKnnVectorsFormat; | ||
|
|
||
| import java.util.Optional; | ||
|
|
||
| /** | ||
| * Class provides per field format implementation for Lucene Knn vector type | ||
| */ | ||
| public class KNN940PerFieldKnnVectorsFormat extends BasePerFieldKnnVectorsFormat { | ||
|
|
||
| public KNN940PerFieldKnnVectorsFormat(final Optional<MapperService> mapperService) { | ||
| super( | ||
| mapperService, | ||
| Lucene94HnswVectorsFormat.DEFAULT_MAX_CONN, | ||
| Lucene94HnswVectorsFormat.DEFAULT_BEAM_WIDTH, | ||
| () -> new Lucene94HnswVectorsFormat(), | ||
| (maxConnm, beamWidth) -> new Lucene94HnswVectorsFormat(maxConnm, beamWidth) | ||
| ); | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You don't need to use String.format here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ack