feat: copy into location support compression. (#13067)

Author: youngsofun

Cargo.lock

@@ -134,7 +134,6 @@ impl Encode for CompressCodec {
 }
 
 impl CompressCodec {
-    #[allow(unused)]
     pub fn compress_all(&mut self, to_compress: &[u8]) -> common_exception::Result<Vec<u8>> {
         let mut compress_bufs = vec![];
         let mut input = PartialBuffer::new(to_compress);

src/common/compress/src/encode.rs

@@ -25,3 +25,4 @@ pub use decode::DecompressCodec;
 pub use decode::DecompressDecoder;
 pub use decode::DecompressReader;
 pub use decode::DecompressState;
+pub use encode::CompressCodec;

src/common/compress/src/lib.rs

@@ -14,6 +14,7 @@ test = false
 [dependencies]
 common-base = { path = "../../../common/base" }
 common-catalog = { path = "../../catalog" }
+common-compress = { path = "../../../common/compress" }
 common-exception = { path = "../../../common/exception" }
 common-expression = { path = "../../expression" }
 common-formats = { path = "../../formats" }

src/query/storages/stage/Cargo.toml

@@ -113,7 +113,13 @@ impl Processor for ParquetFileSink {
     #[async_backtrace::framed]
     async fn async_process(&mut self) -> Result<()> {
         assert!(!self.output_data.is_empty());
-        let path = unload_path(&self.table_info, &self.uuid, self.group_id, self.batch_id);
+        let path = unload_path(
+            &self.table_info,
+            &self.uuid,
+            self.group_id,
+            self.batch_id,
+            None,
+        );
         let data = mem::take(&mut self.output_data);
         self.data_accessor.write(&path, data).await?;
         self.batch_id += 1;

src/query/storages/stage/src/parquet_file/sink_processor.rs

@@ -19,6 +19,7 @@ use common_catalog::table_context::TableContext;
 use common_exception::Result;
 use common_formats::FileFormatOptionsExt;
 use common_pipeline_core::Pipeline;
+use common_pipeline_sources::input_formats::InputContext;
 use opendal::Operator;
 
 use crate::row_based_file::limit_file_size_processor::LimitFileSizeProcessor;
@@ -44,6 +45,11 @@ pub(crate) fn append_data_to_row_based_files(
         table_info.schema(),
         table_info.stage_info.file_format_params.clone(),
     )?;
+    let compression = table_info
+        .stage_info
+        .file_format_params
+        .clone()
+        .compression();
     let prefix = output_format.serialize_prefix()?;
 
     pipeline.add_transform(|input, output| {
@@ -62,6 +68,9 @@ pub(crate) fn append_data_to_row_based_files(
     if max_file_size != usize::MAX {
         pipeline.try_resize(max_threads)?;
     }
+
+    let compression = InputContext::get_compression_alg_copy(compression, "")?;
+
     pipeline.add_sink(|input| {
         let gid = group_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
         RowBasedFileSink::try_create(
@@ -71,6 +80,7 @@ pub(crate) fn append_data_to_row_based_files(
             prefix.clone(),
             uuid.clone(),
             gid,
+            compression,
         )
     })?;
     Ok(())

src/query/storages/stage/src/row_based_file/pipeline.rs

@@ -18,6 +18,8 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use common_catalog::plan::StageTableInfo;
+use common_compress::CompressAlgorithm;
+use common_compress::CompressCodec;
 use common_exception::Result;
 use common_expression::BlockMetaInfoDowncast;
 use common_expression::DataBlock;
@@ -45,6 +47,8 @@ pub struct RowBasedFileSink {
     uuid: String,
     group_id: usize,
     batch_id: usize,
+
+    compression: Option<CompressAlgorithm>,
 }
 
 impl RowBasedFileSink {
@@ -55,6 +59,7 @@ impl RowBasedFileSink {
         prefix: Vec<u8>,
         uuid: String,
         group_id: usize,
+        compression: Option<CompressAlgorithm>,
     ) -> Result<ProcessorPtr> {
         Ok(ProcessorPtr::create(Box::new(RowBasedFileSink {
             table_info,
@@ -66,6 +71,7 @@ impl RowBasedFileSink {
             group_id,
             batch_id: 0,
             output_data: vec![],
+            compression,
         })))
     }
 }
@@ -110,13 +116,22 @@ impl Processor for RowBasedFileSink {
         for b in buffers.buffers {
             output.extend_from_slice(b.as_slice());
         }
+        if let Some(compression) = self.compression {
+            output = CompressCodec::from(compression).compress_all(&output)?;
+        }
         self.output_data = output;
         Ok(())
     }
 
     #[async_backtrace::framed]
     async fn async_process(&mut self) -> Result<()> {
-        let path = unload_path(&self.table_info, &self.uuid, self.group_id, self.batch_id);
+        let path = unload_path(
+            &self.table_info,
+            &self.uuid,
+            self.group_id,
+            self.batch_id,
+            self.compression,
+        );
         let data = mem::take(&mut self.output_data);
         self.data_accessor.write(&path, data).await?;
         self.batch_id += 1;

src/query/storages/stage/src/row_based_file/sink_processor.rs

@@ -31,6 +31,7 @@ use common_catalog::plan::StageTableInfo;
 use common_catalog::table::AppendMode;
 use common_catalog::table::Table;
 use common_catalog::table_context::TableContext;
+use common_compress::CompressAlgorithm;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::BlockThresholds;
@@ -296,24 +297,29 @@ pub fn unload_path(
     uuid: &str,
     group_id: usize,
     batch_id: usize,
+    compression: Option<CompressAlgorithm>,
 ) -> String {
     let format_name = format!(
         "{:?}",
         stage_table_info.stage_info.file_format_params.get_type()
     )
     .to_ascii_lowercase();
 
+    let suffix: &str = &compression
+        .map(|c| format!(".{}", c.extension()))
+        .unwrap_or_default();
+
     let path = &stage_table_info.files_info.path;
 
     if path.ends_with("data_") {
         format!(
-            "{}{}_{:0>4}_{:0>8}.{}",
-            path, uuid, group_id, batch_id, format_name
+            "{}{}_{:0>4}_{:0>8}.{}{}",
+            path, uuid, group_id, batch_id, format_name, suffix
         )
     } else {
         format!(
-            "{}/data_{}_{:0>4}_{:0>8}.{}",
-            path, uuid, group_id, batch_id, format_name
+            "{}/data_{}_{:0>4}_{:0>8}.{}{}",
+            path, uuid, group_id, batch_id, format_name, suffix
         )
     }
 }

src/query/storages/stage/src/stage_table.rs

@@ -0,0 +1,77 @@
+# need to run with '-p 0'
+
+statement ok
+drop stage if exists unload;
+
+statement ok
+create stage unload;
+
+statement ok
+remove @unload;
+
+statement ok
+drop table if exists ii;
+
+statement ok
+create table ii (a int, b int);
+
+statement ok
+insert into ii values (1, 2), (3, 4), (5, 6);
+
+statement ok
+create file format if not exists csv_gzip type=csv compression=gzip;
+
+# test csv
+statement ok
+copy into @unload from ii file_format=(type=csv);
+-----
+
+query 
+select right(name, 4), size from list_stage(location=>'@unload');
+----
+.csv 12
+
+query 
+select $1, $2 from @unload(file_format=>'csv');
+----
+1 2
+3 4
+5 6
+
+# test csv_gzip
+statement ok
+remove @unload;
+
+statement ok
+copy into @unload from ii file_format=(format_name='csv_gzip');
+
+query 
+select right(name, 7), size from list_stage(location=>'@unload');
+----
+.csv.gz 32
+
+query 
+select $1, $2 from @unload(file_format => 'csv_gzip');
+----
+1 2
+3 4
+5 6
+
+# test tsv
+statement ok
+remove @unload;
+
+statement ok
+copy into @unload from ii file_format=(format_name='tsv');
+
+query 
+select right(name, 4), size from list_stage(location=>'@unload');
+----
+.tsv 12
+
+query 
+select $1, $2 from @unload(file_format => 'tsv');
+----
+1 2
+3 4
+5 6