Merge pull request pandas-dev#520 from dimosped/MDP-3095-fix-list-symbols-new

Fix for list_symbols reaching memory limit of MongoDB

Author: dimosped

CHANGES.md

@@ -2,6 +2,7 @@
 
 ### 1.62
   * Bugfix: #517 VersionStore: append does not duplicate data in certain corner cases
+  * Bugfix: #519 VersionStore: list_symbols speed improvement and fix for memory limit exceed
 
 ### 1.61 (2018-3-2)
   * Feature: #288 Mapping reads and writes over chunks in chunkstore

arctic/store/version_store.py

@@ -133,6 +133,9 @@ def list_symbols(self, all_symbols=False, snapshot=None, regex=None, **kwargs):
             query['symbol'] = {'$regex': regex}
         if kwargs:
             for k, v in six.iteritems(kwargs):
+                # TODO: this doesn't work as expected as it ignores the versions with metadata.deleted set
+                #       as a result it will return symbols with matching metadata which have been deleted
+                #       Maybe better add a match step in the pipeline instead of making it part of the query
                 query['metadata.' + k] = v
         if snapshot is not None:
             try:
@@ -148,22 +151,32 @@ def list_symbols(self, all_symbols=False, snapshot=None, regex=None, **kwargs):
             # Match based on user criteria first
             pipeline.append({'$match': query})
         pipeline.extend([
-                        # Id is by insert time which matches version order
-                        {'$sort': bson.SON([('symbol', pymongo.DESCENDING), ('version', pymongo.DESCENDING)])},
-                        # Group by 'symbol'
-                        {'$group': {'_id': '$symbol',
-                                    'deleted': {'$first': '$metadata.deleted'},
-                                    },
-                         },
-                        # Don't include symbols which are part of some snapshot, but really deleted...
-                        {'$match': {'deleted': {'$ne': True}}},
-                        {'$project': {'_id': 0,
-                                      'symbol':  '$_id',
-                                      }
-                         }])
-
-        results = self._versions.aggregate(pipeline)
-        return sorted([x['symbol'] for x in results])
+            # version_custom value is: 2*version + (0 if deleted else 1)
+            # This is used to optimize aggregation query:
+            #  - avoid sorting
+            #  - be able to rely on the latest version (max) for the deleted status
+            #
+            # Be aware of that if you don't use custom sort or if use a sort before $group which utilizes
+            # exactly an existing index, the $group will do best effort to utilize this index:
+            #  - https://jira.mongodb.org/browse/SERVER-4507
+            {'$group': {
+                '_id': '$symbol',
+                'version_custom': {
+                    '$max': {
+                        '$add': [
+                            {'$multiply': ['$version', 2]},
+                            {'$cond': [{'$eq': ['$metadata.deleted', True]}, 1, 0]}
+                        ]
+                    }
+                },
+            }},
+            # Don't include symbols which are part of some snapshot, but really deleted...
+            {'$match': {'version_custom': {'$mod': [2, 0]}}}
+        ])
+
+        # We may hit the group memory limit (100MB), so use allowDiskUse to circumvent this
+        #  - https://docs.mongodb.com/manual/reference/operator/aggregation/group/#group-memory-limit
+        return sorted([x['_id'] for x in self._versions.aggregate(pipeline, allowDiskUse=True)])
 
     @mongo_retry
     def has_symbol(self, symbol, as_of=None):

tests/integration/store/test_version_store.py

@@ -992,6 +992,27 @@ def test_list_symbols_newer_version_with_lower_id(library):
     assert symbol not in library.list_symbols()
 
 
+def test_list_symbols_write_snapshot_write_delete(library):
+    library.write('asdf', {'foo': 'bar'})
+    symbol = 'sym_a'
+    library.write(symbol, {'foo2': 'bar2'}, prune_previous_version=False)
+    library.snapshot('s1')
+    library.write(symbol, {'foo3': 'bar2'}, prune_previous_version=False)
+    library.delete(symbol)
+    # at this point we have one version retained from 's1' and
+    # one more version added with delete (version with foo3 is pruned)
+    # The list_symbols should return only 'asdf'
+    assert library.list_symbols() == ['asdf']
+
+
+def test_list_symbols_delete_write(library):
+    symbol = 'sym_a'
+    library.write(symbol, {'foo': 'bar2'}, prune_previous_version=False)
+    library.delete(symbol)
+    library.write(symbol, {'foo2': 'bar2'}, prune_previous_version=False)
+    assert library.list_symbols() == [symbol]
+
+
 def test_date_range_large(library):
     index = [dt(2017,1,1)]*20000 + [dt(2017,1,2)]*20000
     data = np.random.random((40000, 10))

tests/unit/store/test_version_store.py

@@ -268,12 +268,20 @@ def test_list_symbols_default_pipeline():
     VersionStore.list_symbols(vs)
 
     pipeline = [
-        {'$sort': bson.SON([('symbol', pymongo.DESCENDING), ('version', pymongo.DESCENDING)])},
-        {'$group': {'_id': '$symbol', 'deleted': {'$first': '$metadata.deleted'}}},
-        {'$match': {'deleted': {'$ne': True}}},
-        {'$project': {'_id': 0, 'symbol': '$_id'}}
+        {'$group': {
+            '_id': '$symbol',
+            'version_custom': {
+                '$max': {
+                    '$add': [
+                        {'$multiply': ['$version', 2]},
+                        {'$cond': [{'$eq': ['$metadata.deleted', True]}, 1, 0]}
+                    ]
+                }
+            },
+        }},
+        {'$match': {'version_custom': {'$mod': [2, 0]}}}
     ]
-    versions.aggregate.assert_called_once_with(pipeline)
+    versions.aggregate.assert_called_once_with(pipeline, allowDiskUse=True)
 
 
 def test_snapshot_duplicate_raises_exception():