Merge pull request pandas-dev#513 from dimosped/fix-slow-list-versions

Fix slow list versions

Author: dimosped

CHANGES.md

@@ -3,6 +3,7 @@
 ### 1.61
   * Feature: #288 Mapping reads and writes over chunks in chunkstore
   * Bugfix: #508 VersionStore: list_symbols and read now always returns latest version
+  * Bugfix: #512 Improved performance for list_versions
 
 ### 1.60 (2018-2-13)
   * Bugfix: #503 ChunkStore: speedup check for -1 segments

arctic/store/version_store.py

@@ -247,6 +247,7 @@ def list_versions(self, symbol=None, snapshot=None, latest_only=False):
                 raise NoDataFoundException('No snapshot %s in library %s' % (snapshot, self._arctic_lib.get_name()))
 
         versions = []
+        snapshots = {ss.get('_id'): ss.get('name') for ss in self._snapshots.find()}
         for symbol in symbols:
             query['symbol'] = symbol
             seen_symbols = set()
@@ -259,7 +260,7 @@ def list_versions(self, symbol=None, snapshot=None, latest_only=False):
                                  'deleted': meta.get('deleted', False) if meta else False,
                                  # We return offset-aware datetimes in Local Time.
                                  'date': ms_to_datetime(datetime_to_ms(version['_id'].generation_time)),
-                                 'snapshots': self._find_snapshots(version.get('parent', []))})
+                                 'snapshots': [snapshots[s] for s in version.get('parent', []) if s in snapshots]})
         return versions
 
     def _find_snapshots(self, parent_ids):

tests/unit/store/test_version_store.py

@@ -29,19 +29,41 @@ def test_delete_version_version_not_found():
 def test_list_versions_localTime():
     # Object ID's are stored in UTC. We need to ensure that the returned times
     # for versions are in the local  TimeZone
-    vs = create_autospec(VersionStore, instance=True,
-                         _versions=Mock())
-    vs._find_snapshots.return_value = 'snap'
+    vs = create_autospec(VersionStore, instance=True, _versions=Mock(), _snapshots=Mock())
+    mocked_snap_resp = [{'_id': 'abcde', 'name': 'snap'}]
+    vs._snapshots.find.return_value = mocked_snap_resp
+    vs._snapshots.find_one.return_value = mocked_snap_resp
     date = dt(2013, 4, 1, 9, 0)
     vs._versions.find.return_value = [{'_id': bson.ObjectId.from_datetime(date),
-                                       'symbol': 's', 'version': 10, 'metadata': None}]
-
+                                       'symbol': 's',
+                                       'version': 10,
+                                       'metadata': None,
+                                       'parent': [mocked_snap_resp[0]['_id']]}]
     version = list(VersionStore.list_versions(vs, "symbol"))[0]
     local_date = date.replace(tzinfo=mktz("UTC"))
     assert version == {'symbol': version['symbol'], 'version': version['version'],
                        # We return naive datetimes in 'default' time, which is London for us
                        'date': local_date,
-                       'snapshots': 'snap',
+                       'snapshots': ['snap'],
+                       'deleted': False}
+
+
+def test_list_versions_no_snapshot():
+    vs = create_autospec(VersionStore, instance=True, _versions=Mock(), _snapshots=Mock())
+    vs._snapshots.find.return_value = []
+    vs._snapshots.find_one.return_value = []
+    date = dt(2013, 4, 1, 9, 0)
+    vs._versions.find.return_value = [{'_id': bson.ObjectId.from_datetime(date),
+                                       'symbol': 's',
+                                       'version': 10,
+                                       'metadata': None,
+                                       'parent': []}]
+    version = list(VersionStore.list_versions(vs, "symbol"))[0]
+    local_date = date.replace(tzinfo=mktz("UTC"))
+    assert version == {'symbol': version['symbol'],
+                       'version': version['version'],
+                       'date': local_date,
+                       'snapshots': [],
                        'deleted': False}