From 3888356bf654a3838ffdc8fe5b3c811e18f7d389 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Tue, 21 May 2024 13:53:55 +0200
Subject: [PATCH] wip

---
 lib/logstorage/stats_avg_test.go         |  26 ++
 lib/logstorage/stats_count_empty_test.go | 326 ++++++++++++++++++++++
 lib/logstorage/stats_count_test.go       | 330 +++++++++++++++++++++++
 3 files changed, 682 insertions(+)
 create mode 100644 lib/logstorage/stats_count_empty_test.go
 create mode 100644 lib/logstorage/stats_count_test.go

diff --git a/lib/logstorage/stats_avg_test.go b/lib/logstorage/stats_avg_test.go
index 98c28723f..b4aaf5d2d 100644
--- a/lib/logstorage/stats_avg_test.go
+++ b/lib/logstorage/stats_avg_test.go
@@ -172,6 +172,32 @@ func TestStatsAvg(t *testing.T) {
 		},
 	})
 
+	f("stats by (b) avg(a) if (b:*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"b", "3"},
+		},
+		{
+			{"a", `3`},
+			{"c", `54`},
+		},
+	}, [][]Field{
+		{
+			{"b", "3"},
+			{"x", "1.5"},
+		},
+		{
+			{"b", ""},
+			{"x", "NaN"},
+		},
+	})
+
 	f("stats by (a) avg(b) as x", [][]Field{
 		{
 			{"_msg", `abc`},
diff --git a/lib/logstorage/stats_count_empty_test.go b/lib/logstorage/stats_count_empty_test.go
new file mode 100644
index 000000000..7cde7389f
--- /dev/null
+++ b/lib/logstorage/stats_count_empty_test.go
@@ -0,0 +1,326 @@
+package logstorage
+
+import (
+	"testing"
+)
+
+func TestParseStatsCountEmptySuccess(t *testing.T) {
+	f := func(pipeStr string) {
+		t.Helper()
+		expectParseStatsFuncSuccess(t, pipeStr)
+	}
+
+	f(`count_empty(*)`)
+	f(`count_empty(a)`)
+	f(`count_empty(a, b)`)
+}
+
+func TestParseStatsCountEmptyFailure(t *testing.T) {
+	f := func(pipeStr string) {
+		t.Helper()
+		expectParseStatsFuncFailure(t, pipeStr)
+	}
+
+	f(`count_empty`)
+	f(`count_empty(a b)`)
+	f(`count_empty(x) y`)
+}
+
+func TestStatsCountEmpty(t *testing.T) {
+	f := func(pipeStr string, rows, rowsExpected [][]Field) {
+		t.Helper()
+		expectPipeResults(t, pipeStr, rows, rowsExpected)
+	}
+
+	f("stats count_empty(*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+		},
+		{
+			{"a", `3`},
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "1"},
+		},
+	})
+
+	f("stats count_empty(b) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+		},
+		{
+			{"a", `3`},
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "2"},
+		},
+	})
+
+	f("stats count_empty(a, b) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+		},
+		{
+			{"aa", `3`},
+			{"bb", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "2"},
+		},
+	})
+
+	f("stats count_empty(c) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"a", `3`},
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "3"},
+		},
+	})
+
+	f("stats count_empty(a) if (b:*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count_empty(b) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "1"},
+		},
+		{
+			{"a", "3"},
+			{"x", "0"},
+		},
+	})
+
+	f("stats by (a) count_empty(b) if (!c:foo) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "foo"},
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+			{"c", "bar"},
+		},
+		{
+			{"a", `3`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "0"},
+		},
+		{
+			{"a", "3"},
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count_empty(*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "3"},
+		},
+		{
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", ""},
+			{"x", "1"},
+		},
+		{
+			{"a", "1"},
+			{"x", "0"},
+		},
+		{
+			{"a", "3"},
+			{"x", "0"},
+		},
+	})
+
+	f("stats by (a) count_empty(c) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"a", `3`},
+			{"c", `5`},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "2"},
+		},
+		{
+			{"a", "3"},
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count_empty(a, b, c) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "3"},
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "0"},
+		},
+		{
+			{"a", "3"},
+			{"x", "0"},
+		},
+	})
+
+	f("stats by (a, b) count_empty(a) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "3"},
+		},
+		{
+			{"c", `3`},
+			{"b", `5`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"b", "3"},
+			{"x", "0"},
+		},
+		{
+			{"a", "1"},
+			{"b", ""},
+			{"x", "0"},
+		},
+		{
+			{"a", ""},
+			{"b", "5"},
+			{"x", "1"},
+		},
+	})
+}
diff --git a/lib/logstorage/stats_count_test.go b/lib/logstorage/stats_count_test.go
new file mode 100644
index 000000000..2ff03ef02
--- /dev/null
+++ b/lib/logstorage/stats_count_test.go
@@ -0,0 +1,330 @@
+package logstorage
+
+import (
+	"testing"
+)
+
+func TestParseStatsCountSuccess(t *testing.T) {
+	f := func(pipeStr string) {
+		t.Helper()
+		expectParseStatsFuncSuccess(t, pipeStr)
+	}
+
+	f(`count(*)`)
+	f(`count(a)`)
+	f(`count(a, b)`)
+}
+
+func TestParseStatsCountFailure(t *testing.T) {
+	f := func(pipeStr string) {
+		t.Helper()
+		expectParseStatsFuncFailure(t, pipeStr)
+	}
+
+	f(`count`)
+	f(`count(a b)`)
+	f(`count(x) y`)
+}
+
+func TestStatsCount(t *testing.T) {
+	f := func(pipeStr string, rows, rowsExpected [][]Field) {
+		t.Helper()
+		expectPipeResults(t, pipeStr, rows, rowsExpected)
+	}
+
+	f("stats count(*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+		},
+		{
+			{"a", `3`},
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "4"},
+		},
+	})
+
+	f("stats count(b) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+		},
+		{
+			{"a", `3`},
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "2"},
+		},
+	})
+
+	f("stats count(a, b) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+		},
+		{
+			{"aa", `3`},
+			{"bb", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "2"},
+		},
+	})
+
+	f("stats count(c) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"a", `3`},
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "0"},
+		},
+	})
+
+	f("stats count(a) if (b:*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `2`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"b", `54`},
+		},
+	}, [][]Field{
+		{
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count(b) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "1"},
+		},
+		{
+			{"a", "3"},
+			{"x", "2"},
+		},
+	})
+
+	f("stats by (a) count(b) if (!c:foo) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"b", "aadf"},
+			{"c", "foo"},
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+			{"c", "bar"},
+		},
+		{
+			{"a", `3`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "1"},
+		},
+		{
+			{"a", "3"},
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count(*) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "3"},
+		},
+		{
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+		},
+	}, [][]Field{
+		{
+			{"a", ""},
+			{"x", "1"},
+		},
+		{
+			{"a", "1"},
+			{"x", "2"},
+		},
+		{
+			{"a", "3"},
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count(c) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+		},
+		{
+			{"a", `3`},
+			{"c", `5`},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "0"},
+		},
+		{
+			{"a", "3"},
+			{"x", "1"},
+		},
+	})
+
+	f("stats by (a) count(a, b, c) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "3"},
+		},
+		{
+			{"a", `3`},
+			{"b", `5`},
+		},
+		{
+			{"foo", "bar"},
+		},
+		{
+			{"a", `3`},
+			{"b", `7`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"x", "2"},
+		},
+		{
+			{"a", ""},
+			{"x", "0"},
+		},
+		{
+			{"a", "3"},
+			{"x", "2"},
+		},
+	})
+
+	f("stats by (a, b) count(a) as x", [][]Field{
+		{
+			{"_msg", `abc`},
+			{"a", `1`},
+			{"b", `3`},
+		},
+		{
+			{"_msg", `def`},
+			{"a", `1`},
+			{"c", "3"},
+		},
+		{
+			{"c", `3`},
+			{"b", `5`},
+		},
+	}, [][]Field{
+		{
+			{"a", "1"},
+			{"b", "3"},
+			{"x", "1"},
+		},
+		{
+			{"a", "1"},
+			{"b", ""},
+			{"x", "1"},
+		},
+		{
+			{"a", ""},
+			{"b", "5"},
+			{"x", "0"},
+		},
+	})
+}