Merge branch 'main' into cijothomas/perf-k8s

Author: lquerel

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

@@ -14,11 +14,12 @@ use crate::error::{self, InvalidListArraySnafu};
 use arrow::array::{
     Array, ArrayRef, ArrowPrimitiveType, BinaryArray, BooleanArray, DictionaryArray,
     FixedSizeBinaryArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
-    Int64Array, PrimitiveArray, RecordBatch, StringArray, TimestampNanosecondArray, UInt8Array,
-    UInt16Array, UInt32Array, UInt64Array,
+    Int64Array, PrimitiveArray, RecordBatch, StringArray, StructArray, TimestampNanosecondArray,
+    UInt8Array, UInt16Array, UInt32Array, UInt64Array,
+};
+use arrow::datatypes::{
+    ArrowDictionaryKeyType, ArrowNativeType, DataType, TimeUnit, UInt8Type, UInt16Type,
 };
-use arrow::datatypes::{ArrowDictionaryKeyType, TimeUnit};
-use arrow::datatypes::{ArrowNativeType, DataType, UInt8Type, UInt16Type};
 use paste::paste;
 use snafu::{OptionExt, ensure};
 
@@ -139,9 +140,7 @@ macro_rules! impl_downcast {
 
               pub fn [<get_ $suffix _array> ]<'a>(rb: &'a RecordBatch, name: &str) -> error::Result<&'a $array_type> {
                 use arrow::datatypes::DataType::*;
-                let arr = rb.column_by_name(name).context(error::ColumnNotFoundSnafu {
-            name,
-        })?;
+                let arr = get_required_array(rb, name)?;
 
                  arr.as_any()
                             .downcast_ref::<$array_type>()
@@ -177,6 +176,17 @@ impl_downcast!(
     TimestampNanosecondArray
 );
 
+/// Get reference to array that the caller requires to be in the record batch.
+/// If the column is not in the record batch, returns `ColumnNotFound` error
+pub fn get_required_array<'a>(
+    record_batch: &'a RecordBatch,
+    column_name: &str,
+) -> error::Result<&'a ArrayRef> {
+    record_batch
+        .column_by_name(column_name)
+        .context(error::ColumnNotFoundSnafu { name: column_name })
+}
+
 trait NullableInt64ArrayAccessor {
     fn i64_at(&self, idx: usize) -> error::Result<Option<i64>>;
 }
@@ -243,6 +253,13 @@ pub enum ByteArrayAccessor<'a> {
 }
 
 impl<'a> ByteArrayAccessor<'a> {
+    pub fn try_new_for_column(
+        record_batch: &'a RecordBatch,
+        column_name: &str,
+    ) -> error::Result<Self> {
+        Self::try_new(get_required_array(record_batch, column_name)?)
+    }
+
     pub fn try_new(arr: &'a ArrayRef) -> error::Result<Self> {
         match arr.data_type() {
             DataType::Binary => {
@@ -339,7 +356,11 @@ where
     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()));
+            return Ok(Self::Native(
+                arr.as_any()
+                    .downcast_ref::<T>()
+                    .expect("array can be downcast to it's native datatype"),
+            ));
         }
 
         // determine if the type is a dictionary where the value is the desired datatype
@@ -356,13 +377,13 @@ where
                 DataType::UInt8 => Self::Dictionary8(DictionaryArrayAccessor::new(
                     arr.as_any()
                         .downcast_ref::<DictionaryArray<UInt8Type>>()
-                        .unwrap(),
-                )),
+                        .expect("array can be downcast to DictionaryArray<UInt8Type"),
+                )?),
                 DataType::UInt16 => Self::Dictionary16(DictionaryArrayAccessor::new(
                     arr.as_any()
                         .downcast_ref::<DictionaryArray<UInt16Type>>()
-                        .unwrap(),
-                )),
+                        .expect("array can be downcast to DictionaryArray<UInt16Type>"),
+                )?),
                 _ => {
                     return error::UnsupportedDictionaryKeyTypeSnafu {
                         expect_oneof: vec![DataType::UInt8, DataType::UInt16],
@@ -394,6 +415,13 @@ where
     pub fn try_new(arr: &'a ArrayRef) -> error::Result<Self> {
         Self::try_new_with_datatype(V::DATA_TYPE, arr)
     }
+
+    pub fn try_new_for_column(
+        record_batch: &'a RecordBatch,
+        column_name: &str,
+    ) -> error::Result<Self> {
+        Self::try_new(get_required_array(record_batch, column_name)?)
+    }
 }
 
 impl<'a> MaybeDictArrayAccessor<'a, BinaryArray> {
@@ -412,6 +440,13 @@ 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 fn try_new_for_column(
+        record_batch: &'a RecordBatch,
+        column_name: &str,
+    ) -> error::Result<Self> {
+        Self::try_new(get_required_array(record_batch, column_name)?)
+    }
 }
 
 pub type Int32ArrayAccessor<'a> = MaybeDictArrayAccessor<'a, Int32Array>;
@@ -431,22 +466,109 @@ where
     K: ArrowDictionaryKeyType,
     V: Array + NullableArrayAccessor + 'static,
 {
-    pub fn new(a: &'a DictionaryArray<K>) -> Self {
-        let dict = a.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
-        let value = dict.values().as_any().downcast_ref::<V>().unwrap();
-        Self { inner: dict, value }
+    pub fn new(dict: &'a DictionaryArray<K>) -> error::Result<Self> {
+        let value = dict
+            .values()
+            .as_any()
+            .downcast_ref::<V>()
+            .with_context(|| error::InvalidListArraySnafu {
+                expect_oneof: Vec::new(),
+                actual: dict.values().data_type().clone(),
+            })?;
+        Ok(Self { inner: dict, value })
     }
 
     pub fn value_at(&self, idx: usize) -> Option<V::Native> {
         if self.inner.is_valid(idx) {
-            let offset = self.inner.key(idx).unwrap();
+            let offset = self
+                .inner
+                .key(idx)
+                .expect("dictionary should be valid at index");
             self.value.value_at(offset)
         } else {
             None
         }
     }
 }
 
+/// Helper for accessing columns of a struct array
+///
+/// Methods return various errors into this crate's Error type if
+/// if callers requirments for the struct columns are not met (for
+/// example `ColumnDataTypeMismatchSnafu`)
+pub struct StructColumnAccessor<'a> {
+    inner: &'a StructArray,
+}
+
+impl<'a> StructColumnAccessor<'a> {
+    pub fn new(arr: &'a StructArray) -> Self {
+        Self { inner: arr }
+    }
+
+    pub fn primitive_column<T: ArrowPrimitiveType + 'static>(
+        &self,
+        column_name: &str,
+    ) -> error::Result<&'a PrimitiveArray<T>> {
+        self.primitive_column_op(column_name)?
+            .with_context(|| error::ColumnNotFoundSnafu {
+                name: column_name.to_string(),
+            })
+    }
+
+    pub fn primitive_column_op<T: ArrowPrimitiveType + 'static>(
+        &self,
+        column_name: &str,
+    ) -> error::Result<Option<&'a PrimitiveArray<T>>> {
+        self.inner
+            .column_by_name(column_name)
+            .map(|arr| {
+                arr.as_any()
+                    .downcast_ref::<PrimitiveArray<T>>()
+                    .with_context(|| error::ColumnDataTypeMismatchSnafu {
+                        name: column_name.to_string(),
+                        expect: T::DATA_TYPE,
+                        actual: arr.data_type().clone(),
+                    })
+            })
+            .transpose()
+    }
+
+    pub fn bool_column_op(&self, column_name: &str) -> error::Result<Option<&'a BooleanArray>> {
+        self.inner
+            .column_by_name(column_name)
+            .map(|arr| {
+                arr.as_any()
+                    .downcast_ref()
+                    .with_context(|| error::ColumnDataTypeMismatchSnafu {
+                        name: column_name.to_string(),
+                        expect: DataType::Boolean,
+                        actual: arr.data_type().clone(),
+                    })
+            })
+            .transpose()
+    }
+
+    pub fn string_column_op(
+        &self,
+        column_name: &str,
+    ) -> error::Result<Option<StringArrayAccessor<'a>>> {
+        self.inner
+            .column_by_name(column_name)
+            .map(StringArrayAccessor::try_new)
+            .transpose()
+    }
+
+    pub fn byte_array_column_op(
+        &self,
+        column_name: &str,
+    ) -> error::Result<Option<ByteArrayAccessor<'a>>> {
+        self.inner
+            .column_by_name(column_name)
+            .map(ByteArrayAccessor::try_new)
+            .transpose()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::arrays::{NullableArrayAccessor, StringArrayAccessor};

rust/otel-arrow-rust/src/decode/decoder.rs

@@ -12,9 +12,11 @@
 
 use crate::decode::record_message::RecordMessage;
 use crate::error;
-use crate::otlp::metric::metrics_from;
-use crate::otlp::related_data::RelatedData;
+use crate::otlp::logs::logs_from;
+use crate::otlp::logs::related_data::RelatedData as LogsRelatedData;
+use crate::otlp::metrics::{metrics_from, related_data::RelatedData as MetricsRelatedData};
 use crate::proto::opentelemetry::arrow::v1::{ArrowPayload, ArrowPayloadType, BatchArrowRecords};
+use crate::proto::opentelemetry::collector::logs::v1::ExportLogsServiceRequest;
 use crate::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
 use arrow::array::RecordBatch;
 use arrow::error::ArrowError;
@@ -107,40 +109,55 @@ impl Consumer {
         &mut self,
         records: &mut BatchArrowRecords,
     ) -> error::Result<ExportMetricsServiceRequest> {
-        ensure!(!records.arrow_payloads.is_empty(), error::EmptyBatchSnafu);
-
-        let main_record_type = records.arrow_payloads[0].r#type;
-        let payload_type = ArrowPayloadType::try_from(main_record_type).map_err(|_| {
-            error::UnsupportedPayloadTypeSnafu {
-                actual: main_record_type,
-            }
-            .build()
-        })?;
-        match payload_type {
+        match get_main_payload_type(records)? {
             ArrowPayloadType::UnivariateMetrics => {
                 let record_message = self.consume_bar(records)?;
                 let (mut related_data, metric_record) =
-                    RelatedData::from_record_messages(&record_message)?;
+                    MetricsRelatedData::from_record_messages(&record_message)?;
                 let metric_rec_idx = metric_record.context(error::MetricRecordNotFoundSnafu)?;
                 metrics_from(&record_message[metric_rec_idx].record, &mut related_data)
             }
-
-            ArrowPayloadType::Logs => error::UnsupportedPayloadTypeSnafu {
+            main_record_type => error::UnsupportedPayloadTypeSnafu {
                 actual: main_record_type,
             }
             .fail(),
-            ArrowPayloadType::Spans => error::UnsupportedPayloadTypeSnafu {
-                actual: main_record_type,
+        }
+    }
+
+    pub fn consume_logs_batches(
+        &mut self,
+        records: &mut BatchArrowRecords,
+    ) -> error::Result<ExportLogsServiceRequest> {
+        match get_main_payload_type(records)? {
+            ArrowPayloadType::Logs => {
+                let record_message = self.consume_bar(records)?;
+                let (mut related_data, log_record) =
+                    LogsRelatedData::from_record_messages(&record_message)?;
+                let log_rec_idx = log_record.context(error::LogRecordNotFoundSnafu)?;
+                logs_from(&record_message[log_rec_idx].record, &mut related_data)
             }
-            .fail(),
-            _ => error::UnsupportedPayloadTypeSnafu {
+            main_record_type => error::UnsupportedPayloadTypeSnafu {
                 actual: main_record_type,
             }
             .fail(),
         }
     }
 }
 
+/// Get the main logs, metrics, or traces from a received BatchArrowRecords message.
+fn get_main_payload_type(records: &BatchArrowRecords) -> error::Result<ArrowPayloadType> {
+    ensure!(!records.arrow_payloads.is_empty(), error::EmptyBatchSnafu);
+
+    // Per the specification, the main record type is the first payload
+    let main_record_type = records.arrow_payloads[0].r#type;
+    ArrowPayloadType::try_from(main_record_type).map_err(|_| {
+        error::UnsupportedPayloadTypeSnafu {
+            actual: main_record_type,
+        }
+        .build()
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use crate::test_util::{create_record_batch, create_test_schema};

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

@@ -11,7 +11,7 @@
 // limitations under the License.
 
 use crate::otlp::attributes::store::AttributeValueType;
-use crate::otlp::metric::MetricType;
+use crate::otlp::metrics::MetricType;
 use arrow::datatypes::DataType;
 use arrow::error::ArrowError;
 use num_enum::TryFromPrimitiveError;
@@ -137,6 +137,12 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Log record not found"))]
+    LogRecordNotFound {
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Metric record not found"))]
     MetricRecordNotFound {
         #[snafu(implicit)]

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

@@ -10,9 +10,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod exemplar;
-pub mod metric;
-
 pub mod attributes;
-pub mod data_points;
-pub mod related_data;
+pub mod logs;
+pub mod metrics;
+
+mod common;

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

@@ -11,8 +11,8 @@
 // limitations under the License.
 
 use crate::arrays::{
-    Int64ArrayAccessor, MaybeDictArrayAccessor, NullableArrayAccessor, StringArrayAccessor,
-    get_binary_array_opt, get_bool_array_opt, get_f64_array_opt, get_u8_array,
+    ByteArrayAccessor, Int64ArrayAccessor, MaybeDictArrayAccessor, NullableArrayAccessor,
+    StringArrayAccessor, get_bool_array_opt, get_f64_array_opt, get_u8_array,
 };
 use crate::error;
 use crate::otlp::attributes::parent_id::ParentId;
@@ -80,22 +80,11 @@ where
             .transpose()?;
         let value_type_arr = get_u8_array(rb, consts::ATTRIBUTE_TYPE)?;
 
-        let value_str_arr = StringArrayAccessor::try_new(
-            rb.column_by_name(consts::ATTRIBUTE_STR)
-                .context(error::ColumnNotFoundSnafu {
-                    name: consts::ATTRIBUTE_STR,
-                })?,
-        )?;
-
-        let value_int_arr = Int64ArrayAccessor::try_new(
-            rb.column_by_name(consts::ATTRIBUTE_INT)
-                .context(error::ColumnNotFoundSnafu {
-                    name: consts::ATTRIBUTE_STR,
-                })?,
-        )?;
+        let value_str_arr = StringArrayAccessor::try_new_for_column(rb, consts::ATTRIBUTE_STR)?;
+        let value_int_arr = Int64ArrayAccessor::try_new_for_column(rb, consts::ATTRIBUTE_INT)?;
         let value_double_arr = get_f64_array_opt(rb, consts::ATTRIBUTE_DOUBLE)?;
         let value_bool_arr = get_bool_array_opt(rb, consts::ATTRIBUTE_BOOL)?;
-        let value_bytes_arr = get_binary_array_opt(rb, consts::ATTRIBUTE_BYTES)?;
+        let value_bytes_arr = ByteArrayAccessor::try_new_for_column(rb, consts::ATTRIBUTE_BYTES)?;
 
         for idx in 0..rb.num_rows() {
             let key = key_arr.value_at_or_default(idx);