pr: fix column behavior for short files

Fix a bug where lines of the input file were not correctly distributed
to the columns of the output page if there were fewer lines than the
total number of cells in the output table. For example, before this
commit,

    printf "a\nb\n" | pr -2

would incorrectly produce an output page like this:

    a
    b

After this commit, it produces a more correct output page like this:

    a    b

src/uu/pr/src/pr.rs

@@ -1048,6 +1048,9 @@ fn to_table_merged(
 }
 
 /// Group lines of the file in columns, going top-to-bottom then left-to-right.
+///
+/// This function should be applied when there are more lines than the
+/// total number of cells in the table.
 fn to_table(
     content_lines_per_page: usize,
     columns: usize,
@@ -1062,6 +1065,26 @@ fn to_table(
         .collect()
 }
 
+/// Group lines of the file in columns, going top-to-bottom then left-to-right.
+///
+/// This function should be applied when there are fewer lines than the
+/// total number of cells in the table.
+fn to_table_short_file(
+    content_lines_per_page: usize,
+    columns: usize,
+    lines: &[FileLine],
+) -> Vec<Vec<Option<&FileLine>>> {
+    let num_rows = lines.len() / columns;
+    let mut table: Vec<Vec<_>> = (0..num_rows)
+        .map(|i| (0..columns).map(|j| lines.get(num_rows * j + i)).collect())
+        .collect();
+    // Fill the rest with Nones.
+    for _ in num_rows..content_lines_per_page {
+        table.push(vec![None; columns]);
+    }
+    table
+}
+
 #[allow(clippy::cognitive_complexity)]
 fn write_columns(
     lines: &[FileLine],
@@ -1112,7 +1135,9 @@ fn write_columns(
     // cells, where each row will be printed as a single line in the
     // output.
     let merge = options.merge_files_print.is_some();
-    let table = if across_mode {
+    let table = if !merge && (lines.len() < (content_lines_per_page * columns)) {
+        to_table_short_file(content_lines_per_page, columns, lines)
+    } else if across_mode {
         to_table_across(content_lines_per_page, columns, lines)
     } else if merge {
         to_table_merged(content_lines_per_page, columns, filled_lines)

tests/by-util/test_pr.rs

@@ -667,6 +667,29 @@ fn test_form_feed_followed_by_new_line() {
         .stdout_matches(&regex);
 }
 
+#[test]
+fn test_columns() {
+    let whitespace = " ".repeat(50);
+    let datetime_pattern = r"\d\d\d\d-\d\d-\d\d \d\d:\d\d";
+    let header = format!("\n\n{datetime_pattern}{whitespace}Page 1\n\n\n");
+    // TODO Our output still does not match the behavior of GNU
+    // pr. The correct output should be:
+    //
+    //     "a\t\t\t\t    b\n";
+    //
+    let data = "a                                  \tb                                  \n";
+    let blank_lines_60 = "\n".repeat(60);
+    let pattern = format!("{header}{data}{blank_lines_60}");
+    let regex = Regex::new(&pattern).unwrap();
+
+    // Command line: `printf "a\nb\n" | pr -2`.
+    new_ucmd!()
+        .arg("-2")
+        .pipe_in("a\nb\n")
+        .succeeds()
+        .stdout_matches(&regex);
+}
+
 #[test]
 fn test_merge() {
     // Create the two files to merge.