(constructor)
| 684 | lambda schema, batches: pa.Table.from_batches(batches, schema), |
| 685 | ], ids=['recordbatchreader', 'table']) |
| 686 | def test_roundtrip_reader_capsule(constructor): |
| 687 | batches = make_batches() |
| 688 | schema = batches[0].schema |
| 689 | |
| 690 | gc.collect() # Make sure no Arrow data dangles in a ref cycle |
| 691 | old_allocated = pa.total_allocated_bytes() |
| 692 | |
| 693 | obj = constructor(schema, batches) |
| 694 | |
| 695 | capsule = obj.__arrow_c_stream__() |
| 696 | assert PyCapsule_IsValid(capsule, b"arrow_array_stream") == 1 |
| 697 | imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) |
| 698 | assert imported_reader.schema == schema |
| 699 | imported_batches = list(imported_reader) |
| 700 | assert len(imported_batches) == len(batches) |
| 701 | for batch, expected in zip(imported_batches, batches): |
| 702 | assert batch.equals(expected) |
| 703 | |
| 704 | del obj, imported_reader, batch, expected, imported_batches |
| 705 | |
| 706 | assert pa.total_allocated_bytes() == old_allocated |
| 707 | |
| 708 | obj = constructor(schema, batches) |
| 709 | |
| 710 | bad_schema = pa.schema({'ints': pa.int32()}) |
| 711 | with pytest.raises(pa.lib.ArrowTypeError, match="Field 0 cannot be cast"): |
| 712 | obj.__arrow_c_stream__(bad_schema.__arrow_c_schema__()) |
| 713 | |
| 714 | # Can work with matching schema |
| 715 | matching_schema = pa.schema({'ints': pa.list_(pa.int32())}) |
| 716 | capsule = obj.__arrow_c_stream__(matching_schema.__arrow_c_schema__()) |
| 717 | imported_reader = pa.RecordBatchReader._import_from_c_capsule(capsule) |
| 718 | assert imported_reader.schema == matching_schema |
| 719 | for batch, expected in zip(imported_reader, batches): |
| 720 | assert batch.equals(expected) |
| 721 | |
| 722 | |
| 723 | def test_roundtrip_batch_reader_capsule_requested_schema(): |
nothing calls this directly
no test coverage detected