fix(documents): split import batch sizing into server and client params

Author: tharropoulos

src/typesense/async_/documents.py

@@ -228,6 +228,7 @@ async def import_(
         documents: typing.List[TDoc],
         import_parameters: DocumentImportParametersReturnDocAndId,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[
         typing.Union[ImportResponseWithDocAndId[TDoc], ImportResponseFail[TDoc]]
     ]: ...
@@ -238,6 +239,7 @@ async def import_(
         documents: typing.List[TDoc],
         import_parameters: DocumentImportParametersReturnId,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[typing.Union[ImportResponseWithId, ImportResponseFail[TDoc]]]: ...
 
     @typing.overload
@@ -246,6 +248,7 @@ async def import_(
         documents: typing.List[TDoc],
         import_parameters: typing.Union[DocumentWriteParameters, None] = None,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[typing.Union[ImportResponseSuccess, ImportResponseFail[TDoc]]]: ...
 
     @typing.overload
@@ -254,6 +257,7 @@ async def import_(
         documents: typing.List[TDoc],
         import_parameters: DocumentImportParametersReturnDoc,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[
         typing.Union[ImportResponseWithDoc[TDoc], ImportResponseFail[TDoc]]
     ]: ...
@@ -264,6 +268,7 @@ async def import_(
         documents: typing.List[TDoc],
         import_parameters: _ImportParameters,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[ImportResponse[TDoc]]: ...
 
     @typing.overload
@@ -272,13 +277,15 @@ async def import_(
         documents: typing.Union[bytes, str],
         import_parameters: _ImportParameters = None,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> str: ...
 
     async def import_(
         self,
         documents: typing.Union[bytes, str, typing.List[TDoc]],
         import_parameters: _ImportParameters = None,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.Union[ImportResponse[TDoc], str]:
         """
         Import documents into the collection.
@@ -289,21 +296,33 @@ async def import_(
         Args:
             documents: The documents to import.
             import_parameters: Parameters for the import operation.
-            batch_size: The size of each batch for batch imports.
+            batch_size: Typesense import `batch_size` sent as request query parameter.
+            client_batch_size: Client-side chunk size. When set, this method
+                splits the input list and performs multiple import requests.
 
         Returns:
             The import response, which can be a list of responses or a string.
 
         Raises:
             TypesenseClientError: If an empty list of documents is provided.
         """
+        merged_import_parameters: DocumentImportParameters = {}
+        if import_parameters:
+            merged_import_parameters.update(import_parameters)
+        if batch_size is not None:
+            merged_import_parameters["batch_size"] = batch_size
+
         if isinstance(documents, (str, bytes)):
-            return await self._import_raw(documents, import_parameters)
+            return await self._import_raw(documents, merged_import_parameters)
 
-        if batch_size:
-            return await self._batch_import(documents, import_parameters, batch_size)
+        if client_batch_size:
+            return await self._batch_import(
+                documents,
+                merged_import_parameters,
+                client_batch_size,
+            )
 
-        return await self._bulk_import(documents, import_parameters)
+        return await self._bulk_import(documents, merged_import_parameters)
 
     async def export(
         self,

src/typesense/sync/documents.py

@@ -228,6 +228,7 @@ def import_(
         documents: typing.List[TDoc],
         import_parameters: DocumentImportParametersReturnDocAndId,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[
         typing.Union[ImportResponseWithDocAndId[TDoc], ImportResponseFail[TDoc]]
     ]: ...
@@ -238,6 +239,7 @@ def import_(
         documents: typing.List[TDoc],
         import_parameters: DocumentImportParametersReturnId,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[typing.Union[ImportResponseWithId, ImportResponseFail[TDoc]]]: ...
 
     @typing.overload
@@ -246,6 +248,7 @@ def import_(
         documents: typing.List[TDoc],
         import_parameters: typing.Union[DocumentWriteParameters, None] = None,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[typing.Union[ImportResponseSuccess, ImportResponseFail[TDoc]]]: ...
 
     @typing.overload
@@ -254,6 +257,7 @@ def import_(
         documents: typing.List[TDoc],
         import_parameters: DocumentImportParametersReturnDoc,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[
         typing.Union[ImportResponseWithDoc[TDoc], ImportResponseFail[TDoc]]
     ]: ...
@@ -264,6 +268,7 @@ def import_(
         documents: typing.List[TDoc],
         import_parameters: _ImportParameters,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.List[ImportResponse[TDoc]]: ...
 
     @typing.overload
@@ -272,13 +277,15 @@ def import_(
         documents: typing.Union[bytes, str],
         import_parameters: _ImportParameters = None,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> str: ...
 
     def import_(
         self,
         documents: typing.Union[bytes, str, typing.List[TDoc]],
         import_parameters: _ImportParameters = None,
         batch_size: typing.Union[int, None] = None,
+        client_batch_size: typing.Union[int, None] = None,
     ) -> typing.Union[ImportResponse[TDoc], str]:
         """
         Import documents into the collection.
@@ -289,21 +296,33 @@ def import_(
         Args:
             documents: The documents to import.
             import_parameters: Parameters for the import operation.
-            batch_size: The size of each batch for batch imports.
+            batch_size: Typesense import `batch_size` sent as request query parameter.
+            client_batch_size: Client-side chunk size. When set, this method
+                splits the input list and performs multiple import requests.
 
         Returns:
             The import response, which can be a list of responses or a string.
 
         Raises:
             TypesenseClientError: If an empty list of documents is provided.
         """
+        merged_import_parameters: DocumentImportParameters = {}
+        if import_parameters:
+            merged_import_parameters.update(import_parameters)
+        if batch_size is not None:
+            merged_import_parameters["batch_size"] = batch_size
+
         if isinstance(documents, (str, bytes)):
-            return self._import_raw(documents, import_parameters)
+            return self._import_raw(documents, merged_import_parameters)
 
-        if batch_size:
-            return self._batch_import(documents, import_parameters, batch_size)
+        if client_batch_size:
+            return self._batch_import(
+                documents,
+                merged_import_parameters,
+                client_batch_size,
+            )
 
-        return self._bulk_import(documents, import_parameters)
+        return self._bulk_import(documents, merged_import_parameters)
 
     def export(
         self,

src/typesense/types/document.py

@@ -94,9 +94,12 @@ class DocumentWriteParameters(DirtyValuesParameters):
           If coercion fails, drop the particular field and index the rest of the document.
         - `drop`: Drop the particular field and index the rest of the document.
         - `reject`: Reject the write outright with an error message.
+
+      batch_size (int): Batch size to be sent as query param for the import endpoint.
     """
 
     action: typing.NotRequired[typing.Literal["create", "update", "upsert", "emplace"]]
+    batch_size: typing.NotRequired[int]
 
 
 class UpdateByFilterParameters(typing.TypedDict):

tests/documents_test.py

@@ -296,26 +296,75 @@ def test_import_json_fail(
         actual_documents.import_(generate_companies)
 
 
-def test_import_batch_size(
+def test_import_client_batch_size(
     generate_companies: typing.List[Companies],
     actual_documents: Documents[Companies],
     actual_api_call: ApiCall,
     delete_all: None,
     create_collection: None,
     mocker: MockFixture,
 ) -> None:
-    """Test that the Documents object can import documents in batches."""
-    batch_size = 5
+    """Test that the Documents object can import documents in client-side batches."""
+    client_batch_size = 5
     import_spy = mocker.spy(actual_documents, "import_")
     batch_import_spy = mocker.spy(actual_documents, "_bulk_import")
     request_spy = mocker.spy(actual_api_call, "post")
-    response = actual_documents.import_(generate_companies, batch_size=batch_size)
+    response = actual_documents.import_(
+        generate_companies,
+        client_batch_size=client_batch_size,
+    )
 
     expected = [{"success": True} for _ in generate_companies]
     assert import_spy.call_count == 1
-    assert batch_import_spy.call_count == len(generate_companies) // batch_size
-    assert request_spy.call_count == len(generate_companies) // batch_size
+    assert batch_import_spy.call_count == len(generate_companies) // client_batch_size
+    assert request_spy.call_count == len(generate_companies) // client_batch_size
+    assert response == expected
+
+
+def test_import_batch_size_query_parameter(
+    generate_companies: typing.List[Companies],
+    actual_documents: Documents[Companies],
+    actual_api_call: ApiCall,
+    delete_all: None,
+    create_collection: None,
+    mocker: MockFixture,
+) -> None:
+    """Test that batch_size arg is sent as an import query parameter."""
+    request_spy = mocker.spy(actual_api_call, "post")
+    response = actual_documents.import_(generate_companies, batch_size=42)
+
+    expected = [{"success": True} for _ in generate_companies]
+    assert response == expected
+    request_spy.assert_called_once_with(
+        "/collections/companies/documents/import",
+        body="\n".join([json.dumps(doc) for doc in generate_companies]),
+        params={"batch_size": 42},
+        entity_type=str,
+        as_json=False,
+    )
+
+
+def test_import_batch_size_query_parameter_from_import_parameters(
+    generate_companies: typing.List[Companies],
+    actual_documents: Documents[Companies],
+    actual_api_call: ApiCall,
+    delete_all: None,
+    create_collection: None,
+    mocker: MockFixture,
+) -> None:
+    """Test that import_parameters.batch_size is sent as an import query parameter."""
+    request_spy = mocker.spy(actual_api_call, "post")
+    response = actual_documents.import_(generate_companies, {"batch_size": 42})
+
+    expected = [{"success": True} for _ in generate_companies]
     assert response == expected
+    request_spy.assert_called_once_with(
+        "/collections/companies/documents/import",
+        body="\n".join([json.dumps(doc) for doc in generate_companies]),
+        params={"batch_size": 42},
+        entity_type=str,
+        as_json=False,
+    )
 
 
 def test_import_return_docs(
@@ -524,6 +573,80 @@ async def test_upsert_async(
     assert response == company
 
 
+async def test_import_batch_size_query_parameter_async(
+    generate_companies: typing.List[Companies],
+    actual_async_documents: AsyncDocuments[Companies],
+    actual_async_api_call: AsyncApiCall,
+    delete_all: None,
+    create_collection: None,
+    mocker: MockFixture,
+) -> None:
+    """Test that batch_size arg is sent as an import query parameter."""
+    request_spy = mocker.spy(actual_async_api_call, "post")
+    response = await actual_async_documents.import_(generate_companies, batch_size=42)
+
+    expected = [{"success": True} for _ in generate_companies]
+    assert response == expected
+    request_spy.assert_called_once_with(
+        "/collections/companies/documents/import",
+        body="\n".join([json.dumps(doc) for doc in generate_companies]),
+        params={"batch_size": 42},
+        entity_type=str,
+        as_json=False,
+    )
+
+
+async def test_import_batch_size_query_parameter_from_import_parameters_async(
+    generate_companies: typing.List[Companies],
+    actual_async_documents: AsyncDocuments[Companies],
+    actual_async_api_call: AsyncApiCall,
+    delete_all: None,
+    create_collection: None,
+    mocker: MockFixture,
+) -> None:
+    """Test that import_parameters.batch_size is sent as an import query parameter."""
+    request_spy = mocker.spy(actual_async_api_call, "post")
+    response = await actual_async_documents.import_(
+        generate_companies,
+        {"batch_size": 42},
+    )
+
+    expected = [{"success": True} for _ in generate_companies]
+    assert response == expected
+    request_spy.assert_called_once_with(
+        "/collections/companies/documents/import",
+        body="\n".join([json.dumps(doc) for doc in generate_companies]),
+        params={"batch_size": 42},
+        entity_type=str,
+        as_json=False,
+    )
+
+
+async def test_import_client_batch_size_async(
+    generate_companies: typing.List[Companies],
+    actual_async_documents: AsyncDocuments[Companies],
+    actual_async_api_call: AsyncApiCall,
+    delete_all: None,
+    create_collection: None,
+    mocker: MockFixture,
+) -> None:
+    """Test that AsyncDocuments can import documents in client-side batches."""
+    client_batch_size = 5
+    import_spy = mocker.spy(actual_async_documents, "import_")
+    batch_import_spy = mocker.spy(actual_async_documents, "_bulk_import")
+    request_spy = mocker.spy(actual_async_api_call, "post")
+    response = await actual_async_documents.import_(
+        generate_companies,
+        client_batch_size=client_batch_size,
+    )
+
+    expected = [{"success": True} for _ in generate_companies]
+    assert import_spy.call_count == 1
+    assert batch_import_spy.call_count == len(generate_companies) // client_batch_size
+    assert request_spy.call_count == len(generate_companies) // client_batch_size
+    assert response == expected
+
+
 async def test_export_async(
     actual_async_documents: AsyncDocuments[Companies],
     delete_all: None,