cstrap
6/16/2015 - 9:43 AM

Extending Haystack Elasticsearch backend - multi_match query from http://www.stamkracht.com/extending-haystacks-elasticsearch-backend/

Extending Haystack Elasticsearch backend - multi_match query

from http://www.stamkracht.com/extending-haystacks-elasticsearch-backend/

# -*- coding: utf-8 -*-
from __future__ import absolute_import

from django.conf import settings

from haystack.backends.elasticsearch_backend import (
    ElasticsearchSearchBackend, ElasticsearchSearchEngine, ElasticsearchSearchQuery)
from haystack.query import SearchQuerySet


class ElasticsearchEngineBackendCustom(ElasticsearchSearchBackend):
    DEFAULT_ANALYZER = "snowball"

    def __init__(self, connection_alias, **connection_options):
        super(ElasticsearchEngineBackendCustom, self).__init__(connection_alias, **connection_options)

        user_settings = getattr(settings, 'ELASTICSEARCH_INDEX_SETTINGS', {})
        if user_settings:
            setattr(self, 'DEFAULT_SETTINGS', user_settings)

        user_analyzer = getattr(settings, 'ELASTICSEARCH_DEFAULT_ANALYZER', '')
        if user_analyzer:
            setattr(self, 'DEFAULT_ANALYZER', user_analyzer)

    def build_search_kwargs(self, query_string, sort_by=None, start_offset=0, end_offset=None,
                            fields='', highlight=False, facets=None,
                            date_facets=None, query_facets=None,
                            narrow_queries=None, spelling_query=None,
                            within=None, dwithin=None, distance_point=None,
                            models=None, limit_to_registered_models=None,
                            result_class=None, multi_match=None):

        out = super(ElasticsearchEngineBackendCustom, self).build_search_kwargs(query_string, sort_by, start_offset,
                                                                                end_offset,
                                                                                fields, highlight, facets,
                                                                                date_facets, query_facets,
                                                                                narrow_queries, spelling_query,
                                                                                within, dwithin, distance_point,
                                                                                models, limit_to_registered_models,
                                                                                result_class)

        if multi_match:
             out['query'] = {
                'multi_match': {
                    'query': multi_match['query'],
                    'fields': multi_match['fields'],
                    'tie_breaker': multi_match['tie_breaker'],
                    'minimum_should_match': multi_match['minimum_should_match'],
                }
            }

        return out

    def build_schema(self, fields):
        content_field_name, mapping = super(ElasticsearchEngineBackendCustom, self).build_schema(fields)

        for field_name, field_class in fields.items():
            field_mapping = mapping[field_class.index_fieldname]

            if field_mapping['type'] == 'string' and field_class.indexed:
                if not hasattr(field_class, 'facet_for') or field_class.field_type in ('ngram', 'edge_ngram'):
                    field_mapping['analyzer'] = getattr(field_class, 'analyzer', self.DEFAULT_ANALYZER)
            mapping.update({field_class.index_fieldname: field_mapping})

        return content_field_name, mapping

    def multi_match_run(self, query, fields, minimum_should_match, tie_breaker):
        from elasticsearch_dsl import Search
        from elasticsearch_dsl.query import MultiMatch

        raw = Search().using(self.conn).query(
            MultiMatch(query=u'{}'.format(query), fields=fields, minimum_should_match=minimum_should_match, tie_breaker=tie_breaker)
        ).execute()

        return self._process_results(raw)


class ElasticsearchSearchQueryCustom(ElasticsearchSearchQuery):
    def multi_match(self, query, fields, minimum_should_match, tie_breaker):
        results = self.backend.multi_match_run(query, fields, minimum_should_match, tie_breaker)
        self._results = results.get('results', [])
        self._hit_count = results.get('hits', 0)

    def add_multi_match_query(self, query, fields, minimum_should_match, tie_breaker):
        self.multi_match_query = {
            'query': query,
            'fields': fields,
            'minimum_should_match': minimum_should_match,
            'tie_breaker': tie_breaker
        }

    def build_params(self, spelling_query=None, **kwargs):
        search_kwargs = super(ElasticsearchSearchQueryCustom, self).build_params(spelling_query, **kwargs)
        if self.multi_match_query:
            search_kwargs['multi_match'] = self.multi_match_query

        return search_kwargs


class ElasticsearchSearchQuerySetCustom(SearchQuerySet):
    def multi_match(self, query, fields, minimum_should_match="35%", tie_breaker=0.3):
        clone = self._clone()
        clone.query.add_multi_match_query(query, fields, minimum_should_match, tie_breaker)
        clone.query.multi_match(query, fields, minimum_should_match, tie_breaker)
        return clone


class ElasticsearchEngineCustom(ElasticsearchSearchEngine):
    backend = ElasticsearchEngineBackendCustom
    query = ElasticsearchSearchQueryCustom