From f53a6e2c1c124d6e90eea4853ab994c8c8e08c1e Mon Sep 17 00:00:00 2001
From: Dan Hansen <ohaibbq@gmail.com>
Date: Fri, 1 Mar 2024 16:28:35 -0800
Subject: [PATCH] Ignore nulls when counting window values (#33)

---
 internal/function_window.go | 14 ++++++++++++++
 query_test.go               |  4 ++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/internal/function_window.go b/internal/function_window.go
index ad1a00f..33f2e01 100644
--- a/internal/function_window.go
+++ b/internal/function_window.go
@@ -72,6 +72,20 @@ func (f *WINDOW_AVG) Done(agg *WindowFuncAggregatedStatus) (Value, error) {
 type WINDOW_COUNT struct {
 }
 
+func (f *WINDOW_COUNT) Step(values []Value, agg *WindowFuncAggregatedStatus) error {
+	if len(values) == 0 {
+		return nil
+	}
+
+	value := values[0]
+	if value == nil {
+		return nil
+	}
+
+	agg.Values = append(agg.Values, value)
+	return nil
+}
+
 func (f *WINDOW_COUNT) Done(agg *WindowFuncAggregatedStatus) (Value, error) {
 	values, err := agg.RelevantValues()
 	if err != nil {
diff --git a/query_test.go b/query_test.go
index de841a6..b6c85e4 100644
--- a/query_test.go
+++ b/query_test.go
@@ -766,8 +766,8 @@ SELECT ARRAY_CONCAT_AGG(x) AS array_concat_agg FROM (
 		},
 		{
 			name:         "count with null",
-			query:        `SELECT COUNT(x) FROM UNNEST([NULL]) AS x`,
-			expectedRows: [][]interface{}{{int64(0)}},
+			query:        `WITH toks AS (SELECT 1 AS x UNION ALL SELECT null) SELECT COUNT(x) FROM toks`,
+			expectedRows: [][]interface{}{{int64(1)}},
 		},
 		{
 			name:         "count with if",