feat: compress with zstd for loading data

cli/src/main.rs

@@ -66,24 +66,20 @@ impl InputFormat {
                 options.insert("field_delimiter", ",");
                 options.insert("quote", "\"");
                 options.insert("skip_header", "0");
-                options.insert("compression", "NONE");
             }
             InputFormat::TSV => {
                 options.insert("type", "TSV");
                 options.insert("record_delimiter", "\n");
                 options.insert("field_delimiter", "\t");
-                options.insert("compression", "NONE");
             }
             InputFormat::NDJSON => {
                 options.insert("type", "NDJSON");
-                options.insert("compression", "NONE");
             }
             InputFormat::Parquet => {
                 options.insert("type", "Parquet");
             }
             InputFormat::XML => {
                 options.insert("type", "XML");
-                options.insert("compression", "NONE");
                 options.insert("row_tag", "row");
             }
         }

driver/Cargo.toml

@@ -30,6 +30,7 @@ databend-client = { workspace = true }
 databend-driver-macros = { workspace = true }
 databend-sql = { workspace = true }
 
+async-compression = { version = "0.4", features = ["tokio", "zstd"] }
 async-trait = "0.1"
 chrono = { version = "0.4.35", default-features = false, features = ["clock"] }
 csv = "1.3"

driver/src/rest_api.rs

@@ -14,14 +14,17 @@
 
 use std::collections::{BTreeMap, VecDeque};
 use std::future::Future;
+use std::io::Cursor;
 use std::path::Path;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+use async_compression::tokio::write::ZstdEncoder;
 use async_trait::async_trait;
 use log::info;
 use tokio::fs::File;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio_stream::Stream;
 
 use databend_client::error::Error as ClientError;
@@ -113,10 +116,32 @@ impl Connection for RestAPIConnection {
             .timestamp_nanos_opt()
             .ok_or_else(|| Error::IO("Failed to get current timestamp".to_string()))?;
         let stage = format!("@~/client/load/{}", now);
-        self.upload_to_stage(&stage, data, size).await?;
-        let file_format_options =
+
+        let mut file_format_options =
             file_format_options.unwrap_or_else(Self::default_file_format_options);
         let copy_options = copy_options.unwrap_or_else(Self::default_copy_options);
+
+        let compressed = match file_format_options.get("compression") {
+            None => {
+                let mut buffer = Vec::new();
+                let mut data = data;
+                let real_size = data.read_to_end(&mut buffer).await?;
+                if real_size != size as usize && size != 0 {
+                    return Err(Error::IO(format!(
+                        "Failed to read all data, expected: {}, read: {}",
+                        size, real_size
+                    )));
+                }
+                let mut encoder = ZstdEncoder::new(Vec::new());
+                encoder.write_all(&buffer).await?;
+                encoder.shutdown().await?;
+                file_format_options.insert("compression", "ZSTD");
+                Box::new(Cursor::new(encoder.into_inner()))
+            }
+            _ => data,
+        };
+
+        self.upload_to_stage(&stage, compressed, size).await?;
         let resp = self
             .client
             .insert_with_stage(sql, &stage, file_format_options, copy_options)