[arrow-datafusion] branch master updated: Add print format param with support for tsv print format to datafusion cli (#292)

[dh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dheres pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new dcf6c11  Add print format param with support for tsv print format to datafusion cli (#292)
dcf6c11 is described below

commit dcf6c110e7309e867e4b2fc1ac75cf9313095667
Author: Jiayu Liu <Ji...@users.noreply.github.com>
AuthorDate: Tue May 11 05:49:33 2021 +0800

    Add print format param with support for tsv print format to datafusion cli (#292)
    
    * add csv mode to datafusion cli
    
    * adding tsv format
    
    * use Self whereas possible
    
    * prune import
---
 datafusion-cli/src/format/print_format.rs | 41 ++++++++++++++++++-------------
 datafusion-cli/src/main.rs                | 10 ++++----
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/datafusion-cli/src/format/print_format.rs b/datafusion-cli/src/format/print_format.rs
index 921e29f..a9fc56b 100644
--- a/datafusion-cli/src/format/print_format.rs
+++ b/datafusion-cli/src/format/print_format.rs
@@ -26,38 +26,45 @@ use std::str::FromStr;
 #[derive(Debug, Clone)]
 pub enum PrintFormat {
     Csv,
+    Tsv,
     Table,
 }
 
 impl FromStr for PrintFormat {
     type Err = ();
-    fn from_str(s: &str) -> std::result::Result<PrintFormat, ()> {
+    fn from_str(s: &str) -> std::result::Result<Self, ()> {
         match s {
-            "csv" => Ok(PrintFormat::Csv),
-            "table" => Ok(PrintFormat::Table),
+            "csv" => Ok(Self::Csv),
+            "tsv" => Ok(Self::Tsv),
+            "table" => Ok(Self::Table),
             _ => Err(()),
         }
     }
 }
 
+fn print_batches_with_sep(batches: &[RecordBatch], delimiter: u8) -> Result<String> {
+    let mut bytes = vec![];
+    {
+        let builder = WriterBuilder::new()
+            .has_headers(true)
+            .with_delimiter(delimiter);
+        let mut writer = builder.build(&mut bytes);
+        for batch in batches {
+            writer.write(batch)?;
+        }
+    }
+    let formatted = String::from_utf8(bytes)
+        .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+    Ok(formatted)
+}
+
 impl PrintFormat {
     /// print the batches to stdout using the specified format
     pub fn print_batches(&self, batches: &[RecordBatch]) -> Result<()> {
         match self {
-            PrintFormat::Csv => {
-                let mut bytes = vec![];
-                {
-                    let builder = WriterBuilder::new().has_headers(true);
-                    let mut writer = builder.build(&mut bytes);
-                    for batch in batches {
-                        writer.write(batch)?;
-                    }
-                }
-                let csv = String::from_utf8(bytes)
-                    .map_err(|e| DataFusionError::Execution(e.to_string()))?;
-                println!("{}", csv);
-            }
-            PrintFormat::Table => pretty::print_batches(batches)?,
+            Self::Csv => println!("{}", print_batches_with_sep(batches, b',')?),
+            Self::Tsv => println!("{}", print_batches_with_sep(batches, b'\t')?),
+            Self::Table => pretty::print_batches(batches)?,
         }
         Ok(())
     }
diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs
index f15c1d8..52d3ccc 100644
--- a/datafusion-cli/src/main.rs
+++ b/datafusion-cli/src/main.rs
@@ -66,7 +66,7 @@ pub async fn main() {
         )
         .arg(
             Arg::with_name("format")
-                .help("Output format (possible values: table, csv)")
+                .help("Output format (possible values: table, csv, tsv)")
                 .long("format")
                 .default_value("table")
                 .validator(is_valid_format)
@@ -183,10 +183,10 @@ async fn exec_from_repl(execution_config: ExecutionConfig, print_format: PrintFo
 }
 
 fn is_valid_format(format: String) -> std::result::Result<(), String> {
-    match format.to_lowercase().as_str() {
-        "csv" => Ok(()),
-        "table" => Ok(()),
-        _ => Err(format!("Format '{}' not supported", format)),
+    if format.parse::<PrintFormat>().is_ok() {
+        Ok(())
+    } else {
+        Err(format!("Format '{}' not supported", format))
     }
 }