feat: support specifying compression when unloading to parquet. (#17664)

youngsofun · web-flow · commit 73862939fec7 · 2025-03-28T10:39:04.000+08:00
compression = zstd | snappy
default is zstd.
diff --git a/src/meta/app/src/principal/file_format.rs b/src/meta/app/src/principal/file_format.rs
@@ -109,7 +109,7 @@ impl FileFormatParams {
             FileFormatParams::NdJson(v) => v.compression,
             FileFormatParams::Json(v) => v.compression,
             FileFormatParams::Xml(v) => v.compression,
-            FileFormatParams::Parquet(_) => StageFileCompression::None,
+            FileFormatParams::Parquet(v) => v.compression,
             FileFormatParams::Orc(_) => StageFileCompression::None,
             FileFormatParams::Avro(_) => StageFileCompression::None,
         }
@@ -153,18 +153,18 @@ impl FileFormatParams {
             StageFileFormatType::Xml => {
                 let default = XmlFileFormatParams::default();
                 let row_tag = reader.take_string(OPT_ROW_TAG, default.row_tag);
-                let compression = reader.take_compression()?;
+                let compression = reader.take_compression_default_none()?;
                 FileFormatParams::Xml(XmlFileFormatParams {
                     compression,
                     row_tag,
                 })
             }
             StageFileFormatType::Json => {
-                let compression = reader.take_compression()?;
+                let compression = reader.take_compression_default_none()?;
                 FileFormatParams::Json(JsonFileFormatParams { compression })
             }
             StageFileFormatType::NdJson => {
-                let compression = reader.take_compression()?;
+                let compression = reader.take_compression_default_none()?;
                 let missing_field_as = reader.options.remove(MISSING_FIELD_AS);
                 let null_field_as = reader.options.remove(NULL_FIELD_AS);
                 let null_if = parse_null_if(reader.options.remove(NULL_IF))?;
@@ -176,7 +176,7 @@ impl FileFormatParams {
                 )?)
             }
             StageFileFormatType::Avro => {
-                let compression = reader.take_compression()?;
+                let compression = reader.take_compression_default_none()?;
                 let missing_field_as = reader.options.remove(MISSING_FIELD_AS);
                 let null_if = parse_null_if(reader.options.remove(NULL_IF))?;
                 FileFormatParams::Avro(AvroFileFormatParams::try_create(
@@ -186,9 +186,11 @@ impl FileFormatParams {
                 )?)
             }
             StageFileFormatType::Parquet => {
+                let compression = reader.take_compression(StageFileCompression::Zstd)?;
                 let missing_field_as = reader.options.remove(MISSING_FIELD_AS);
                 let null_if = parse_null_if(reader.options.remove(NULL_IF))?;
                 FileFormatParams::Parquet(ParquetFileFormatParams::try_create(
+                    compression,
                     missing_field_as.as_deref(),
                     null_if,
                 )?)
@@ -201,7 +203,7 @@ impl FileFormatParams {
             }
             StageFileFormatType::Csv => {
                 let default = CsvFileFormatParams::default();
-                let compression = reader.take_compression()?;
+                let compression = reader.take_compression_default_none()?;
                 let headers = reader.take_u64(OPT_SKIP_HEADER, default.headers)?;
                 let field_delimiter =
                     reader.take_string(OPT_FIELD_DELIMITER, default.field_delimiter);
@@ -246,7 +248,7 @@ impl FileFormatParams {
             }
             StageFileFormatType::Tsv => {
                 let default = TsvFileFormatParams::default();
-                let compression = reader.take_compression()?;
+                let compression = reader.take_compression_default_none()?;
                 let headers = reader.take_u64(OPT_SKIP_HEADER, default.headers)?;
                 let field_delimiter =
                     reader.take_string(OPT_FIELD_DELIMITER, default.field_delimiter);
@@ -329,6 +331,7 @@ impl FileFormatParams {
 impl Default for FileFormatParams {
     fn default() -> Self {
         FileFormatParams::Parquet(ParquetFileFormatParams {
+            compression: StageFileCompression::Zstd,
             missing_field_as: NullAs::Error,
             null_if: vec![],
         })
@@ -384,10 +387,13 @@ impl FileFormatOptionsReader {
         }
     }
 
-    fn take_compression(&mut self) -> Result<StageFileCompression> {
+    fn take_compression_default_none(&mut self) -> Result<StageFileCompression> {
+        self.take_compression(StageFileCompression::None)
+    }
+    fn take_compression(&mut self, default: StageFileCompression) -> Result<StageFileCompression> {
         match self.options.remove("compression") {
             Some(c) => StageFileCompression::from_str(&c).map_err(ErrorCode::IllegalFileFormat),
-            None => Ok(StageFileCompression::None),
+            None => Ok(default),
         }
     }
 
@@ -752,16 +758,41 @@ impl AvroFileFormatParams {
     }
 }
 
-#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ParquetFileFormatParams {
+    // used only for unload
+    pub compression: StageFileCompression,
     pub missing_field_as: NullAs,
     pub null_if: Vec<String>,
 }
 
+impl Default for ParquetFileFormatParams {
+    fn default() -> Self {
+        Self {
+            compression: StageFileCompression::Zstd,
+            missing_field_as: Default::default(),
+            null_if: Default::default(),
+        }
+    }
+}
+
 impl ParquetFileFormatParams {
-    pub fn try_create(missing_field_as: Option<&str>, null_if: Vec<String>) -> Result<Self> {
+    pub fn try_create(
+        compression: StageFileCompression,
+        missing_field_as: Option<&str>,
+        null_if: Vec<String>,
+    ) -> Result<Self> {
+        if !matches!(
+            compression,
+            StageFileCompression::Zstd | StageFileCompression::Snappy
+        ) {
+            return Err(ErrorCode::InvalidArgument(format!(
+                "compression algorithm {compression} not supported, only support Zstd and Snappy."
+            )));
+        }
         let missing_field_as = NullAs::parse(missing_field_as, MISSING_FIELD_AS, NullAs::Error)?;
         Ok(Self {
+            compression,
             missing_field_as,
             null_if,
         })
diff --git a/src/meta/proto-conv/src/file_format_from_to_protobuf_impl.rs b/src/meta/proto-conv/src/file_format_from_to_protobuf_impl.rs
@@ -21,6 +21,7 @@ use databend_common_io::GeometryDataType;
 use databend_common_meta_app as mt;
 use databend_common_meta_app::principal::BinaryFormat;
 use databend_common_meta_app::principal::EmptyFieldAs;
+use databend_common_meta_app::principal::StageFileCompression;
 use databend_common_protos::pb;
 use num::FromPrimitive;
 
@@ -335,14 +336,29 @@ impl FromToProto for mt::principal::ParquetFileFormatParams {
     fn from_pb(p: pb::ParquetFileFormatParams) -> Result<Self, Incompatible>
     where Self: Sized {
         reader_check_msg(p.ver, p.min_reader_ver)?;
-        mt::principal::ParquetFileFormatParams::try_create(p.missing_field_as.as_deref(), p.null_if)
-            .map_err(|e| Incompatible::new(format!("{e}")))
+        let mut compression = mt::principal::StageFileCompression::from_pb_enum(
+            FromPrimitive::from_i32(p.compression).ok_or_else(|| {
+                Incompatible::new(format!("invalid StageFileCompression: {}", p.compression))
+            })?,
+        )?;
+        if compression == StageFileCompression::Auto {
+            compression = StageFileCompression::Zstd;
+        };
+        mt::principal::ParquetFileFormatParams::try_create(
+            compression,
+            p.missing_field_as.as_deref(),
+            p.null_if,
+        )
+        .map_err(|e| Incompatible::new(format!("{e}")))
     }
 
     fn to_pb(&self) -> Result<pb::ParquetFileFormatParams, Incompatible> {
+        let compression =
+            mt::principal::StageFileCompression::to_pb_enum(&self.compression)? as i32;
         Ok(pb::ParquetFileFormatParams {
             ver: VER,
             min_reader_ver: MIN_READER_VER,
+            compression,
             missing_field_as: Some(self.missing_field_as.to_string()),
             null_if: self.null_if.clone(),
         })
diff --git a/src/meta/proto-conv/src/util.rs b/src/meta/proto-conv/src/util.rs
@@ -152,6 +152,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (120, "2025-02-11: Add: Add new UserPrivilege CreateWarehouse and new OwnershipObject::Warehouse"),
     (121, "2025-03-03: Add: Add new FileFormat AvroFileFormatParams"),
     (122, "2025-03-11: Add: table_meta and virtual_data_schema"),
+    (123, "2025-03-27: Add: add compression in user.proto/ParquetFileFormatParam"),
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)
diff --git a/src/meta/proto-conv/tests/it/main.rs b/src/meta/proto-conv/tests/it/main.rs
@@ -117,3 +117,4 @@ mod v119_virtual_column;
 mod v120_warehouse_ownershipobject;
 mod v121_avro_format_params;
 mod v122_virtual_schema;
+mod v123_parquet_format_params;
diff --git a/src/meta/proto-conv/tests/it/v032_file_format_params.rs b/src/meta/proto-conv/tests/it/v032_file_format_params.rs
@@ -138,6 +138,7 @@ fn test_decode_v32_parquet_file_format_params() -> anyhow::Result<()> {
 
     let want = || {
         mt::principal::FileFormatParams::Parquet(ParquetFileFormatParams {
+            compression: StageFileCompression::Zstd,
             missing_field_as: Default::default(),
             null_if: vec![],
         })
diff --git a/src/meta/proto-conv/tests/it/v066_stage_create_on.rs b/src/meta/proto-conv/tests/it/v066_stage_create_on.rs
@@ -15,6 +15,7 @@
 use chrono::DateTime;
 use chrono::Utc;
 use databend_common_meta_app as mt;
+use databend_common_meta_app::principal::StageFileCompression;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::storage::StorageParams;
 use databend_common_meta_app::storage::StorageS3Config;
@@ -55,6 +56,7 @@ fn test_decode_v66_stage() -> anyhow::Result<()> {
         is_temporary: false,
         file_format_params: mt::principal::FileFormatParams::Parquet(
             mt::principal::ParquetFileFormatParams {
+                compression: StageFileCompression::Zstd,
                 missing_field_as: Default::default(),
                 null_if: vec![],
             },
diff --git a/src/meta/proto-conv/tests/it/v093_parquet_format_params.rs b/src/meta/proto-conv/tests/it/v093_parquet_format_params.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use databend_common_meta_app::principal::ParquetFileFormatParams;
+use databend_common_meta_app::principal::StageFileCompression;
 use fastrace::func_name;
 
 use crate::common;
@@ -31,6 +32,7 @@ use crate::common;
 fn test_decode_v93_parquet_file_format_params() -> anyhow::Result<()> {
     let parquet_file_format_params_v93 = vec![34, 0, 34, 1, 97, 160, 6, 93, 168, 6, 24];
     let want = || ParquetFileFormatParams {
+        compression: StageFileCompression::Zstd,
         missing_field_as: Default::default(),
         null_if: vec!["".to_string(), "a".to_string()],
     };
diff --git a/src/meta/proto-conv/tests/it/v099_parquet_format_params.rs b/src/meta/proto-conv/tests/it/v099_parquet_format_params.rs
@@ -14,6 +14,7 @@
 
 use databend_common_meta_app::principal::NullAs;
 use databend_common_meta_app::principal::ParquetFileFormatParams;
+use databend_common_meta_app::principal::StageFileCompression;
 use fastrace::func_name;
 
 use crate::common;
@@ -35,6 +36,7 @@ fn test_decode_v99_parquet_file_format_params() -> anyhow::Result<()> {
         168, 6, 24,
     ];
     let want = || ParquetFileFormatParams {
+        compression: StageFileCompression::Zstd,
         missing_field_as: NullAs::FieldDefault,
         null_if: vec!["".to_string(), "a".to_string()],
     };
diff --git a/src/meta/proto-conv/tests/it/v123_parquet_format_params.rs b/src/meta/proto-conv/tests/it/v123_parquet_format_params.rs
@@ -0,0 +1,52 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_meta_app::principal::NullAs;
+use databend_common_meta_app::principal::ParquetFileFormatParams;
+use databend_common_meta_app::principal::StageFileCompression;
+use fastrace::func_name;
+
+use crate::common;
+
+// These bytes are built when a new version in introduced,
+
+// and are kept for backward compatibility test.
+//
+// *************************************************************
+// * These messages should never be updated,                   *
+// * only be added when a new version is added,                *
+// * or be removed when an old version is no longer supported. *
+// *************************************************************
+//
+#[test]
+fn test_decode_v123_parquet_file_format_params() -> anyhow::Result<()> {
+    let parquet_file_format_params_v123 = vec![
+        10, 13, 70, 73, 69, 76, 68, 95, 68, 69, 70, 65, 85, 76, 84, 16, 8, 34, 0, 34, 1, 97, 160,
+        6, 123, 168, 6, 24,
+    ];
+
+    let want = || ParquetFileFormatParams {
+        compression: StageFileCompression::Snappy,
+        missing_field_as: NullAs::FieldDefault,
+        null_if: vec!["".to_string(), "a".to_string()],
+    };
+    common::test_load_old(
+        func_name!(),
+        parquet_file_format_params_v123.as_slice(),
+        123,
+        want(),
+    )?;
+    common::test_pb_from_to(func_name!(), want())?;
+    Ok(())
+}
diff --git a/src/meta/protos/proto/file_format.proto b/src/meta/protos/proto/file_format.proto
@@ -93,6 +93,7 @@ message ParquetFileFormatParams {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;
   optional string missing_field_as = 1;
+  StageFileCompression compression = 2;
   repeated string null_if = 4;
 }
 
diff --git a/src/query/ast/src/parser/stage.rs b/src/query/ast/src/parser/stage.rs
@@ -103,7 +103,7 @@ pub fn format_options(i: Input) -> IResult<FileFormatOptions> {
 
     let option_compression = map(
         rule! {
-            COMPRESSION ~ "=" ~ ( AUTO | NONE | GZIP | BZ2 | BROTLI | ZSTD | DEFLATE | RAWDEFLATE | XZ )
+            COMPRESSION ~ "=" ~ ( AUTO | NONE | GZIP | BZ2 | BROTLI | ZSTD | DEFLATE | RAWDEFLATE | XZ | SNAPPY)
         },
         |(_, _, v)| {
             (
diff --git a/src/query/storages/stage/src/append/parquet_file/writer_processor.rs b/src/query/storages/stage/src/append/parquet_file/writer_processor.rs
diff --git a/tests/sqllogictests/suites/stage/formats/parquet/options/parquet_compression.test b/tests/sqllogictests/suites/stage/formats/parquet/options/parquet_compression.test
diff --git a/tests/sqllogictests/suites/stage/formats/parquet/options/select_options.test b/tests/sqllogictests/suites/stage/formats/parquet/options/select_options.test

Original file line number	Diff line number	Diff line change
`@@ -93,6 +93,7 @@ message ParquetFileFormatParams {`
`93`	`93`	`uint64 ver = 100;`
`94`	`94`	`uint64 min_reader_ver = 101;`
`95`	`95`	`optional string missing_field_as = 1;`
	`96`	`+ StageFileCompression compression = 2;`
`96`	`97`	`repeated string null_if = 4;`
`97`	`98`	`}`
`98`	`99`