Skip to content

Commit d67e5fc

Browse files
committed
Further refactoring
Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
1 parent e79b84f commit d67e5fc

File tree

2 files changed

+42
-58
lines changed

2 files changed

+42
-58
lines changed

src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java

+42-54
Original file line numberDiff line numberDiff line change
@@ -112,8 +112,12 @@ public String buildShape(
112112
Boolean showFieldType,
113113
Set<Index> successfulSearchShardIndices
114114
) {
115-
Index firstIndex = successfulSearchShardIndices.iterator().next();
116-
Map<String, Object> propertiesAsMap = getPropertiesMapForIndex(firstIndex);
115+
Index firstIndex = null;
116+
Map<String, Object> propertiesAsMap = null;
117+
if (successfulSearchShardIndices != null) {
118+
firstIndex = successfulSearchShardIndices.iterator().next();
119+
propertiesAsMap = getPropertiesMapForIndex(firstIndex);
120+
}
117121

118122
StringBuilder shape = new StringBuilder();
119123
shape.append(buildQueryShape(source.query(), showFieldName, showFieldType, propertiesAsMap, firstIndex));
@@ -345,79 +349,63 @@ String buildFieldDataString(
345349
}
346350

347351
String getFieldType(String fieldName, Map<String, Object> propertiesAsMap, Index index) {
348-
String fieldType = getFieldTypeFromCache(fieldName, index);
349-
if (fieldType != null) {
350-
if (fieldType.equals(NO_FIELD_TYPE_VALUE)) {
351-
return null;
352-
} else {
353-
return fieldType;
354-
}
352+
if (propertiesAsMap == null || index == null) {
353+
return null;
355354
}
355+
// Attempt to get field type from cache
356+
String fieldType = getFieldTypeFromCache(fieldName, index);
356357

357-
fieldType = getFieldTypeFromMapping(index, fieldName, propertiesAsMap);
358358
if (fieldType != null) {
359-
String finalFieldType = fieldType;
360-
fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()).computeIfAbsent(fieldName, k -> finalFieldType);
361359
return fieldType;
362360
}
363-
return null;
364-
}
365361

366-
String getFieldTypeFromMapping(Index index, String fieldName, Map<String, Object> propertiesAsMap) {
367-
// Recursively find the field type from properties
368-
if (propertiesAsMap != null) {
369-
String fieldType = findFieldTypeInProperties(propertiesAsMap, fieldName.split("\\."), 0);
362+
// Retrieve field type from mapping and cache it if found
363+
fieldType = getFieldTypeFromProperties(fieldName, propertiesAsMap);
370364

371-
// Cache the NO_FIELD_TYPE_VALUE result if no field type is found in this index
372-
if (fieldType == null) {
373-
fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()).computeIfAbsent(fieldName, k -> fieldType);
374-
}
375-
return fieldType;
376-
}
377-
return null;
365+
// Cache field type or NO_FIELD_TYPE_VALUE if not found
366+
fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>())
367+
.putIfAbsent(fieldName, fieldType != null ? fieldType : NO_FIELD_TYPE_VALUE);
368+
369+
return fieldType;
378370
}
379371

380-
private String findFieldTypeInProperties(Map<String, ?> properties, String[] fieldParts, int index) {
381-
if (index >= fieldParts.length) {
372+
String getFieldTypeFromProperties(String fieldName, Map<String, Object> propertiesAsMap) {
373+
if (propertiesAsMap == null) {
382374
return null;
383375
}
384376

385-
String currentField = fieldParts[index];
386-
Object currentMapping = properties.get(currentField);
377+
String[] fieldParts = fieldName.split("\\.");
378+
Map<String, ?> currentProperties = propertiesAsMap;
387379

388-
// Check if current mapping is a map and contains further nested properties or a type
389-
if (currentMapping instanceof Map) {
390-
Map<String, ?> currentMappingMap = (Map<String, ?>) currentMapping;
380+
for (int depth = 0; depth < fieldParts.length; depth++) {
381+
Object currentMapping = currentProperties.get(fieldParts[depth]);
391382

392-
// If it has a 'properties' key, go into the nested object
393-
if (currentMappingMap.containsKey("properties")) {
394-
Map<String, ?> nestedProperties = (Map<String, ?>) currentMappingMap.get("properties");
395-
return findFieldTypeInProperties(nestedProperties, fieldParts, index + 1);
396-
}
383+
if (currentMapping instanceof Map) {
384+
Map<String, ?> currentMap = (Map<String, ?>) currentMapping;
397385

398-
// If it has a 'fields' key, handle multifields (e.g., title.raw) and is not the parent field (parent field case handled below)
399-
if (currentMappingMap.containsKey("fields") && index + 1 < fieldParts.length) {
400-
Map<String, ?> fieldsMap = (Map<String, ?>) currentMappingMap.get("fields");
401-
return findFieldTypeInProperties(fieldsMap, fieldParts, index + 1); // Recursively check subfield
402-
}
403-
404-
// If it has a 'type' key, return the type. Also handles parent field case in multifields
405-
if (currentMappingMap.containsKey("type")) {
406-
return currentMappingMap.get("type").toString();
386+
// Navigate into nested properties if available
387+
if (currentMap.containsKey("properties")) {
388+
currentProperties = (Map<String, ?>) currentMap.get("properties");
389+
}
390+
// Handle multifields (e.g., title.raw)
391+
else if (currentMap.containsKey("fields") && depth + 1 < fieldParts.length) {
392+
currentProperties = (Map<String, ?>) currentMap.get("fields");
393+
}
394+
// Return type if found
395+
else if (currentMap.containsKey("type")) {
396+
return (String) currentMap.get("type");
397+
} else {
398+
return null;
399+
}
400+
} else {
401+
return null;
407402
}
408403
}
409404

410405
return null;
411406
}
412407

413408
String getFieldTypeFromCache(String fieldName, Index index) {
414-
Map<String, String> indexMap = fieldTypeMap.get(index);
415-
if (indexMap != null) {
416-
String fieldType = indexMap.get(fieldName);
417-
if (fieldType != null) {
418-
return fieldType;
419-
}
420-
}
421-
return null;
409+
return fieldTypeMap.getOrDefault(index, new ConcurrentHashMap<>()).get(fieldName);
422410
}
423411
}

src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java

-4
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
import static org.mockito.Mockito.atLeastOnce;
1414
import static org.mockito.Mockito.mock;
1515
import static org.mockito.Mockito.spy;
16-
import static org.mockito.Mockito.times;
1716
import static org.mockito.Mockito.verify;
1817
import static org.mockito.Mockito.when;
1918
import static org.opensearch.index.query.QueryBuilders.boolQuery;
@@ -597,11 +596,8 @@ public void testFieldTypeCachingForNonExistentField() throws IOException {
597596

598597
verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), any(Index.class));
599598

600-
verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromMapping(any(Index.class), eq(nonExistentField), any());
601-
602599
queryShapeGeneratorSpy.buildShape(sourceBuilder, true, true, successfulSearchShardIndices);
603600

604-
verify(queryShapeGeneratorSpy, times(2)).getFieldTypeFromMapping(any(Index.class), eq(nonExistentField), any());
605601
verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), any(Index.class));
606602
}
607603

0 commit comments

Comments
 (0)