(tempdir)
| 4603 | |
| 4604 | @pytest.mark.parquet |
| 4605 | def test_write_dataset_max_open_files(tempdir): |
| 4606 | directory = tempdir / 'ds' |
| 4607 | file_format = "parquet" |
| 4608 | partition_column_id = 1 |
| 4609 | column_names = ['c1', 'c2'] |
| 4610 | record_batch_1 = pa.record_batch(data=[[1, 2, 3, 4, 0, 10], |
| 4611 | ['a', 'b', 'c', 'd', 'e', 'a']], |
| 4612 | names=column_names) |
| 4613 | record_batch_2 = pa.record_batch(data=[[5, 6, 7, 8, 0, 1], |
| 4614 | ['a', 'b', 'c', 'd', 'e', 'c']], |
| 4615 | names=column_names) |
| 4616 | record_batch_3 = pa.record_batch(data=[[9, 10, 11, 12, 0, 1], |
| 4617 | ['a', 'b', 'c', 'd', 'e', 'd']], |
| 4618 | names=column_names) |
| 4619 | record_batch_4 = pa.record_batch(data=[[13, 14, 15, 16, 0, 1], |
| 4620 | ['a', 'b', 'c', 'd', 'e', 'b']], |
| 4621 | names=column_names) |
| 4622 | |
| 4623 | table = pa.Table.from_batches([record_batch_1, record_batch_2, |
| 4624 | record_batch_3, record_batch_4]) |
| 4625 | |
| 4626 | partitioning = ds.partitioning( |
| 4627 | pa.schema([(column_names[partition_column_id], pa.string())]), |
| 4628 | flavor="hive") |
| 4629 | |
| 4630 | data_source_1 = directory / "default" |
| 4631 | |
| 4632 | ds.write_dataset(data=table, base_dir=data_source_1, |
| 4633 | partitioning=partitioning, format=file_format) |
| 4634 | |
| 4635 | # Here we consider the number of unique partitions created when |
| 4636 | # partitioning column contains duplicate records. |
| 4637 | # Returns: (number_of_files_generated, number_of_partitions) |
| 4638 | def _get_compare_pair(data_source, record_batch, file_format, col_id): |
| 4639 | num_of_files_generated = _get_num_of_files_generated( |
| 4640 | base_directory=data_source, file_format=file_format) |
| 4641 | number_of_partitions = len(pa.compute.unique(record_batch[col_id])) |
| 4642 | return num_of_files_generated, number_of_partitions |
| 4643 | |
| 4644 | # CASE 1: when max_open_files=default & max_open_files >= num_of_partitions |
| 4645 | # In case of a writing to disk via partitioning based on a |
| 4646 | # particular column (considering row labels in that column), |
| 4647 | # the number of unique rows must be equal |
| 4648 | # to the number of files generated |
| 4649 | |
| 4650 | num_of_files_generated, number_of_partitions \ |
| 4651 | = _get_compare_pair(data_source_1, record_batch_1, file_format, |
| 4652 | partition_column_id) |
| 4653 | assert num_of_files_generated == number_of_partitions |
| 4654 | |
| 4655 | # CASE 2: when max_open_files > 0 & max_open_files < num_of_partitions |
| 4656 | # the number of files generated must be greater than the number of |
| 4657 | # partitions |
| 4658 | |
| 4659 | data_source_2 = directory / "max_1" |
| 4660 | |
| 4661 | max_open_files = 3 |
| 4662 |
nothing calls this directly
no test coverage detected