⚡️ Speed up function get_exported_parquet_files by 23%

[codeflash-ai]

📄 23% (0.23x) speedup for get_exported_parquet_files in src/datasets/utils/_dataset_viewer.py

⏱️ Runtime : 8.84 milliseconds → 7.20 milliseconds (best of 61 runs)

📝 Explanation and details

The optimization achieves a 22% runtime improvement by introducing a small LRU cache for authentication header generation, which is a frequently called operation in the dataset loading path.

Key Changes:

1. Cached Header Building: Added @lru_cache(maxsize=128) decorator to a new _cached_build_hf_headers() function that wraps the call to huggingface_hub.utils.build_hf_headers(). The cache key is based on the token parameter (which is hashable as None/str/bool).

2. Safe Cache Usage: Returns dict(_cached_build_hf_headers(token)) - creating a shallow copy of the cached dictionary to prevent callers from mutating the cached object, maintaining thread-safety and correctness.

Why This Is Faster:

The line profiler reveals that in the original code, get_authentication_headers_for_url() spent 98.8% of its time (2.72ms out of 2.75ms) calling huggingface_hub.utils.build_hf_headers(). In the optimized version, this drops to 89.1% (198μs out of 223μs) - a ~13.7x speedup for this specific function.

The build_hf_headers() call involves:

• Version string formatting
• Library metadata construction
• Dictionary allocation and population

With caching, repeated calls with the same token (common in batch operations) become simple dictionary lookups and shallow copies rather than full header reconstruction.

Impact on Workloads:

Based on function_references, get_exported_parquet_files() is called from src/datasets/load.py during dataset module initialization - a hot path executed for every dataset load operation. The optimization particularly benefits:

• Batch dataset loading: When loading multiple datasets or configurations with the same authentication token, subsequent calls hit the cache
• Dataset iteration: The test results show 31-34% speedup for successful parquet retrieval cases, indicating real-world benefit
• Authenticated workflows: Most valuable when token parameter is consistent across calls (the common case)

Test Results Pattern:

The annotated tests show consistent 31-34% speedup for successful retrieval scenarios (tests with matching commit hashes, valid responses), while error/exception paths show minimal difference (~0-2%). This is expected since the optimization targets the header building bottleneck, which only matters in successful API call paths.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 37 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

from unittest.mock import \
    MagicMock  # used to create lightweight callable mocks

# imports
import pytest  # used for our unit tests
# Import the function and error class from the real module under test.
# We import the module object too so we can monkeypatch attributes that were imported
# at module import time (e.g., get_session).
from src.datasets.utils import _dataset_viewer as dv
from src.datasets.utils._dataset_viewer import (DatasetViewerError,
                                                get_exported_parquet_files)

# Helper response class to simulate HTTP responses returned by get_session().get(...)
class DummyResponse:
    """
    A tiny response object that mimics the interface used by get_exported_parquet_files:
    - .headers (dict)
    - .json() -> returns parsed JSON body
    - .raise_for_status() -> may raise an exception to simulate HTTP errors
    """

    def __init__(self, headers=None, json_body=None, raise_exc=None):
        # store headers, defaulting to empty dict
        self.headers = headers or {}
        # store JSON body for json() to return
        self._json = json_body or {}
        # optionally store an exception instance to be raised by raise_for_status()
        self._raise_exc = raise_exc

    def raise_for_status(self):
        # If configured to raise, raise the provided exception (simulates HTTP errors)
        if self._raise_exc:
            raise self._raise_exc
        # else do nothing (200 OK like)
        return None

    def json(self):
        # Return the preconfigured JSON body
        return self._json

def test_basic_success_exact_commit_hash(monkeypatch):
    # Basic success path:
    # - The response contains X-Revision equal to the commit_hash argument
    # - The JSON indicates the parquet export is ready (partial=False, pending=False, failed=False)
    # - The JSON contains "parquet_files" and those are returned by the function

    dataset = "some-user/some-dataset"  # dataset identifier string
    commit_hash = "abc123"  # expected commit hash to match header
    parquet_files_expected = [{"path": "file-1.parquet"}, {"path": "file-2.parquet"}]

    # Build a dummy response that matches the "ready" conditions in the function
    resp = DummyResponse(
        headers={"X-Revision": commit_hash},
        json_body={"partial": False, "pending": False, "failed": False, "parquet_files": parquet_files_expected},
    )

    # Create a session-like object whose .get(...) returns our dummy response.
    # We also capture the arguments it was called with to assert URL/header construction.
    called = {}

    class Session:
        def get(self, url, headers, timeout):
            # record the arguments for assertions below
            called['url'] = url
            called['headers'] = headers
            called['timeout'] = timeout
            return resp

    # Monkeypatch the function's get_session name in the module under test to return our Session.
    monkeypatch.setattr(dv, "get_session", lambda: Session())

    # Monkeypatch get_authentication_headers_for_url to avoid relying on huggingface_hub for headers.
    # Return a deterministic header dict so we can assert it is forwarded to the session.get call.
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {"Authorization": "token123"})

    # Call the function under test with the exact commit hash
    codeflash_output = get_exported_parquet_files(dataset=dataset, commit_hash=commit_hash, token="token123"); result = codeflash_output # 6.82μs -> 6.36μs (7.19% faster)

def test_success_with_commit_hash_none_accepts_any_revision(monkeypatch):
    # When commit_hash is None, the function should accept any X-Revision header value and return files.
    dataset = "owner/dataset-2"
    # header has some revision value that differs from what would be passed if not None
    resp = DummyResponse(
        headers={"X-Revision": "some-other-revision"},
        json_body={"partial": False, "pending": False, "failed": False, "parquet_files": [{"p": 1}]},
    )

    # session that returns our response
    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=lambda url, headers, timeout: resp))
    # Avoid external header building
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    # Passing commit_hash=None should bypass revision matching and return parquet files
    codeflash_output = get_exported_parquet_files(dataset=dataset, commit_hash=None, token=None); result = codeflash_output # 181μs -> 182μs (0.344% slower)

@pytest.mark.parametrize(
    "json_body, note",
    [
        ({"partial": True, "pending": False, "failed": False, "parquet_files": [{"a": 1}]}, "partial is True"),
        ({"partial": False, "pending": True, "failed": False, "parquet_files": [{"a": 1}]}, "pending is True"),
        ({"partial": False, "pending": False, "failed": True, "parquet_files": [{"a": 1}]}, "failed is True"),
        ({"partial": False, "pending": False, "failed": False}, "parquet_files missing"),
    ],
)
def test_not_ready_variants_raise_dataset_viewer_error(monkeypatch, json_body, note):
    # Edge cases where the export is not considered "completely ready":
    # - partial == True
    # - pending == True
    # - failed == True
    # - missing parquet_files key entirely
    # Each should lead to a DatasetViewerError being raised.

    dataset = "owner/dataset-not-ready"
    commit_hash = "hash-ready"

    resp = DummyResponse(headers={"X-Revision": commit_hash}, json_body=json_body)

    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=lambda url, headers, timeout: resp))
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    # The function is expected to raise DatasetViewerError in all these not-ready scenarios.
    with pytest.raises(DatasetViewerError):
        get_exported_parquet_files(dataset=dataset, commit_hash=commit_hash, token="tok") # 728μs -> 734μs (0.841% slower)

def test_outdated_revision_logs_and_raises(monkeypatch):
    # If the X-Revision header exists but does not match the provided commit_hash,
    # the function should not return files and should raise DatasetViewerError.

    dataset = "owner/outdated"
    header_revision = "newer"
    provided_commit = "older"

    resp = DummyResponse(
        headers={"X-Revision": header_revision},
        json_body={"partial": False, "pending": False, "failed": False, "parquet_files": [{"x": "y"}]},
    )

    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=lambda url, headers, timeout: resp))
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    # Because revisions don't match, DatasetViewerError should be raised.
    with pytest.raises(DatasetViewerError):
        get_exported_parquet_files(dataset=dataset, commit_hash=provided_commit, token=None) # 180μs -> 181μs (0.595% slower)

def test_missing_revision_header_raises(monkeypatch):
    # If the HTTP response lacks the X-Revision header entirely, the function should raise DatasetViewerError.

    dataset = "owner/no-revision"
    resp = DummyResponse(
        headers={},  # no X-Revision
        json_body={"partial": False, "pending": False, "failed": False, "parquet_files": [{"ok": True}]},
    )

    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=lambda url, headers, timeout: resp))
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    with pytest.raises(DatasetViewerError):
        get_exported_parquet_files(dataset=dataset, commit_hash="whatever", token=None) # 177μs -> 180μs (1.52% slower)

def test_raise_for_status_exception_leads_to_dataset_viewer_error(monkeypatch):
    # If response.raise_for_status() throws an exception (e.g., HTTP error),
    # the function should catch it and ultimately raise DatasetViewerError.

    dataset = "owner/http-error"
    # Simulate raise_for_status throwing a requests-like HTTPError (we can use a plain Exception)
    resp = DummyResponse(
        headers={"X-Revision": "hash"},
        json_body={"partial": False, "pending": False, "failed": False, "parquet_files": [{"ok": True}]},
        raise_exc=Exception("HTTP 401 Unauthorized"),
    )

    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=lambda url, headers, timeout: resp))
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    with pytest.raises(DatasetViewerError):
        get_exported_parquet_files(dataset=dataset, commit_hash="hash", token="bad-token") # 180μs -> 182μs (0.912% slower)

def test_get_raises_exception_propagates_to_dataset_viewer_error(monkeypatch):
    # If the session.get(...) call itself raises (network failure), the function catches it
    # and raises DatasetViewerError.

    dataset = "owner/get-exception"

    def raising_get(url, headers, timeout):
        raise ConnectionError("connection refused")

    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=raising_get))
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    with pytest.raises(DatasetViewerError):
        get_exported_parquet_files(dataset=dataset, commit_hash="irrelevant", token=None) # 182μs -> 184μs (0.848% slower)

def test_large_scale_parquet_files_returned(monkeypatch):
    # Large-scale test to assert that the function handles a reasonably large list of parquet files.
    # We keep element count under the instructed limit (1000); use 500 entries.

    dataset = "owner/large-dataset"
    commit_hash = "big-hash"

    # Create 500 dummy parquet file descriptors
    large_parquet_files = [{"path": f"part-{i}.parquet", "size": i} for i in range(500)]

    resp = DummyResponse(
        headers={"X-Revision": commit_hash},
        json_body={"partial": False, "pending": False, "failed": False, "parquet_files": large_parquet_files},
    )

    # Ensure the session returns our large payload
    monkeypatch.setattr(dv, "get_session", lambda: MagicMock(get=lambda url, headers, timeout: resp))
    monkeypatch.setattr(dv, "get_authentication_headers_for_url", lambda url, token: {})

    codeflash_output = get_exported_parquet_files(dataset=dataset, commit_hash=commit_hash, token="tok"); result = codeflash_output # 181μs -> 182μs (0.616% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import json
from unittest.mock import MagicMock, Mock, patch

import pytest
from requests import Response
from requests.exceptions import ConnectionError, Timeout
from src.datasets import config
from src.datasets.utils._dataset_viewer import get_exported_parquet_files
from src.datasets.utils.file_utils import get_authentication_headers_for_url

def test_successful_parquet_export_with_matching_commit():
    """
    Test successful retrieval of parquet files when commit hash matches.
    This is the happy path scenario.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123def456"
    expected_files = [
        {"filename": "data-00000-of-00001.parquet", "size": 1024},
        {"filename": "data-00001-of-00001.parquet", "size": 2048},
    ]
    
    # Mock the response from the API
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 306μs -> 229μs (33.4% faster)

def test_successful_parquet_export_with_none_commit():
    """
    Test successful retrieval when commit_hash is None.
    The function should accept any commit hash in this case.
    """
    dataset_name = "test/dataset"
    expected_files = [{"filename": "data.parquet", "size": 5000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": "any_commit_hash"}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, None, token=None); result = codeflash_output # 301μs -> 225μs (33.9% faster)

def test_successful_parquet_export_with_token():
    """
    Test successful retrieval with authentication token provided.
    Verifies token is passed to authentication headers.
    """
    dataset_name = "private/dataset"
    commit_hash = "def789"
    token = "hf_test_token_123"
    expected_files = [{"filename": "private_data.parquet", "size": 3000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        with patch("src.datasets.utils._dataset_viewer.get_authentication_headers_for_url") as mock_auth:
            mock_auth.return_value = {"Authorization": f"Bearer {token}"}
            mock_session.return_value.get.return_value = mock_response
            codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=token); result = codeflash_output

def test_single_parquet_file():
    """
    Test retrieval of a single parquet file.
    """
    dataset_name = "simple/dataset"
    commit_hash = "hash123"
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 302μs -> 227μs (32.7% faster)

def test_mismatched_commit_hash():
    """
    Test that function raises error when commit hash doesn't match.
    The API returns a different revision than requested.
    """
    dataset_name = "test/dataset"
    requested_commit = "abc123"
    actual_commit = "def456"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": actual_commit}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": [{"filename": "data.parquet", "size": 1000}],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, requested_commit, token=None)

def test_parquet_export_pending():
    """
    Test that function raises error when parquet export is still pending.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": True,  # Export is still pending
        "failed": False,
        "parquet_files": [],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_parquet_export_failed():
    """
    Test that function raises error when parquet export has failed.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": True,  # Export has failed
        "parquet_files": [],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_parquet_export_partial():
    """
    Test that function raises error when parquet export is only partial.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": True,  # Export is only partial
        "pending": False,
        "failed": False,
        "parquet_files": [{"filename": "data.parquet", "size": 1000}],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_missing_parquet_files_key():
    """
    Test that function raises error when parquet_files key is missing from response.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        # Missing "parquet_files" key
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_missing_x_revision_header():
    """
    Test that function raises error when X-Revision header is missing.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {}  # Missing X-Revision header
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": [{"filename": "data.parquet", "size": 1000}],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_http_error_response():
    """
    Test that function raises error when HTTP request fails with status error.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.raise_for_status.side_effect = Exception("404 Not Found")
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_timeout_error():
    """
    Test that function raises error when request times out.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.side_effect = Timeout("Request timed out")
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_connection_error():
    """
    Test that function raises error when network connection fails.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.side_effect = ConnectionError("Connection failed")
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_invalid_json_response():
    """
    Test that function raises error when API response contains invalid JSON.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_empty_parquet_files_list():
    """
    Test that function raises error when parquet_files list is empty.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": [],  # Empty list
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        # Empty list is still valid; it should return the empty list
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 296μs -> 224μs (32.2% faster)

def test_dataset_name_with_special_characters():
    """
    Test with dataset names containing special characters and slashes.
    """
    dataset_name = "user-123/my_dataset.v2"
    commit_hash = "abc123"
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 296μs -> 225μs (31.3% faster)

def test_boolean_token():
    """
    Test with token passed as boolean (False or True).
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        with patch("src.datasets.utils._dataset_viewer.get_authentication_headers_for_url") as mock_auth:
            mock_auth.return_value = {}
            mock_session.return_value.get.return_value = mock_response
            codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=False); result = codeflash_output

def test_pending_default_true():
    """
    Test that missing pending field defaults to True (treats as pending).
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        # "pending" is missing - should default to True
        "failed": False,
        "parquet_files": [{"filename": "data.parquet", "size": 1000}],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_failed_default_true():
    """
    Test that missing failed field defaults to True (treats as failed).
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        # "failed" is missing - should default to True
        "parquet_files": [{"filename": "data.parquet", "size": 1000}],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        with pytest.raises(Exception) as exc_info:
            get_exported_parquet_files(dataset_name, commit_hash, token=None)

def test_many_parquet_files():
    """
    Test retrieval of a large number of parquet files.
    Ensures function handles many files efficiently.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    
    # Create a list of 500 parquet files
    expected_files = [
        {"filename": f"data-{i:05d}-of-00500.parquet", "size": 1000 + i}
        for i in range(500)
    ]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 302μs -> 228μs (32.6% faster)

def test_large_file_sizes():
    """
    Test parquet files with large file sizes (gigabytes).
    """
    dataset_name = "test/large-dataset"
    commit_hash = "abc123"
    
    # Create files with large sizes (in bytes)
    expected_files = [
        {"filename": f"data-{i:02d}.parquet", "size": (i + 1) * 1000000000}  # 1GB, 2GB, 3GB, etc.
        for i in range(5)
    ]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 298μs -> 224μs (33.4% faster)

def test_long_dataset_name():
    """
    Test with very long dataset name.
    """
    dataset_name = "user/" + "x" * 200 + "/dataset"
    commit_hash = "abc123"
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 297μs -> 224μs (32.3% faster)

def test_long_commit_hash():
    """
    Test with very long commit hash string.
    """
    dataset_name = "test/dataset"
    commit_hash = "a" * 256  # Very long commit hash
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 297μs -> 226μs (31.0% faster)

def test_response_with_extra_fields():
    """
    Test that function handles responses with unexpected extra fields.
    Should ignore extra fields and process normally.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
        "extra_field_1": "some_value",
        "extra_field_2": {"nested": "data"},
        "extra_field_3": [1, 2, 3],
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 296μs -> 225μs (31.6% faster)

def test_parquet_file_with_extra_metadata():
    """
    Test that parquet files in response can have additional metadata fields.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    expected_files = [
        {
            "filename": "data.parquet",
            "size": 1000,
            "sha256": "abc123def456",
            "num_rows": 5000,
            "num_bytes": 1000,
            "dataset": "test/dataset",
            "config": "default",
            "split": "train",
        }
    ]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        mock_session.return_value.get.return_value = mock_response
        codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=None); result = codeflash_output # 296μs -> 224μs (32.2% faster)

def test_long_token_string():
    """
    Test with very long authentication token.
    """
    dataset_name = "test/dataset"
    commit_hash = "abc123"
    token = "hf_" + "x" * 500  # Very long token
    expected_files = [{"filename": "data.parquet", "size": 1000}]
    
    mock_response = Mock(spec=Response)
    mock_response.headers = {"X-Revision": commit_hash}
    mock_response.json.return_value = {
        "partial": False,
        "pending": False,
        "failed": False,
        "parquet_files": expected_files,
    }
    
    with patch("src.datasets.utils._dataset_viewer.get_session") as mock_session:
        with patch("src.datasets.utils._dataset_viewer.get_authentication_headers_for_url") as mock_auth:
            mock_auth.return_value = {"Authorization": f"Bearer {token}"}
            mock_session.return_value.get.return_value = mock_response
            codeflash_output = get_exported_parquet_files(dataset_name, commit_hash, token=token); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-get_exported_parquet_files-mlcukeob and push.