Fix: New Datafusion-cli streaming printing way should handle corner case for only one small batch which lines are less than max_rows (#14921)

* BUG: fix datafusion-cli missing case with only one batch and smaller than max_rows size

* Fix

* Add test case

Author: zhuqi-lucas

datafusion-cli/src/print_options.rs

@@ -151,7 +151,6 @@ impl PrintOptions {
                             print_options
                                 .format
                                 .print_header(&schema, &widths, writer)?;
-                            header_printed = true;
                         }
                         for preview_batch in preview_batches.drain(..) {
                             print_options.format.print_batch_with_widths(
@@ -198,6 +197,13 @@ impl PrintOptions {
                         writer,
                     )?;
                 }
+                for preview_batch in preview_batches.drain(..) {
+                    print_options.format.print_batch_with_widths(
+                        &preview_batch,
+                        precomputed_widths.as_ref().unwrap(),
+                        writer,
+                    )?;
+                }
             }
             if let Some(ref widths) = precomputed_widths {
                 print_options.format.print_bottom_border(widths, writer)?;

datafusion-cli/tests/cli_integration.rs

@@ -47,6 +47,12 @@ fn init() {
     ["--command", "show datafusion.execution.batch_size", "--format", "json", "-q", "-b", "1"],
     "[{\"name\":\"datafusion.execution.batch_size\",\"value\":\"1\"}]\n"
 )]
+
+/// Add case fixed issue: https://github.com/apache/datafusion/issues/14920
+#[case::exec_from_commands(
+    ["--command", "SELECT * FROM generate_series(1, 5) t1(v1) ORDER BY v1 DESC;", "--format", "table", "-q"],
+    "+----+\n| v1 |\n+----+\n| 5  |\n| 4  |\n| 3  |\n| 2  |\n| 1  |\n+----+\n"
+)]
 #[test]
 fn cli_quick_test<'a>(
     #[case] args: impl IntoIterator<Item = &'a str>,