Skip to content

Instantly share code, notes, and snippets.

@acdha acdha/collapsed.py
Last active Dec 11, 2016

Embed
What would you like to do?
Django Haystack backend support for Solr's Collapsing query parser and Expand component
# 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
You can’t perform that action at this time.