pitrou commented on code in PR #13041:
URL: https://github.com/apache/arrow/pull/13041#discussion_r869084288
##########
python/pyarrow/_flight.pyx:
##########
@@ -977,22 +977,28 @@ cdef class MetadataRecordBatchWriter(_CRecordBatchWriter):
check_flight_status(
self._writer().WriteMetadata(c_buf))
- def write_batch(self, RecordBatch batch):
+ def write_batch(self, RecordBatch batch, custom_metadata=None):
"""
Write RecordBatch to stream.
Parameters
----------
batch : RecordBatch
"""
+
# Override superclass method to use check_flight_status so we
# can generate FlightWriteSizeExceededError. We don't do this
# for write_table as callers who intend to handle the error
# and retry with a smaller batch should be working with
# individual batches to have control.
+
+ if not (custom_metadata is None or isinstance(custom_metadata,
KeyValueMetadata)):
+ custom_metadata = KeyValueMetadata(custom_metadata)
Review Comment:
I think you can use `ensure_metadata` to abstract these checks away.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
