Use os.walk to speedup reading from DirectoryStore

usage of os.walk for tree (os.scandir for folders) is faster than
listdir as it avoids many stats call.

This Should make the DirectoryStore faster.

Author: Carreau

docs/release.rst

@@ -5,6 +5,8 @@ Release notes
 Next release
 ------------
 
+* `DirectoryStore` now uses `os.scandir`, which should make listing large store
+  faster, :issue:`563`
 * Fix minor bug in `N5Store`. 
   By :user:`gsakkis`, :issue:`550`.
 

zarr/storage.py

@@ -845,15 +845,27 @@ def __eq__(self, other):
 
     def keys(self):
         if os.path.exists(self.path):
-            directories = [(self.path, '')]
-            while directories:
-                dir_name, prefix = directories.pop()
-                for name in os.listdir(dir_name):
-                    path = os.path.join(dir_name, name)
-                    if os.path.isfile(path):
-                        yield prefix + name
-                    elif os.path.isdir(path):
-                        directories.append((path, prefix + name + '/'))
+            yield from self._keys_fast(self.path)
+
+    @staticmethod
+    def _keys_fast(path, walker=os.walk):
+        """
+
+        Faster logic on platform where the separator is `/` and using
+        `os.walk()` to decrease the number of stats.call.
+
+        """
+        it = iter(walker(path))
+        d0, dirnames, filenames = next(it)
+        if d0.endswith('/'):
+            root_len = len(d0)
+        else:
+            root_len = len(d0)+1
+        for f in filenames:
+            yield f
+        for dirpath, _, filenames in it:
+            for f in filenames:
+                yield dirpath[root_len:].replace('\\', '/')+'/'+f
 
     def __iter__(self):
         return self.keys()

zarr/tests/test_storage.py

@@ -827,6 +827,34 @@ def test_normalize_keys(self):
         assert 'FOO' in store
         assert 'foo' in store
 
+    def test_listing_keys_slash(self):
+
+        def mock_walker_slash(_path):
+            yield from [
+                # trailing slash in first key
+                ('root_with_slash/', ['d1', 'g1'], ['.zgroup']),
+                ('root_with_slash/d1', [], ['.zarray']),
+                ('root_with_slash/g1', [], ['.zgroup'])
+            ]
+
+        res = set(DirectoryStore._keys_fast('root_with_slash/', walker=mock_walker_slash))
+        assert res == {'.zgroup', 'g1/.zgroup', 'd1/.zarray'}
+
+    def test_listing_keys_no_slash(self):
+
+        def mock_walker_no_slash(_path):
+            yield from [
+                # no trainling slash in first key
+                ('root_with_no_slash', ['d1', 'g1'], ['.zgroup']),
+                ('root_with_no_slash/d1', [], ['.zarray']),
+                ('root_with_no_slash/g1', [], ['.zgroup'])
+            ]
+
+        res = set(
+            DirectoryStore._keys_fast('root_with_no_slash', mock_walker_no_slash)
+                )
+        assert res == {'.zgroup', 'g1/.zgroup', 'd1/.zarray'}
+
 
 @pytest.mark.skipif(have_fsspec is False, reason="needs fsspec")
 class TestFSStore(StoreTests, unittest.TestCase):