c

Author: JakeDern

rust/otap-dataflow/crates/pdata/src/otap.rs

@@ -184,7 +184,7 @@ impl From<Traces> for OtapArrowRecords {
 /// storing and retrieving Arrow record batches in a type-safe manner. It is
 /// implemented by various structs that represent each signal type and provides
 /// methods to efficiently set and get record batches.
-pub trait OtapBatchStore: Default + Clone {
+pub trait OtapBatchStore: Default + Clone + Into<OtapArrowRecords> {
     /// Internally, implementers should use a bitmask for the types they support.
     /// The offsets in the bitmask should correspond to the ArrowPayloadType enum values.
     const TYPE_MASK: u64;

rust/otap-dataflow/crates/pdata/src/otap/testing.rs

@@ -26,15 +26,18 @@ use crate::schema::schema::{DataType as OtapDataType, SimpleType};
 /// Known ID column paths that need plain encoding metadata.
 pub const ID_COLUMN_PATHS: &[&str] = &["id", "resource.id", "scope.id", "parent_id"];
 
-/// Build a `[Option<RecordBatch>; Logs::COUNT]` from payload‑type / column specs.
+/// Build a [`Logs`] store from payload‑type / column specs.
 ///
 /// Missing required columns are automatically filled in using the schema spec.
+/// Use `.into()` to convert to [`OtapArrowRecords`], or `.into_batches()` to
+/// get the raw `[Option<RecordBatch>; Logs::COUNT]` array.
 ///
 /// ```ignore
-/// logs!(
+/// let logs: Logs = logs!(
 ///     (Logs, ("id", UInt16, vec![0u16, 1, 2])),
 ///     (LogAttrs, ("parent_id", UInt16, vec![0u16, 1, 2])),
-/// )
+/// );
+/// let records: OtapArrowRecords = logs.into();
 /// ```
 #[macro_export]
 macro_rules! logs {
@@ -54,9 +57,11 @@ macro_rules! logs {
 }
 pub use logs;
 
-/// Build a `[Option<RecordBatch>; Traces::COUNT]` from payload‑type / column specs.
+/// Build a [`Traces`] store from payload‑type / column specs.
 ///
 /// Missing required columns are automatically filled in using the schema spec.
+/// Use `.into()` to convert to [`OtapArrowRecords`], or `.into_batches()` to
+/// get the raw `[Option<RecordBatch>; Traces::COUNT]` array.
 #[macro_export]
 macro_rules! traces {
     ($(($payload:ident, $($record_batch_args:tt)*)),* $(,)?) => {
@@ -75,9 +80,11 @@ macro_rules! traces {
 }
 pub use traces;
 
-/// Build a `[Option<RecordBatch>; Metrics::COUNT]` from payload‑type / column specs.
+/// Build a [`Metrics`] store from payload‑type / column specs.
 ///
 /// Missing required columns are automatically filled in using the schema spec.
+/// Use `.into()` to convert to [`OtapArrowRecords`], or `.into_batches()` to
+/// get the raw `[Option<RecordBatch>; Metrics::COUNT]` array.
 #[macro_export]
 macro_rules! metrics {
     ($(($payload:ident, $($record_batch_args:tt)*)),* $(,)?) => {
@@ -180,15 +187,15 @@ fn normalize_test_batch(
     batch
 }
 
-/// Create a `[Option<RecordBatch>; N]` from a list of (payload_type, batch) pairs.
+/// Create an `OtapBatchStore` (`Logs`, `Metrics`, or `Traces`) from payload-type / column specs.
 ///
 /// Each batch is normalized (dotted columns -> structs), completed (missing columns
 /// filled from spec), and marked with encoding metadata. For metrics, missing data
 /// point tables are auto-generated based on metric_type.
 #[must_use]
 pub fn make_test_batch<S: OtapBatchStore, const N: usize>(
     inputs: Vec<(ArrowPayloadType, RecordBatch)>,
-) -> [Option<RecordBatch>; N] {
+) -> S {
     let allowed = S::allowed_payload_types();
     let mut result: [Option<RecordBatch>; N] = std::array::from_fn(|_| None);
     let all_payload_types: Vec<ArrowPayloadType> = inputs.iter().map(|(pt, _)| *pt).collect();
@@ -217,7 +224,13 @@ pub fn make_test_batch<S: OtapBatchStore, const N: usize>(
     // Auto-generate missing data point tables for metrics
     generate_missing_dp_tables(&mut result);
 
-    result
+    // Move the completed batches into a new store
+    let mut store = S::new();
+    let batches_mut = store.batches_mut();
+    for (idx, batch) in result.into_iter().enumerate() {
+        batches_mut[idx] = batch;
+    }
+    store
 }
 
 /// Fill in missing columns for a test batch using the OTAP schema spec.