Last active
February 20, 2023 14:39
-
-
Save acdha/0a66ca23984bc8d607936fecd9c29941 to your computer and use it in GitHub Desktop.
Django Haystack backend support for Solr's Collapsing query parser and Expand component
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# encoding: utf-8 | |
""" | |
Experimental SearchQuerySet exposing Solr's Collapse filter and Expand component | |
See https://cwiki.apache.org/confluence/display/solr/Collapse+and+Expand+Results | |
Usage:: | |
sqs = sqs.collapse('item_grouping', sort='"score DESC, wdl_id ASC"') | |
sqs = sqs.expand(rows=2, fq='django_ct:core.file') | |
Note that _process_results has some non-portable code to handle inconsistencies | |
in the response types for numeric item IDs in JSON. | |
""" | |
from __future__ import absolute_import, division, print_function | |
from haystack.backends.solr_backend import SolrEngine, SolrSearchBackend, SolrSearchQuery | |
from haystack.query import SearchQuerySet | |
class CollapsedSearchQuerySet(SearchQuerySet): | |
def __init__(self, *args, **kwargs): | |
super(CollapsedSearchQuerySet, self).__init__(*args, **kwargs) | |
def collapse(self, field_name, **kwargs): | |
# To implement collapsing we'll need to add a filter-query like this: | |
# &fq={!collapse field=item_grouping sort="django_ct desc, score desc"} | |
# FIXME: handle escaping, etc. | |
fq = ['field=%s' % field_name] | |
for k, v in kwargs.items(): | |
fq.append('%s=%s' % (k, v)) | |
return self.narrow('{!collapse %s}' % ' '.join(fq)) | |
def expand(self, **kwargs): | |
clone = self._clone() | |
assert isinstance(clone.query, CollapsedSearchQuery) | |
clone.query.add_expand(**kwargs) | |
return clone | |
class CollapsedSearchQuery(SolrSearchQuery): | |
def __init__(self, *args, **kwargs): | |
super(CollapsedSearchQuery, self).__init__(*args, **kwargs) | |
self.expand = None | |
def add_expand(self, **kwargs): | |
self.expand = kwargs | |
def build_params(self, *args, **kwargs): | |
params = super(CollapsedSearchQuery, self).build_params(*args, **kwargs) | |
if self.expand is not None: | |
params['expand'] = 'true' | |
for k, v in self.expand.items(): | |
params['expand.%s' % k] = v | |
return params | |
class CollapsedSolrSearchBackend(SolrSearchBackend): | |
def _process_results(self, raw_results, *args, **kwargs): | |
res = super(CollapsedSolrSearchBackend, self)._process_results(raw_results, *args, **kwargs) | |
if 'expanded' in raw_results.raw_response: | |
# The keys should be numeric (WDL ID or Item Grouping) but they'll be strings in the raw | |
# JSON response since JSON inherits JavaScript's lack of integer safety. We'll convert: | |
expanded = {int(k): v for k, v in raw_results.raw_response['expanded'].items()} | |
for i in res['results']: | |
extra = expanded.get(i.wdl_id) or expanded.get(i.item_grouping) | |
setattr(i, 'expanded_results', extra) | |
return res | |
class CollapsedSolrEngine(SolrEngine): | |
backend = CollapsedSolrSearchBackend | |
query = CollapsedSearchQuery |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment