Exception when not fully consuming generator/iterator that feeds into ArrowStreamDataset

[kszlim]

Repro code:

import tensorflow as tf
import tensorflow_io.arrow as arrow_io
import pyarrow as pa

if __name__ == "__main__":
    data = pa.array(list(range(5000)), type=pa.float32())
    names = ["data"]
    batches = (pa.RecordBatch.from_arrays([data], names=names) for _ in range(100))
    ds = arrow_io.ArrowStreamDataset.from_record_batches(batches, output_types=(tf.float32,), batch_size=5000, batch_mode="drop_remainder")
    for sample in ds.take(10):
        print(sample)

Exception in thread Thread-1 (run_server):
Traceback (most recent call last):
File ".../venv/lib/python3.10/site-packages/tensorflow_io/python/ops/arrow_dataset_ops.py", line 572, in run_server
writer.write_batch(batch)
File "pyarrow/ipc.pxi", line 484, in pyarrow.lib._CRecordBatchWriter.write_batch
File "/usr/lib/python3.10/socket.py", line 723, in write
return self._sock.send(b)
BrokenPipeError: [Errno 32] Broken pipe

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/usr/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
self.run()
File "/usr/lib/python3.10/threading.py", line 953, in run
self._target(*self._args, **self._kwargs)
File ".../venv/lib/python3.10/site-packages/tensorflow_io/python/ops/arrow_dataset_ops.py", line 578, in run_server
outfile.close()
File "/usr/lib/python3.10/socket.py", line 723, in write
return self._sock.send(b)
BrokenPipeError: [Errno 32] Broken pipe

[kszlim]

Anyone run into this?