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