From acf57c803448983def12dd1926bfd92099913d44 Mon Sep 17 00:00:00 2001 From: zhuqi-lucas <821684824@qq.com> Date: Fri, 28 Feb 2025 00:38:55 +0800 Subject: [PATCH 1/3] BUG: fix datafusion-cli missing case with only one batch and smaller than max_rows size --- datafusion-cli/src/print_options.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index 092483faed63f..f0bd7fa474179 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -126,6 +126,7 @@ impl PrintOptions { while let Some(batch) = stream.next().await { let batch = batch?; + println!("batch: {:?}", batch); let batch_rows = batch.num_rows(); if !max_rows_reached && total_count < max_rows { @@ -190,13 +191,19 @@ impl PrintOptions { let widths = print_options .format .compute_column_widths(&preview_batches, schema.clone())?; - precomputed_widths = Some(widths); + precomputed_widths = Some(widths.clone()); if !header_printed { print_options.format.print_header( &schema, precomputed_widths.as_ref().unwrap(), writer, )?; + header_printed = true; + } + 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 { From 1d6b50f5bfbe9ed59897a86055bd85deac1c8021 Mon Sep 17 00:00:00 2001 From: zhuqi-lucas <821684824@qq.com> Date: Fri, 28 Feb 2025 08:24:27 +0800 Subject: [PATCH 2/3] Fix --- datafusion-cli/src/print_options.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index f0bd7fa474179..8dd7ca9c8081d 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -126,7 +126,6 @@ impl PrintOptions { while let Some(batch) = stream.next().await { let batch = batch?; - println!("batch: {:?}", batch); let batch_rows = batch.num_rows(); if !max_rows_reached && total_count < max_rows { @@ -152,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( @@ -191,19 +189,20 @@ impl PrintOptions { let widths = print_options .format .compute_column_widths(&preview_batches, schema.clone())?; - precomputed_widths = Some(widths.clone()); + precomputed_widths = Some(widths); if !header_printed { print_options.format.print_header( &schema, precomputed_widths.as_ref().unwrap(), writer, )?; - header_printed = true; } for preview_batch in preview_batches.drain(..) { - print_options - .format - .print_batch_with_widths(&preview_batch, precomputed_widths.as_ref().unwrap(), writer)?; + print_options.format.print_batch_with_widths( + &preview_batch, + precomputed_widths.as_ref().unwrap(), + writer, + )?; } } if let Some(ref widths) = precomputed_widths { From 3ba5943f98bc39a979c8a44ebc6aa0327e3c2c1d Mon Sep 17 00:00:00 2001 From: zhuqi-lucas <821684824@qq.com> Date: Fri, 28 Feb 2025 12:57:09 +0800 Subject: [PATCH 3/3] Add test case --- datafusion-cli/tests/cli_integration.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 27cabf15afecb..55c5d53e80984 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/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,