()
| 112 | |
| 113 | |
| 114 | def test_write_with_features(): |
| 115 | output = pa.BufferOutputStream() |
| 116 | features = Features({"labels": ClassLabel(names=["neg", "pos"])}) |
| 117 | with ArrowWriter(stream=output, features=features) as writer: |
| 118 | writer.write({"labels": 0}) |
| 119 | writer.write({"labels": 1}) |
| 120 | num_examples, num_bytes = writer.finalize() |
| 121 | assert num_examples == 2 |
| 122 | assert num_bytes > 0 |
| 123 | assert writer._schema == features.arrow_schema |
| 124 | assert writer._schema.metadata == features.arrow_schema.metadata |
| 125 | stream = pa.BufferReader(output.getvalue()) |
| 126 | f = pa.ipc.open_stream(stream) |
| 127 | pa_table: pa.Table = f.read_all() |
| 128 | schema = pa_table.schema |
| 129 | assert pa_table.num_rows == 2 |
| 130 | assert schema == features.arrow_schema |
| 131 | assert schema.metadata == features.arrow_schema.metadata |
| 132 | assert features == Features.from_arrow_schema(schema) |
| 133 | |
| 134 | |
| 135 | @pytest.mark.parametrize("writer_batch_size", [None, 2, 10]) |
nothing calls this directly
no test coverage detected