AVRO-3901: [Rust] Unit tests and impl for better union support (#2583)

* AVRO-3901: Unit tests and impl for better union support * Apply suggestions from code review Rename test cases Co-authored-by: Martin Grigorov <[email protected]> * Cleanup from review * AVRO-3901: Fix formatting, clippy errors and improve naming Signed-off-by: Martin Tzvetanov Grigorov <[email protected]> --------- Signed-off-by: Martin Tzvetanov Grigorov <[email protected]> Co-authored-by: Simon Gittins <[email protected]> Co-authored-by: Martin Grigorov <[email protected]> Co-authored-by: Martin Tzvetanov Grigorov <[email protected]>
apache · Nov 14, 2023 · 85ddfcd · 85ddfcd
1 parent da98719
commit 85ddfcd
Show file tree

Hide file tree

Showing 2 changed files with 372 additions and 3 deletions.
diff --git a/lang/rust/avro/src/encode.rs b/lang/rust/avro/src/encode.rs
@@ -19,7 +19,7 @@ use crate::{
     decimal::serialize_big_decimal,
     schema::{
         DecimalSchema, EnumSchema, FixedSchema, Name, Namespace, RecordSchema, ResolvedSchema,
-        Schema, SchemaKind,
+        Schema, SchemaKind, UnionSchema,
     },
     types::{Value, ValueKind},
     util::{zig_i32, zig_i64},
@@ -71,7 +71,19 @@ pub(crate) fn encode_internal<S: Borrow<Schema>>(
     }
 
     match value {
-        Value::Null => (),
+        Value::Null => {
+            if let Schema::Union(union) = schema {
+                match union.schemas.iter().position(|sch| *sch == Schema::Null) {
+                    None => {
+                        return Err(Error::EncodeValueAsSchemaError {
+                            value_kind: ValueKind::Null,
+                            supported_schema: vec![SchemaKind::Null, SchemaKind::Union],
+                        })
+                    }
+                    Some(p) => encode_long(p as i64, buffer),
+                }
+            } // else {()}
+        }
         Value::Boolean(b) => buffer.push(u8::from(*b)),
         // Pattern | Pattern here to signify that these _must_ have the same encoding.
         Value::Int(i) | Value::Date(i) | Value::TimeMillis(i) => encode_int(*i, buffer),
@@ -242,11 +254,26 @@ pub(crate) fn encode_internal<S: Borrow<Schema>>(
                         ));
                     }
                 }
+            } else if let Schema::Union(UnionSchema { schemas, .. }) = schema {
+                let original_size = buffer.len();
+                for (index, schema) in schemas.iter().enumerate() {
+                    encode_long(index as i64, buffer);
+                    match encode_internal(value, schema, names, enclosing_namespace, buffer) {
+                        Ok(_) => return Ok(()),
+                        Err(_) => {
+                            buffer.truncate(original_size); //undo any partial encoding
+                        }
+                    }
+                }
+                return Err(Error::EncodeValueAsSchemaError {
+                    value_kind: ValueKind::Record,
+                    supported_schema: vec![SchemaKind::Record, SchemaKind::Union],
+                });
             } else {
                 error!("invalid schema type for Record: {:?}", schema);
                 return Err(Error::EncodeValueAsSchemaError {
                     value_kind: ValueKind::Record,
-                    supported_schema: vec![SchemaKind::Record],
+                    supported_schema: vec![SchemaKind::Record, SchemaKind::Union],
                 });
             }
         }