Ignore nulls when counting window values (#33)

Recidiviz · Mar 2, 2024 · f53a6e2 · f53a6e2
1 parent 4be95a3
commit f53a6e2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git 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
@@ -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",