| 1288 | |
| 1289 | |
| 1290 | def test_make_csv_fragment_from_buffer(dataset_reader, pickle_module): |
| 1291 | content = textwrap.dedent(""" |
| 1292 | alpha,num,animal |
| 1293 | a,12,dog |
| 1294 | b,11,cat |
| 1295 | c,10,rabbit |
| 1296 | """) |
| 1297 | buffer = pa.py_buffer(content.encode('utf-8')) |
| 1298 | |
| 1299 | csv_format = ds.CsvFileFormat() |
| 1300 | fragment = csv_format.make_fragment(buffer) |
| 1301 | |
| 1302 | # When buffer, fragment open returns a BufferReader, not NativeFile |
| 1303 | assert isinstance(fragment.open(), pa.BufferReader) |
| 1304 | |
| 1305 | expected = pa.table([['a', 'b', 'c'], |
| 1306 | [12, 11, 10], |
| 1307 | ['dog', 'cat', 'rabbit']], |
| 1308 | names=['alpha', 'num', 'animal']) |
| 1309 | assert dataset_reader.to_table(fragment).equals(expected) |
| 1310 | |
| 1311 | pickled = pickle_module.loads(pickle_module.dumps(fragment)) |
| 1312 | assert dataset_reader.to_table(pickled).equals(fragment.to_table()) |
| 1313 | |
| 1314 | |
| 1315 | def test_make_json_fragment_from_buffer(dataset_reader, pickle_module): |