Added multiprocessing support to update_index! Thanks to CMGdigital

for funding development of this feature.

Author: toastdriven

AUTHORS

@@ -51,3 +51,4 @@ Thanks to
     * Luke Hatcher (lukeman) for documentation patches.
     * Trey Hunner (treyhunner) for a Whoosh field boosting patch.
     * Kent Gormat of Retail Catalyst for funding the development of multiple index support.
+    * CMGdigital for funding the development on multiprocessing update_index.

docs/management_commands.rst

@@ -53,6 +53,9 @@ arguments::
     ``--remove``:
         Remove objects from the index that are no longer present in the
         database.
+    ``--workers``:
+        Allows for the use multiple workers to parallelize indexing. Requires
+        ``multiprocessing``.
     ``--verbosity``:
         If provided, dumps out more information about what's being done.
         
@@ -68,9 +71,9 @@ arguments::
 .. note::
 
     This command *ONLY* updates records in the index. It does *NOT* handle
-    deletions, so you may need to write a separate script that handles deleted
-    models, such as a queue consumer or something that runs through all records
-    and tries to load the model for it. Alternatively, you can use the
+    deletions unless the ``--remove`` flag is provided. You might consider
+    a queue consumer if the memory requirements for ``--remove`` don't
+    fit your needs. Alternatively, you can use the
     ``RealTimeSearchIndex``, which will automatically handle deletions.
 
 

haystack/management/commands/update_index.py

@@ -1,11 +1,12 @@
 import datetime
+import os
 from optparse import make_option
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.core.management.base import AppCommand
 from django.db import reset_queries
 from django.utils.encoding import smart_str
-from haystack import connections, connection_router
+from haystack import connections as haystack_connections
 from haystack.constants import DEFAULT_ALIAS
 from haystack.query import SearchQuerySet
 
@@ -14,6 +15,93 @@
 DEFAULT_AGE = None
 
 
+def worker(bits):
+    # We need to reset the connections, otherwise the different processes
+    # will try to share the connection, which causes things to blow up.
+    from django.db import connections
+    
+    for alias, info in connections.databases.items():
+        # We need to also tread lightly with SQLite, because blindly wiping
+        # out connections (via ``... = {}``) destroys in-memory DBs.
+        if not 'sqlite3' in info['ENGINE']:
+            try:
+                del(connections._connections[alias])
+            except KeyError:
+                pass
+    
+    if bits[0] == 'do_update':
+        func, model, start, end, total, using, age, verbosity = bits
+    elif bits[0] == 'do_remove':
+        func, model, pks_seen, start, upper_bound, using, verbosity = bits
+    else:
+        return
+    
+    unified_index = haystack_connections[using].get_unified_index()
+    index = unified_index.get_index(model)
+    backend = haystack_connections[using].get_backend()
+    
+    if func == 'do_update':
+        qs = build_queryset(index, model, age=age, verbosity=verbosity)
+        do_update(backend, index, qs, start, end, total, verbosity=verbosity)
+    elif bits[0] == 'do_remove':
+        do_remove(backend, index, model, pks_seen, start, upper_bound, verbosity=verbosity)
+
+
+def build_queryset(index, model, age=DEFAULT_AGE, verbosity=1):
+    extra_lookup_kwargs = {}
+    updated_field = index.get_updated_field()
+    
+    if age:
+        if updated_field:
+            extra_lookup_kwargs['%s__gte' % updated_field] = datetime.datetime.now() - datetime.timedelta(hours=age)
+        else:
+            if verbosity >= 2:
+                print "No updated date field found for '%s' - not restricting by age." % model.__name__
+    
+    if not hasattr(index.index_queryset(), 'filter'):
+        raise ImproperlyConfigured("The '%r' class must return a 'QuerySet' in the 'get_queryset' method." % index)
+    
+    # `.select_related()` seems like a good idea here but can fail on
+    # nullable `ForeignKey` as well as what seems like other cases.
+    return index.index_queryset().filter(**extra_lookup_kwargs).order_by(model._meta.pk.name)
+
+
+def do_update(backend, index, qs, start, end, total, verbosity=1):
+    # Get a clone of the QuerySet so that the cache doesn't bloat up
+    # in memory. Useful when reindexing large amounts of data.
+    small_cache_qs = qs.all()
+    current_qs = small_cache_qs[start:end]
+    
+    if verbosity >= 2:
+        if os.getpid() == os.getppid():
+            print "  indexed %s - %d of %d." % (start+1, end, total)
+        else:
+            print "  indexed %s - %d of %d (by %s)." % (start+1, end, total, os.getpid())
+    
+    # FIXME: Get the right backend.
+    backend.update(index, current_qs)
+    
+    # Clear out the DB connections queries because it bloats up RAM.
+    reset_queries()
+
+
+def do_remove(backend, index, model, pks_seen, start, upper_bound, verbosity=1):
+    # Fetch a list of results.
+    # Can't do pk range, because id's are strings (thanks comments
+    # & UUIDs!).
+    stuff_in_the_index = SearchQuerySet().models(model)[start:upper_bound]
+    
+    # Iterate over those results.
+    for result in stuff_in_the_index:
+        # Be careful not to hit the DB.
+        if not smart_str(result.pk) in pks_seen:
+            # The id is NOT in the small_cache_qs, issue a delete.
+            if verbosity >= 2:
+                print "  removing %s." % result.pk
+            
+            backend.remove(".".join([result.app_label, result.model_name, str(result.pk)]))
+
+
 class Command(AppCommand):
     help = "Freshens the index for the given app(s)."
     base_options = (
@@ -31,6 +119,10 @@ class Command(AppCommand):
         make_option("-u", "--using", action="store", type="string", dest="using", default=None,
             help='If provided, chooses a connection to work with.'
         ),
+        make_option('-k', '--workers', action='store', dest='workers',
+            default=0, type='int', 
+            help='Allows for the use multiple workers to parallelize indexing. Requires multiprocessing.'
+        ),
     )
     option_list = AppCommand.option_list + base_options
 
@@ -51,11 +143,12 @@ class Command(AppCommand):
 
     def handle(self, *apps, **options):
         self.verbosity = int(options.get('verbosity', 1))
+        self.batchsize = options.get('batchsize', DEFAULT_BATCH_SIZE)
         self.age = options.get('age', DEFAULT_AGE)
         self.remove = options.get('remove', False)
         self.using = options.get('using') or DEFAULT_ALIAS
-        
-        self.backend = connections[self.using].get_backend()
+        self.workers = int(options.get('workers', 0))
+        self.backend = haystack_connections[self.using].get_backend()
 
         if not apps:
             from django.db.models import get_app
@@ -77,7 +170,10 @@ def handle_app(self, app, **options):
         from django.db.models import get_models
         from haystack.exceptions import NotHandled
 
-        unified_index = connections[self.using].get_unified_index()
+        unified_index = haystack_connections[self.using].get_unified_index()
+        
+        if self.workers > 0:
+            import multiprocessing
 
         for model in get_models(app):
             try:
@@ -87,73 +183,49 @@ def handle_app(self, app, **options):
                     print "Skipping '%s' - no index." % model
                 continue
 
-            extra_lookup_kwargs = {}
-            updated_field = index.get_updated_field()
-            
-            if self.age:
-                if updated_field:
-                    extra_lookup_kwargs['%s__gte' % updated_field] = datetime.datetime.now() - datetime.timedelta(hours=self.age)
-                else:
-                    if self.verbosity >= 2:
-                        print "No updated date field found for '%s' - not restricting by age." % model.__name__
-            
-            if not hasattr(index.index_queryset(), 'filter'):
-                raise ImproperlyConfigured("The '%r' class must return a 'QuerySet' in the 'get_queryset' method." % index)
-            
-            # `.select_related()` seems like a good idea here but can fail on
-            # nullable `ForeignKey` as well as what seems like other cases.
-            qs = index.index_queryset().filter(**extra_lookup_kwargs).order_by(model._meta.pk.name)
+            qs = build_queryset(index, model, age=self.age, verbosity=self.verbosity)
             total = qs.count()
 
             if self.verbosity >= 1:
                 print "Indexing %d %s." % (total, smart_str(model._meta.verbose_name_plural))
 
-            pks_seen = set()
+            pks_seen = set([smart_str(pk) for pk in qs.values_list('pk', flat=True)])
+            batch_size = self.batchsize or self.backend.batch_size
 
-            for start in range(0, total, self.backend.batch_size):
-                end = min(start + self.backend.batch_size, total)
-                
-                # Get a clone of the QuerySet so that the cache doesn't bloat up
-                # in memory. Useful when reindexing large amounts of data.
-                small_cache_qs = qs.all()
-                current_qs = small_cache_qs[start:end]
-                
-                for obj in current_qs:
-                    pks_seen.add(smart_str(obj.pk))
-                
-                if self.verbosity >= 2:
-                    print "  indexing %s - %d of %d." % (start+1, end, total)
-                
-                self.backend.update(index, current_qs)
+            if self.workers > 0:
+                ghetto_queue = []
+            
+            for start in range(0, total, batch_size):
+                end = min(start + batch_size, total)
 
-                # Clear out the DB connections queries because it bloats up RAM.
-                reset_queries()
+                if self.workers == 0:
+                    do_update(self.backend, index, qs, start, end, total, self.verbosity)
+                else:
+                    ghetto_queue.append(('do_update', model, start, end, total, self.using, self.age, self.verbosity))
+            
+            if self.workers > 0:
+                pool = multiprocessing.Pool(self.workers)
+                pool.map(worker, ghetto_queue)
 
             if self.remove:
                 if self.age or total <= 0:
                     # They're using a reduced set, which may not incorporate
                     # all pks. Rebuild the list with everything.
-                    pks_seen = set()
                     qs = index.index_queryset().values_list('pk', flat=True)
-                    total = qs.count()
-                    
-                    for pk in qs:
-                        pks_seen.add(smart_str(pk))
+                    pks_seen = set([smart_str(pk) for pk in qs])
+                    total = len(pks_seen)
 
-                for start in range(0, total, self.backend.batch_size):
-                    upper_bound = start + self.backend.batch_size
-                    
-                    # Fetch a list of results.
-                    # Can't do pk range, because id's are strings (thanks comments
-                    # & UUIDs!).
-                    stuff_in_the_index = SearchQuerySet().models(model)[start:upper_bound]
+                if self.workers > 0:
+                    ghetto_queue = []
+                
+                for start in range(0, total, batch_size):
+                    upper_bound = start + batch_size
 
-                    # Iterate over those results.
-                    for result in stuff_in_the_index:
-                        # Be careful not to hit the DB.
-                        if not smart_str(result.pk) in pks_seen:
-                            # The id is NOT in the small_cache_qs, issue a delete.
-                            if self.verbosity >= 2:
-                                print "  removing %s." % result.pk
-                            
-                            self.backend.remove(".".join([result.app_label, result.model_name, str(result.pk)]))
+                    if self.workers == 0:
+                        do_remove(self.backend, index, model, pks_seen, start, upper_bound)
+                    else:
+                        ghetto_queue.append(('do_remove', model, pks_seen, start, upper_bound, self.using, self.verbosity))
+                
+                if self.workers > 0:
+                    pool = multiprocessing.Pool(self.workers)
+                    pool.map(worker, ghetto_queue)

tests/settings.py

@@ -1,6 +1,10 @@
 # Haystack settings for running tests.
-DATABASE_ENGINE = 'sqlite3'
-DATABASE_NAME = 'haystack_tests.db'
+DATABASES = {
+    'default': {
+        'ENGINE': 'sqlite3',
+        'NAME': 'haystack_tests.db',
+    }
+}
 
 INSTALLED_APPS = [
     'django.contrib.admin',

tests/solr_tests/tests/__init__.py

@@ -2,6 +2,7 @@
 warnings.simplefilter('ignore', Warning)
 
 from solr_tests.tests.admin import *
+from solr_tests.tests.management_commands import *
 from solr_tests.tests.solr_query import *
 from solr_tests.tests.solr_backend import *
 from solr_tests.tests.templatetags import *

tests/solr_tests/tests/management_commands.py

@@ -0,0 +1,76 @@
+import pysolr
+from django.conf import settings
+from django.core.management import call_command
+from django.test import TestCase
+from haystack import connections
+from haystack import indexes
+from haystack.utils.loading import UnifiedIndex
+from core.models import MockModel
+
+
+class SolrMockSearchIndex(indexes.SearchIndex):
+    text = indexes.CharField(document=True, use_template=True)
+    name = indexes.CharField(model_attr='author', faceted=True)
+    pub_date = indexes.DateField(model_attr='pub_date')
+    
+    def get_model(self):
+        return MockModel
+
+
+class ManagementCommandTestCase(TestCase):
+    fixtures = ['bulk_data.json']
+    
+    def setUp(self):
+        super(ManagementCommandTestCase, self).setUp()
+        self.solr = pysolr.Solr(settings.HAYSTACK_CONNECTIONS['default']['URL'])
+        
+        # Stow.
+        self.old_ui = connections['default'].get_unified_index()
+        self.ui = UnifiedIndex()
+        self.smmi = SolrMockSearchIndex()
+        self.ui.build(indexes=[self.smmi])
+        connections['default']._index = self.ui
+    
+    def tearDown(self):
+        connections['default']._index = self.old_ui
+        super(ManagementCommandTestCase, self).tearDown()
+    
+    def test_basic_commands(self):
+        call_command('clear_index', interactive=False, verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 0)
+        
+        call_command('update_index', verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 23)
+        
+        call_command('clear_index', interactive=False, verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 0)
+        
+        call_command('rebuild_index', interactive=False, verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 23)
+    
+    def test_remove(self):
+        call_command('clear_index', interactive=False, verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 0)
+        
+        call_command('update_index', verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 23)
+        
+        # Remove a model instance.
+        MockModel.objects.get(pk=1).delete()
+        self.assertEqual(self.solr.search('*:*').hits, 23)
+        
+        # Plain ``update_index`` doesn't fix it.
+        call_command('update_index', verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 23)
+        
+        # With the remove flag, it's gone.
+        call_command('update_index', remove=True, verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 22)
+    
+    def test_multiprocessing(self):
+        call_command('clear_index', interactive=False, verbosity=0)
+        self.assertEqual(self.solr.search('*:*').hits, 0)
+        
+        # Watch the output, make sure there are multiple pids.
+        call_command('update_index', verbosity=2, workers=2, batchsize=5)
+        self.assertEqual(self.solr.search('*:*').hits, 23)