Address review feedback

Author: cameronwaterman

README.rst

@@ -34,6 +34,16 @@ To install **nisystemlink-clients**, use one of the following methods:
 
    $ python -m easy_install nisystemlink-clients
 
+Optional Arrow (pyarrow) Support
+--------------------------------
+The base install does not pull in ``pyarrow``. Install the optional extra if you
+plan to use Arrow RecordBatch ingestion with ``DataFrameClient.append_table_data``::
+
+   $ python -m pip install "nisystemlink-clients[pyarrow]"
+
+Without the extra, Arrow-specific code paths will raise a ``RuntimeError`` when
+attempting to append ``pyarrow.RecordBatch`` instances.
+
 .. _usage_section:
 
 Usage

docs/api_reference/dataframe.rst

@@ -25,3 +25,27 @@ nisystemlink.clients.dataframe
 .. automodule:: nisystemlink.clients.dataframe.models
    :members:
    :imported-members:
+
+Arrow / JSON Ingestion Notes
+----------------------------
+``append_table_data`` accepts multiple data forms:
+
+* ``AppendTableDataRequest`` (JSON)
+* ``DataFrame`` model (JSON)
+* Single ``pyarrow.RecordBatch`` (Arrow IPC)
+* Iterable of ``pyarrow.RecordBatch`` (Arrow IPC)
+* ``None`` with ``end_of_data`` (flush only)
+
+Arrow support is optional and requires installing the ``pyarrow`` extra::
+
+   pip install "nisystemlink-clients[pyarrow]"
+
+If ``pyarrow`` is not installed and a RecordBatch (or iterable) is passed, a
+``RuntimeError`` is raised. When Arrow is used, the batches are serialized into
+an IPC stream and sent with content type
+``application/vnd.apache.arrow.stream``; the ``end_of_data`` flag is sent as a
+query parameter. JSON ingestion places ``endOfData`` in the request body.
+
+If the target SystemLink DataFrame Service does not yet support Arrow and
+responds with HTTP 400, the client raises an explanatory ``ApiException``
+advising to upgrade or fall back to JSON ingestion.

examples/dataframe/create_write_data.py

@@ -1,6 +1,10 @@
 import random
 from datetime import datetime
 
+try:
+    import pyarrow as pa  # type: ignore
+except Exception:
+    pa = None
 from nisystemlink.clients.dataframe import DataFrameClient
 from nisystemlink.clients.dataframe.models import (
     AppendTableDataRequest,
@@ -25,12 +29,48 @@
     )
 )
 
-# Generate example data
-frame = DataFrame(
-    data=[[i, random.random(), datetime.now().isoformat()] for i in range(100)]
+# Append via explicit AppendTableDataRequest (JSON)
+frame_request = DataFrame(
+    data=[[i, random.random(), datetime.now().isoformat()] for i in range(3)]
 )
+client.append_table_data(table_id, AppendTableDataRequest(frame=frame_request))
 
-# Write example data to table
-client.append_table_data(
-    table_id, data=AppendTableDataRequest(frame=frame, endOfData=True)
+# Append via DataFrame model directly (JSON)
+frame_direct = DataFrame(
+    data=[[i + 3, random.random(), datetime.now().isoformat()] for i in range(3)]
 )
+client.append_table_data(table_id, frame_direct)
+
+if pa is not None:
+    # Append via single RecordBatch (Arrow)
+    batch_single = pa.record_batch(
+        [
+            pa.array([6, 7, 8]),
+            pa.array([0.1, 0.2, 0.3]),
+            pa.array([datetime.now().isoformat()] * 3),
+        ],
+        names=["ix", "Float_Column", "Timestamp_Column"],
+    )
+    client.append_table_data(table_id, batch_single)
+
+    # Append via iterable of RecordBatches (Arrow)
+    batch_list = [
+        pa.record_batch(
+            [
+                pa.array([9, 10]),
+                pa.array([0.4, 0.5]),
+                pa.array([datetime.now().isoformat(), datetime.now().isoformat()]),
+            ],
+            names=["ix", "Float_Column", "Timestamp_Column"],
+        )
+    ]
+    client.append_table_data(table_id, batch_list)
+
+    # Flush only (no additional rows, mark end_of_data)
+    client.append_table_data(table_id, None, end_of_data=True)
+
+    # Empty iterable (must supply end_of_data)
+    client.append_table_data(table_id, [], end_of_data=True)
+else:
+    # If pyarrow not installed, flush via JSON path
+    client.append_table_data(table_id, None, end_of_data=True)

nisystemlink/clients/dataframe/_data_frame_client.py

@@ -4,7 +4,10 @@
 from io import BytesIO
 from typing import List, Optional, Union
 
-import pyarrow as pa  # type: ignore
+try:
+    import pyarrow as pa  # type: ignore
+except Exception:
+    pa = None
 from nisystemlink.clients import core
 from nisystemlink.clients.core._uplink._base_client import BaseClient
 from nisystemlink.clients.core._uplink._methods import (
@@ -265,7 +268,7 @@ def _append_table_data_json(
         content_type="application/vnd.apache.arrow.stream",
     )
     def _append_table_data_arrow(
-        self, id: str, data: bytes, end_of_data: Optional[bool] = None
+        self, id: str, data: Iterable[bytes], end_of_data: Optional[bool] = None
     ) -> None:
         """Internal uplink-implemented Arrow (binary) append call."""
         ...
@@ -277,6 +280,7 @@ def append_table_data(
             Union[
                 models.AppendTableDataRequest,
                 models.DataFrame,
+                "pa.RecordBatch",  # type: ignore[name-defined]
                 Iterable["pa.RecordBatch"],  # type: ignore[name-defined]
             ]
         ],
@@ -291,9 +295,12 @@ def append_table_data(
                 * AppendTableDataRequest: Sent as-is via JSON; ``end_of_data`` must be ``None``.
                 * DataFrame (service model): Wrapped into an AppendTableDataRequest (``end_of_data``
                     optional) and sent as JSON.
+                * Single pyarrow.RecordBatch: Treated the same as an iterable containing one batch
+                    and streamed as Arrow IPC. ``end_of_data`` (if provided) is sent as a query
+                    parameter.
                 * Iterable[pyarrow.RecordBatch]: Streamed as Arrow IPC. ``end_of_data`` (if provided)
-                    is sent as a query parameter. If the iterator yields no batches, it is treated as
-                    ``None`` and requires ``end_of_data``.
+                    is sent as a query parameter. If the iterator yields no batches, it is treated
+                    as ``None`` and requires ``end_of_data``.
                 * None: ``end_of_data`` must be provided; sends JSON containing only the
                     ``endOfData`` flag.
             end_of_data: Whether additional rows may be appended in future requests. Required when
@@ -323,6 +330,9 @@ def append_table_data(
             self._append_table_data_json(id, request_model)
             return
 
+        if pa is not None and isinstance(data, pa.RecordBatch):
+            data = [data]
+
         if isinstance(data, Iterable):
             iterator = iter(data)
             try:
@@ -348,22 +358,36 @@ def append_table_data(
                     "Iterable provided to data must yield pyarrow.RecordBatch objects."
                 )
 
-            def _build_body() -> bytes:
+            def _generate_body() -> Iterable[memoryview]:
+                data_iter = iter(data)
+                try:
+                    batch = next(data_iter)
+                except StopIteration:
+                    return
                 with BytesIO() as buf:
                     options = pa.ipc.IpcWriteOptions(compression="zstd")
-                    with pa.ipc.new_stream(
-                        buf, first_batch.schema, options=options
-                    ) as writer:
-                        writer.write_batch(first_batch)
-                        for batch in iterator:
-                            writer.write_batch(batch)
-                    return buf.getvalue()
+                    writer = pa.ipc.new_stream(buf, batch.schema, options=options)
+
+                    while True:
+                        writer.write_batch(batch)
+                        with buf.getbuffer() as view, view[0 : buf.tell()] as slice:
+                            yield slice
+                        buf.seek(0)
+                        try:
+                            batch = next(data_iter)
+                        except StopIteration:
+                            break
+
+                    writer.close()
+                    with buf.getbuffer() as view:
+                        with view[0 : buf.tell()] as slice:
+                            yield slice
 
             try:
                 self._append_table_data_arrow(
                     id,
-                    _build_body(),
-                    (end_of_data if end_of_data is not None else None),
+                    _generate_body(),
+                    end_of_data,
                 )
             except core.ApiException as ex:
                 if ex.http_status_code == 400:

poetry.lock

@@ -43,7 +43,10 @@ uplink   = [
 pydantic = "^2.11.3"
 pyyaml = "^6.0.1"
 pandas = "^2.1.0"
-pyarrow = "^21.0.0"
+pyarrow = { version = "^21.0.0", optional = true }
+
+[tool.poetry.extras]
+pyarrow = ["pyarrow"]
 
 [tool.poetry.group.dev.dependencies]
 black               = ">=22.10,<25.0"

pyproject.toml

@@ -698,6 +698,16 @@ def test__append_table_data__arrow_ingestion_success(
         with pytest.raises(ApiException):
             client.append_table_data(table_id, None, end_of_data=True)
 
+    def test__append_table_data__single_recordbatch_success(
+        self, client: DataFrameClient, create_table
+    ):
+        pa = pytest.importorskip("pyarrow")
+        table_id = self._new_single_int_table(create_table)
+        batch = pa.record_batch([pa.array([1, 2, 3])], names=["a"])
+        client.append_table_data(table_id, batch, end_of_data=True)
+        with pytest.raises(ApiException):
+            client.append_table_data(table_id, None, end_of_data=True)
+
     def test__append_table_data__arrow_ingestion_with_end_of_data_query_param_false(
         self, client: DataFrameClient, create_table
     ):