Repository: PennyLaneAI/pennylane
Pull request #4681 (2023-10-16 19:09:40)

Fix full dataset download after partial
**Context:**
[Link to shortcut story](https://app.shortcut.com/xanaduai/story/42605/download-dataset-works-after-download-partial-is-called-on-some-dataset-without-force-set)

If you call load and download only a partial dataset (eg. load("qchem", molname="H2", attributes=["hf_state"])), then try to load the full dataset (call load() again but without the attributes kwarg), it will not download the rest of the dataset. This is because if _download_dataset sees that there is already a file for some dataset, it assumes it has the data it needs. You can add force=True to ensure it's re-downloaded, but you shouldn't need to.

**Description of the Change:**
If all attributes are requested in `qml.load()` (e.g `attributes` is `None`), and the dataset already exists locally, `_download_partial()` will be used to download the missing attributes. If `force` is specified, the existing attributes in the local dataset will be overwritten.

`_download_full()` will only be called if the dataset does not exist at all locally.

**Benefits:**
Mixed calls to `qml.load()` work as expected.

**Possible Drawbacks:**

**Related GitHub Issues:**


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
brownj85: Heads up, this will be slow on large datasets until #4674 is merged
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4681?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`4c3ab72`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/4c3ab729c0715c0a78a0e9b3454a650f0e7e7096?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`4911c1f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4681?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.
> Report is 1 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #4681   +/-   ##
=======================================
  Coverage   99.64%   99.64%           
=======================================
  Files         376      376           
  Lines       33828    33835    +7     
=======================================
+ Hits        33707    33714    +7     
  Misses        121      121           
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4681?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/data/data\_manager/\_\_init\_\_.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4681?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RhdGEvZGF0YV9tYW5hZ2VyL19faW5pdF9fLnB5) | `98.62% <100.00%> (+0.06%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4681?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9394ac111..0b139fee1 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -258,6 +258,10 @@
   of the computed qubit operator, if imaginary components are smaller than a threshold. 
   [(#4639)](https://github.com/PennyLaneAI/pennylane/pull/4639)
 
+* `qml.data.load` correctly performs a full download of the dataset after a partial download of the
+  same dataset has already been performed.
+  [(#4681)](https://github.com/PennyLaneAI/pennylane/pull/4681)
+  
 * Improved performance of `qml.data.load()` when partially loading a dataset
   [(#4674)](https://github.com/PennyLaneAI/pennylane/pull/4674)
 
@@ -273,7 +277,6 @@
   [(#4675)](https://github.com/PennyLaneAI/pennylane/pull/4675)
 
 
-
 <h3>Breaking changes 💔</h3>
 
 * The device test suite now converts device kwargs to integers or floats if they can be converted to integers or floats.
diff --git a/pennylane/data/data_manager/__init__.py b/pennylane/data/data_manager/__init__.py
index 877fb37df..cfde2ed5e 100644
--- a/pennylane/data/data_manager/__init__.py
+++ b/pennylane/data/data_manager/__init__.py
@@ -56,17 +56,57 @@ def _get_data_struct():
 
 
 def _download_partial(
-    s3_url: str, dest: Path, attributes: typing.Iterable[str], overwrite: bool, block_size: int
+    s3_url: str,
+    dest: Path,
+    attributes: Optional[typing.Iterable[str]],
+    overwrite: bool,
+    block_size: int,
 ) -> None:
-    """Download only the requested attributes of the Dataset at ``s3_path`` into ``dest``.
-    If a dataset already exists at ``dest``, the attributes will be loaded into
-    the existing dataset.
+    """Download the requested attributes of the Dataset at ``s3_path`` into ``dest``.
+
+    Args:
+        s3_url: URL of the remote dataset
+        dest: Destination dataset path
+        attributes: Requested attributes to download. Passing ``None`` is equivalent
+            to requesting all attributes of the remote dataset.
+        overwrite: If True, overwrite attributes that already exist at ``dest``. Otherwise,
+            only download attributes that do not exist at ``dest``.
     """
 
-    remote_dataset = Dataset(open_hdf5_s3(s3_url, block_size=block_size))
-    remote_dataset.write(dest, "a", attributes, overwrite=overwrite)
+    dest_dataset = Dataset.open(dest, mode="a")
+    remote_dataset = None
+
+    attributes_to_fetch = set()
+
+    if attributes is not None:
+        attributes_to_fetch.update(attributes)
+    else:
+        remote_dataset = Dataset(open_hdf5_s3(s3_url, block_size=block_size))
+        attributes_to_fetch.update(remote_dataset.attrs)
 
+    if not overwrite:
+        attributes_to_fetch.difference_update(dest_dataset.attrs)
+
+    if len(attributes_to_fetch) > 0:
+        remote_dataset = remote_dataset or Dataset(open_hdf5_s3(s3_url, block_size=block_size))
+        remote_dataset.write(dest_dataset, "a", attributes, overwrite=overwrite)
+
+    if remote_dataset:
+        remote_dataset.close()
+
+    dest_dataset.close()
     del remote_dataset
+    del dest_dataset
+
+
+def _download_full(s3_url: str, dest: Path):
+    """Download the full dataset file at ``s3_url`` to ``path``."""
+
+    with open(dest, "wb") as f:
+        resp = get(s3_url, timeout=5.0)
+        resp.raise_for_status()
+
+        f.write(resp.content)
 
 
 def _download_dataset(
@@ -77,25 +117,23 @@ def _download_dataset(
     force: bool = False,
 ) -> None:
     """Downloads the dataset at ``data_path`` to ``dest``, optionally downloading
-    only requested attributes."""
+    only requested attributes. If ``attributes`` is not provided, every attribute
+    will be requested.
+
+    If any of the attributes of the remote dataset are already downloaded locally,
+    they will not be overwritten unless ``force`` is True.
+    """
 
     # URL-escape special characters like '+', '$', and '%' in the data path
     url_safe_datapath = urllib.parse.quote(str(data_path))
+    s3_url = f"{S3_URL}/{url_safe_datapath}"
 
-    s3_path = f"{S3_URL}/{url_safe_datapath}"
-
-    if attributes is not None:
-        _download_partial(s3_path, dest, attributes, overwrite=force, block_size=block_size)
-        return
-
-    if dest.exists() and not force:
-        return
-
-    with open(dest, "wb") as f:
-        resp = get(s3_path, timeout=5.0)
-        resp.raise_for_status()
-
-        f.write(resp.content)
+    if attributes is not None or dest.exists():
+        _download_partial(
+            s3_url, dest=dest, attributes=attributes, overwrite=force, block_size=block_size
+        )
+    else:
+        _download_full(s3_url, dest=dest)
 
 
 def _validate_attributes(data_struct: dict, data_name: str, attributes: typing.Iterable[str]):


Test patch:
diff --git a/tests/data/data_manager/test_dataset_access.py b/tests/data/data_manager/test_dataset_access.py
index 97a55cfee..e7fdcea43 100644
--- a/tests/data/data_manager/test_dataset_access.py
+++ b/tests/data/data_manager/test_dataset_access.py
@@ -24,7 +24,7 @@ import requests
 import pennylane as qml
 import pennylane.data.data_manager
 from pennylane.data import Dataset
-from pennylane.data.data_manager import DataPath, S3_URL, _validate_attributes
+from pennylane.data.data_manager import S3_URL, DataPath, _validate_attributes
 
 # pylint:disable=protected-access,redefined-outer-name
 
@@ -283,60 +283,170 @@ def test_load_except(monkeypatch, tmp_path):
         )
 
 
+@patch("pennylane.data.data_manager._download_partial")
+@patch("pennylane.data.data_manager._download_full")
+@pytest.mark.parametrize("force", (True, False))
+@pytest.mark.parametrize(
+    "attributes, dest_exists, called_partial",
+    [(["x"], True, True), (["x"], False, True), (None, True, True), (None, False, False)],
+)
+def test_download_dataset_full_or_partial(
+    download_full, download_partial, attributes, dest_exists, force, called_partial
+):  # pylint: disable=too-many-arguments
+    """Test that _download_dataset calls ``_download_partial()`` if ``attributes`` is not None,
+    or the dataset already exists at ``dest``, and that it only calls ``_download_full()`` if
+    the dataset does not exist at ``dest`` and ``attributes`` is None.
+    """
+
+    dest = MagicMock()
+    dest.exists.return_value = dest_exists
+
+    pennylane.data.data_manager._download_dataset(
+        "dataset/path", attributes=attributes, dest=dest, force=force, block_size=1
+    )
+
+    assert download_partial.called is called_partial
+    assert download_full.called is not called_partial
+
+
+@pytest.mark.parametrize("force", (True, False))
+@patch("pennylane.data.data_manager._download_full")
+def test_download_dataset_full_call(download_full, force):
+    """Test that ``_download_dataset()`` passes the correct parameters
+    to ``_download_full()``
+    """
+    dest = MagicMock()
+    dest.exists.return_value = False
+
+    pennylane.data.data_manager._download_dataset(
+        "dataset/path", attributes=None, dest=dest, force=force, block_size=1
+    )
+
+    download_full.assert_called_once_with(f"{S3_URL}/dataset/path", dest=dest)
+
+
+@pytest.mark.parametrize("attributes", [None, ["x"]])
+@pytest.mark.parametrize("force", (True, False))
+@patch("pennylane.data.data_manager._download_partial")
+def test_download_dataset_partial_call(download_partial, attributes, force):
+    """Test that ``_download_dataset()`` passes the correct parameters
+    to ``_download_partial()``
+    """
+    dest = MagicMock()
+    dest.exists.return_value = True
+
+    pennylane.data.data_manager._download_dataset(
+        "dataset/path", attributes=attributes, dest=dest, force=force, block_size=1
+    )
+
+    download_partial.assert_called_once_with(
+        f"{S3_URL}/dataset/path", dest=dest, attributes=attributes, overwrite=force, block_size=1
+    )
+
+
 @pytest.mark.usefixtures("mock_requests_get")
 @pytest.mark.parametrize("mock_requests_get", [b"This is binary data"], indirect=True)
-def test_download_dataset_full(tmp_path):
+def test_download_full(tmp_path):
     """Tests that _download_dataset will fetch the dataset file
-    using requests if all attributes are requested."""
+    at ``s3_url`` into ``dest``."""
 
-    pennylane.data.data_manager._download_dataset(
-        "dataset/path", tmp_path / "dataset", attributes=None, block_size=1
+    pennylane.data.data_manager._download_full(
+        "dataset/path",
+        tmp_path / "dataset",
     )
 
     with open(tmp_path / "dataset", "rb") as f:
         assert f.read() == b"This is binary data"
 
 
-@pytest.mark.usefixtures("mock_requests_get")
-@pytest.mark.parametrize("mock_requests_get", [b"This is downloaded data"], indirect=True)
+@pytest.mark.parametrize("overwrite", [True, False])
 @pytest.mark.parametrize(
-    "force, expect_data", [(False, b"This is local data"), (True, b"This is downloaded data")]
+    "attributes, expect_attrs",
+    [(None, {"x": 1, "y": 2}), (["x"], {"x": 1}), (["x", "y"], {"x": 1, "y": 2})],
 )
-def test_download_dataset_full_already_exists(tmp_path, force, expect_data):
-    """Tests that _download_dataset will only overwrite an exists
-    dataset in the same path if `force` is True."""
+def test_download_partial_dest_not_exists(
+    tmp_path, monkeypatch, attributes, expect_attrs, overwrite
+):
+    """Tests that _download_dataset will fetch only the requested attributes
+    of a dataset when the destination does not exist."""
+    remote_dataset = Dataset()
+    remote_dataset.x = 1
+    remote_dataset.y = 2
 
-    with open(tmp_path / "dataset", "wb") as f:
-        f.write(b"This is local data")
+    monkeypatch.setattr(
+        pennylane.data.data_manager, "open_hdf5_s3", MagicMock(return_value=remote_dataset.bind)
+    )
 
-    pennylane.data.data_manager._download_dataset(
-        "dataset/path", tmp_path / "dataset", attributes=None, force=force, block_size=1
+    pennylane.data.data_manager._download_partial(
+        "dataset/path",
+        dest=tmp_path / "dataset",
+        attributes=attributes,
+        overwrite=overwrite,
+        block_size=1,
     )
 
-    with open(tmp_path / "dataset", "rb") as f:
-        assert f.read() == expect_data
+    local = Dataset.open(tmp_path / "dataset")
 
+    assert local.attrs == expect_attrs
 
-def test_download_dataset_partial(tmp_path, monkeypatch):
-    """Tests that _download_dataset will fetch the dataset file
-    using requests if all attributes are requested."""
 
+@pytest.mark.parametrize(
+    "attributes, overwrite, expect_attrs",
+    [
+        (None, False, {"x": "local_value", "y": 2, "z": 3}),
+        (["x", "y", "z"], False, {"x": "local_value", "y": 2, "z": 3}),
+        (None, True, {"x": "remote_value", "y": 2, "z": 3}),
+        (["x", "y", "z"], True, {"x": "remote_value", "y": 2, "z": 3}),
+        (["x"], False, {"x": "local_value"}),
+        (["x"], True, {"x": "remote_value"}),
+        (["y"], True, {"x": "local_value", "y": 2}),
+        (["y"], False, {"x": "local_value", "y": 2}),
+    ],
+)
+def test_download_partial_dest_exists(tmp_path, monkeypatch, attributes, overwrite, expect_attrs):
+    """Tests that _download_dataset will fetch only the requested attributes
+    of a dataset when the dataset does not exist locally."""
     remote_dataset = Dataset()
-    remote_dataset.x = 1
+    remote_dataset.x = "remote_value"
     remote_dataset.y = 2
+    remote_dataset.z = 3
 
     monkeypatch.setattr(
         pennylane.data.data_manager, "open_hdf5_s3", MagicMock(return_value=remote_dataset.bind)
     )
 
-    pennylane.data.data_manager._download_dataset(
-        "dataset/path", tmp_path / "dataset", attributes=["x"], block_size=1
+    local_dataset = Dataset.open(tmp_path / "dataset", "w")
+    local_dataset.x = "local_value"
+    local_dataset.close()
+
+    pennylane.data.data_manager._download_partial(
+        "dataset/path",
+        dest=tmp_path / "dataset",
+        attributes=attributes,
+        overwrite=overwrite,
+        block_size=1,
     )
 
     local = Dataset.open(tmp_path / "dataset")
 
-    assert local.x == 1
-    assert not hasattr(local, "y")
+    assert local.attrs == expect_attrs
+
+
+@patch("pennylane.data.data_manager.open_hdf5_s3")
+def test_download_partial_no_check_remote(open_hdf5_s3, tmp_path):
+    """Test that `_download_partial()` will not open the remote dataset if all requested attributes
+    are present in the local dataset, and ``overwrite`` is ``False``."""
+    local_dataset = Dataset.open(tmp_path / "dataset", "w")
+    local_dataset.x = 1
+    local_dataset.y = 2
+
+    local_dataset.close()
+
+    pennylane.data.data_manager._download_partial(
+        "dataset_url", tmp_path / "dataset", ["x", "y"], overwrite=False, block_size=1
+    )
+
+    open_hdf5_s3.assert_not_called()
 
 
 @patch("builtins.open")
@@ -363,7 +473,7 @@ def test_download_dataset_escapes_url(_, mock_get_args, datapath, escaped):
     "datapath, escaped",
     [("data/NH3+/data.h5", "data/NH3%2B/data.h5"), ("data/CA$H/money.h5", "data/CA%24H/money.h5")],
 )
-def test_download_dataset_escapes_url_partial(mock_download_partial, datapath, escaped):
+def test_download_dataset_escapes_url_partial(download_partial, datapath, escaped):
     """Tests that _download_dataset escapes special characters in a URL when doing a partial
     download."""
     dest = Path("dest")
@@ -374,8 +484,8 @@ def test_download_dataset_escapes_url_partial(mock_download_partial, datapath, e
         DataPath(datapath), dest=dest, attributes=attributes, force=force, block_size=1
     )
 
-    mock_download_partial.assert_called_once_with(
-        f"{S3_URL}/{escaped}", dest, attributes, overwrite=force, block_size=1
+    download_partial.assert_called_once_with(
+        f"{S3_URL}/{escaped}", dest=dest, attributes=attributes, overwrite=force, block_size=1
     )