aboutsummaryrefslogtreecommitdiffstats
path: root/container-search/src/main/java/com/yahoo/search/searchchain/testutil/DocumentSourceSearcher.java
blob: 06fa3982b186d402ad25f5cca623f50d37d210bd (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.search.searchchain.testutil;


import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.Map;
import java.util.HashMap;
import java.util.List;

import com.yahoo.net.URI;
import com.yahoo.search.Query;
import com.yahoo.search.Result;
import com.yahoo.search.Searcher;
import com.yahoo.search.result.Hit;
import com.yahoo.search.searchchain.Execution;

/**
 * <p>Implements a document source.  You pass in a query and a Result
 * set.  When this Searcher is called with that query it will return
 * that result set.</p>
 *
 * <p>This supports multi-phase search.</p>
 *
 * <p>To avoid having to add type information for the fields, a quick hack is used to
 * support testing of attribute prefetching.
 * Any field in the configured hits which has a name starting by attribute
 * will be returned when attribute prefetch filling is requested.</p>
 *
 * @author bratseth
 */
public class DocumentSourceSearcher extends Searcher {

    // using null as name in the API would just be a horrid headache
    public static final String DEFAULT_SUMMARY_CLASS = "default";

    // TODO: update tests to explicitly set hits, so that the default results can be removed entirely.
    private Result defaultFilledResult;

    private final Map<Query, Result> completelyFilledResults = new HashMap<>();
    private final Map<Query, Result> unFilledResults = new HashMap<>();
    private final Map<String, Set<String>> summaryClasses = new HashMap<>();

    private int queryCount;

    public DocumentSourceSearcher() {
        addDefaultResults();
    }

    /**
     * Adds a result which can be searched for and filled.
     * Summary fields starting by "a" are attributes, others are not.
     *
     * @return true when replacing an existing &lt;query, result&gt; pair.
     */
    public boolean addResult(Query query, Result fullResult) {
        Result emptyResult = new Result(query.clone());
        emptyResult.setTotalHitCount(fullResult.getTotalHitCount());
        for (Hit fullHit : fullResult.hits().asList()) {
            Hit emptyHit = fullHit.clone();
            emptyHit.clearFields();
            emptyHit.setFillable();
            emptyHit.setRelevance(fullHit.getRelevance());

            emptyResult.hits().add(emptyHit);
        }
        unFilledResults.put(getQueryKeyClone(query), emptyResult);

        if (completelyFilledResults.put(getQueryKeyClone(query), fullResult.clone()) != null) {
           // TODO: throw exception if the key exists from before, change the method to void
           return true;
        }
        return false;
    }

    public void addSummaryClass(String name, Set<String> fields) {
        summaryClasses.put(name,fields);
    }

    public void addSummaryClassByCopy(String name, Collection<String> fields) {
        addSummaryClass(name, new HashSet<>(fields));
    }

    private void addDefaultResults() {
        Query q = new Query("?query=default");
        Result r = new Result(q);
        r.hits().add(new Hit("http://default-1.html", 0));
        r.hits().add(new Hit("http://default-2.html", 0));
        r.hits().add(new Hit("http://default-3.html", 0));
        r.hits().add(new Hit("http://default-4.html", 0));
        defaultFilledResult = r;
        addResult(q, r);
    }

    @Override
    public Result search(Query query, Execution execution)  {
        queryCount++;
        Result r = unFilledResults.get(getQueryKeyClone(query));
        if (r == null)
            r = defaultFilledResult.clone();
        else
            r = r.clone();

        r.setQuery(query);
        r.hits().trim(query.getOffset(), query.getHits());
        query.setOffset(0);
        return r;
    }

    /**
     * Returns a query clone which has sourcr, offset and hits set to null. This is used by access to
     * the maps using the query as key to achieve lookup independent of offset/hits value
     */
    private Query getQueryKeyClone(Query query) {
        Query key = query.clone();
        key.setWindow(0,0);
        key.getModel().setSources("");
        return key;
    }

    @Override
    public void fill(Result result, String summaryClass, Execution execution) {
        Result filledResult;
        filledResult = completelyFilledResults.get(getQueryKeyClone(result.getQuery()));

        if (filledResult == null) {
            filledResult = defaultFilledResult;
        }
        fillHits(filledResult,result,summaryClass);
    }

    private void fillHits(Result filledHits, Result hitsToFill, String summaryClass) {
        Set<String> fieldsToFill = summaryClasses.get(summaryClass);

        if (fieldsToFill == null ) {
            fieldsToFill = summaryClasses.get(DEFAULT_SUMMARY_CLASS);
        }

        for (Hit hitToFill : hitsToFill.hits()) {
            Hit filledHit = getMatchingFilledHit(hitToFill.getId(), filledHits);

            if (filledHit != null) {
                if (fieldsToFill != null) {
                    copyFieldValuesThatExist(filledHit,hitToFill,fieldsToFill);
                } else {
                    // TODO: remove this block and update fieldsToFill above to throw an exception if no appropriate summary class is found
                    for (var iter = filledHit.fieldIterator(); iter.hasNext();) {
                        var propertyEntry = iter.next();
                        hitToFill.setField(propertyEntry.getKey(), propertyEntry.getValue());
                    }
                }
                hitToFill.setFilled(summaryClass == null ? DEFAULT_SUMMARY_CLASS : summaryClass);
            }
        }
        hitsToFill.analyzeHits();
    }

    private Hit getMatchingFilledHit(URI hitToFillId, Result filledHits) {
        Hit filledHit = null;

        for ( Hit filledHitCandidate : filledHits.hits()) {
            if ( hitToFillId == filledHitCandidate.getId() ) {
                filledHit = filledHitCandidate;
                break;
            }
        }
        return filledHit;
    }

    private void copyFieldValuesThatExist(Hit filledHit, Hit hitToFill, Set<String> fieldsToFill) {
        for (String fieldToFill : fieldsToFill ) {
            if ( filledHit.getField(fieldToFill) != null ) {
                hitToFill.setField(fieldToFill, filledHit.getField(fieldToFill));
            }
        }
    }

    /**
     * Returns the number of queries made to this searcher since the last
     * reset. For testing - not reliable if multiple threads makes
     * queries simultaneously
     */
    public int getQueryCount() { return queryCount; }

    public void resetQueryCount() { queryCount = 0; }

}