[Window] Support LOGICAL_OR and LOGICAL_AND (#42)

* [Window] Support `LOGICAL_OR` and `LOGICAL_AND` * lint * fix LOGICAL_OR default value; add test for no rows
Recidiviz · Apr 7, 2024 · debde73 · debde73
1 parent 2ca16b0
commit debde73
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 0 deletions.
diff --git a/internal/function_bind.go b/internal/function_bind.go
@@ -3707,6 +3707,18 @@ func bindWindowCountIf() func() *WindowAggregator {
 	}
 }
 
+func bindWindowLogicalAnd() func() *WindowAggregator {
+	return func() *WindowAggregator {
+		return newSingleItemWindowAggregator(&WINDOW_LOGICAL_AND{})
+	}
+}
+
+func bindWindowLogicalOr() func() *WindowAggregator {
+	return func() *WindowAggregator {
+		return newSingleItemWindowAggregator(&WINDOW_LOGICAL_OR{})
+	}
+}
+
 func bindWindowMax() func() *WindowAggregator {
 	return func() *WindowAggregator {
 		return newSingleItemWindowAggregator(&WINDOW_MAX{})

diff --git a/internal/function_register.go b/internal/function_register.go
@@ -332,6 +332,8 @@ var windowFuncs = []*WindowFuncInfo{
 	{Name: "count", BindFunc: bindWindowCount},
 	{Name: "count_star", BindFunc: bindWindowCountStar},
 	{Name: "countif", BindFunc: bindWindowCountIf},
+	{Name: "logical_and", BindFunc: bindWindowLogicalAnd},
+	{Name: "logical_or", BindFunc: bindWindowLogicalOr},
 	{Name: "max", BindFunc: bindWindowMax},
 	{Name: "min", BindFunc: bindWindowMin},
 	{Name: "string_agg", BindFunc: bindWindowStringAgg},

diff --git a/internal/function_window.go b/internal/function_window.go
@@ -406,6 +406,52 @@ func (f *WINDOW_LAG) Done(agg *WindowFuncAggregatedStatus) (Value, error) {
 	return agg.Values[len(agg.Values)-f.offset-1], nil
 }
 
+type WINDOW_LOGICAL_AND struct {
+}
+
+func (f *WINDOW_LOGICAL_AND) Done(agg *WindowFuncAggregatedStatus) (Value, error) {
+	values, err := agg.RelevantValues()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, cond := range values {
+		b, err := cond.ToBool()
+		if err != nil {
+			return nil, err
+		}
+
+		if !b {
+			return BoolValue(false), nil
+		}
+	}
+
+	return BoolValue(true), nil
+}
+
+type WINDOW_LOGICAL_OR struct {
+}
+
+func (f *WINDOW_LOGICAL_OR) Done(agg *WindowFuncAggregatedStatus) (Value, error) {
+	values, err := agg.RelevantValues()
+	if err != nil {
+		return nil, err
+	}
+
+	for _, cond := range values {
+		b, err := cond.ToBool()
+		if err != nil {
+			return nil, err
+		}
+
+		if b {
+			return BoolValue(true), nil
+		}
+	}
+
+	return BoolValue(false), nil
+}
+
 type WINDOW_PERCENTILE_CONT struct {
 	percentile Value
 }

diff --git a/query_test.go b/query_test.go
@@ -832,6 +832,59 @@ SELECT LOGICAL_AND(x) AS logical_or FROM toks`,
 SELECT LOGICAL_OR(x) AS logical_or FROM toks`,
 			expectedRows: [][]interface{}{{false}},
 		},
+		{name: "logical_and default",
+			query: `SELECT LOGICAL_AND(a) OVER (ORDER BY b), a, b
+FROM (
+  SELECT CAST(a AS BOOL) AS a, b FROM UNNEST([
+      STRUCT(False AS a, 1 AS b),
+      STRUCT(False AS a, 2 AS b),
+      STRUCT(false AS a, 3 AS b),
+      STRUCT(False AS a, 4 AS b)
+  ])
+);`,
+			expectedRows: [][]interface{}{
+				{false, false, int64(1)},
+				{false, false, int64(2)},
+				{false, false, int64(3)},
+				{false, false, int64(4)},
+			}},
+		{name: "logical_and no rows",
+			query:        `SELECT LOGICAL_AND(a) OVER (ORDER BY a) FROM ( SELECT CAST(a AS BOOL) AS a FROM UNNEST([]) a );`,
+			expectedRows: make([][]interface{}, 0),
+		},
+		{
+			name: "logical_and with window",
+			query: `WITH toks AS ( SELECT true AS x UNION ALL SELECT false UNION ALL SELECT true)
+							SELECT LOGICAL_AND(x) OVER (ORDER BY x) FROM toks`,
+			expectedRows: [][]interface{}{{false}, {false}, {false}},
+		},
+
+		{name: "logical_or window default",
+			query: `SELECT LOGICAL_OR(a) OVER (ORDER BY b), a, b
+FROM (
+  SELECT CAST(a AS BOOL) AS a, b FROM UNNEST([
+      STRUCT(False AS a, 1 AS b),
+      STRUCT(False AS a, 2 AS b),
+      STRUCT(True AS a, 3 AS b),
+      STRUCT(False AS a, 4 AS b)
+  ])
+);`,
+			expectedRows: [][]interface{}{
+				{false, false, int64(1)},
+				{false, false, int64(2)},
+				{true, true, int64(3)},
+				{true, false, int64(4)},
+			}},
+		{name: "logical_or window no rows",
+			query:        `SELECT LOGICAL_OR(a) OVER (ORDER BY a) FROM ( SELECT CAST(a AS BOOL) AS a FROM UNNEST([]) a );`,
+			expectedRows: make([][]interface{}, 0),
+		},
+		{
+			name: "logical_or with window",
+			query: `WITH toks AS ( SELECT true AS x UNION ALL SELECT false UNION ALL SELECT true)
+							SELECT LOGICAL_OR(x) OVER (ORDER BY x) FROM toks`,
+			expectedRows: [][]interface{}{{false}, {true}, {true}},
+		},
 		{
 			name:         "max from int group",
 			query:        `SELECT MAX(x) AS max FROM UNNEST([8, 37, 4, 55]) AS x`,