-
Notifications
You must be signed in to change notification settings - Fork 144
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: John Mazanec <jmazane@amazon.com>
- Loading branch information
1 parent
f006b33
commit b721e83
Showing
5 changed files
with
264 additions
and
38 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
68 changes: 68 additions & 0 deletions
68
src/main/java/org/opensearch/knn/index/codec/derivedsource/DerivedSourceMapHelper.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.derivedsource; | ||
|
|
||
| import lombok.extern.log4j.Log4j2; | ||
| import org.opensearch.common.xcontent.support.XContentMapValues; | ||
|
|
||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
|
|
||
| /** | ||
| * Utility class for manipulating the source map | ||
| */ | ||
| @Log4j2 | ||
| public class DerivedSourceMapHelper { | ||
|
|
||
| /** | ||
| * Removes all fields in the array from the source | ||
| * | ||
| * @param fields Fields to remove | ||
| * @param source Source map to remove from | ||
| * @return Map with filtered fields | ||
| */ | ||
| public static Map<String, Object> filterFields(String[] fields, Map<String, Object> source) { | ||
| return XContentMapValues.filter(null, fields).apply(source); | ||
| } | ||
|
|
||
| /** | ||
| * Check whether field exists in the document | ||
| * | ||
| * @param source source document | ||
| * @param fieldName field to check. Field should be flattened. i.e. my.path.field | ||
| * @return whether or not the field exists in the object | ||
| */ | ||
| public static boolean fieldExists(Map<String, Object> source, String fieldName) { | ||
| return XContentMapValues.extractValue(fieldName, source, NullValue.INSTANCE) != null; | ||
| } | ||
|
|
||
| /** | ||
| * Injects vector into source, handling object field path if necessary | ||
| * | ||
| * @param sourceAsMap source to be injected into | ||
| * @param vector vector to inject | ||
| * @param fieldName field name injecting at | ||
| */ | ||
| public static void injectVector(Map<String, Object> sourceAsMap, Object vector, String fieldName) { | ||
| // If a field contains ".", we need to ensure that we properly nest it. | ||
| String[] fields = ParentChildHelper.splitPath(fieldName); | ||
| if (fields.length < 2) { | ||
| sourceAsMap.put(fieldName, vector); | ||
| } | ||
|
|
||
| Map<String, Object> currentMap = sourceAsMap; | ||
| for (int i = 0; i < fields.length - 1; i++) { | ||
| String field = fields[i]; | ||
| currentMap = (Map<String, Object>) currentMap.computeIfAbsent(field, k -> new HashMap<>()); | ||
| } | ||
| currentMap.put(fields[fields.length - 1], vector); | ||
|
|
||
| } | ||
|
|
||
| private static class NullValue { | ||
| private static final NullValue INSTANCE = new NullValue(); | ||
| } | ||
| } |
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
153 changes: 153 additions & 0 deletions
153
src/test/java/org/opensearch/knn/index/codec/derivedsource/DerivedSourceMapHelperTests.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,153 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.index.codec.derivedsource; | ||
|
|
||
| import lombok.SneakyThrows; | ||
| import org.opensearch.common.xcontent.XContentFactory; | ||
| import org.opensearch.common.xcontent.XContentHelper; | ||
| import org.opensearch.core.common.bytes.BytesReference; | ||
| import org.opensearch.core.xcontent.XContentBuilder; | ||
| import org.opensearch.knn.KNNTestCase; | ||
|
|
||
| import java.util.Collections; | ||
| import java.util.HashMap; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
|
|
||
| public class DerivedSourceMapHelperTests extends KNNTestCase { | ||
|
|
||
| @SneakyThrows | ||
| public void testFilterFields() { | ||
| // Basic filter | ||
| String[] filterFields = new String[] { "field1", "field2", "field3" }; | ||
| XContentBuilder builder = XContentFactory.jsonBuilder() | ||
| .startObject() | ||
| .field("field1", "value") | ||
| .field("field2", "value") | ||
| .field("field3", "value") | ||
| .endObject(); | ||
| Map<String, Object> source = DerivedSourceMapHelper.filterFields(filterFields, xContentToMap(builder)); | ||
| assertEquals(Collections.emptyMap(), source); | ||
| source = DerivedSourceMapHelper.filterFields(new String[] { "field1" }, xContentToMap(builder)); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(source, "field2")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(source, "field3")); | ||
| assertFalse(DerivedSourceMapHelper.fieldExists(source, "field1")); | ||
|
|
||
| // Object case | ||
| filterFields = new String[] { "level1.level2.test", "field.nonexist" }; | ||
| builder = XContentFactory.jsonBuilder() | ||
| .startObject() | ||
| .startObject("level1") | ||
| .startObject("level2") | ||
| .field("test", "test") | ||
| .endObject() | ||
| .endObject() | ||
| .endObject(); | ||
| source = DerivedSourceMapHelper.filterFields(filterFields, xContentToMap(builder)); | ||
| assertEquals(Map.of("level1", Map.of("level2", Collections.emptyMap())), source); | ||
|
|
||
| // Nested case. In this case, if filtering out a value leads it to be an empty map, then it will be removed | ||
| // from the array | ||
| filterFields = new String[] { "nested.deep.value", "nested.vector" }; | ||
| builder = XContentFactory.jsonBuilder() | ||
| .startObject() | ||
| .startArray("nested") | ||
| .startObject() | ||
| .field("text", "text1") | ||
| .startArray("deep") | ||
| .startObject() | ||
| .field("value", "text2") | ||
| .endObject() | ||
| .endArray() | ||
| .endObject() | ||
| .startObject() | ||
| .field("vector", "text1") | ||
| .endObject() | ||
| .startObject() | ||
| .field("text", "text1") | ||
| .field("vector", "text1") | ||
| .endObject() | ||
| .endArray() | ||
| .endObject(); | ||
| source = DerivedSourceMapHelper.filterFields(filterFields, xContentToMap(builder)); | ||
| assertTrue(source.containsKey("nested")); | ||
| Object nested = source.get("nested"); | ||
| assertTrue(nested instanceof List<?>); | ||
| List<?> nestedList = (List<?>) nested; | ||
| assertEquals(2, nestedList.size()); | ||
| assertTrue(nestedList.get(0) instanceof Map); | ||
| Map<String, Object> nestedMap = (Map<String, Object>) nestedList.get(0); | ||
| assertTrue(nestedMap.containsKey("deep")); | ||
| Object deep = nestedMap.get("deep"); | ||
| assertTrue(deep instanceof List<?>); | ||
| assertEquals(0, ((List<?>) deep).size()); | ||
| } | ||
|
|
||
| @SneakyThrows | ||
| public void testFieldExists() { | ||
| // Basic, flat validation | ||
| XContentBuilder builder = XContentFactory.jsonBuilder().startObject().field("field", "value").endObject(); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "field")); | ||
| assertFalse(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "non-existent-field")); | ||
|
|
||
| // Object type | ||
| builder = XContentFactory.jsonBuilder().startObject().startObject("field").field("test", "test").endObject().endObject(); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "field")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "field.test")); | ||
| assertFalse(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "test")); | ||
|
|
||
| // Complex nested object | ||
| builder = XContentFactory.jsonBuilder() | ||
| .startObject() | ||
| .startArray("nested1") | ||
| .startObject() | ||
| .field("field1", "test") | ||
| .endObject() | ||
| .startObject() | ||
| .field("field1", "test") | ||
| .field("field2", "test") | ||
| .endObject() | ||
| .startObject() | ||
| .field("field3", "test") | ||
| .endObject() | ||
| .endArray() | ||
| .startArray("nested2") | ||
| .startObject() | ||
| .field("field1", "test") | ||
| .endObject() | ||
| .startObject() | ||
| .field("field1", "test") | ||
| .field("field2", "test") | ||
| .endObject() | ||
| .startObject() | ||
| .field("field3", "test") | ||
| .endObject() | ||
| .endArray() | ||
| .field("notnested", "test") | ||
| .endObject(); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested1")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested2")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "notnested")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested1.field1")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested1.field2")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested1.field3")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested2.field1")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested2.field2")); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "nested2.field3")); | ||
| assertFalse(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "field1")); | ||
| assertFalse(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "field2")); | ||
| assertFalse(DerivedSourceMapHelper.fieldExists(xContentToMap(builder), "field3")); | ||
|
|
||
| // Null test | ||
| Map<String, Object> map = new HashMap<>(); | ||
| map.put("test", null); | ||
| assertTrue(DerivedSourceMapHelper.fieldExists(map, "test")); | ||
| } | ||
|
|
||
| private Map<String, Object> xContentToMap(XContentBuilder xContentBuilder) { | ||
| return XContentHelper.convertToMap(BytesReference.bytes(xContentBuilder), true, xContentBuilder.contentType()).v2(); | ||
| } | ||
| } |