From debde73c0a0de2164ef15456c2dd5d2e8ae86f6a Mon Sep 17 00:00:00 2001 From: Dan Hansen Date: Sat, 6 Apr 2024 23:35:25 -0700 Subject: [PATCH] [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 --- internal/function_bind.go | 12 ++++++++ internal/function_register.go | 2 ++ internal/function_window.go | 46 ++++++++++++++++++++++++++++++ query_test.go | 53 +++++++++++++++++++++++++++++++++++ 4 files changed, 113 insertions(+) diff --git a/internal/function_bind.go b/internal/function_bind.go index 6660806..344a6bc 100644 --- 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 index 4b2c272..eb5938e 100644 --- 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 index 33f2e01..00889f7 100644 --- 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 index d848d4e..4905f7c 100644 --- 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`,