Add download_sample for a dataset's public sample

docs/index.md

@@ -113,6 +113,34 @@ dataset = download_dataset("your-dataset-id")
 
 The files will be stored under `MDC_DOWNLOAD_PATH` (default `~/.mozdata/datasets`).
 
+### Download a dataset's sample
+
+Some publishers provide a **sample**: a small excerpt of the dataset for evaluation before
+downloading the full archive. Use `download_sample` to fetch it:
+
+```python
+from datacollective import download_sample
+
+sample_path = download_sample("your-dataset-id")
+```
+
+The sample archive is saved under `samples/<dataset-id>/` inside `MDC_DOWNLOAD_PATH`,
+apart from the full dataset archives. A sample may share a filename with the full archive,
+so keeping them in separate directories ensures one is never mistaken for the other.
+
+The function returns the path to the sample archive. The archive is not extracted; use
+`load_dataset` for the full dataset if you want a loaded DataFrame.
+
+If the sample is already present, it is returned without downloading again. Pass
+`overwrite_existing=True` to download it again and replace the existing file.
+
+!!! note
+    Downloading a sample requires an API key whose organization has agreed to the dataset's
+    terms on the MDC platform, just like downloading the full dataset. Datasets that
+    publish no sample raise `FileNotFoundError`. To find datasets with a sample, pass
+    `has_sample=True` to `list_datasets` (see
+    [Browse and search the dataset catalog](#browse-and-search-the-dataset-catalog)).
+
 ## Programmatic submissions and uploads
 
 The SDK supports creating dataset submissions and uploading files with resumable uploads. 
@@ -368,6 +396,10 @@ or `load_dataset` again.
 - If only partial files exist without proper tracking data, the SDK safely starts
   a fresh download.
 
+!!! note
+    Sample downloads (`download_sample`) do not resume. The API returns no checksum for
+    samples, so an interrupted sample download starts again from the beginning.
+
 
 ### Automatically check for extracted archives
 

src/datacollective/__init__.py

@@ -2,6 +2,7 @@
 
 from datacollective.datasets import (
     download_dataset,
+    download_sample,
     get_dataset_details,
     list_dataset_filters,
     list_datasets,
@@ -33,6 +34,7 @@
 
 __all__ = [
     "download_dataset",
+    "download_sample",
     "save_dataset_to_disk",
     "load_dataset",
     "get_dataset_details",

src/datacollective/datasets.py

@@ -23,8 +23,10 @@
 from datacollective.download import (
     DOWNLOAD_SOURCE_DOWNLOAD,
     DOWNLOAD_SOURCE_LOAD,
+    DOWNLOAD_SOURCE_SAMPLE,
     DOWNLOAD_SOURCE_SAVE,
     _download_dataset,
+    _download_sample,
 )
 from datacollective.hf_utils import _convert_to_hf, _require_datasets
 from datacollective.logging_utils import (
@@ -123,6 +125,63 @@ def download_dataset(
     )
 
 
+def download_sample(
+    dataset_id: str,
+    download_directory: str | None = None,
+    show_progress: bool = True,
+    overwrite_existing: bool = False,
+    enable_logging: bool = False,
+) -> Path:
+    """
+    Download the dataset's sample archive to a local directory and return the archive path.
+    The sample is a small excerpt of the dataset that its publisher provides for evaluation.
+
+    The archive is saved under `samples/<dataset_id>/` inside the download directory, so it
+    can never be mistaken for the full dataset archive even when the two share a filename.
+    Skips download if the sample the API currently serves is already there, unless
+    `overwrite_existing=True`. A sample the publisher has since replaced is downloaded under
+    its new filename, leaving the earlier file in place.
+
+    Unlike `download_dataset`, an interrupted sample download starts again rather than
+    resuming: the API returns no checksum for samples to match a partial file against.
+
+    Requires an API key whose organization has agreed to the dataset's terms on the MDC
+    platform. Buying the dataset is not required.
+
+    Args:
+        dataset_id: The dataset ID (as shown in MDC platform) or slug.
+        download_directory: Directory under which to save the sample archive.
+            If None or empty, falls back to env MDC_DOWNLOAD_PATH or default.
+        show_progress: Whether to show a progress bar during download.
+        overwrite_existing: Whether to download the sample again, replacing the existing one.
+        enable_logging: Whether to enable SDK logging to console and a local log file.
+
+    Returns:
+        Path to the downloaded sample archive.
+
+    Raises:
+        ValueError: If dataset_id is empty.
+        FileNotFoundError: If the dataset does not exist, or publishes no sample (404).
+        ResourceRemovedError: If the dataset is no longer available (410).
+        PermissionError: If the dataset's terms have not been agreed to or the dataset is
+            private to another organization (403), or download directory is not writable.
+        RuntimeError: If rate limit is exceeded (429) or unexpected response format.
+        requests.HTTPError: For other non-2xx responses.
+    """
+    _enable_logging(enable_logging)
+    logger.info(f"Downloading sample of dataset {dataset_id}")
+
+    dataset_details = get_dataset_details(dataset_id)
+
+    return _download_sample(
+        dataset_id=dataset_details.id,
+        download_directory=download_directory,
+        show_progress=show_progress,
+        overwrite_existing=overwrite_existing,
+        download_source=DOWNLOAD_SOURCE_SAMPLE,
+    )
+
+
 def _download_archive(
     dataset_id: str,
     download_directory: str | None,

src/datacollective/download.py

@@ -20,16 +20,21 @@
 DOWNLOAD_SOURCE_DOWNLOAD = "download_dataset"
 DOWNLOAD_SOURCE_SAVE = "save_dataset_to_disk"
 DOWNLOAD_SOURCE_LOAD = "load_dataset"
+DOWNLOAD_SOURCE_SAMPLE = "download_sample"
+
+# Samples live under <download directory>/samples/<dataset id>/ rather than next to the full
+# archives: a sample may carry the same filename as its dataset's archive, or as another
+# dataset's sample, and sharing a directory would let one be mistaken for the other.
+SAMPLES_DIRNAME = "samples"
 
 
 class DownloadPlan(NonEmptyStrModel):
     download_url: str
     target_filepath: Path
     tmp_filepath: Path
     size_bytes: int = Field(..., gt=0)
-    checksum: (
-        str | None
-    )  # We allow None if the API does not return a checksum, but that's generally unexpected.
+    # The API returns no checksum for a sample, so a sample download can never be resumed.
+    checksum: str | None
     checksum_filepath: Path
 
 
@@ -80,6 +85,93 @@ def _download_dataset(
         download_source=download_source,
     )
 
+    return _run_download_plan(
+        download_plan=download_plan,
+        show_progress=show_progress,
+        overwrite_existing=overwrite_existing,
+        download_source=download_source,
+    )
+
+
+def _download_sample(
+    dataset_id: str,
+    download_directory: str | None,
+    show_progress: bool,
+    overwrite_existing: bool,
+    download_source: str | None = None,
+) -> Path:
+    """
+    Download a dataset's sample archive from MDC to a local directory.
+
+    Follows the same flow as `_download_dataset`, with two differences:
+
+    - The sample's filename is only known once the API is called, so the API is asked first
+      and the existence check is done on the exact filename it returns. Samples are kept in a
+      directory of their own, `<download directory>/samples/<dataset id>/`, since a sample's
+      filename may collide with the full archive's or with another dataset's sample.
+    - The API returns no checksum for a sample, so an interrupted download cannot be matched
+      to the current sample and starts again instead of resuming.
+
+    Args:
+        dataset_id: The unique dataset ID.
+        download_directory: The base directory to download into.
+        show_progress: Show a progress bar when downloading the sample.
+        overwrite_existing: If True, download the sample again and replace any existing one.
+        download_source: Optional context appended to the User-Agent for download analytics.
+    Returns:
+        The full path to the downloaded sample archive.
+    """
+    sample_dir = (
+        _resolve_download_dir(download_directory) / SAMPLES_DIRNAME / dataset_id
+    )
+
+    download_plan = _get_sample_download_plan(
+        dataset_id=dataset_id,
+        sample_dir=sample_dir,
+        download_source=download_source,
+    )
+    target_filepath = download_plan.target_filepath
+
+    if target_filepath.exists() and not overwrite_existing:
+        logger.info(
+            f"Skipping download. Sample archive already exists at `{target_filepath}`"
+        )
+        return target_filepath
+
+    # Created only now, so a refused request leaves no empty directory behind.
+    sample_dir.mkdir(parents=True, exist_ok=True)
+
+    # With no checksum in the plan, `_determine_resume_state` discards any leftover
+    # .part file and the download starts again, which is all a sample can do.
+    return _run_download_plan(
+        download_plan=download_plan,
+        show_progress=show_progress,
+        overwrite_existing=overwrite_existing,
+        download_source=download_source,
+    )
+
+
+def _run_download_plan(
+    download_plan: DownloadPlan,
+    show_progress: bool,
+    overwrite_existing: bool,
+    download_source: str | None = None,
+) -> Path:
+    """
+    Carry out a download plan: clean up when overwriting, resume where possible, download to
+    the temporary file, then move it into place.
+
+    Args:
+        download_plan: The DownloadPlan object with download details.
+        show_progress: Whether to show a progress bar during download.
+        overwrite_existing: If True, remove any existing complete or partial download first.
+        download_source: Optional context appended to the User-Agent for download analytics.
+
+    Returns:
+        The full path to the downloaded archive.
+    """
+    target_filepath = download_plan.target_filepath
+
     # If overwriting, clean up any existing complete or partial download files
     if overwrite_existing:
         logger.info(
@@ -138,6 +230,62 @@ def _get_download_plan(
     )
 
     payload: dict[str, Any] = resp.json()
+    return _build_download_plan(payload, target_filepath)
+
+
+def _get_sample_download_plan(
+    dataset_id: str,
+    sample_dir: Path,
+    download_source: str | None = None,
+) -> DownloadPlan:
+    """
+    Send a POST request to the API to receive the download details for a dataset's sample.
+
+    Args:
+        dataset_id: The dataset ID (as shown in MDC platform).
+        sample_dir: Resolved, absolute directory to save the sample in.
+        download_source: Optional context appended to the User-Agent for download analytics.
+
+    Returns:
+        a DownloadPlan object
+
+    Raises:
+        FileNotFoundError: If the dataset does not exist or publishes no sample (404).
+        ResourceRemovedError: If the dataset is no longer available (410).
+        PermissionError: If access is denied (403).
+        RuntimeError: If rate limit is exceeded (429) or unexpected response format.
+        requests.HTTPError: For other non-2xx responses.
+    """
+    sample_url = f"{_get_api_url()}/datasets/{dataset_id}/sample-download"
+    resp = _send_api_request(
+        method="POST", url=sample_url, source_function=download_source
+    )
+
+    payload: dict[str, Any] = resp.json()
+    filename = str(payload.get("filename") or "")
+    # The filename becomes a path on disk, so accept a bare, visible name only.
+    if Path(filename).name != filename or not filename or filename.startswith("."):
+        raise RuntimeError(f"Unexpected sample filename in response: {filename!r}")
+
+    return _build_download_plan(payload, sample_dir / filename)
+
+
+def _build_download_plan(
+    payload: dict[str, Any], target_filepath: Path
+) -> DownloadPlan:
+    """
+    Build a DownloadPlan from a download response of the MDC API.
+
+    Args:
+        payload: The JSON body of the response.
+        target_filepath: Resolved, absolute full file path to save the download.
+
+    Returns:
+        a DownloadPlan object
+
+    Raises:
+        RuntimeError: If the response lacks a download URL or size.
+    """
     download_url = payload.get("downloadUrl")
     size_bytes = payload.get("sizeBytes")
     checksum = payload.get("checksum")
@@ -172,6 +320,8 @@ def _determine_resume_state(download_plan: DownloadPlan) -> str | None:
     A download is resumed only when both the .part and .checksum files exist
     and the stored checksum matches the current one. In every other case the
     leftover .part / .checksum files are removed and the download starts fresh.
+    A plan without a checksum (a sample) can never be resumed, so any leftover
+    .part file is simply discarded.
 
     Args:
         download_plan: The DownloadPlan object with download details.
@@ -180,6 +330,16 @@ def _determine_resume_state(download_plan: DownloadPlan) -> str | None:
         The checksum to use for resumption, or None if starting fresh.
     """
     part_exists = download_plan.tmp_filepath.exists()
+
+    if download_plan.checksum is None:
+        if part_exists:
+            logger.info(
+                "Partial download found but the download cannot be resumed. "
+                "Starting fresh download..."
+            )
+        _cleanup_partial_download(download_plan)
+        return None
+
     checksum_filepath = download_plan.checksum_filepath
     stored_checksum = (
         checksum_filepath.read_text().strip() if checksum_filepath.exists() else None

tests/test_datasets.py

@@ -73,3 +73,33 @@ def test_save_dataset_to_disk_warns_and_reports_its_own_source(
             "download_source": "save_dataset_to_disk",
         }
     ]
+
+
+def test_download_sample_resolves_id_and_reports_its_own_source(
+    monkeypatch: MonkeyPatch,
+) -> None:
+    calls: list[dict[str, object]] = []
+
+    def fake_download_sample(**kwargs: object) -> Path:
+        calls.append(kwargs)
+        return Path("sample.tar.gz")
+
+    monkeypatch.setattr(
+        datasets,
+        "get_dataset_details",
+        lambda dataset_id: DatasetDetails(id="canonical-id"),
+    )
+    monkeypatch.setattr(datasets, "_download_sample", fake_download_sample)
+
+    result = datasets.download_sample("a-dataset-slug", download_directory="dir")
+
+    assert result == Path("sample.tar.gz")
+    assert calls == [
+        {
+            "dataset_id": "canonical-id",
+            "download_directory": "dir",
+            "show_progress": True,
+            "overwrite_existing": False,
+            "download_source": "download_sample",
+        }
+    ]