feat(async-index): restore batch_size and show_progress on AsyncIndex.upsert

jhamon · jhamon · commit 5c588d4fde6a · 2026-05-04T03:08:32.000Z
## Purpose

BCG-090 (severity H): the async `upsert()` method was missing the
`batch_size` and `show_progress` parameters available on the sync
`Index.upsert()`, leaving async users without automatic batching.

## Solution

- Added `batch_size: int | None = None` and `show_progress: bool = True`
  parameters to `AsyncIndex.upsert()`.
- Extracted the single-request logic into a private
  `async def _upsert_one_batch(...)` method.
- When `batch_size` is set, vectors are split via `chunked()` and each
  batch is sent **sequentially** (one `await` at a time) to match sync
  semantics. Users who want parallel emission can fan out with
  `asyncio.gather()` themselves.
- `response_info` is set to `None` on aggregated multi-batch responses;
  preserved on single-request responses.
- Added 12 unit tests in `tests/unit/test_async_index_upsert_batching.py`
  mirroring the sync test suite, plus a sequential-emission contract test.

## Follow-ups

`AsyncIndex.upsert_from_dataframe` remains `NotImplementedError` —
that is a separate additive feature outside this task's scope.
diff --git a/pinecone/async_client/async_index.py b/pinecone/async_client/async_index.py
@@ -13,6 +13,7 @@
 
 from pinecone._internal.adapters.imports_adapter import ImportsAdapter
 from pinecone._internal.adapters.vectors_adapter import VectorsAdapter, extract_response_info
+from pinecone._internal.batching import chunked, validate_batch_size, with_progress
 from pinecone._internal.config import PineconeConfig
 from pinecone._internal.constants import DATA_PLANE_API_VERSION
 from pinecone._internal.data_plane_helpers import _validate_host, _vector_to_dict
@@ -221,6 +222,8 @@ async def upsert(
             | dict[str, Any]
         ],
         namespace: str = "",
+        batch_size: int | None = None,
+        show_progress: bool = True,
         timeout: float | None = None,
     ) -> UpsertResponse:
         """Upsert a batch of vectors into a namespace.
@@ -235,13 +238,27 @@ async def upsert(
                 and optional ``sparse_values`` / ``metadata`` keys.
             namespace (str): Target namespace. Defaults to the default
                 (empty-string) namespace.
+            batch_size (int | None): Split *vectors* into chunks of this size
+                and send one request per chunk. Default ``None`` sends a single
+                request (current behaviour). Must be a positive integer if
+                provided.
+            show_progress (bool): When ``True`` and ``tqdm`` is installed,
+                display a progress bar across batches. Has no effect when
+                ``batch_size`` is ``None`` or ``tqdm`` is not installed.
+                Defaults to ``True``.
+            timeout (float | None): Per-request timeout in seconds. Overrides
+                the client-level default for this call only.
 
         Returns:
             :class:`UpsertResponse` with the count of vectors upserted.
+            When ``batch_size`` triggers multiple requests, ``response_info``
+            on the aggregated response is ``None`` — there is no single
+            response to attribute it to.
 
         Raises:
             :exc:`PineconeTypeError`: If a vector element is not a recognized format.
             :exc:`PineconeValueError`: If a vector element is malformed.
+            :exc:`PineconeValueError`: If *batch_size* is not a positive integer.
             :exc:`ApiError`: If the API returns an error response.
             :exc:`PineconeConnectionError`: If a network-level connection
                 fails (DNS, refused, transport error).
@@ -266,10 +283,18 @@ async def upsert(
                 )
                 print(response.upserted_count)
 
+                # Upsert 1000 vectors in batches of 100
+                response = await idx.upsert(
+                    vectors=large_vector_list,
+                    batch_size=100,
+                    show_progress=True,
+                )
+                print(response.upserted_count)
+
         .. note::
-           All vectors are sent in a single request. For large datasets,
-           batch your calls (recommended batch size: 100–500 vectors).
-           For sync clients, use :meth:`upsert_from_dataframe` for automatic batching.
+           Batches are sent **sequentially** (one ``await`` at a time). For
+           parallel emission, fan out at the call site with
+           ``asyncio.gather(*[idx.upsert(vectors=b) for b in batches])``.
            For very large datasets (millions of vectors), consider
            :meth:`start_import` for bulk import from cloud storage.
 
@@ -279,6 +304,33 @@ async def upsert(
            - :meth:`start_import` — for bulk loading millions of vectors
              from cloud storage (S3, GCS).
         """
+        if batch_size is None:
+            return await self._upsert_one_batch(
+                vectors=vectors, namespace=namespace, timeout=timeout
+            )
+        validate_batch_size(batch_size)
+        built = list(vectors)
+        batches = chunked(built, batch_size)
+        total_count = 0
+        for batch in with_progress(batches, show_progress=show_progress):
+            result = await self._upsert_one_batch(
+                vectors=batch, namespace=namespace, timeout=timeout
+            )
+            total_count += result.upserted_count
+        return UpsertResponse(upserted_count=total_count)
+
+    async def _upsert_one_batch(
+        self,
+        *,
+        vectors: Sequence[
+            Vector
+            | tuple[str, list[float]]
+            | tuple[str, list[float], dict[str, Any]]
+            | dict[str, Any]
+        ],
+        namespace: str,
+        timeout: float | None,
+    ) -> UpsertResponse:
         built = [VectorFactory.build(v) for v in vectors]
         body: dict[str, Any] = {
             "vectors": [_vector_to_dict(v) for v in built],
diff --git a/tests/unit/test_async_index_upsert_batching.py b/tests/unit/test_async_index_upsert_batching.py
@@ -0,0 +1,245 @@
+"""Unit tests for AsyncIndex.upsert() batch_size and show_progress parameters."""
+
+from __future__ import annotations
+
+import sys
+from unittest.mock import patch
+
+import httpx
+import pytest
+import respx
+
+from pinecone import AsyncIndex
+from pinecone.errors.exceptions import PineconeValueError
+from pinecone.models.vectors.responses import UpsertResponse
+from pinecone.models.vectors.vector import Vector
+
+INDEX_HOST = "test-index-abc1234.svc.us-east1-gcp.pinecone.io"
+INDEX_HOST_HTTPS = f"https://{INDEX_HOST}"
+UPSERT_URL = f"{INDEX_HOST_HTTPS}/vectors/upsert"
+
+
+def _make_upsert_response(*, upserted_count: int) -> dict[str, object]:
+    return {"upsertedCount": upserted_count}
+
+
+def _make_async_index() -> AsyncIndex:
+    return AsyncIndex(host=INDEX_HOST, api_key="test-key")
+
+
+def _make_vectors(n: int) -> list[Vector]:
+    return [Vector(id=f"v{i}", values=[float(i), float(i + 1)]) for i in range(n)]
+
+
+class TestAsyncUpsertNoBatchSize:
+    """batch_size=None (default) sends a single request."""
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_no_batch_size_sends_single_request(self) -> None:
+        route = respx.post(UPSERT_URL).mock(
+            return_value=httpx.Response(200, json=_make_upsert_response(upserted_count=100))
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=_make_vectors(100), batch_size=None)
+        assert len(route.calls) == 1
+        assert result.upserted_count == 100
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_response_info_preserved_when_not_batched(self) -> None:
+        respx.post(UPSERT_URL).mock(
+            return_value=httpx.Response(
+                200,
+                json=_make_upsert_response(upserted_count=5),
+                headers={"X-Pinecone-Request-Id": "req-abc123"},
+            )
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=_make_vectors(5))
+        assert result.response_info is not None
+
+
+class TestAsyncUpsertWithBatchSize:
+    """batch_size=N splits vectors and sends one request per chunk."""
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_with_batch_size_sends_multiple_requests(self) -> None:
+        route = respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=100)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=100)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=50)),
+            ]
+        )
+        idx = _make_async_index()
+        await idx.upsert(vectors=_make_vectors(250), batch_size=100, show_progress=False)
+        assert len(route.calls) == 3
+
+        import orjson
+
+        first_body = orjson.loads(route.calls[0].request.content)
+        second_body = orjson.loads(route.calls[1].request.content)
+        third_body = orjson.loads(route.calls[2].request.content)
+        assert len(first_body["vectors"]) == 100
+        assert len(second_body["vectors"]) == 100
+        assert len(third_body["vectors"]) == 50
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_with_batch_size_aggregates_response(self) -> None:
+        respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=100)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=100)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=50)),
+            ]
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=_make_vectors(250), batch_size=100, show_progress=False)
+        assert isinstance(result, UpsertResponse)
+        assert result.upserted_count == 250
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_response_info_is_none_when_batched(self) -> None:
+        respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=100)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=100)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=50)),
+            ]
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=_make_vectors(250), batch_size=100, show_progress=False)
+        assert result.response_info is None
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_namespace_forwarded_per_batch(self) -> None:
+        route = respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+            ]
+        )
+        idx = _make_async_index()
+        await idx.upsert(
+            vectors=_make_vectors(10), batch_size=5, namespace="my-ns", show_progress=False
+        )
+
+        import orjson
+
+        for call in route.calls:
+            body = orjson.loads(call.request.content)
+            assert body.get("namespace") == "my-ns"
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_timeout_forwarded_per_batch(self) -> None:
+        route = respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+            ]
+        )
+        idx = _make_async_index()
+        await idx.upsert(vectors=_make_vectors(10), batch_size=5, timeout=5.0, show_progress=False)
+        assert len(route.calls) == 2
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_empty_vectors_with_batch_size(self) -> None:
+        route = respx.post(UPSERT_URL).mock(
+            return_value=httpx.Response(200, json=_make_upsert_response(upserted_count=0))
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=[], batch_size=100, show_progress=False)
+        assert len(route.calls) == 0
+        assert result.upserted_count == 0
+
+
+class TestAsyncUpsertInvalidBatchSize:
+    """Invalid batch_size values raise PineconeValueError."""
+
+    @pytest.mark.asyncio
+    async def test_upsert_invalid_batch_size_raises(self) -> None:
+        idx = _make_async_index()
+        with pytest.raises(PineconeValueError):
+            await idx.upsert(vectors=_make_vectors(5), batch_size=0)
+
+    @pytest.mark.asyncio
+    async def test_upsert_invalid_batch_size_negative(self) -> None:
+        idx = _make_async_index()
+        with pytest.raises(PineconeValueError):
+            await idx.upsert(vectors=_make_vectors(5), batch_size=-1)
+
+    @pytest.mark.asyncio
+    async def test_upsert_invalid_batch_size_float(self) -> None:
+        idx = _make_async_index()
+        with pytest.raises(PineconeValueError):
+            await idx.upsert(vectors=_make_vectors(5), batch_size=1.5)  # type: ignore[arg-type]
+
+
+class TestAsyncUpsertShowProgress:
+    """show_progress behavior with and without tqdm."""
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_show_progress_false_does_not_import_tqdm(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setitem(sys.modules, "tqdm", None)  # type: ignore[arg-type]
+        monkeypatch.setitem(sys.modules, "tqdm.auto", None)  # type: ignore[arg-type]
+        respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+            ]
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=_make_vectors(10), batch_size=5, show_progress=False)
+        assert result.upserted_count == 10
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_upsert_show_progress_true_works_without_tqdm(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setitem(sys.modules, "tqdm", None)  # type: ignore[arg-type]
+        monkeypatch.setitem(sys.modules, "tqdm.auto", None)  # type: ignore[arg-type]
+        respx.post(UPSERT_URL).mock(
+            side_effect=[
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+                httpx.Response(200, json=_make_upsert_response(upserted_count=5)),
+            ]
+        )
+        idx = _make_async_index()
+        result = await idx.upsert(vectors=_make_vectors(10), batch_size=5, show_progress=True)
+        assert result.upserted_count == 10
+
+
+class TestAsyncUpsertSequential:
+    """Batches are emitted sequentially, not in parallel."""
+
+    @pytest.mark.asyncio
+    async def test_upsert_batches_are_sequential_not_parallel(self) -> None:
+        call_order: list[int] = []
+
+        async def mock_upsert_one_batch(
+            *,
+            vectors: object,
+            namespace: str,
+            timeout: object,
+        ) -> UpsertResponse:
+            batch_index = len(call_order)
+            call_order.append(batch_index)
+            return UpsertResponse(upserted_count=5)
+
+        idx = _make_async_index()
+        with patch.object(idx, "_upsert_one_batch", side_effect=mock_upsert_one_batch):
+            result = await idx.upsert(vectors=_make_vectors(15), batch_size=5, show_progress=False)
+
+        assert call_order == [0, 1, 2]
+        assert result.upserted_count == 15
