Reindex implementation

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

@@ -195,6 +195,20 @@ pub enum Error {
     #[error("Mixed signals")]
     MixedSignals,
 
+    #[error(
+        "Too many items. signal: {:?}, size: {}, max: {}, message: {}",
+        payload_type,
+        count,
+        max,
+        message
+    )]
+    TooManyItems {
+        payload_type: ArrowPayloadType,
+        count: usize,
+        max: u64,
+        message: String,
+    },
+
     #[error("Encoding error: {}", error)]
     Encoding {
         #[from]

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

@@ -158,6 +158,24 @@ impl OtapArrowRecords {
     }
 }
 
+impl From<Logs> for OtapArrowRecords {
+    fn from(logs: Logs) -> Self {
+        Self::Logs(logs)
+    }
+}
+
+impl From<Metrics> for OtapArrowRecords {
+    fn from(metrics: Metrics) -> Self {
+        Self::Metrics(metrics)
+    }
+}
+
+impl From<Traces> for OtapArrowRecords {
+    fn from(traces: Traces) -> Self {
+        Self::Traces(traces)
+    }
+}
+
 /// The ArrowBatchStore helper trait is used to define a common interface for
 /// storing and retrieving Arrow record batches in a type-safe manner. It is
 /// implemented by various structs that represent each signal type and provides
@@ -187,6 +205,12 @@ pub trait OtapBatchStore: Default + Clone {
     /// Return a list of the allowed payload types associated with this type of batch
     fn allowed_payload_types() -> &'static [ArrowPayloadType];
 
+    /// Returns the payload type at the given index in [BatchArray]
+    #[must_use]
+    fn payload_type_at_idx(index: usize) -> ArrowPayloadType {
+        Self::allowed_payload_types()[index]
+    }
+
     /// Decode the delta-encoded and quasi-delta encoded IDs & parent IDs on each Arrow
     /// Record Batch contained in this Otap Batch
     fn decode_transport_optimized_ids(otap_batch: &mut OtapArrowRecords) -> Result<()>;
@@ -340,9 +364,9 @@ impl OtapBatchStore for Logs {
 
     fn allowed_payload_types() -> &'static [ArrowPayloadType] {
         &[
-            ArrowPayloadType::Logs,
             ArrowPayloadType::ResourceAttrs,
             ArrowPayloadType::ScopeAttrs,
+            ArrowPayloadType::Logs,
             ArrowPayloadType::LogAttrs,
         ]
     }
@@ -414,6 +438,9 @@ const DATA_POINTS_TYPES: [ArrowPayloadType; 4] = [
 ];
 
 /// Fetch the number of items as defined by the batching system
+/// TODO [JD] we don't need a unifying function like this once we remove
+/// all the usages in to groups.rs. Instead we can have each batch store
+/// define these.
 #[must_use]
 fn num_items<const N: usize>(batches: &[Option<RecordBatch>; N]) -> usize {
     match N {
@@ -482,10 +509,9 @@ impl OtapBatchStore for Metrics {
 
     fn allowed_payload_types() -> &'static [ArrowPayloadType] {
         &[
-            ArrowPayloadType::UnivariateMetrics,
-            ArrowPayloadType::MultivariateMetrics,
             ArrowPayloadType::ResourceAttrs,
             ArrowPayloadType::ScopeAttrs,
+            ArrowPayloadType::UnivariateMetrics,
             ArrowPayloadType::NumberDataPoints,
             ArrowPayloadType::SummaryDataPoints,
             ArrowPayloadType::HistogramDataPoints,
@@ -500,6 +526,7 @@ impl OtapBatchStore for Metrics {
             ArrowPayloadType::NumberDpExemplarAttrs,
             ArrowPayloadType::HistogramDpExemplarAttrs,
             ArrowPayloadType::ExpHistogramDpExemplarAttrs,
+            ArrowPayloadType::MultivariateMetrics,
             ArrowPayloadType::MetricAttrs,
         ]
     }
@@ -717,9 +744,9 @@ impl OtapBatchStore for Traces {
 
     fn allowed_payload_types() -> &'static [ArrowPayloadType] {
         &[
-            ArrowPayloadType::Spans,
             ArrowPayloadType::ResourceAttrs,
             ArrowPayloadType::ScopeAttrs,
+            ArrowPayloadType::Spans,
             ArrowPayloadType::SpanAttrs,
             ArrowPayloadType::SpanEvents,
             ArrowPayloadType::SpanLinks,
@@ -830,6 +857,7 @@ impl OtapBatchStore for Traces {
 }
 
 /// Return the child payload types for the given payload type
+/// TODO [JD]: This is pretty much made obsolete by payload_relations
 #[must_use]
 pub const fn child_payload_types(payload_type: ArrowPayloadType) -> &'static [ArrowPayloadType] {
     match payload_type {
@@ -3374,4 +3402,17 @@ mod test {
             }
         }
     }
+
+    #[test]
+    fn test_idx_to_payload_type() {
+        test_idx_to_payload_type_impl::<Logs>();
+        test_idx_to_payload_type_impl::<Metrics>();
+        test_idx_to_payload_type_impl::<Traces>();
+    }
+
+    fn test_idx_to_payload_type_impl<T: OtapBatchStore>() {
+        for i in 0..T::allowed_payload_types().len() {
+            assert_eq!(T::payload_type_at_idx(i), T::allowed_payload_types()[i]);
+        }
+    }
 }

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

@@ -102,19 +102,18 @@ fn generate_metrics_batching_test_cases() -> Vec<MetricsBatchingTestCase> {
     }
 
     // All four metric types present - temporarily disabled due to schema compatibility issues
-    // TODO: Investigate and fix schema merging for mixed metric types
-    // for limit in [10, 20, 50] {
-    //     add_case(
-    //         &format!("all_types_small_limit_{}", limit),
-    //         MetricsConfig::new()
-    //             .with_gauges(vec![2, 3])
-    //             .with_sums(vec![2])
-    //             .with_histograms(vec![1, 2])
-    //             .with_summaries(vec![1]),
-    //         limit as u64,
-    //         3,
-    //     );
-    // }
+    for limit in [10, 20, 50] {
+        add_case(
+            &format!("all_types_small_limit_{}", limit),
+            MetricsConfig::new()
+                .with_gauges(vec![2, 3])
+                .with_sums(vec![2])
+                .with_histograms(vec![1, 2])
+                .with_summaries(vec![1]),
+            limit as u64,
+            3,
+        );
+    }
 
     // Mixed sizes: some fit, some oversized
     for limit in [10, 15, 20] {
@@ -143,19 +142,18 @@ fn generate_metrics_batching_test_cases() -> Vec<MetricsBatchingTestCase> {
     }
 
     // Complex multi-type scenarios - temporarily disabled
-    // TODO: Investigate and fix schema merging for mixed metric types
-    // for limit in [25, 50, 100] {
-    //     add_case(
-    //         &format!("complex_all_types_limit_{}", limit),
-    //         MetricsConfig::new()
-    //             .with_gauges(vec![3, 5, 7, 2])
-    //             .with_sums(vec![4, 6, 3])
-    //             .with_histograms(vec![2, 8])
-    //             .with_summaries(vec![5, 3, 2]),
-    //         limit as u64,
-    //         6,
-    //     );
-    // }
+    for limit in [25, 50, 100] {
+        add_case(
+            &format!("complex_all_types_limit_{}", limit),
+            MetricsConfig::new()
+                .with_gauges(vec![3, 5, 7, 2])
+                .with_sums(vec![4, 6, 3])
+                .with_histograms(vec![2, 8])
+                .with_summaries(vec![5, 3, 2]),
+            limit as u64,
+            6,
+        );
+    }
 
     // Stress test: many metrics of varying sizes - single type only
     for limit in [50, 100, 200] {
@@ -220,7 +218,8 @@ fn generate_metrics_batching_test_cases() -> Vec<MetricsBatchingTestCase> {
     }
 
     // Test with varying attributes enabled - temporarily disabled
-    // TODO: Investigate parent_id column type issue with varying attributes
+    // TODO: Investigate parent_id column type issue with varying attributes,
+    // there is a problem with splitting
     // for limit in [20, 50] {
     //     add_case(
     //         &format!("with_attrs_gauges_limit_{}", limit),