(self, array_feature, shape_1, shape_2)
| 205 | self.assertTupleEqual(dataset[:2]["matrix"].to_numpy().shape, (2, *shape_2)) |
| 206 | |
| 207 | def test_write(self, array_feature, shape_1, shape_2): |
| 208 | with tempfile.TemporaryDirectory() as tmp_dir: |
| 209 | my_features = self.get_features(array_feature, shape_1, shape_2) |
| 210 | my_examples = [ |
| 211 | (0, self.get_dict_example_0(shape_1, shape_2)), |
| 212 | (1, self.get_dict_example_1(shape_1, shape_2)), |
| 213 | ] |
| 214 | with ArrowWriter(features=my_features, path=os.path.join(tmp_dir, "beta.arrow")) as writer: |
| 215 | for key, record in my_examples: |
| 216 | example = my_features.encode_example(record) |
| 217 | writer.write(example) |
| 218 | num_examples, num_bytes = writer.finalize() |
| 219 | dataset = datasets.Dataset.from_file(os.path.join(tmp_dir, "beta.arrow")) |
| 220 | self._check_getitem_output_type(dataset, shape_1, shape_2, my_examples[0][1]["matrix"]) |
| 221 | del dataset |
| 222 | |
| 223 | def test_write_batch(self, array_feature, shape_1, shape_2): |
| 224 | with tempfile.TemporaryDirectory() as tmp_dir: |
nothing calls this directly
no test coverage detected