diff --git a/src/main/java/org/opensearch/knn/index/fetch/KNNFetchSubPhase.java b/src/main/java/org/opensearch/knn/index/fetch/KNNFetchSubPhase.java index 1def7e43d7..edb7f8ff58 100644 --- a/src/main/java/org/opensearch/knn/index/fetch/KNNFetchSubPhase.java +++ b/src/main/java/org/opensearch/knn/index/fetch/KNNFetchSubPhase.java @@ -106,7 +106,7 @@ public void process(HitContext hitContext) throws IOException { } if (hasNested) { - syntheticNestedDocValues(mapperService, hitContext, maps); + syntheticNestedFieldWithDocValues(mapperService, hitContext, maps); } for (DocValueField f : fields) { if (maps.containsKey(f.field)) { @@ -123,11 +123,10 @@ public void process(HitContext hitContext) throws IOException { hitContext.hit().sourceRef(BytesReference.bytes(builder)); } - protected void syntheticNestedDocValues(MapperService mapperService, HitContext hitContext, Map maps) + protected void syntheticNestedFieldWithDocValues(MapperService mapperService, HitContext hitContext, Map sourceMaps) throws IOException { DocumentMapper documentMapper = mapperService.documentMapper(); Map mapperMap = documentMapper.objectMappers(); - SearchHit hit = hitContext.hit(); for (ObjectMapper objectMapper : mapperMap.values()) { if (objectMapper == null) { @@ -138,19 +137,11 @@ protected void syntheticNestedDocValues(MapperService mapperService, HitContext } String path = objectMapper.fullPath(); for (DocValueField f : fields) { - if (!f.field.startsWith(path)) { - continue; - } - if (!maps.containsKey(path)) { - continue; - } - - // path to nested field: - Object nestedObj = maps.get(path); - if (!(nestedObj instanceof ArrayList)) { + if (!checkNestedField(path, f, sourceMaps)) { continue; } // nested array in one nested path + Object nestedObj = sourceMaps.get(path); ArrayList nestedDocList = (ArrayList) nestedObj; log.debug( @@ -162,64 +153,97 @@ protected void syntheticNestedDocValues(MapperService mapperService, HitContext + f.field ); - Query parentFilter = Queries.newNonNestedFilter(); - Query childFilter = objectMapper.nestedTypeFilter(); - ContextIndexSearcher searcher = fetchContext.searcher(); - final Weight childWeight = searcher.createWeight(searcher.rewrite(childFilter), ScoreMode.COMPLETE_NO_SCORES, 1f); - - LeafReaderContext subReaderContext = hitContext.readerContext(); - Scorer childScorer = childWeight.scorer(subReaderContext); - DocIdSetIterator childIter = childScorer.iterator(); - BitSet parentBits = fetchContext.getQueryShardContext().bitsetFilter(parentFilter).getBitSet(subReaderContext); - - int currentParent = hit.docId() - subReaderContext.docBase; - int previousParent = parentBits.prevSetBit(currentParent - 1); - int childDocId = childIter.advance(previousParent + 1); - SourceLookup nestedVecSourceLookup = new SourceLookup(); - - // when nested field only have vector field and exclude source, list is empty - boolean isEmpty = nestedDocList.isEmpty(); - - for (int offset = 0; childDocId < currentParent && childDocId != DocIdSetIterator.NO_MORE_DOCS; childDocId = childIter - .nextDoc(), offset++) { - nestedVecSourceLookup.setSegmentAndDocument(subReaderContext, childDocId); - List nestedVecDocValuesSource = f.fetcher.fetchValues(nestedVecSourceLookup); - if (nestedVecDocValuesSource == null || nestedVecDocValuesSource.isEmpty()) { - continue; - } - if (isEmpty) { - nestedDocList.add(new HashMap()); - } - if (offset < nestedDocList.size()) { - Object o2 = nestedDocList.get(offset); - if (o2 instanceof Map) { - Map o2map = (Map) o2; - String suffix = f.field.substring(path.length() + 1); - o2map.put(suffix, nestedVecDocValuesSource.get(0)); - } - } else { - /** - * TODO nested field partial doc only have vector and source exclude - * this source map nestedDocList would out-of-order, can not fill the vector into right offset - * "nested_field" : [ - * {"nested_vector": [2.6, 2.6]}, - * {"nested_numeric": 2, "nested_vector": [3.1, 2.3]} - * ] - */ - - throw new UnsupportedOperationException( - String.format( - "\"Nested Path \"%s\" in Field \"%s\" with _ID \"%s\" can not be empty\"", - path, - f.field, - hit.getId() - ) - ); - } + innerProcessOneNestedField(objectMapper, hitContext, nestedDocList, f, path); + } + } + } + + private void innerProcessOneNestedField( + ObjectMapper objectMapper, + HitContext hitContext, + ArrayList nestedDocList, + DocValueField f, + String path + ) throws IOException { + + BitSet parentBits = getParentDocBitSet(hitContext); + DocIdSetIterator childIter = getChildDocIdSetIterator(objectMapper, hitContext); + LeafReaderContext subReaderContext = hitContext.readerContext(); + + SearchHit hit = hitContext.hit(); + int currentParent = hit.docId() - subReaderContext.docBase; + int previousParent = parentBits.prevSetBit(currentParent - 1); + int childDocId = childIter.advance(previousParent + 1); + SourceLookup nestedVecSourceLookup = new SourceLookup(); + + // when nested field only have vector field and exclude source, list is empty + boolean isEmpty = nestedDocList.isEmpty(); + + for (int offset = 0; childDocId < currentParent && childDocId != DocIdSetIterator.NO_MORE_DOCS; childDocId = childIter + .nextDoc(), offset++) { + nestedVecSourceLookup.setSegmentAndDocument(subReaderContext, childDocId); + List nestedVecDocValuesSource = f.fetcher.fetchValues(nestedVecSourceLookup); + if (nestedVecDocValuesSource == null || nestedVecDocValuesSource.isEmpty()) { + continue; + } + if (isEmpty) { + nestedDocList.add(new HashMap()); + } + if (offset < nestedDocList.size()) { + Object sourceObj = nestedDocList.get(offset); + if (sourceObj instanceof Map) { + Map sourceMap = (Map) sourceObj; + String suffix = f.field.substring(path.length() + 1); + sourceMap.put(suffix, nestedVecDocValuesSource.get(0)); } + } else { + /** + * TODO nested field partial doc only have vector and source exclude + * this source map nestedDocList would out-of-order, can not fill the vector into right offset + * "nested_field" : [ + * {"nested_vector": [2.6, 2.6]}, + * {"nested_numeric": 2, "nested_vector": [3.1, 2.3]} + * ] + */ + throw new UnsupportedOperationException( + String.format("\"Nested Path \"%s\" in Field \"%s\" with _ID \"%s\" can not be empty\"", path, f.field, hit.getId()) + ); } } } + + private BitSet getParentDocBitSet(HitContext hitContext) throws IOException { + Query parentFilter = Queries.newNonNestedFilter(); + LeafReaderContext subReaderContext = hitContext.readerContext(); + BitSet parentBits = fetchContext.getQueryShardContext().bitsetFilter(parentFilter).getBitSet(subReaderContext); + return parentBits; + } + + private DocIdSetIterator getChildDocIdSetIterator(ObjectMapper objectMapper, HitContext hitContext) throws IOException { + Query childFilter = objectMapper.nestedTypeFilter(); + ContextIndexSearcher searcher = fetchContext.searcher(); + LeafReaderContext subReaderContext = hitContext.readerContext(); + final Weight childWeight = searcher.createWeight(searcher.rewrite(childFilter), ScoreMode.COMPLETE_NO_SCORES, 1f); + Scorer childScorer = childWeight.scorer(subReaderContext); + DocIdSetIterator childIter = childScorer.iterator(); + return childIter; + } + + private boolean checkNestedField(String path, DocValueField f, Map sourceMaps) { + if (!f.field.startsWith(path)) { + return false; + } + if (!sourceMaps.containsKey(path)) { + return false; + } + + // path to nested field: + Object nestedObj = sourceMaps.get(path); + if (!(nestedObj instanceof ArrayList)) { + return false; + } + return true; + } } @Getter