(constructor)
| 651 | lambda schema, batches: pa.Table.from_batches(batches, schema), |
| 652 | ], ids=['recordbatchreader', 'table']) |
| 653 | def test_roundtrip_reader_capsule(constructor): |
| 654 | batches = make_batches() |
| 655 | schema = batches[0].schema |
| 656 | |
| 657 | gc.collect() # Make sure no Arrow data dangles in a ref cycle |
| 658 | old_allocated = pa.total_allocated_bytes() |
| 659 | |
| 660 | obj = constructor(schema, batches) |
| 661 | |
| 662 | capsule = obj.__arrow_c_stream__() |
| 663 | assert PyCapsule_IsValid(capsule, b"arrow_array_stream") == 1 |
| 664 | imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) |
| 665 | assert imported_reader.schema == schema |
| 666 | imported_batches = list(imported_reader) |
| 667 | assert len(imported_batches) == len(batches) |
| 668 | for batch, expected in zip(imported_batches, batches): |
| 669 | assert batch.equals(expected) |
| 670 | |
| 671 | del obj, imported_reader, batch, expected, imported_batches |
| 672 | |
| 673 | assert pa.total_allocated_bytes() == old_allocated |
| 674 | |
| 675 | obj = constructor(schema, batches) |
| 676 | |
| 677 | bad_schema = pa.schema({'ints': pa.int32()}) |
| 678 | with pytest.raises(pa.lib.ArrowTypeError, match="Field 0 cannot be cast"): |
| 679 | obj.__arrow_c_stream__(bad_schema.__arrow_c_schema__()) |
| 680 | |
| 681 | # Can work with matching schema |
| 682 | matching_schema = pa.schema({'ints': pa.list_(pa.int32())}) |
| 683 | capsule = obj.__arrow_c_stream__(matching_schema.__arrow_c_schema__()) |
| 684 | imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) |
| 685 | assert imported_reader.schema == matching_schema |
| 686 | for batch, expected in zip(imported_reader, batches): |
| 687 | assert batch.equals(expected) |
| 688 | |
| 689 | |
| 690 | def test_roundtrip_batch_reader_capsule_requested_schema(): |
nothing calls this directly
no test coverage detected