vibhatha commented on a change in pull request #11911:
URL: https://github.com/apache/arrow/pull/11911#discussion_r771765893
##########
File path: python/pyarrow/tests/test_dataset.py
##########
@@ -3621,6 +3621,191 @@ def compare_tables_ignoring_order(t1, t2):
assert not extra_file.exists()
+def _generate_random_int_array(size=4, min=1, max=10):
+ return np.random.randint(min, max, size)
+
+
+def _generate_data_and_columns(num_of_columns, num_of_records):
+ data = []
+ column_names = []
+ for i in range(num_of_columns):
+ data.append(_generate_random_int_array(size=num_of_records,
+ min=1,
+ max=num_of_records))
+ column_names.append("c" + str(i))
+ record_batch = pa.record_batch(data=data, names=column_names)
+ return record_batch
+
+
+def _get_num_of_files_generated(base_directory, file_format):
+ return len(list(pathlib.Path(base_directory).glob(f'**/*.{file_format}')))
+
+
+def test_write_dataset_max_rows_per_file(tempdir):
+ directory = tempdir / 'ds'
+ max_rows_per_file = 10
+ max_rows_per_group = 10
+ num_of_columns = 2
+ num_of_records = 35
+
+ record_batch = _generate_data_and_columns(num_of_columns,
+ num_of_records)
+
+ ds.write_dataset(record_batch, directory, format="parquet",
+ max_rows_per_file=max_rows_per_file,
+ max_rows_per_group=max_rows_per_group)
+
+ files_in_dir = os.listdir(directory)
+
+ # number of partitions with max_rows and the partition with the remainder
+ expected_partitions = num_of_records // max_rows_per_file + 1
+
+ # test whether the expected amount of files are written
+ assert len(files_in_dir) == expected_partitions
+
+ # compute the number of rows per each file written
+ result_row_combination = []
+ for _, f_file in enumerate(files_in_dir):
+ f_path = directory / str(f_file)
+ dataset = ds.dataset(f_path, format="parquet")
Review comment:
Thanks a lot for taking time and reviewing my code very thoroughly. I
appreciate that 🙏
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
