Merge pull request #2194 from pared/bulk_unused_links

state: bulk sql for unused links deletion

dvc/remote/base.py

@@ -636,7 +636,7 @@ class RemoteBASE(object):
         if self.no_traverse and hasattr(self, "batch_exists"):
             with ThreadPoolExecutor(max_workers=self.JOBS) as executor:
                 path_infos = [self.checksum_to_path_info(x) for x in checksums]
-                chunks = to_chunks(path_infos, self.JOBS)
+                chunks = to_chunks(path_infos, num_chunks=self.JOBS)
                 results = executor.map(exists_with_progress, chunks)
                 in_remote = itertools.chain.from_iterable(results)
                 ret = list(itertools.compress(checksums, in_remote))

dvc/remote/local/__init__.py

@@ -360,9 +360,9 @@ class RemoteLOCAL(RemoteBASE):
             from_infos = cache
 
         return (
-            to_chunks(from_infos, jobs),
-            to_chunks(to_infos, jobs),
-            to_chunks(names, jobs),
+            to_chunks(from_infos, num_chunks=jobs),
+            to_chunks(to_infos, num_chunks=jobs),
+            to_chunks(names, num_chunks=jobs),
         )
 
     def _process(

dvc/remote/ssh/__init__.py

@@ -157,7 +157,7 @@ class RemoteSSH(RemoteBASE):
 
                 with ThreadPoolExecutor(max_workers=max_workers) as executor:
                     paths = [path_info.path for path_info in path_infos]
-                    chunks = to_chunks(paths, max_workers)
+                    chunks = to_chunks(paths, num_chunks=max_workers)
                     chunks_and_channels = zip(chunks, channels)
                     outcome = executor.map(_exists, chunks_and_channels)
                     results = list(itertools.chain.from_iterable(outcome))

dvc/state.py

@@ -7,11 +7,12 @@ import sqlite3
 import logging
 
 from dvc.config import Config
-from dvc.utils import remove, current_timestamp, relpath
+from dvc.utils import remove, current_timestamp, relpath, to_chunks
 from dvc.exceptions import DvcException
 from dvc.utils.fs import get_mtime_and_size, get_inode
 from dvc.utils.compat import fspath_py35
 
+SQLITE_MAX_VARIABLES_NUMBER = 999
 
 logger = logging.getLogger(__name__)
 
@@ -449,6 +450,10 @@ class State(StateBase):  # pylint: disable=too-many-instance-attributes
                 remove(path)
                 unused.append(relpath)
 
-        for relpath in unused:
-            cmd = "DELETE FROM {} WHERE path = ?".format(self.LINK_STATE_TABLE)
-            self._execute(cmd, (relpath,))
+        for chunk_unused in to_chunks(
+            unused, chunk_size=SQLITE_MAX_VARIABLES_NUMBER
+        ):
+            cmd = "DELETE FROM {} WHERE path IN ({})".format(
+                self.LINK_STATE_TABLE, ",".join(["?"] * len(chunk_unused))
+            )
+            self._execute(cmd, tuple(chunk_unused))

dvc/utils/__init__.py

@@ -183,16 +183,33 @@ def remove(path):
             raise
 
 
-def to_chunks(l, jobs):
-    n = int(math.ceil(len(l) / jobs))
+def _split(list_to_split, chunk_size):
+    return [
+        list_to_split[i : i + chunk_size]
+        for i in range(0, len(list_to_split), chunk_size)
+    ]
 
-    if len(l) == 1:
-        return [l]
 
-    if n == 0:
-        n = 1
+def _to_chunks_by_chunks_number(list_to_split, num_chunks):
+    chunk_size = int(math.ceil(float(len(list_to_split)) / num_chunks))
 
-    return [l[x : x + n] for x in range(0, len(l), n)]
+    if len(list_to_split) == 1:
+        return [list_to_split]
+
+    if chunk_size == 0:
+        chunk_size = 1
+
+    return _split(list_to_split, chunk_size)
+
+
+def to_chunks(list_to_split, num_chunks=None, chunk_size=None):
+    if (num_chunks and chunk_size) or (not num_chunks and not chunk_size):
+        raise ValueError(
+            "One and only one of 'num_chunks', 'chunk_size' must be defined"
+        )
+    if chunk_size:
+        return _split(list_to_split, chunk_size)
+    return _to_chunks_by_chunks_number(list_to_split, num_chunks)
 
 
 # NOTE: Check if we are in a bundle

tests/func/test_state.py

@@ -8,72 +8,85 @@ from dvc.utils import file_md5
 from dvc.main import main
 from dvc.path_info import PathInfo
 
-from tests.basic_env import TestDvc
 
+def test_state(dvc_repo, repo_dir):
+    path = os.path.join(dvc_repo.root_dir, repo_dir.FOO)
+    path_info = PathInfo(path)
+    md5 = file_md5(path)[0]
+
+    state = State(dvc_repo, dvc_repo.config.config)
+
+    with state:
+        state.save(path_info, md5)
+        entry_md5 = state.get(path_info)
+        assert entry_md5 == md5
+
+        os.unlink(path)
+        with open(path, "a") as fd:
+            fd.write("1")
+
+        entry_md5 = state.get(path_info)
+        assert entry_md5 is None
 
-class TestState(TestDvc):
-    def test_update(self):
-        path = os.path.join(self.dvc.root_dir, self.FOO)
-        path_info = PathInfo(path)
         md5 = file_md5(path)[0]
+        state.save(path_info, md5)
 
-        state = State(self.dvc, self.dvc.config.config)
-
-        with state:
-            state.save(path_info, md5)
-            entry_md5 = state.get(path_info)
-            self.assertEqual(entry_md5, md5)
-
-            os.unlink(path)
-            with open(path, "a") as fd:
-                fd.write("1")
-
-            entry_md5 = state.get(path_info)
-            self.assertTrue(entry_md5 is None)
-
-            md5 = file_md5(path)[0]
-            state.save(path_info, md5)
-
-            entry_md5 = state.get(path_info)
-            self.assertEqual(entry_md5, md5)
+        entry_md5 = state.get(path_info)
+        assert entry_md5 == md5
 
 
-class TestStateOverflow(TestDvc):
-    def test(self):
-        # NOTE: trying to add more entries than state can handle,
-        # to see if it will clean up and vacuum successfully
-        ret = main(["config", "state.row_limit", "10"])
-        self.assertEqual(ret, 0)
+def test_state_overflow(dvc_repo):
+    # NOTE: trying to add more entries than state can handle,
+    # to see if it will clean up and vacuum successfully
+    ret = main(["config", "state.row_limit", "10"])
+    assert ret == 0
 
-        dname = "dir"
-        os.mkdir(dname)
-        for i in range(20):
-            with open(os.path.join(dname, str(i)), "w+") as fobj:
-                fobj.write(str(i))
+    dname = "dir"
+    os.mkdir(dname)
+    for i in range(20):
+        with open(os.path.join(dname, str(i)), "w+") as fobj:
+            fobj.write(str(i))
 
-        ret = main(["add", "dir"])
-        self.assertEqual(ret, 0)
+    ret = main(["add", "dir"])
+    assert ret == 0
 
 
-class TestGetStateRecordForInode(TestDvc):
-    @staticmethod
-    def mock_get_inode(inode):
-        def get_inode_mocked(path):
-            return inode
+def mock_get_inode(inode):
+    def get_inode_mocked(path):
+        return inode
 
-        return get_inode_mocked
+    return get_inode_mocked
 
-    @mock.patch("dvc.state.get_inode", autospec=True)
-    def test_transforms_inode(self, get_inode_mock):
-        state = State(self.dvc, self.dvc.config.config)
-        inode = state.MAX_INT + 2
-        self.assertNotEqual(inode, state._to_sqlite(inode))
 
-        path = os.path.join(self.dvc.root_dir, self.FOO)
-        md5 = file_md5(path)[0]
-        get_inode_mock.side_effect = self.mock_get_inode(inode)
+@mock.patch("dvc.state.get_inode", autospec=True)
+def test_get_state_record_for_inode(get_inode_mock, dvc_repo, repo_dir):
+    state = State(dvc_repo, dvc_repo.config.config)
+    inode = state.MAX_INT + 2
+    assert inode != state._to_sqlite(inode)
 
-        with state:
-            state.save(PathInfo(path), md5)
-            ret = state.get_state_record_for_inode(inode)
-            self.assertIsNotNone(ret)
+    path = os.path.join(dvc_repo.root_dir, repo_dir.FOO)
+    md5 = file_md5(path)[0]
+    get_inode_mock.side_effect = mock_get_inode(inode)
+
+    with state:
+        state.save(PathInfo(path), md5)
+        ret = state.get_state_record_for_inode(inode)
+        assert ret is not None
+
+
+def test_remove_unused_links(repo_dir, dvc_repo):
+    stages = dvc_repo.add(repo_dir.FOO)
+    assert len(stages) == 1
+
+    stages = dvc_repo.add(repo_dir.BAR)
+    assert len(stages) == 1
+
+    cmd_count_links = "SELECT count(*) FROM {}".format(State.LINK_STATE_TABLE)
+    with dvc_repo.state:
+        result = dvc_repo.state._execute(cmd_count_links).fetchone()[0]
+        assert result == 2
+
+        dvc_repo.state.remove_unused_links([])
+
+        result = dvc_repo.state._execute(cmd_count_links).fetchone()[0]
+        assert result == 0

tests/unit/utils/test_utils.py

@@ -0,0 +1,30 @@
+import pytest
+
+from dvc.utils import to_chunks
+
+
+@pytest.mark.parametrize(
+    "chunk_size, expected_chunks",
+    [(1, [[1], [2], [3], [4]]), (2, [[1, 2], [3, 4]]), (3, [[1, 2, 3], [4]])],
+)
+def test_to_chunks_chunk_size(chunk_size, expected_chunks):
+    list_to_chunk = [1, 2, 3, 4]
+    result = list(to_chunks(list_to_chunk, chunk_size=chunk_size))
+    assert result == expected_chunks
+
+
+@pytest.mark.parametrize("num_chunks, chunk_size", [(1, 2), (None, None)])
+def test_to_chunks_should_raise(num_chunks, chunk_size):
+    list_to_chunk = [1, 2, 3]
+    with pytest.raises(ValueError):
+        to_chunks(list_to_chunk, num_chunks, chunk_size)
+
+
+@pytest.mark.parametrize(
+    "num_chunks, expected_chunks",
+    [(4, [[1], [2], [3], [4]]), (3, [[1, 2], [3, 4]]), (2, [[1, 2], [3, 4]])],
+)
+def test_to_chunks_num_chunks(num_chunks, expected_chunks):
+    list_to_chunk = [1, 2, 3, 4]
+    result = to_chunks(list_to_chunk, num_chunks=num_chunks)
+    assert result == expected_chunks