diff --git a/CHANGELOG.md b/CHANGELOG.md index b0bcfb50..821c1060 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ Unreleased ---------- - Added a `content_settings` parameter to `open()` (write modes), `pipe_file()` and `put_file()` to set blob content settings (content type, content disposition, cache control, ...) on write. Accepts a `dict` (recommended) or an `azure.storage.blob.ContentSettings` instance. [#554](https://github.com/fsspec/adlfs/pull/554) +- Optimized single-block writes to use a single `upload_blob` call instead of staging and committing a block list 2026.5.0 -------- diff --git a/adlfs/spec.py b/adlfs/spec.py index fba2e761..b8425400 100644 --- a/adlfs/spec.py +++ b/adlfs/spec.py @@ -2281,12 +2281,27 @@ async def _async_upload_chunk(self, final: bool = False, **kwargs): """ data = self.buffer.getvalue() length = len(data) - block_id = self._get_block_id() commit_kw = {} if self.mode == "xb": commit_kw["headers"] = {"If-None-Match": "*"} if self.mode in {"wb", "xb"}: try: + if final and not self._block_list: + async with self.container_client.get_blob_client( + blob=self.blob + ) as bc: + response = await bc.upload_blob( + data=data, + overwrite=(self.mode == "wb"), + metadata=self.metadata, + content_settings=self._content_settings, + max_concurrency=self.fs.max_concurrency or 1, + ) + if self.fs.version_aware: + self.version_id = response.get("version_id") + return + + block_id = self._get_block_id() max_concurrency = self.fs.max_concurrency or 1 semaphore = asyncio.Semaphore(max_concurrency) tasks = [] @@ -2314,24 +2329,8 @@ async def _async_upload_chunk(self, final: bool = False, **kwargs): except ResourceExistsError as e: raise FileExistsError(self.path) from e except Exception as e: - # This step handles the situation where data="" and length=0 - # which is throws an InvalidHeader error from Azure, so instead - # of staging a block, we directly upload the empty blob - # This isn't actually tested, since Azureite behaves differently. - if not self._block_list and length == 0 and final: - async with self.container_client.get_blob_client( - blob=self.blob - ) as bc: - response = await bc.upload_blob( - data=data, - metadata=self.metadata, - content_settings=self._content_settings, - overwrite=(self.mode == "wb"), - ) - if self.fs.version_aware: - self.version_id = response.get("version_id") - elif length == 0 and final: - # just finalize + # Finalize a multi-block upload whose last chunk is empty + if length == 0 and final and self._block_list: block_list = [BlobBlock(_id) for _id in self._block_list] async with self.container_client.get_blob_client( blob=self.blob diff --git a/adlfs/tests/test_spec.py b/adlfs/tests/test_spec.py index d98ab94c..9b9f4b88 100644 --- a/adlfs/tests/test_spec.py +++ b/adlfs/tests/test_spec.py @@ -2256,16 +2256,16 @@ def test_write_populates_version_id_when_version_aware(storage, mocker): skip_instance_cache=True, ) - mock_commit_block_list = mocker.patch.object( + mock_upload_blob = mocker.patch.object( BlobClient, - "commit_block_list", + "upload_blob", return_value={"version_id": "test-version-id"}, ) with fs.open("data/version-aware-write.bin", "wb") as f: f.write(b"hello world") - assert mock_commit_block_list.called + mock_upload_blob.assert_called_once() assert f.version_id == "test-version-id" @@ -2279,16 +2279,16 @@ def test_write_does_not_populate_version_id_when_not_version_aware(storage, mock skip_instance_cache=True, ) - mock_commit_block_list = mocker.patch.object( + mock_upload_blob = mocker.patch.object( BlobClient, - "commit_block_list", + "upload_blob", return_value={"version_id": "test-version-id"}, ) with fs.open("data/not-version-aware-wb.bin", "wb") as f: f.write(b"hello world") - assert mock_commit_block_list.called + mock_upload_blob.assert_called_once() assert f.version_id is None @@ -2707,3 +2707,76 @@ def test_etag_normalized_form(storage): ) def test_striping_etag(input_etag, expected_etag): assert _normalize_etag_quotes(input_etag) == expected_etag + + +def test_small_write_uses_single_upload_blob(storage, mocker): + from azure.storage.blob.aio import BlobClient + + fs = AzureBlobFileSystem( + account_name=storage.account_name, + connection_string=CONN_STR, + skip_instance_cache=True, + ) + + mock_stage_block = mocker.patch.object(BlobClient, "stage_block") + mock_commit_block_list = mocker.patch.object(BlobClient, "commit_block_list") + mock_upload_blob = mocker.patch.object(BlobClient, "upload_blob", return_value={}) + + with fs.open("data/small-file.txt", "wb") as f: + f.write(b"test content") + + mock_upload_blob.assert_called_once() + mock_stage_block.assert_not_called() + mock_commit_block_list.assert_not_called() + + call_kwargs = mock_upload_blob.call_args.kwargs + assert call_kwargs["data"] == b"test content" + + +@pytest.mark.xfail( + parse_version(fsspec.__version__) <= parse_version("2024.10.0"), + reason="not supported upstream yet", +) +def test_small_write_uses_single_upload_blob_x(storage, mocker): + from azure.storage.blob.aio import BlobClient + + fs = AzureBlobFileSystem( + account_name=storage.account_name, + connection_string=CONN_STR, + skip_instance_cache=True, + ) + + mock_stage_block = mocker.patch.object(BlobClient, "stage_block") + mock_commit_block_list = mocker.patch.object(BlobClient, "commit_block_list") + mock_upload_blob = mocker.patch.object(BlobClient, "upload_blob", return_value={}) + + with fs.open("data/small-file-xb.txt", "xb") as f: + f.write(b"test content") + + mock_upload_blob.assert_called_once() + mock_stage_block.assert_not_called() + mock_commit_block_list.assert_not_called() + + call_kwargs = mock_upload_blob.call_args.kwargs + assert call_kwargs["data"] == b"test content" + assert call_kwargs["overwrite"] is False + + +def test_small_write_upload_failure_raises_runtime_error(storage, mocker): + from azure.core.exceptions import HttpResponseError + from azure.storage.blob.aio import BlobClient + + fs = AzureBlobFileSystem( + account_name=storage.account_name, + connection_string=CONN_STR, + skip_instance_cache=True, + ) + + mocker.patch.object(BlobClient, "upload_blob", side_effect=HttpResponseError) + mock_commit_block_list = mocker.patch.object(BlobClient, "commit_block_list") + + with pytest.raises(RuntimeError): + with fs.open("data/small-file.txt", "wb") as f: + f.write(b"test content") + + mock_commit_block_list.assert_not_called()