()
| 619 | |
| 620 | |
| 621 | def test_table_c_stream_interface(): |
| 622 | class StreamWrapper: |
| 623 | def __init__(self, batches): |
| 624 | self.batches = batches |
| 625 | |
| 626 | def __arrow_c_stream__(self, requested_schema=None): |
| 627 | reader = pa.RecordBatchReader.from_batches( |
| 628 | self.batches[0].schema, self.batches) |
| 629 | return reader.__arrow_c_stream__(requested_schema) |
| 630 | |
| 631 | data = [ |
| 632 | pa.record_batch([pa.array([1, 2, 3], type=pa.int64())], names=['a']), |
| 633 | pa.record_batch([pa.array([4, 5, 6], type=pa.int64())], names=['a']) |
| 634 | ] |
| 635 | wrapper = StreamWrapper(data) |
| 636 | |
| 637 | # Can roundtrip through the wrapper. |
| 638 | result = pa.table(wrapper) |
| 639 | expected = pa.Table.from_batches(data) |
| 640 | assert result == expected |
| 641 | |
| 642 | # Passing schema works if already that schema |
| 643 | result = pa.table(wrapper, schema=data[0].schema) |
| 644 | assert result == expected |
| 645 | |
| 646 | # Passing a different schema will cast |
| 647 | good_schema = pa.schema([pa.field('a', pa.int32())]) |
| 648 | result = pa.table(wrapper, schema=good_schema) |
| 649 | assert result == expected.cast(good_schema) |
| 650 | |
| 651 | # If schema doesn't match, raises NotImplementedError |
| 652 | with pytest.raises( |
| 653 | pa.lib.ArrowTypeError, match="Field 0 cannot be cast" |
| 654 | ): |
| 655 | pa.table( |
| 656 | wrapper, schema=pa.schema([pa.field('a', pa.list_(pa.int32()))]) |
| 657 | ) |
| 658 | |
| 659 | |
| 660 | def test_recordbatch_itercolumns(): |
nothing calls this directly
no test coverage detected