Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BUG FIX] Fix knn index shard to get bwc engine paths #309

Merged
merged 6 commits into from
Mar 8, 2022
Merged

[BUG FIX] Fix knn index shard to get bwc engine paths #309

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
066f8b9
Fix index shard to get bwc engine paths
jmazanec15 Mar 8, 2022
6efcdff
Remove type from KNNSingleNodeTestCase
jmazanec15 Mar 8, 2022
4817493
Remove unused import
jmazanec15 Mar 8, 2022
c8b1678
Fix path in test
jmazanec15 Mar 8, 2022
6636c70
Modify stream logic
jmazanec15 Mar 8, 2022
d0d96ca
Apply spotless
jmazanec15 Mar 8, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 15 additions & 10 deletions src/main/java/org/opensearch/knn/index/KNNIndexShard.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,16 @@

import java.io.IOException;
import java.nio.file.Path;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFileName;
import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFilePrefix;
import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFileSuffix;

/**
* KNNIndexShard wraps IndexShard and adds methods to perform k-NN related operations against the shard
Expand Down Expand Up @@ -126,17 +127,21 @@ private Map<String, SpaceType> getEnginePaths(IndexReader indexReader, KNNEngine
// was L2. So, if Space Type is not present, just fall back to L2
String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue());
SpaceType spaceType = SpaceType.getSpace(spaceTypeName);
String engineFileName = buildEngineFileName(reader.getSegmentInfo().info.name,
knnEngine.getLatestBuildVersion(), fieldInfo.name, fileExtension);

engineFiles.putAll(reader.getSegmentInfo().files().stream()
.filter(fileName -> fileName.equals(engineFileName))
.map(fileName -> shardPath.resolve(fileName).toString())
.filter(Objects::nonNull)
.collect(Collectors.toMap(fileName -> fileName, fileName -> spaceType)));

engineFiles.putAll(getEnginePaths(reader.getSegmentInfo().files(),
reader.getSegmentInfo().info.name, fieldInfo.name, fileExtension, shardPath, spaceType));
}
}
}
return engineFiles;
}

protected Map<String, SpaceType> getEnginePaths(Collection<String> files, String segmentName, String fieldName,
String fileExtension, Path shardPath, SpaceType spaceType) {
return files.stream()
.filter(fileName -> fileName.startsWith(buildEngineFilePrefix(segmentName)))
Copy link
Member

@VijayanB VijayanB Mar 8, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't have to compute buildEngineFilePrefix for every file. This can be called once outside streams and used inside filter.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense. Updated.

.filter(fileName -> fileName.endsWith(buildEngineFileSuffix(fieldName, fileExtension)))
.map(fileName -> shardPath.resolve(fileName).toString())
.collect(Collectors.toMap(fileName -> fileName, fileName -> spaceType));
}
}
11 changes: 10 additions & 1 deletion src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,15 @@ public static KNNCodecUtil.Pair getFloats(BinaryDocValues values) throws IOExcep

public static String buildEngineFileName(String segmentName, String latestBuildVersion, String fieldName,
String extension) {
return String.format("%s_%s_%s%s", segmentName, latestBuildVersion, fieldName, extension);
return String.format("%s%s%s", buildEngineFilePrefix(segmentName), latestBuildVersion,
buildEngineFileSuffix(fieldName, extension));
}

public static String buildEngineFilePrefix(String segmentName) {
return String.format("%s_", segmentName);
}

public static String buildEngineFileSuffix(String fieldName, String extension) {
return String.format("_%s%s", fieldName, extension);
}
}
2 changes: 1 addition & 1 deletion src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ protected IndexService createKNNIndex(String indexName) {
* Create simple k-NN mapping
*/
protected void createKnnIndexMapping(String indexName, String fieldName, Integer dimensions) {
PutMappingRequest request = new PutMappingRequest(indexName).type("_doc");
PutMappingRequest request = new PutMappingRequest(indexName);
request.source(fieldName, "type=knn_vector,dimension="+dimensions);
OpenSearchAssertions.assertAcked(client().admin().indices().putMapping(request).actionGet());
}
Expand Down
38 changes: 38 additions & 0 deletions src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,24 @@

package org.opensearch.knn.index;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import org.opensearch.knn.KNNSingleNodeTestCase;
import org.opensearch.index.IndexService;
import org.opensearch.index.engine.Engine;
import org.opensearch.index.shard.IndexShard;
import org.opensearch.knn.index.memory.NativeMemoryCacheManager;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static org.opensearch.knn.index.memory.NativeMemoryCacheManager.GRAPH_COUNT;

Expand Down Expand Up @@ -116,4 +123,35 @@ public void testGetHNSWPaths() throws IOException, ExecutionException, Interrupt
assertTrue(paths.get(0).contains("hnsw") || paths.get(0).contains("hnswc"));
searcher.close();
}

public void testGetEnginePaths() {
// Check that the correct engine paths are being returned by the KNNIndexShard
String segmentName = "_0";
String fieldName = "test_field";
String fileExt = ".test";
SpaceType spaceType = SpaceType.L2;

Set<String> includedFileNames = ImmutableSet.of(
String.format("%s_111_%s%s", segmentName, fieldName, fileExt),
String.format("%s_7_%s%s", segmentName, fieldName, fileExt),
String.format("%s_53_%s%s", segmentName, fieldName, fileExt)
);

List<String> excludedFileNames = ImmutableList.of(
String.format("_111_%s%s", fieldName, fileExt), // missing segment name
String.format("%s_111_%s", segmentName, fileExt), // missing field name
String.format("%s_111_%s.invalid", segmentName, fieldName) // missing extension
);

List<String> files = Stream.concat(includedFileNames.stream(), excludedFileNames.stream())
.collect(Collectors.toList());

KNNIndexShard knnIndexShard = new KNNIndexShard(null);

Path path = Paths.get("");
Map<String, SpaceType> included = knnIndexShard.getEnginePaths(files, segmentName, fieldName, fileExt, path, spaceType);

assertEquals(includedFileNames.size(), included.size());
included.keySet().forEach(o -> assertTrue(includedFileNames.contains(o)));
}
}