poisson upscaling without storing count in the sample types

Author: realFlowControl

Diff for: profiling/src/api.rs

@@ -196,6 +196,12 @@ pub enum UpscalingInfo {
         count_value_offset: usize,
         sampling_distance: u64,
     },
+    PoissonCount {
+        // sum_value_offset is an offset in the profile values type array
+        sum_value_offset: usize,
+        count_value: u64,
+        sampling_distance: u64,
+    },
     Proportional {
         scale: f64,
     },
@@ -213,6 +219,15 @@ impl std::fmt::Display for UpscalingInfo {
                 "Poisson = sum_value_offset: {}, count_value_offset: {}, sampling_distance: {}",
                 sum_value_offset, count_value_offset, sampling_distance
             ),
+            UpscalingInfo::PoissonCount {
+                sum_value_offset,
+                count_value,
+                sampling_distance,
+            } => write!(
+                f,
+                "Poisson = sum_value_offset: {}, count_value: {}, sampling_distance: {}",
+                sum_value_offset, count_value, sampling_distance
+            ),
             UpscalingInfo::Proportional { scale } => {
                 write!(f, "Proportional = scale: {}", scale)
             }
@@ -241,6 +256,28 @@ impl UpscalingInfo {
                     sampling_distance
                 )
             }
+            UpscalingInfo::PoissonCount {
+                sum_value_offset,
+                count_value,
+                sampling_distance,
+            } => {
+                anyhow::ensure!(
+                    sum_value_offset < &number_of_values,
+                    "sum_value_offset {} must be strictly less than {}",
+                    sum_value_offset,
+                    number_of_values
+                );
+                anyhow::ensure!(
+                    count_value != &0,
+                    "count_value {} must be greater than 0",
+                    count_value
+                );
+                anyhow::ensure!(
+                    sampling_distance != &0,
+                    "sampling_distance {} must be greater than 0",
+                    sampling_distance
+                )
+            }
             UpscalingInfo::Proportional { scale: _ } => (),
         }
         anyhow::Ok(())

Diff for: profiling/src/internal/profile/mod.rs

@@ -1410,6 +1410,38 @@ mod api_tests {
         assert_eq!(first.values, vec![1, 298, 29]);
     }
 
+    #[test]
+    fn test_upscaling_by_value_on_one_value_with_poisson_count() {
+        let sample_types = create_samples_types();
+
+        let mut profile = Profile::new(SystemTime::now(), &sample_types, None);
+
+        let sample1 = api::Sample {
+            locations: vec![],
+            values: vec![1, 16, 29],
+            labels: vec![],
+        };
+
+        profile.add_sample(sample1, None).expect("add to success");
+
+        let upscaling_info = UpscalingInfo::PoissonCount {
+            sum_value_offset: 1,
+            count_value: 29,
+            sampling_distance: 10,
+        };
+        let values_offset: Vec<usize> = vec![1];
+        profile
+            .add_upscaling_rule(values_offset.as_slice(), "", "", upscaling_info)
+            .expect("Rule added");
+
+        let serialized_profile = pprof::roundtrip_to_pprof(profile).unwrap();
+
+        assert_eq!(serialized_profile.samples.len(), 1);
+        let first = serialized_profile.samples.first().expect("one sample");
+
+        assert_eq!(first.values, vec![1, 298, 29]);
+    }
+
     #[test]
     fn test_upscaling_by_value_on_zero_value_with_poisson() {
         let sample_types = create_samples_types();

Diff for: profiling/src/internal/upscaling.rs

@@ -28,6 +28,20 @@ impl UpscalingRule {
                 let avg = values[sum_value_offset] as f64 / values[count_value_offset] as f64;
                 1_f64 / (1_f64 - (-avg / sampling_distance as f64).exp())
             }
+            UpscalingInfo::PoissonCount {
+                sum_value_offset,
+                count_value,
+                sampling_distance,
+            } => {
+                // This should not happen, but if it happens,
+                // do not upscale
+                if values[sum_value_offset] == 0 || count_value == 0 {
+                    return 1_f64;
+                }
+
+                let avg = values[sum_value_offset] as f64 / count_value as f64;
+                1_f64 / (1_f64 - (-avg / sampling_distance as f64).exp())
+            }
             UpscalingInfo::Proportional { scale } => scale,
         }
     }