(fields, writer_batch_size)
| 97 | ], |
| 98 | ) |
| 99 | def test_write(fields, writer_batch_size): |
| 100 | output = pa.BufferOutputStream() |
| 101 | schema = pa.schema(fields) if fields else None |
| 102 | with ArrowWriter(stream=output, schema=schema, writer_batch_size=writer_batch_size) as writer: |
| 103 | writer.write({"col_1": "foo", "col_2": 1}) |
| 104 | writer.write({"col_1": "bar", "col_2": 2}) |
| 105 | num_examples, num_bytes = writer.finalize() |
| 106 | assert num_examples == 2 |
| 107 | assert num_bytes > 0 |
| 108 | if not fields: |
| 109 | fields = {"col_1": pa.string(), "col_2": pa.int64()} |
| 110 | assert writer._schema == pa.schema(fields, metadata=writer._schema.metadata) |
| 111 | _check_output(output.getvalue(), expected_num_chunks=num_examples if writer_batch_size == 1 else 1) |
| 112 | |
| 113 | |
| 114 | def test_write_with_features(): |
nothing calls this directly
no test coverage detected