Add print format param with support for tsv print format to datafusio…

…n cli (#292) * add csv mode to datafusion cli * adding tsv format * use Self whereas possible * prune import
apache · May 10, 2021 · dcf6c11 · dcf6c11
1 parent 11634f3
commit dcf6c11
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 22 deletions.
diff --git 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
@@ -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))
     }
 }