(tempdir)
| 4958 | |
| 4959 | |
| 4960 | def test_write_dataset_csv(tempdir): |
| 4961 | table = pa.table([ |
| 4962 | pa.array(range(20)), pa.array(random.random() for _ in range(20)), |
| 4963 | pa.array(['a'] * 10 + ['b'] * 10) |
| 4964 | ], names=["f1", "f2", "chr1"]) |
| 4965 | |
| 4966 | base_dir = tempdir / 'csv_dataset' |
| 4967 | ds.write_dataset(table, base_dir, format="csv") |
| 4968 | # check that all files are present |
| 4969 | file_paths = list(base_dir.rglob("*")) |
| 4970 | expected_paths = [base_dir / "part-0.csv"] |
| 4971 | assert set(file_paths) == set(expected_paths) |
| 4972 | # check Table roundtrip |
| 4973 | result = ds.dataset(base_dir, format="csv").to_table() |
| 4974 | assert result.equals(table) |
| 4975 | |
| 4976 | # using custom options |
| 4977 | format = ds.CsvFileFormat(read_options=pyarrow.csv.ReadOptions( |
| 4978 | column_names=table.schema.names)) |
| 4979 | opts = format.make_write_options(include_header=False) |
| 4980 | base_dir = tempdir / 'csv_dataset_noheader' |
| 4981 | ds.write_dataset(table, base_dir, format=format, file_options=opts) |
| 4982 | result = ds.dataset(base_dir, format=format).to_table() |
| 4983 | assert result.equals(table) |
| 4984 | |
| 4985 | |
| 4986 | @pytest.mark.parquet |
nothing calls this directly
no test coverage detected