| 4753 | |
| 4754 | |
| 4755 | def test_write_table(tempdir): |
| 4756 | table = pa.table([ |
| 4757 | pa.array(range(20)), pa.array(random.random() for _ in range(20)), |
| 4758 | pa.array(['a'] * 10 + ['b'] * 10) |
| 4759 | ], names=["f1", "f2", "part"]) |
| 4760 | |
| 4761 | base_dir = tempdir / 'single' |
| 4762 | ds.write_dataset(table, base_dir, |
| 4763 | basename_template='dat_{i}.arrow', format="feather") |
| 4764 | # check that all files are present |
| 4765 | file_paths = list(base_dir.rglob("*")) |
| 4766 | expected_paths = [base_dir / "dat_0.arrow"] |
| 4767 | assert set(file_paths) == set(expected_paths) |
| 4768 | # check Table roundtrip |
| 4769 | result = ds.dataset(base_dir, format="ipc").to_table() |
| 4770 | assert result.equals(table) |
| 4771 | |
| 4772 | # with partitioning |
| 4773 | base_dir = tempdir / 'partitioned' |
| 4774 | expected_paths = [ |
| 4775 | base_dir / "part=a", base_dir / "part=a" / "dat_0.arrow", |
| 4776 | base_dir / "part=b", base_dir / "part=b" / "dat_0.arrow" |
| 4777 | ] |
| 4778 | |
| 4779 | visited_paths = [] |
| 4780 | visited_sizes = [] |
| 4781 | |
| 4782 | def file_visitor(written_file): |
| 4783 | visited_paths.append(written_file.path) |
| 4784 | visited_sizes.append(written_file.size) |
| 4785 | |
| 4786 | partitioning = ds.partitioning( |
| 4787 | pa.schema([("part", pa.string())]), flavor="hive") |
| 4788 | ds.write_dataset(table, base_dir, format="feather", |
| 4789 | basename_template='dat_{i}.arrow', |
| 4790 | partitioning=partitioning, file_visitor=file_visitor) |
| 4791 | file_paths = list(base_dir.rglob("*")) |
| 4792 | assert set(file_paths) == set(expected_paths) |
| 4793 | actual_sizes = [os.path.getsize(path) for path in visited_paths] |
| 4794 | assert visited_sizes == actual_sizes |
| 4795 | result = ds.dataset(base_dir, format="ipc", partitioning=partitioning) |
| 4796 | assert result.to_table().equals(table) |
| 4797 | assert len(visited_paths) == 2 |
| 4798 | for visited_path in visited_paths: |
| 4799 | assert pathlib.Path(visited_path) in expected_paths |
| 4800 | |
| 4801 | |
| 4802 | def test_write_table_multiple_fragments(tempdir): |