suggester_completion.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. // Copyright 2012-present Oliver Eilhard. All rights reserved.
  2. // Use of this source code is governed by a MIT-license.
  3. // See http://olivere.mit-license.org/license.txt for details.
  4. package elastic
  5. import "errors"
  6. // CompletionSuggester is a fast suggester for e.g. type-ahead completion.
  7. // See https://www.elastic.co/guide/en/elasticsearch/reference/5.2/search-suggesters-completion.html
  8. // for more details.
  9. type CompletionSuggester struct {
  10. Suggester
  11. name string
  12. text string
  13. field string
  14. analyzer string
  15. size *int
  16. shardSize *int
  17. contextQueries []SuggesterContextQuery
  18. }
  19. // Creates a new completion suggester.
  20. func NewCompletionSuggester(name string) *CompletionSuggester {
  21. return &CompletionSuggester{
  22. name: name,
  23. contextQueries: make([]SuggesterContextQuery, 0),
  24. }
  25. }
  26. func (q *CompletionSuggester) Name() string {
  27. return q.name
  28. }
  29. func (q *CompletionSuggester) Text(text string) *CompletionSuggester {
  30. q.text = text
  31. return q
  32. }
  33. func (q *CompletionSuggester) Field(field string) *CompletionSuggester {
  34. q.field = field
  35. return q
  36. }
  37. func (q *CompletionSuggester) Analyzer(analyzer string) *CompletionSuggester {
  38. q.analyzer = analyzer
  39. return q
  40. }
  41. func (q *CompletionSuggester) Size(size int) *CompletionSuggester {
  42. q.size = &size
  43. return q
  44. }
  45. func (q *CompletionSuggester) ShardSize(shardSize int) *CompletionSuggester {
  46. q.shardSize = &shardSize
  47. return q
  48. }
  49. func (q *CompletionSuggester) ContextQuery(query SuggesterContextQuery) *CompletionSuggester {
  50. q.contextQueries = append(q.contextQueries, query)
  51. return q
  52. }
  53. func (q *CompletionSuggester) ContextQueries(queries ...SuggesterContextQuery) *CompletionSuggester {
  54. q.contextQueries = append(q.contextQueries, queries...)
  55. return q
  56. }
  57. // completionSuggesterRequest is necessary because the order in which
  58. // the JSON elements are routed to Elasticsearch is relevant.
  59. // We got into trouble when using plain maps because the text element
  60. // needs to go before the completion element.
  61. type completionSuggesterRequest struct {
  62. Text string `json:"text"`
  63. Completion interface{} `json:"completion"`
  64. }
  65. // Creates the source for the completion suggester.
  66. func (q *CompletionSuggester) Source(includeName bool) (interface{}, error) {
  67. cs := &completionSuggesterRequest{}
  68. if q.text != "" {
  69. cs.Text = q.text
  70. }
  71. suggester := make(map[string]interface{})
  72. cs.Completion = suggester
  73. if q.analyzer != "" {
  74. suggester["analyzer"] = q.analyzer
  75. }
  76. if q.field != "" {
  77. suggester["field"] = q.field
  78. }
  79. if q.size != nil {
  80. suggester["size"] = *q.size
  81. }
  82. if q.shardSize != nil {
  83. suggester["shard_size"] = *q.shardSize
  84. }
  85. switch len(q.contextQueries) {
  86. case 0:
  87. case 1:
  88. src, err := q.contextQueries[0].Source()
  89. if err != nil {
  90. return nil, err
  91. }
  92. suggester["context"] = src
  93. default:
  94. ctxq := make(map[string]interface{})
  95. for _, query := range q.contextQueries {
  96. src, err := query.Source()
  97. if err != nil {
  98. return nil, err
  99. }
  100. // Merge the dictionary into ctxq
  101. m, ok := src.(map[string]interface{})
  102. if !ok {
  103. return nil, errors.New("elastic: context query is not a map")
  104. }
  105. for k, v := range m {
  106. ctxq[k] = v
  107. }
  108. }
  109. suggester["contexts"] = ctxq
  110. }
  111. // TODO(oe) Add completion-suggester specific parameters here
  112. if !includeName {
  113. return cs, nil
  114. }
  115. source := make(map[string]interface{})
  116. source[q.name] = cs
  117. return source, nil
  118. }