feat(codecs): Support more avro types

changelog.d/24773_support_more_avro_types.fix.md

@@ -0,0 +1,3 @@
+The Avro codec now supports encoding and decoding Date, Fixed, TimeMillis, TimestampNanos, LocalTimestampNanos values.
+
+authors: omwbennett

lib/codecs/Cargo.toml

@@ -63,6 +63,7 @@ vector-config-macros = { path = "../vector-config-macros", default-features = fa
 vector-core = { path = "../vector-core", default-features = false, features = ["vrl"] }
 vector-vrl-functions.workspace = true
 toml = { version = "0.9.8", optional = true }
+uuid.workspace = true
 
 [dev-dependencies]
 criterion.workspace = true
@@ -74,7 +75,6 @@ similar-asserts = "1.7.0"
 vector-core = { path = "../vector-core", default-features = false, features = ["vrl", "test"] }
 rstest = "0.26.1"
 tracing-test = "0.2.6"
-uuid.workspace = true
 vrl.workspace = true
 
 [features]

lib/codecs/src/decoding/format/avro.rs

@@ -91,11 +91,8 @@ impl From<&AvroDeserializerOptions> for AvroSerializerOptions {
 pub struct AvroDeserializerOptions {
     /// The Avro schema definition.
     /// **Note**: The following [`apache_avro::types::Value`] variants are *not* supported:
-    /// * `Date`
     /// * `Decimal`
     /// * `Duration`
-    /// * `Fixed`
-    /// * `TimeMillis`
     #[configurable(metadata(
         docs::examples = r#"{ "type": "record", "name": "log", "fields": [{ "name": "message", "type": "string" }] }"#,
         docs::additional_props_description = r#"Supports most avro data types, unsupported data types includes
@@ -189,10 +186,8 @@ pub fn try_from(value: AvroValue) -> vector_common::Result<VrlValue> {
             Ok(VrlValue::Array(vector))
         }
         AvroValue::Boolean(boolean) => Ok(VrlValue::from(boolean)),
-        AvroValue::Bytes(bytes) => Ok(VrlValue::from(bytes)),
-        AvroValue::Date(_) => Err(vector_common::Error::from(
-            "AvroValue::Date is not supported",
-        )),
+        AvroValue::Bytes(bytes) => Ok(VrlValue::Bytes(Bytes::from(bytes))),
+        AvroValue::Date(days) => Ok(VrlValue::from(days)),
         AvroValue::Decimal(_) => Err(vector_common::Error::from(
             "AvroValue::Decimal is not supported",
         )),
@@ -201,9 +196,7 @@ pub fn try_from(value: AvroValue) -> vector_common::Result<VrlValue> {
             "AvroValue::Duration is not supported",
         )),
         AvroValue::Enum(_, string) => Ok(VrlValue::from(string)),
-        AvroValue::Fixed(_, _) => Err(vector_common::Error::from(
-            "AvroValue::Fixed is not supported",
-        )),
+        AvroValue::Fixed(_, bytes) => Ok(VrlValue::Bytes(Bytes::from(bytes))),
         AvroValue::Float(float) => Ok(VrlValue::from_f64_or_zero(float as f64)),
         AvroValue::Int(int) => Ok(VrlValue::from(int)),
         AvroValue::Long(long) => Ok(VrlValue::from(long)),
@@ -220,9 +213,7 @@ pub fn try_from(value: AvroValue) -> vector_common::Result<VrlValue> {
             .map(|v| VrlValue::Object(v.into_iter().collect())),
         AvroValue::String(string) => Ok(VrlValue::from(string)),
         AvroValue::TimeMicros(time_micros) => Ok(VrlValue::from(time_micros)),
-        AvroValue::TimeMillis(_) => Err(vector_common::Error::from(
-            "AvroValue::TimeMillis is not supported",
-        )),
+        AvroValue::TimeMillis(time_millis) => Ok(VrlValue::from(time_millis)),
         AvroValue::TimestampMicros(ts_micros) => Ok(VrlValue::from(ts_micros)),
         AvroValue::TimestampMillis(ts_millis) => Ok(VrlValue::from(ts_millis)),
         AvroValue::Union(_, v) => try_from(*v),
@@ -232,12 +223,8 @@ pub fn try_from(value: AvroValue) -> vector_common::Result<VrlValue> {
         AvroValue::BigDecimal(_) => Err(vector_common::Error::from(
             "AvroValue::BigDecimal is not supported",
         )),
-        AvroValue::TimestampNanos(_) => Err(vector_common::Error::from(
-            "AvroValue::TimestampNanos is not supported",
-        )),
-        AvroValue::LocalTimestampNanos(_) => Err(vector_common::Error::from(
-            "AvroValue::LocalTimestampNanos is not supported",
-        )),
+        AvroValue::TimestampNanos(ts_nanos) => Ok(VrlValue::from(ts_nanos)),
+        AvroValue::LocalTimestampNanos(ts_nanos) => Ok(VrlValue::from(ts_nanos)),
     }
 }
 

lib/codecs/src/encoding/format/avro.rs

@@ -6,6 +6,109 @@ use vector_core::{config::DataType, event::Event, schema};
 
 use crate::encoding::BuildError;
 
+type AvroValue = apache_avro::types::Value;
+
+/// `apache_avro::to_value` may serialize VRL values into Avro types which later
+/// cannot be resolved against certain Avro types
+/// (e.g. VRL integer (i64) -> Avro `Long` which cannot be resolved to Avro `Date`).
+/// `coerce_logical_types` does a recursive pre-pass to fix such cases.
+fn coerce_logical_types(
+    value: AvroValue,
+    schema: &apache_avro::Schema,
+) -> vector_common::Result<AvroValue> {
+    use apache_avro::Schema;
+    match (value, schema) {
+        (AvroValue::Long(days), Schema::Date) => {
+            i32::try_from(days).map(AvroValue::Date).map_err(|_| {
+                vector_common::Error::from(format!(
+                    "Avro date value {days} is out of range for i32"
+                ))
+            })
+        }
+        (AvroValue::Long(millis), Schema::TimeMillis) => i32::try_from(millis)
+            .map(AvroValue::TimeMillis)
+            .map_err(|_| {
+                vector_common::Error::from(format!(
+                    "Avro time-millis value {millis} is out of range for i32"
+                ))
+            }),
+        (AvroValue::Record(fields), Schema::Record(record_schema)) => {
+            let fields = fields
+                .into_iter()
+                .map(|(name, value)| {
+                    let value = match record_schema.lookup.get(&name) {
+                        Some(index) => {
+                            let field_schema = &record_schema.fields[*index].schema;
+                            coerce_logical_types(value, field_schema)?
+                        }
+                        None => value,
+                    };
+                    Ok((name, value))
+                })
+                .collect::<vector_common::Result<Vec<_>>>()?;
+            Ok(AvroValue::Record(fields))
+        }
+        (AvroValue::Map(entries), Schema::Record(record_schema)) => {
+            let entries = entries
+                .into_iter()
+                .map(|(name, value)| {
+                    let value = match record_schema.lookup.get(&name) {
+                        Some(index) => {
+                            let field_schema = &record_schema.fields[*index].schema;
+                            coerce_logical_types(value, field_schema)?
+                        }
+                        None => value,
+                    };
+                    Ok((name, value))
+                })
+                .collect::<vector_common::Result<_>>()?;
+            Ok(AvroValue::Map(entries))
+        }
+        (AvroValue::Array(items), Schema::Array(array_schema)) => items
+            .into_iter()
+            .map(|item| coerce_logical_types(item, &array_schema.items))
+            .collect::<Result<Vec<_>, _>>()
+            .map(AvroValue::Array),
+        (AvroValue::Map(entries), Schema::Map(map_schema)) => entries
+            .into_iter()
+            .map(|(key, value)| {
+                coerce_logical_types(value, &map_schema.types).map(|value| (key, value))
+            })
+            .collect::<vector_common::Result<_>>()
+            .map(AvroValue::Map),
+        (AvroValue::Union(index, value), Schema::Union(union_schema)) => {
+            let schema = union_schema
+                .variants()
+                .get(index as usize)
+                .unwrap_or(schema);
+            coerce_logical_types(*value, schema)
+                .map(|value| AvroValue::Union(index, Box::new(value)))
+        }
+        (value, Schema::Union(union_schema)) => {
+            if let Ok(resolved) = value.clone().resolve(schema) {
+                return Ok(resolved);
+            }
+
+            let mut last_err = None;
+            for (index, variant) in union_schema.variants().iter().enumerate() {
+                match coerce_logical_types(value.clone(), variant) {
+                    Ok(coerced) if coerced.clone().resolve(variant).is_ok() => {
+                        return Ok(AvroValue::Union(index as u32, Box::new(coerced)));
+                    }
+                    Ok(_) => {}
+                    Err(err) => last_err = Some(err),
+                }
+            }
+
+            match last_err {
+                Some(err) => Err(err),
+                None => Ok(value),
+            }
+        }
+        (value, _) => Ok(value),
+    }
+}
+
 /// Config used to build a `AvroSerializer`.
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct AvroSerializerConfig {
@@ -71,6 +174,7 @@ impl Encoder<Event> for AvroSerializer {
     fn encode(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Self::Error> {
         let log = event.into_log();
         let value = apache_avro::to_value(log)?;
+        let value = coerce_logical_types(value, &self.schema)?;
         let value = value.resolve(&self.schema)?;
         let bytes = apache_avro::to_avro_datum(&self.schema, value)?;
         buffer.put_slice(&bytes);
@@ -113,4 +217,158 @@ mod tests {
 
         assert_eq!(bytes.freeze(), b"\0\x06bar".as_slice());
     }
+
+    #[test]
+    fn coerce_date_fields_recursively() {
+        let schema = apache_avro::Schema::parse_str(indoc! {r#"
+            {
+                "type": "record",
+                "name": "Outer",
+                "fields": [
+                    {
+                        "name": "direct_date",
+                        "type": {"type": "int", "logicalType": "date"}
+                    },
+                    {
+                        "name": "inner",
+                        "type": {
+                            "type": "record",
+                            "name": "Inner",
+                            "fields": [
+                                {
+                                    "name": "date",
+                                    "type": {"type": "int", "logicalType": "date"}
+                                }
+                            ]
+                        }
+                    },
+                    {
+                        "name": "record_as_map",
+                        "type": {
+                            "type": "record",
+                            "name": "MapBackedInner",
+                            "fields": [
+                                {
+                                    "name": "date",
+                                    "type": {"type": "int", "logicalType": "date"}
+                                }
+                            ]
+                        }
+                    },
+                    {
+                        "name": "date_array",
+                        "type": {
+                            "type": "array",
+                            "items": {"type": "int", "logicalType": "date"}
+                        }
+                    },
+                    {
+                        "name": "date_map",
+                        "type": {
+                            "type": "map",
+                            "values": {"type": "int", "logicalType": "date"}
+                        }
+                    },
+                    {
+                        "name": "union_date",
+                        "type": ["null", {"type": "int", "logicalType": "date"}]
+                    },
+                    {
+                        "name": "fallback_union_date",
+                        "type": [
+                            "null",
+                            {"type": "int", "logicalType": "date"},
+                            "long"
+                        ]
+                    },
+                    {
+                        "name": "logical_only_union_date",
+                        "type": [
+                            "null",
+                            {"type": "int", "logicalType": "date"}
+                        ]
+                    }
+                ]
+            }
+        "#})
+        .unwrap();
+        let value = AvroValue::Record(vec![
+            ("direct_date".to_owned(), AvroValue::Long(20_000)),
+            (
+                "inner".to_owned(),
+                AvroValue::Record(vec![("date".to_owned(), AvroValue::Long(20_001))]),
+            ),
+            (
+                "record_as_map".to_owned(),
+                AvroValue::Map(
+                    [("date".to_owned(), AvroValue::Long(20_002))]
+                        .into_iter()
+                        .collect(),
+                ),
+            ),
+            (
+                "date_array".to_owned(),
+                AvroValue::Array(vec![AvroValue::Long(20_003), AvroValue::Long(20_004)]),
+            ),
+            (
+                "date_map".to_owned(),
+                AvroValue::Map(
+                    [
+                        ("first".to_owned(), AvroValue::Long(20_005)),
+                        ("second".to_owned(), AvroValue::Long(20_006)),
+                    ]
+                    .into_iter()
+                    .collect(),
+                ),
+            ),
+            ("union_date".to_owned(), AvroValue::Long(20_007)),
+            ("fallback_union_date".to_owned(), AvroValue::Long(20_009)),
+            (
+                "logical_only_union_date".to_owned(),
+                AvroValue::Long(20_008),
+            ),
+        ]);
+
+        let value = coerce_logical_types(value, &schema).unwrap();
+        let value = value.resolve(&schema).unwrap();
+
+        assert!(matches!(
+            value,
+            AvroValue::Record(fields) if {
+                matches!(fields[0].1, AvroValue::Date(20_000))
+                    && matches!(
+                        &fields[1].1,
+                        AvroValue::Record(inner) if matches!(inner[0].1, AvroValue::Date(20_001))
+                    )
+                    && matches!(
+                        &fields[2].1,
+                        AvroValue::Record(inner) if matches!(inner[0].1, AvroValue::Date(20_002))
+                    )
+                    && matches!(
+                        &fields[3].1,
+                        AvroValue::Array(items)
+                            if matches!(items.as_slice(), [AvroValue::Date(20_003), AvroValue::Date(20_004)])
+                    )
+                    && matches!(
+                        &fields[4].1,
+                        AvroValue::Map(entries)
+                            if matches!(entries.get("first"), Some(AvroValue::Date(20_005)))
+                                && matches!(entries.get("second"), Some(AvroValue::Date(20_006)))
+                    )
+                    && matches!(
+                        &fields[5].1,
+                        AvroValue::Union(1, value) if matches!(value.as_ref(), AvroValue::Date(20_007))
+                    )
+                    && matches!(
+                        &fields[6].1,
+                        AvroValue::Union(2, value)
+                            if matches!(value.as_ref(), AvroValue::Long(20_009))
+                    )
+                    && matches!(
+                        &fields[7].1,
+                        AvroValue::Union(1, value) if matches!(value.as_ref(), AvroValue::Date(20_008))
+                    )
+            }
+        ));
+    }
 }

lib/codecs/tests/avro.rs

@@ -21,9 +21,7 @@ use vector_core::{config::LogNamespace, event::Event};
 #[case(true)]
 #[case(false)]
 fn roundtrip_avro_fixtures(
-    #[files("tests/data/avro/generated/*.avro")]
-    #[exclude(".*(date|fixed|time_millis).avro")]
-    path: PathBuf,
+    #[files("tests/data/avro/generated/*.avro")] path: PathBuf,
     #[case] reserialize: bool,
 ) {
     let schema_path = path.as_path().with_extension("avsc");