From 448ce648bec90e882d20eebda910347c261b5735 Mon Sep 17 00:00:00 2001
From: Daniel Hillman <hillmandj@gmail.com>
Date: Fri, 18 Oct 2024 11:03:43 -0400
Subject: [PATCH] enhancement(sample transform): add stratified sampling
 capability (#21274)

* (enhancement sample transform): add stratified sampling capability

* Convert group_by key to Template

* Generate documentation

* Add changelog

* Fix typo in docs

* more documentation clean up

* Fix formatting to pass linter, fix punctuation in docs

---------

Co-authored-by: Jesse Szwedko <jesse.szwedko@datadoghq.com>
---
 ...ampling_to_sample_transform.enhancement.md |  3 +
 src/transforms/sample/config.rs               | 14 ++++
 src/transforms/sample/transform.rs            | 82 +++++++++++++++++--
 .../components/transforms/base/sample.cue     | 13 +++
 4 files changed, 106 insertions(+), 6 deletions(-)
 create mode 100644 changelog.d/20921_add_stratified_sampling_to_sample_transform.enhancement.md

diff --git a/changelog.d/20921_add_stratified_sampling_to_sample_transform.enhancement.md b/changelog.d/20921_add_stratified_sampling_to_sample_transform.enhancement.md
new file mode 100644
index 0000000000000..5fad60cda3d94
--- /dev/null
+++ b/changelog.d/20921_add_stratified_sampling_to_sample_transform.enhancement.md
@@ -0,0 +1,3 @@
+The `sample` transform can now take in a `group_by` configuration option that will allow logs with unique values for the patterns passed in to be sampled independently. This can reduce the complexity of the topology, since users would no longer need to create separate samplers with similar configuration to handle different log streams.
+
+authors: hillmandj
diff --git a/src/transforms/sample/config.rs b/src/transforms/sample/config.rs
index d05338493ef80..532bc2b5d0bbd 100644
--- a/src/transforms/sample/config.rs
+++ b/src/transforms/sample/config.rs
@@ -10,6 +10,7 @@ use crate::{
         TransformOutput,
     },
     schema,
+    template::Template,
     transforms::Transform,
 };
 
@@ -44,6 +45,16 @@ pub struct SampleConfig {
     #[configurable(metadata(docs::examples = "message"))]
     pub key_field: Option<String>,
 
+    /// The value to group events into separate buckets to be sampled independently.
+    ///
+    /// If left unspecified, or if the event doesn't have `group_by`, then the event is not
+    /// sampled separately.
+    #[configurable(metadata(
+        docs::examples = "{{ service }}",
+        docs::examples = "{{ hostname }}-{{ service }}"
+    ))]
+    pub group_by: Option<Template>,
+
     /// A logical condition used to exclude events from sampling.
     pub exclude: Option<AnyCondition>,
 }
@@ -53,6 +64,7 @@ impl GenerateConfig for SampleConfig {
         toml::Value::try_from(Self {
             rate: 10,
             key_field: None,
+            group_by: None,
             exclude: None::<AnyCondition>,
         })
         .unwrap()
@@ -67,6 +79,7 @@ impl TransformConfig for SampleConfig {
             Self::NAME.to_string(),
             self.rate,
             self.key_field.clone(),
+            self.group_by.clone(),
             self.exclude
                 .as_ref()
                 .map(|condition| condition.build(&context.enrichment_tables))
@@ -126,6 +139,7 @@ mod tests {
             let config = SampleConfig {
                 rate: 1,
                 key_field: None,
+                group_by: None,
                 exclude: None,
             };
             let (tx, rx) = mpsc::channel(1);
diff --git a/src/transforms/sample/transform.rs b/src/transforms/sample/transform.rs
index d9c7434d15566..d789950799721 100644
--- a/src/transforms/sample/transform.rs
+++ b/src/transforms/sample/transform.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use vector_lib::config::LegacyKey;
 use vrl::event_path;
 
@@ -5,6 +6,8 @@ use crate::{
     conditions::Condition,
     event::Event,
     internal_events::SampleEventDiscarded,
+    sinks::prelude::TemplateRenderingError,
+    template::Template,
     transforms::{FunctionTransform, OutputBuffer},
 };
 
@@ -13,26 +16,29 @@ pub struct Sample {
     name: String,
     rate: u64,
     key_field: Option<String>,
+    group_by: Option<Template>,
     exclude: Option<Condition>,
-    count: u64,
+    counter: HashMap<Option<String>, u64>,
 }
 
 impl Sample {
     // This function is dead code when the feature flag `transforms-impl-sample` is specified but not
     // `transforms-sample`.
     #![allow(dead_code)]
-    pub const fn new(
+    pub fn new(
         name: String,
         rate: u64,
         key_field: Option<String>,
+        group_by: Option<Template>,
         exclude: Option<Condition>,
     ) -> Self {
         Self {
             name,
             rate,
             key_field,
+            group_by,
             exclude,
-            count: 0,
+            counter: HashMap::new(),
         }
     }
 }
@@ -69,13 +75,42 @@ impl FunctionTransform for Sample {
             })
             .map(|v| v.to_string_lossy());
 
+        // Fetch actual field value if group_by option is set.
+        let group_by_key = self.group_by.as_ref().and_then(|group_by| match &event {
+            Event::Log(event) => group_by
+                .render_string(event)
+                .map_err(|error| {
+                    emit!(TemplateRenderingError {
+                        error,
+                        field: Some("group_by"),
+                        drop_event: false,
+                    })
+                })
+                .ok(),
+            Event::Trace(event) => group_by
+                .render_string(event)
+                .map_err(|error| {
+                    emit!(TemplateRenderingError {
+                        error,
+                        field: Some("group_by"),
+                        drop_event: false,
+                    })
+                })
+                .ok(),
+            Event::Metric(_) => panic!("component can never receive metric events"),
+        });
+
+        let counter_value: u64 = *self.counter.entry(group_by_key.clone()).or_default();
+
         let num = if let Some(value) = value {
             seahash::hash(value.as_bytes())
         } else {
-            self.count
+            counter_value
         };
 
-        self.count = (self.count + 1) % self.rate;
+        // reset counter for particular key, or default key if group_by option isn't provided
+        let increment: u64 = (counter_value + 1) % self.rate;
+        self.counter.insert(group_by_key.clone(), increment);
 
         if num % self.rate == 0 {
             match event {
@@ -134,6 +169,7 @@ mod tests {
             "sample".to_string(),
             2,
             log_schema().message_key().map(ToString::to_string),
+            None,
             Some(condition_contains(
                 log_schema().message_key().unwrap().to_string().as_str(),
                 "na",
@@ -156,6 +192,7 @@ mod tests {
             "sample".to_string(),
             25,
             log_schema().message_key().map(ToString::to_string),
+            None,
             Some(condition_contains(
                 log_schema().message_key().unwrap().to_string().as_str(),
                 "na",
@@ -181,6 +218,7 @@ mod tests {
             "sample".to_string(),
             2,
             log_schema().message_key().map(ToString::to_string),
+            None,
             Some(condition_contains(
                 log_schema().message_key().unwrap().to_string().as_str(),
                 "na",
@@ -216,6 +254,7 @@ mod tests {
                 "sample".to_string(),
                 0,
                 key_field.clone(),
+                None,
                 Some(condition_contains(
                     log_schema().message_key().unwrap().to_string().as_str(),
                     "important",
@@ -232,6 +271,33 @@ mod tests {
         }
     }
 
+    #[test]
+    fn handles_group_by() {
+        for group_by in &[None, Some(Template::try_from("{{ other_field }}").unwrap())] {
+            let mut event = Event::Log(LogEvent::from("nananana"));
+            let log = event.as_mut_log();
+            log.insert("other_field", "foo");
+            let mut sampler = Sample::new(
+                "sample".to_string(),
+                0,
+                log_schema().message_key().map(ToString::to_string),
+                group_by.clone(),
+                Some(condition_contains(
+                    log_schema().message_key().unwrap().to_string().as_str(),
+                    "na",
+                )),
+            );
+            let iterations = 0..1000;
+            let total_passed = iterations
+                .filter_map(|_| {
+                    transform_one(&mut sampler, event.clone())
+                        .map(|result| assert_eq!(result, event))
+                })
+                .count();
+            assert_eq!(total_passed, 1000);
+        }
+    }
+
     #[test]
     fn handles_key_field() {
         for key_field in &[None, Some("other_field".into())] {
@@ -242,6 +308,7 @@ mod tests {
                 "sample".to_string(),
                 0,
                 key_field.clone(),
+                None,
                 Some(condition_contains("other_field", "foo")),
             );
             let iterations = 0..1000;
@@ -264,6 +331,7 @@ mod tests {
                 "sample".to_string(),
                 10,
                 key_field.clone(),
+                None,
                 Some(condition_contains(&message_key, "na")),
             );
             let passing = events
@@ -278,6 +346,7 @@ mod tests {
                 "sample".to_string(),
                 25,
                 key_field.clone(),
+                None,
                 Some(condition_contains(&message_key, "na")),
             );
             let passing = events
@@ -292,6 +361,7 @@ mod tests {
                 "sample".to_string(),
                 25,
                 key_field.clone(),
+                None,
                 Some(condition_contains(&message_key, "na")),
             );
             let event = Event::Log(LogEvent::from("nananana"));
@@ -304,7 +374,7 @@ mod tests {
     fn handles_trace_event() {
         let event: TraceEvent = LogEvent::from("trace").into();
         let trace = Event::Trace(event);
-        let mut sampler = Sample::new("sample".to_string(), 2, None, None);
+        let mut sampler = Sample::new("sample".to_string(), 2, None, None, None);
         let iterations = 0..2;
         let total_passed = iterations
             .filter_map(|_| transform_one(&mut sampler, trace.clone()))
diff --git a/website/cue/reference/components/transforms/base/sample.cue b/website/cue/reference/components/transforms/base/sample.cue
index 865f6e370aca2..d6b573e409334 100644
--- a/website/cue/reference/components/transforms/base/sample.cue
+++ b/website/cue/reference/components/transforms/base/sample.cue
@@ -6,6 +6,19 @@ base: components: transforms: sample: configuration: {
 		required:    false
 		type: condition: {}
 	}
+	group_by: {
+		description: """
+			The value to group events into separate buckets to be sampled independently.
+
+			If left unspecified, or if the event doesn't have `group_by`, then the event is not
+			sampled separately.
+			"""
+		required: false
+		type: string: {
+			examples: ["{{ service }}", "{{ hostname }}-{{ service }}"]
+			syntax: "template"
+		}
+	}
 	key_field: {
 		description: """
 			The name of the field whose value is hashed to determine if the event should be