def test_file_reader_writer():
data = [
pa.array([1, 2, 3, 4]),
pa.array(['foo', 'bar', 'baz', None]),
pa.array([True, None, False, True])
]
batch = pa.RecordBatch.from_arrays(data, ['f0', 'f1', 'f2'])
sink = pa.BufferOutputStream()
with pytest.warns(FutureWarning):
stream_writer = pa.StreamWriter(sink, batch.schema)
assert isinstance(stream_writer, pa.RecordBatchStreamWriter)
sink2 = pa.BufferOutputStream()
with pytest.warns(FutureWarning):
file_writer = pa.FileWriter(sink2, batch.schema)
assert isinstance(file_writer, pa.RecordBatchFileWriter)
file_writer.write_batch(batch)
stream_writer.write_batch(batch)
file_writer.close()
stream_writer.close()
buf = sink.get_result()
buf2 = sink2.get_result()
with pytest.warns(FutureWarning):
stream_reader = pa.StreamReader(buf)
assert isinstance(stream_reader, pa.RecordBatchStreamReader)
with pytest.warns(FutureWarning):
file_reader = pa.FileReader(buf2)
assert isinstance(file_reader, pa.RecordBatchFileReader)
def _get_writer(self, sink, schema):
return pa.FileWriter(sink, schema)
def write_file(batch, sink):
writer = pa.FileWriter(sink, batch.schema)
writer.write_batch(batch)
writer.close()