summaryrefslogtreecommitdiff
path: root/core/src/main/java/org/elasticsearch/search/suggest/SuggestBuilder.java
blob: fec2b0f3bf31d0d96ea3ddeedddaafc6dd11cb4d (plain)
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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.elasticsearch.search.suggest;

import org.elasticsearch.action.support.ToXContentToBytes;
import org.elasticsearch.common.Nullable;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.ParsingException;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.lucene.BytesRefs;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;

/**
 * Defines how to perform suggesting. This builders allows a number of global options to be specified and
 * an arbitrary number of {@link SuggestionBuilder} instances.
 * <p>
 * Suggesting works by suggesting terms/phrases that appear in the suggest text that are similar compared
 * to the terms in provided text. These suggestions are based on several options described in this class.
 */
public class SuggestBuilder extends ToXContentToBytes implements Writeable {
    protected static final ParseField GLOBAL_TEXT_FIELD = new ParseField("text");

    private String globalText;
    private final Map<String, SuggestionBuilder<?>> suggestions = new HashMap<>();

    /**
     * Build an empty SuggestBuilder.
     */
    public SuggestBuilder() {
    }

    /**
     * Read from a stream.
     */
    public SuggestBuilder(StreamInput in) throws IOException {
        globalText = in.readOptionalString();
        final int size = in.readVInt();
        for (int i = 0; i < size; i++) {
            suggestions.put(in.readString(), in.readNamedWriteable(SuggestionBuilder.class));
        }
    }

    @Override
    public void writeTo(StreamOutput out) throws IOException {
        out.writeOptionalString(globalText);
        final int size = suggestions.size();
        out.writeVInt(size);
        for (Entry<String, SuggestionBuilder<?>> suggestion : suggestions.entrySet()) {
            out.writeString(suggestion.getKey());
            out.writeNamedWriteable(suggestion.getValue());
        }
    }

    /**
     * Sets the text to provide suggestions for. The suggest text is a required option that needs
     * to be set either via this setter or via the {@link org.elasticsearch.search.suggest.SuggestionBuilder#text(String)} method.
     * <p>
     * The suggest text gets analyzed by the suggest analyzer or the suggest field search analyzer.
     * For each analyzed token, suggested terms are suggested if possible.
     */
    public SuggestBuilder setGlobalText(@Nullable String globalText) {
        this.globalText = globalText;
        return this;
    }

    /**
     * Gets the global suggest text
     */
    @Nullable
    public String getGlobalText() {
        return globalText;
    }

    /**
     * Adds an {@link org.elasticsearch.search.suggest.SuggestionBuilder} instance under a user defined name.
     * The order in which the <code>Suggestions</code> are added, is the same as in the response.
     * @throws IllegalArgumentException if two suggestions added have the same name
     */
    public SuggestBuilder addSuggestion(String name, SuggestionBuilder<?> suggestion) {
        Objects.requireNonNull(name, "every suggestion needs a name");
        if (suggestions.get(name) == null) {
            suggestions.put(name, suggestion);
        } else {
            throw new IllegalArgumentException("already added another suggestion with name [" + name + "]");
        }
        return this;
    }

    /**
     * Get all the <code>Suggestions</code> that were added to the global {@link SuggestBuilder},
     * together with their names
     */
    public Map<String, SuggestionBuilder<?>> getSuggestions() {
        return suggestions;
    }

    @Override
    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
        builder.startObject();
        if (globalText != null) {
            builder.field("text", globalText);
        }
        for (Entry<String, SuggestionBuilder<?>> suggestion : suggestions.entrySet()) {
            builder.startObject(suggestion.getKey());
            suggestion.getValue().toXContent(builder, params);
            builder.endObject();
        }
        builder.endObject();
        return builder;
    }

    public static SuggestBuilder fromXContent(XContentParser parser) throws IOException {
        SuggestBuilder suggestBuilder = new SuggestBuilder();
        String fieldName = null;

        if (parser.currentToken() == null) {
            // when we parse from RestSuggestAction the current token is null, advance the token
            parser.nextToken();
        }
        assert parser.currentToken() == XContentParser.Token.START_OBJECT : "current token must be a start object";
        XContentParser.Token token;
        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
            if (token == XContentParser.Token.FIELD_NAME) {
                fieldName = parser.currentName();
            } else if (token.isValue()) {
                if (GLOBAL_TEXT_FIELD.match(fieldName)) {
                    suggestBuilder.setGlobalText(parser.text());
                } else {
                    throw new IllegalArgumentException("[suggest] does not support [" + fieldName + "]");
                }
            } else if (token == XContentParser.Token.START_OBJECT) {
                String suggestionName = fieldName;
                if (suggestionName == null) {
                    throw new IllegalArgumentException("suggestion must have name");
                }
                suggestBuilder.addSuggestion(suggestionName, SuggestionBuilder.fromXContent(parser));
            } else {
                throw new ParsingException(parser.getTokenLocation(), "unexpected token [" + token + "] after [" + fieldName + "]");
            }
        }
        return suggestBuilder;
    }

    public SuggestionSearchContext build(QueryShardContext context) throws IOException {
        SuggestionSearchContext suggestionSearchContext = new SuggestionSearchContext();
        for (Entry<String, SuggestionBuilder<?>> suggestion : suggestions.entrySet()) {
            SuggestionContext suggestionContext = suggestion.getValue().build(context);
            if (suggestionContext.getText() == null) {
                if (globalText == null) {
                    throw new IllegalArgumentException("The required text option is missing");
                }
                suggestionContext.setText(BytesRefs.toBytesRef(globalText));
            }
            suggestionSearchContext.addSuggestion(suggestion.getKey(), suggestionContext);
        }
        return suggestionSearchContext;
    }

    @Override
    public boolean equals(Object other) {
        if (this == other) {
            return true;
        }
        if (other == null || getClass() != other.getClass()) {
            return false;
        }
        SuggestBuilder o = (SuggestBuilder) other;
        return Objects.equals(globalText, o.globalText) &&
               Objects.equals(suggestions, o.suggestions);
    }

    @Override
    public int hashCode() {
        return Objects.hash(globalText, suggestions);
    }
}