fix some incompatibilities between Go and Rust schemas (#390)

Adds capability to the Rust implementation to handle some of the Arrow
schema incompatibilities identified in
#353

The incompatibilities identified were of two types:
- The Go implementation using dictionary encoding, and Rust not
expecting this arrow type
- The Go implementation using a one binary type (FixedSizeBinary) for
some fields (exemplar trace/span id) and Rust expecting a different one
(Binary).

To flexibly handle the different types, this PR introduces two wrappers
around Arrow Arrays:
- `MaybeDictionaryArrayAccessor` is an enum that can contain either the
native array type, or a dictionary where the value is this type (Note:
we originally had `StringArrayAccessor` that did this for strings, and
this PR just makes the implementation generic so it also covers
primitive arrays, FixedSizeBinary and Binary types).
- `BinaryArrayAccessor` that can contain either `FixedSizeBinary` or
`Binary` array type (or dictionaries containing these types) and be
transparently treated as a `NullableArrayAccessor` that returns
`Vec<u8>`

Note tested using a protobuf message that was created by golang and
consuming using rust
```rs
use otel_arrow_rust::opentelemetry::BatchArrowRecords;
use otel_arrow_rust::Consumer;

use prost::Message;
use std::fs::File;
use std::io::Read;

fn main() {
    // otap_metrics.pb contains a BatchArrowRecords message serialized as protobuf.
    // this was generated from golang including fixes to field nullability from this PR
    // #374
    let mut file = File::open("otap_metrics.pb").unwrap();
    let mut contents = vec![];
    file.read_to_end(&mut contents).unwrap();
    let mut bar = BatchArrowRecords::decode(contents.as_ref()).unwrap();
    let mut consumer = Consumer::default();
    let result = consumer.consume_metrics_batches(&mut bar);
    println!("Result is OK {}", result.is_ok()) // prints "Result is OK true"
}
```

---------

Co-authored-by: Joshua MacDonald <[email protected]>
Co-authored-by: Drew Relmas <[email protected]>

Author: 3 people

rust/otel-arrow-rust/src/arrays.rs

@@ -10,11 +10,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::error;
+use crate::error::{self, InvalidListArraySnafu};
 use arrow::array::{
-    Array, ArrayRef, ArrowPrimitiveType, BinaryArray, BooleanArray, DictionaryArray, Float32Array,
-    Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, PrimitiveArray, RecordBatch,
-    StringArray, TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
+    Array, ArrayRef, ArrowPrimitiveType, BinaryArray, BooleanArray, DictionaryArray,
+    FixedSizeBinaryArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
+    Int64Array, PrimitiveArray, RecordBatch, StringArray, TimestampNanosecondArray, UInt8Array,
+    UInt16Array, UInt32Array, UInt64Array,
 };
 use arrow::datatypes::{ArrowDictionaryKeyType, TimeUnit};
 use arrow::datatypes::{ArrowNativeType, DataType, UInt8Type, UInt16Type};
@@ -95,6 +96,18 @@ impl NullableArrayAccessor for BinaryArray {
     }
 }
 
+impl NullableArrayAccessor for FixedSizeBinaryArray {
+    type Native = Vec<u8>;
+
+    fn value_at(&self, idx: usize) -> Option<Self::Native> {
+        if self.is_valid(idx) {
+            Some(self.value(idx).to_vec())
+        } else {
+            None
+        }
+    }
+}
+
 impl NullableArrayAccessor for StringArray {
     type Native = String;
 
@@ -221,76 +234,190 @@ where
     }
 }
 
-pub type DictionaryStringArrayAccessor<'a, K> = DictionaryArrayAccessor<'a, K, StringArray>;
+/// Wrapper around various arrays that may return a byte slice. Note that
+/// this delegates to the underlying NullableArrayAccessor implementation
+/// for the Arrow array which copies the bytes when value_at is called
+pub enum ByteArrayAccessor<'a> {
+    Binary(MaybeDictArrayAccessor<'a, BinaryArray>),
+    FixedSizeBinary(MaybeDictArrayAccessor<'a, FixedSizeBinaryArray>),
+}
 
-/// [StringArrayAccessor] allows to access string values from [StringArray]s and [DictionaryArray]s.
-pub enum StringArrayAccessor<'a> {
-    /// Plain StringArray
-    String(&'a StringArray),
-    /// DictionaryArray with UInt8 keys and String values.
-    Dictionary8(DictionaryStringArrayAccessor<'a, UInt8Type>),
-    /// DictionaryArray with UInt16 keys and String values.
-    Dictionary16(DictionaryStringArrayAccessor<'a, UInt16Type>),
+impl<'a> ByteArrayAccessor<'a> {
+    pub fn try_new(arr: &'a ArrayRef) -> error::Result<Self> {
+        match arr.data_type() {
+            DataType::Binary => {
+                MaybeDictArrayAccessor::<BinaryArray>::try_new(arr).map(Self::Binary)
+            }
+            DataType::FixedSizeBinary(dims) => {
+                MaybeDictArrayAccessor::<FixedSizeBinaryArray>::try_new(arr, *dims)
+                    .map(Self::FixedSizeBinary)
+            }
+            DataType::Dictionary(_, val) => match **val {
+                DataType::Binary => {
+                    MaybeDictArrayAccessor::<BinaryArray>::try_new(arr).map(Self::Binary)
+                }
+                DataType::FixedSizeBinary(dims) => {
+                    MaybeDictArrayAccessor::<FixedSizeBinaryArray>::try_new(arr, dims)
+                        .map(Self::FixedSizeBinary)
+                }
+                _ => error::UnsupportedDictionaryValueTypeSnafu {
+                    expect_oneof: vec![DataType::Binary, DataType::FixedSizeBinary(-1)],
+                    actual: (**val).clone(),
+                }
+                .fail(),
+            },
+            _ => InvalidListArraySnafu {
+                expect_oneof: vec![
+                    DataType::Binary,
+                    DataType::FixedSizeBinary(-1),
+                    DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Binary)),
+                    DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Binary)),
+                    DataType::Dictionary(
+                        Box::new(DataType::UInt8),
+                        Box::new(DataType::FixedSizeBinary(-1)),
+                    ),
+                    DataType::Dictionary(
+                        Box::new(DataType::UInt16),
+                        Box::new(DataType::FixedSizeBinary(-1)),
+                    ),
+                ],
+                actual: arr.data_type().clone(),
+            }
+            .fail(),
+        }
+    }
 }
 
-impl NullableArrayAccessor for StringArrayAccessor<'_> {
-    type Native = String;
+impl NullableArrayAccessor for ByteArrayAccessor<'_> {
+    type Native = Vec<u8>;
 
     fn value_at(&self, idx: usize) -> Option<Self::Native> {
         match self {
-            StringArrayAccessor::String(s) => s.value_at(idx),
-            StringArrayAccessor::Dictionary8(d) => d.value_at(idx),
-            StringArrayAccessor::Dictionary16(d) => d.value_at(idx),
+            Self::Binary(b) => b.value_at(idx),
+            Self::FixedSizeBinary(b) => b.value_at(idx),
         }
     }
 }
 
-impl<'a> StringArrayAccessor<'a> {
-    pub fn new(a: &'a ArrayRef) -> error::Result<Self> {
-        let result = match a.data_type() {
-            DataType::Utf8 => {
-                // safety: we've checked array data type
-                Self::String(a.as_any().downcast_ref::<StringArray>().unwrap())
-            }
-            DataType::Dictionary(key, v) => {
-                ensure!(
-                    **v == DataType::Utf8,
-                    error::UnsupportedStringColumnTypeSnafu {
-                        data_type: (**v).clone()
-                    }
-                );
-                match **key {
-                    DataType::UInt8 => Self::Dictionary8(DictionaryArrayAccessor::new(
-                        // safety: we've checked the key type
-                        a.as_any()
-                            .downcast_ref::<DictionaryArray<UInt8Type>>()
-                            .unwrap(),
-                    )),
-                    DataType::UInt16 => Self::Dictionary16(DictionaryArrayAccessor::new(
-                        // safety: we've checked the key type
-                        a.as_any()
-                            .downcast_ref::<DictionaryArray<UInt16Type>>()
-                            .unwrap(),
-                    )),
-                    _ => {
-                        return error::UnsupportedStringDictKeyTypeSnafu {
-                            data_type: a.data_type().clone(),
-                        }
-                        .fail();
-                    }
+/// Wrapper around an array that might be a dictionary or it might just be an unencoded
+/// array of the base type
+pub enum MaybeDictArrayAccessor<'a, V> {
+    Native(&'a V),
+    Dictionary8(DictionaryArrayAccessor<'a, UInt8Type, V>),
+    Dictionary16(DictionaryArrayAccessor<'a, UInt16Type, V>),
+}
+
+impl<'a, T> NullableArrayAccessor for MaybeDictArrayAccessor<'a, T>
+where
+    T: Array + NullableArrayAccessor + 'static,
+{
+    type Native = T::Native;
+
+    fn value_at(
+        &self,
+        idx: usize,
+    ) -> Option<<MaybeDictArrayAccessor<'a, T> as NullableArrayAccessor>::Native> {
+        match self {
+            Self::Native(s) => s.value_at(idx),
+            Self::Dictionary8(d) => d.value_at(idx),
+            Self::Dictionary16(d) => d.value_at(idx),
+        }
+    }
+}
+
+impl<'a, T> MaybeDictArrayAccessor<'a, T>
+where
+    T: Array + NullableArrayAccessor + 'static,
+{
+    /// Inspects the given array to determine whether it can be treated as an array
+    /// of the specified data type. The array must either:
+    /// - Directly have the expected data type, or
+    /// - Be a dictionary array whose value type matches the expected data type.
+    ///
+    /// Returns a wrapped native array if the type matches.
+    /// Returns an error if the array type can't be treated as this datatype
+    fn try_new_with_datatype(data_type: DataType, arr: &'a ArrayRef) -> error::Result<Self> {
+        // if the type isn't a dictionary, we treat it as an unencoded array
+        if *arr.data_type() == data_type {
+            return Ok(Self::Native(arr.as_any().downcast_ref::<T>().unwrap()));
+        }
+
+        // determine if the type is a dictionary where the value is the desired datatype
+        if let DataType::Dictionary(key, v) = arr.data_type() {
+            ensure!(
+                **v == data_type,
+                error::UnsupportedDictionaryValueTypeSnafu {
+                    expect_oneof: vec![data_type],
+                    actual: (**v).clone()
                 }
-            }
-            _ => {
-                return error::UnsupportedStringColumnTypeSnafu {
-                    data_type: a.data_type().clone(),
+            );
+
+            let result = match **key {
+                DataType::UInt8 => Self::Dictionary8(DictionaryArrayAccessor::new(
+                    arr.as_any()
+                        .downcast_ref::<DictionaryArray<UInt8Type>>()
+                        .unwrap(),
+                )),
+                DataType::UInt16 => Self::Dictionary16(DictionaryArrayAccessor::new(
+                    arr.as_any()
+                        .downcast_ref::<DictionaryArray<UInt16Type>>()
+                        .unwrap(),
+                )),
+                _ => {
+                    return error::UnsupportedDictionaryKeyTypeSnafu {
+                        expect_oneof: vec![DataType::UInt8, DataType::UInt16],
+                        actual: (**key).clone(),
+                    }
+                    .fail();
                 }
-                .fail();
-            }
-        };
-        Ok(result)
+            };
+
+            return Ok(result);
+        }
+
+        InvalidListArraySnafu {
+            expect_oneof: vec![
+                data_type.clone(),
+                DataType::Dictionary(Box::new(DataType::UInt8), Box::new(data_type.clone())),
+                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(data_type.clone())),
+            ],
+            actual: arr.data_type().clone(),
+        }
+        .fail()
     }
 }
 
+impl<'a, V> MaybeDictArrayAccessor<'a, PrimitiveArray<V>>
+where
+    V: ArrowPrimitiveType,
+{
+    pub fn try_new(arr: &'a ArrayRef) -> error::Result<Self> {
+        Self::try_new_with_datatype(V::DATA_TYPE, arr)
+    }
+}
+
+impl<'a> MaybeDictArrayAccessor<'a, BinaryArray> {
+    pub fn try_new(arr: &'a ArrayRef) -> error::Result<Self> {
+        Self::try_new_with_datatype(BinaryArray::DATA_TYPE, arr)
+    }
+}
+
+impl<'a> MaybeDictArrayAccessor<'a, FixedSizeBinaryArray> {
+    pub fn try_new(arr: &'a ArrayRef, dims: i32) -> error::Result<Self> {
+        Self::try_new_with_datatype(DataType::FixedSizeBinary(dims), arr)
+    }
+}
+
+impl<'a> MaybeDictArrayAccessor<'a, StringArray> {
+    pub fn try_new(arr: &'a ArrayRef) -> error::Result<Self> {
+        Self::try_new_with_datatype(StringArray::DATA_TYPE, arr)
+    }
+}
+
+pub type Int32ArrayAccessor<'a> = MaybeDictArrayAccessor<'a, Int32Array>;
+pub type Int64ArrayAccessor<'a> = MaybeDictArrayAccessor<'a, Int64Array>;
+pub type StringArrayAccessor<'a> = MaybeDictArrayAccessor<'a, StringArray>;
+
 pub struct DictionaryArrayAccessor<'a, K, V>
 where
     K: ArrowDictionaryKeyType,
@@ -311,8 +438,12 @@ where
     }
 
     pub fn value_at(&self, idx: usize) -> Option<V::Native> {
-        let offset = self.inner.key(idx).unwrap();
-        self.value.value_at(offset)
+        if self.inner.is_valid(idx) {
+            let offset = self.inner.key(idx).unwrap();
+            self.value.value_at(offset)
+        } else {
+            None
+        }
     }
 }
 
@@ -327,7 +458,7 @@ mod tests {
     fn test_dictionary_accessor() {
         let expected: DictionaryArray<UInt16Type> = vec!["a", "a", "b", "c"].into_iter().collect();
         let dict = Arc::new(expected) as ArrayRef;
-        let accessor = StringArrayAccessor::new(&dict).unwrap();
+        let accessor = StringArrayAccessor::try_new(&dict).unwrap();
         assert_eq!("a", accessor.value_at(0).unwrap());
         assert_eq!("a", accessor.value_at(1).unwrap());
         assert_eq!("b", accessor.value_at(2).unwrap());

rust/otel-arrow-rust/src/error.rs

@@ -96,9 +96,13 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display("Invalid List array data type, expect {}, actual {}", expect, actual))]
+    #[snafu(display(
+        "Invalid List array data type, expect one of {:?}, actual {}",
+        expect_oneof,
+        actual
+    ))]
     InvalidListArray {
-        expect: DataType,
+        expect_oneof: Vec<DataType>,
         actual: DataType,
         #[snafu(implicit)]
         location: Location,
@@ -139,6 +143,30 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display(
+        "Unsupported dictionary key type, expect one of {:?}, actual {}",
+        expect_oneof,
+        actual
+    ))]
+    UnsupportedDictionaryKeyType {
+        expect_oneof: Vec<DataType>,
+        actual: DataType,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display(
+        "Unsupported dictionary value type. expect {:?}, actual {}",
+        expect_oneof,
+        actual
+    ))]
+    UnsupportedDictionaryValueType {
+        expect_oneof: Vec<DataType>,
+        actual: DataType,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Unsupported string column type, given: {}", data_type))]
     UnsupportedStringColumnType {
         data_type: DataType,

rust/otel-arrow-rust/src/otlp/attributes/parent_id.rs

@@ -10,42 +10,30 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::arrays::NullableArrayAccessor;
 use crate::otlp::attributes::decoder::{
     Attrs16ParentIdDecoder, Attrs32ParentIdDecoder, AttrsParentIdDecoder,
 };
-use arrow::array::{UInt16Array, UInt32Array};
-use arrow::datatypes::DataType;
+use arrow::datatypes::{UInt16Type, UInt32Type};
 use num_enum::TryFromPrimitive;
 use std::hash::Hash;
 use std::ops::{Add, AddAssign};
 
 pub trait ParentId: Copy + Hash + Eq + Default + Add<Output = Self> + AddAssign {
-    type Array: NullableArrayAccessor<Native = Self> + 'static;
-
-    fn arrow_data_type() -> DataType;
+    type ArrayType;
 
     fn new_decoder() -> AttrsParentIdDecoder<Self>;
 }
 
 impl ParentId for u16 {
-    type Array = UInt16Array;
-
-    fn arrow_data_type() -> DataType {
-        DataType::UInt16
-    }
+    type ArrayType = UInt16Type;
 
     fn new_decoder() -> AttrsParentIdDecoder<Self> {
         Attrs16ParentIdDecoder::default()
     }
 }
 
 impl ParentId for u32 {
-    type Array = UInt32Array;
-
-    fn arrow_data_type() -> DataType {
-        DataType::UInt32
-    }
+    type ArrayType = UInt32Type;
 
     fn new_decoder() -> AttrsParentIdDecoder<Self> {
         Attrs32ParentIdDecoder::default()