From 47e9789e61014e0ff800fd4f2fc0c1881862720e Mon Sep 17 00:00:00 2001
From: Ian <i3149@users.noreply.github.com>
Date: Tue, 10 Sep 2024 22:35:52 -0700
Subject: [PATCH] Allowing rollups to keep undefined terms if needed (#749)

---
 cmd/ktranslate/main.go |  7 +++++++
 config.go              | 14 ++++++++------
 pkg/rollup/rollup.go   | 18 +++++++++++++++---
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/cmd/ktranslate/main.go b/cmd/ktranslate/main.go
index 5b0e5a2f..bcac2798 100644
--- a/cmd/ktranslate/main.go
+++ b/cmd/ktranslate/main.go
@@ -592,6 +592,13 @@ func applyFlags(cfg *ktranslate.Config) error {
 				cfg.Rollup.TopK = v
 			case "rollups":
 				cfg.Rollup.Formats = strings.Split(val, filter.AndToken)
+			case "rollup_keep_undefined":
+				v, err := strconv.ParseBool(val)
+				if err != nil {
+					errCh <- err
+					return
+				}
+				cfg.Rollup.KeepUndefined = v
 			// pkg/eggs/kmux
 			case "dir":
 				cfg.KMux.Dir = val
diff --git a/config.go b/config.go
index 5665def9..65ee743a 100644
--- a/config.go
+++ b/config.go
@@ -138,9 +138,10 @@ type DDogSinkConfig struct {
 
 // RollupConfig is the config for rollups
 type RollupConfig struct {
-	JoinKey string
-	TopK    int
-	Formats []string
+	JoinKey       string
+	TopK          int
+	Formats       []string
+	KeepUndefined bool
 }
 
 // KMuxConfig is the config for the mux server
@@ -451,9 +452,10 @@ func DefaultConfig() *Config {
 			RelayURL: "",
 		},
 		Rollup: &RollupConfig{
-			JoinKey: "^",
-			TopK:    10,
-			Formats: []string{},
+			JoinKey:       "^",
+			TopK:          10,
+			Formats:       []string{},
+			KeepUndefined: false,
 		},
 		KMux: &KMuxConfig{
 			Dir: ".",
diff --git a/pkg/rollup/rollup.go b/pkg/rollup/rollup.go
index 503244ad..acbcc8ba 100644
--- a/pkg/rollup/rollup.go
+++ b/pkg/rollup/rollup.go
@@ -26,12 +26,14 @@ const (
 	Percentile            = "entropy"
 
 	KENTIK_EVENT_TYPE = "KFlow:%s:%s"
+	UndefinedKey      = "undefined"
 )
 
 var (
-	rollups RollupFlag
-	keyJoin string
-	topK    int
+	rollups       RollupFlag
+	keyJoin       string
+	topK          int
+	keepUndefined bool
 )
 
 type RollupFlag []string
@@ -48,6 +50,7 @@ func init() {
 	flag.Var(&rollups, "rollups", "Any rollups to use. Format: type, name, metric, dimension 1, dimension 2, ..., dimension n: sum,bytes,in_bytes,dst_addr")
 	flag.StringVar(&keyJoin, "rollup_key_join", "^", "Token to use to join dimension keys together")
 	flag.IntVar(&topK, "rollup_top_k", 10, "Export only these top values")
+	flag.BoolVar(&keepUndefined, "rollup_keep_undefined", false, "If set, mark undefined values with the string undefined.")
 }
 
 type Roller interface {
@@ -146,6 +149,9 @@ func GetRollups(log logger.Underlying, cfg *ktranslate.RollupConfig) ([]Roller,
 		}
 	}
 
+	// If true, don't drop the rollups which don't get matched.
+	keepUndefined = cfg.KeepUndefined
+
 	return rolls, nil
 }
 
@@ -230,6 +236,9 @@ func (r *rollupBase) getKey(mapr map[string]interface{}) string {
 				// Skip?
 			}
 		}
+		if keepUndefined && keyPts[i] == "" {
+			keyPts[i] = UndefinedKey
+		}
 	}
 	next := len(r.dims)
 	for _, d := range r.multiDims { // Now handle the 2 level deep maps
@@ -250,6 +259,9 @@ func (r *rollupBase) getKey(mapr map[string]interface{}) string {
 				keyPts[next] = strconv.Itoa(int(dd[d[1]]))
 			}
 		}
+		if keepUndefined && keyPts[next] == "" {
+			keyPts[next] = UndefinedKey
+		}
 		next++
 	}