Address second round of review feedback

Author: cameronwaterman

examples/dataframe/create_write_data.py

@@ -66,10 +66,10 @@
     ]
     client.append_table_data(table_id, batch_list)
 
-    # Flush only (no additional rows, mark end_of_data)
+    # Mark end_of_data for the table
+    # Supply `None` and `end_of_data=True`
     client.append_table_data(table_id, None, end_of_data=True)
-
-    # Empty iterable (must supply end_of_data)
+    # OR supply an empty Iterable and `end_of_data=True`
     client.append_table_data(table_id, [], end_of_data=True)
 else:
     # If pyarrow not installed, flush via JSON path

nisystemlink/clients/dataframe/_data_frame_client.py

@@ -359,11 +359,7 @@ def append_table_data(
                 )
 
             def _generate_body() -> Iterable[memoryview]:
-                data_iter = iter(data)
-                try:
-                    batch = next(data_iter)
-                except StopIteration:
-                    return
+                batch = first_batch
                 with BytesIO() as buf:
                     options = pa.ipc.IpcWriteOptions(compression="zstd")
                     writer = pa.ipc.new_stream(buf, batch.schema, options=options)
@@ -374,14 +370,13 @@ def _generate_body() -> Iterable[memoryview]:
                             yield slice
                         buf.seek(0)
                         try:
-                            batch = next(data_iter)
+                            batch = next(iterator)
                         except StopIteration:
                             break
 
                     writer.close()
-                    with buf.getbuffer() as view:
-                        with view[0 : buf.tell()] as slice:
-                            yield slice
+                    with buf.getbuffer() as view, view[0 : buf.tell()] as slice:
+                        yield slice
 
             try:
                 self._append_table_data_arrow(

tests/integration/dataframe/test_dataframe.py

@@ -684,9 +684,10 @@ def test__append_table_data__flush_only_with_none(
         self, client: DataFrameClient, create_table
     ):
         table_id = self._new_single_int_table(create_table)
-        frame = DataFrame(columns=["a"], data=[["1"]])
-        client.append_table_data(table_id, frame)
         client.append_table_data(table_id, None, end_of_data=True)
+        metadata = client.get_table_metadata(table_id)
+        assert metadata.row_count == 0
+        assert metadata.supports_append is False
 
     def test__append_table_data__arrow_ingestion_success(
         self, client: DataFrameClient, create_table
@@ -823,22 +824,14 @@ def test__append_table_data__arrow_ingestion_non_400_passthrough(
         self, client: DataFrameClient
     ):
         pa = pytest.importorskip("pyarrow")
-        table_id = "mock_table_id"
         batch = pa.record_batch([pa.array([1, 2, 3])], names=["a"])
-        with responses.RequestsMock() as rsps:
-            rsps.add(
-                responses.POST,
-                f"{client.session.base_url}tables/{table_id}/data",
-                status=409,
-            )
-            with pytest.raises(ApiException) as excinfo:
-                client.append_table_data(table_id, [batch], end_of_data=True)
+        with pytest.raises(ApiException) as excinfo:
+            client.append_table_data("111111111111111111111111", [batch], end_of_data=True)
         assert "Arrow ingestion request was rejected" not in str(excinfo.value)
 
     def test__append_table_data__unsupported_type_raises(
         self, client: DataFrameClient, create_table
     ):
         table_id = create_table(basic_table_model)
         with pytest.raises(ValueError, match="Unsupported type"):
-            # cast to Any to satisfy type checker while still exercising runtime path
-            client.append_table_data(table_id, cast(Any, 123))
+            client.append_table_data(table_id, 123)  # type: ignore[arg-type]