feat: support load Avro Files. (#17548)

Author: youngsofun

Cargo.lock

@@ -220,6 +220,7 @@ ahash = "0.8"
 aho-corasick = { version = "1.0.1" } #
 anyerror = { version = "=0.1.13" }
 anyhow = { version = "1.0.65" }
+apache-avro = { version = "0.17.0", features = ["snappy", "zstandard", "xz", "snappy", "bzip"] }
 approx = "0.5.1"
 arrow = { version = "53" }
 arrow-array = { version = "53" }

Cargo.toml

@@ -61,6 +61,7 @@ pub enum FileFormatParams {
     Xml(XmlFileFormatParams),
     Parquet(ParquetFileFormatParams),
     Orc(OrcFileFormatParams),
+    Avro(AvroFileFormatParams),
 }
 
 impl FileFormatParams {
@@ -73,6 +74,7 @@ impl FileFormatParams {
             FileFormatParams::Xml(_) => StageFileFormatType::Xml,
             FileFormatParams::Parquet(_) => StageFileFormatType::Parquet,
             FileFormatParams::Orc(_) => StageFileFormatType::Orc,
+            FileFormatParams::Avro(_) => StageFileFormatType::Avro,
         }
     }
 
@@ -90,6 +92,9 @@ impl FileFormatParams {
                 Ok(FileFormatParams::Json(JsonFileFormatParams::default()))
             }
             StageFileFormatType::Orc => Ok(FileFormatParams::Orc(OrcFileFormatParams::default())),
+            StageFileFormatType::Avro => {
+                Ok(FileFormatParams::Avro(AvroFileFormatParams::default()))
+            }
             _ => Err(ErrorCode::IllegalFileFormat(format!(
                 "Unsupported file format type: {:?}",
                 format_type
@@ -106,6 +111,7 @@ impl FileFormatParams {
             FileFormatParams::Xml(v) => v.compression,
             FileFormatParams::Parquet(_) => StageFileCompression::None,
             FileFormatParams::Orc(_) => StageFileCompression::None,
+            FileFormatParams::Avro(_) => StageFileCompression::None,
         }
     }
 
@@ -169,6 +175,16 @@ impl FileFormatParams {
                     null_if,
                 )?)
             }
+            StageFileFormatType::Avro => {
+                let compression = reader.take_compression()?;
+                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(
+                    compression,
+                    missing_field_as.as_deref(),
+                    null_if,
+                )?)
+            }
             StageFileFormatType::Parquet => {
                 let missing_field_as = reader.options.remove(MISSING_FIELD_AS);
                 let null_if = parse_null_if(reader.options.remove(NULL_IF))?;
@@ -690,6 +706,52 @@ impl NdJsonFileFormatParams {
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AvroFileFormatParams {
+    pub compression: StageFileCompression,
+    pub missing_field_as: NullAs,
+    pub null_if: Vec<String>,
+}
+
+impl AvroFileFormatParams {
+    pub fn try_create(
+        compression: StageFileCompression,
+        missing_field_as: Option<&str>,
+        null_if: Vec<String>,
+    ) -> Result<Self> {
+        let missing_field_as = NullAs::parse(missing_field_as, MISSING_FIELD_AS, NullAs::Error)?;
+        if matches!(missing_field_as, NullAs::Null) {
+            return Err(ErrorCode::InvalidArgument(
+                "Invalid option value for Avro: NULL_FIELD_AS is set to NULL. The valid values are ERROR | FIELD_DEFAULT.",
+            ));
+        }
+        Ok(Self {
+            compression,
+            missing_field_as,
+            null_if,
+        })
+    }
+}
+
+impl Default for crate::principal::AvroFileFormatParams {
+    fn default() -> Self {
+        crate::principal::AvroFileFormatParams {
+            compression: StageFileCompression::None,
+            missing_field_as: NullAs::Error,
+            null_if: vec![],
+        }
+    }
+}
+
+impl AvroFileFormatParams {
+    pub fn downcast_unchecked(params: &FileFormatParams) -> &AvroFileFormatParams {
+        match params {
+            FileFormatParams::Avro(p) => p,
+            _ => unreachable!(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ParquetFileFormatParams {
     pub missing_field_as: NullAs,
@@ -776,6 +838,9 @@ impl Display for FileFormatParams {
                     params.compression, params.missing_field_as, params.null_field_as
                 )
             }
+            FileFormatParams::Avro(params) => {
+                write!(f, "TYPE = AVRO, NULL_FIELDS_AS = {}", params.compression,)
+            }
             FileFormatParams::Parquet(params) => {
                 write!(
                     f,

src/meta/app/src/principal/file_format.rs

@@ -192,11 +192,9 @@ impl FromStr for StageFileFormatType {
             "XML" => Ok(StageFileFormatType::Xml),
             "JSON" => Ok(StageFileFormatType::Json),
             "ORC" => Ok(StageFileFormatType::Orc),
-            "AVRO" => Err(format!(
-                "File format type '{s}' not implemented yet', must be one of ( CSV | TSV | NDJSON | PARQUET | ORC)"
-            )),
+            "AVRO" => Ok(StageFileFormatType::Avro),
             _ => Err(format!(
-                "Unknown file format type '{s}', must be one of ( CSV | TSV | NDJSON | PARQUET | ORC)"
+                "Unknown file format type '{s}', must be one of ( CSV | TSV | NDJSON | PARQUET | ORC | AVRO | JSON )"
             )),
         }
     }

src/meta/app/src/principal/user_stage.rs

@@ -247,6 +247,11 @@ impl FromToProto for mt::principal::FileFormatParams {
                     mt::principal::XmlFileFormatParams::from_pb(p)?,
                 ))
             }
+            Some(pb::file_format_params::Format::Avro(p)) => {
+                Ok(mt::principal::FileFormatParams::Avro(
+                    mt::principal::AvroFileFormatParams::from_pb(p)?,
+                ))
+            }
             None => Err(Incompatible::new(
                 "FileFormatParams.format cannot be None".to_string(),
             )),
@@ -275,6 +280,11 @@ impl FromToProto for mt::principal::FileFormatParams {
                     mt::principal::JsonFileFormatParams::to_pb(p)?,
                 )),
             }),
+            Self::Avro(p) => Ok(Self::PB {
+                format: Some(pb::file_format_params::Format::Avro(
+                    mt::principal::AvroFileFormatParams::to_pb(p)?,
+                )),
+            }),
             Self::Tsv(p) => Ok(Self::PB {
                 format: Some(pb::file_format_params::Format::Tsv(
                     mt::principal::TsvFileFormatParams::to_pb(p)?,
@@ -557,3 +567,39 @@ impl FromToProto for mt::principal::TsvFileFormatParams {
         })
     }
 }
+
+impl FromToProto for mt::principal::AvroFileFormatParams {
+    type PB = pb::AvroFileFormatParams;
+    fn get_pb_ver(p: &Self::PB) -> u64 {
+        p.ver
+    }
+
+    fn from_pb(p: pb::AvroFileFormatParams) -> Result<Self, Incompatible>
+    where Self: Sized {
+        reader_check_msg(p.ver, p.min_reader_ver)?;
+        let compression = mt::principal::StageFileCompression::from_pb_enum(
+            FromPrimitive::from_i32(p.compression).ok_or_else(|| {
+                Incompatible::new(format!("invalid StageFileCompression: {}", p.compression))
+            })?,
+        )?;
+
+        mt::principal::AvroFileFormatParams::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::AvroFileFormatParams, Incompatible> {
+        let compression =
+            mt::principal::StageFileCompression::to_pb_enum(&self.compression)? as i32;
+        Ok(pb::AvroFileFormatParams {
+            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(),
+        })
+    }
+}

src/meta/proto-conv/src/file_format_from_to_protobuf_impl.rs

@@ -150,6 +150,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (118, "2025-01-22: Add: config.proto: add user_name in WebhdfsConfig"),
     (119, "2025-01-25: Add: virtual_column add alias_names and auto_generated field"),
     (120, "2025-02-11: Add: Add new UserPrivilege CreateWarehouse and new OwnershipObject::Warehouse"),
+    (121, "2025-03-03: Add: Add new FileFormat AvroFileFormatParams"),
     // 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`)

src/meta/proto-conv/src/util.rs

@@ -118,3 +118,4 @@ mod v117_webhdfs_add_disable_list_batch;
 mod v118_webhdfs_add_user_name;
 mod v119_virtual_column;
 mod v120_warehouse_ownershipobject;
+mod v121_avro_format_params;

src/meta/proto-conv/tests/it/main.rs

@@ -0,0 +1,71 @@
+// 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::AvroFileFormatParams;
+use databend_common_meta_app::principal::FileFormatParams;
+use databend_common_meta_app::principal::NullAs;
+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_v121_avro_file_format_params() -> anyhow::Result<()> {
+    let avro_file_format_params_v121 = vec![
+        8, 4, 18, 13, 70, 73, 69, 76, 68, 95, 68, 69, 70, 65, 85, 76, 84, 26, 4, 110, 117, 108,
+        108, 26, 4, 78, 85, 76, 76, 160, 6, 121, 168, 6, 24,
+    ];
+
+    let want = || AvroFileFormatParams {
+        compression: StageFileCompression::Zstd,
+        missing_field_as: NullAs::FieldDefault,
+        null_if: vec!["null".to_string(), "NULL".to_string()],
+    };
+    common::test_load_old(
+        func_name!(),
+        avro_file_format_params_v121.as_slice(),
+        121,
+        want(),
+    )?;
+    common::test_pb_from_to(func_name!(), want())?;
+    Ok(())
+}
+
+#[test]
+fn test_decode_v121_file_format_params() -> anyhow::Result<()> {
+    let file_format_params_v121 = vec![
+        66, 35, 8, 4, 18, 13, 70, 73, 69, 76, 68, 95, 68, 69, 70, 65, 85, 76, 84, 26, 4, 110, 117,
+        108, 108, 26, 4, 78, 85, 76, 76, 160, 6, 121, 168, 6, 24,
+    ];
+    let want = || {
+        FileFormatParams::Avro(AvroFileFormatParams {
+            compression: StageFileCompression::Zstd,
+            missing_field_as: NullAs::FieldDefault,
+            null_if: vec!["null".to_string(), "NULL".to_string()],
+        })
+    };
+    common::test_load_old(func_name!(), file_format_params_v121.as_slice(), 0, want())?;
+    common::test_pb_from_to(func_name!(), want())?;
+    Ok(())
+}

src/meta/proto-conv/tests/it/v121_avro_format_params.rs

@@ -85,6 +85,7 @@ message FileFormatParams {
     NdJsonFileFormatParams nd_json = 5;
     XmlFileFormatParams xml = 6;
     OrcFileFormatParams orc = 7;
+    AvroFileFormatParams avro = 8;
   }
 }
 
@@ -155,4 +156,12 @@ message OrcFileFormatParams {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;
   optional string missing_field_as = 1;
+}
+
+message AvroFileFormatParams {
+  uint64 ver = 100;
+  uint64 min_reader_ver = 101;
+  StageFileCompression compression = 1;
+  optional string missing_field_as = 2;
+  repeated string null_if = 3;
 }

src/meta/protos/proto/file_format.proto

@@ -120,12 +120,12 @@ pub fn copy_into(i: Input) -> IResult<Statement> {
          #copy_into_location:"`COPY
                 INTO { internalStage | externalStage | externalLocation }
                 FROM { [<database_name>.]<table_name> | ( <query> ) }
-                [ FILE_FORMAT = ( { TYPE = { CSV | JSON | PARQUET | TSV } [ formatTypeOptions ] } ) ]
+                [ FILE_FORMAT = ( { TYPE = { CSV | NDJSON | PARQUET | TSV | AVRO } [ formatTypeOptions ] } ) ]
                 [ copyOptions ]`"
          | #copy_into_table: "`COPY
                 INTO { [<database_name>.]<table_name> { ( <columns> ) } }
                 FROM { internalStage | externalStage | externalLocation | ( <query> ) }
-                [ FILE_FORMAT = ( { TYPE = { CSV | JSON | PARQUET | TSV } [ formatTypeOptions ] } ) ]
+                [ FILE_FORMAT = ( { TYPE = { CSV | NDJSON | PARQUET | TSV | AVRO } [ formatTypeOptions ] } ) ]
                 [ FILES = ( '<file_name>' [ , '<file_name>' ] [ , ... ] ) ]
                 [ PATTERN = '<regex_pattern>' ]
                 [ VALIDATION_MODE = RETURN_ROWS ]