@@ -8,10 +8,18 @@ package sql
88import (
99 "context"
1010
11+ "github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
1112 "github.com/cockroachdb/cockroach/pkg/sql/execinfra/execagg"
1213 "github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
14+ "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
15+ "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
16+ "github.com/cockroachdb/cockroach/pkg/sql/rowenc"
1317 "github.com/cockroachdb/cockroach/pkg/sql/rowexec"
18+ "github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
19+ "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
20+ "github.com/cockroachdb/cockroach/pkg/sql/sem/tree/treewindow"
1421 "github.com/cockroachdb/cockroach/pkg/sql/types"
22+ "github.com/cockroachdb/cockroach/pkg/util/duration"
1523 "github.com/cockroachdb/errors"
1624)
1725
@@ -49,12 +57,207 @@ func createWindowFnSpec(
4957 }
5058 if funcInProgress .frame != nil {
5159 // funcInProgress has a custom window frame.
52- frameSpec := execinfrapb.WindowerSpec_Frame {}
53- if err := frameSpec . InitFromAST (ctx , funcInProgress .frame , planCtx .EvalContext ()); err != nil {
60+ frameSpec := & execinfrapb.WindowerSpec_Frame {}
61+ if err := initFrameFromAST (ctx , frameSpec , funcInProgress .frame , planCtx .EvalContext ()); err != nil {
5462 return execinfrapb.WindowerSpec_WindowFn {}, outputType , err
5563 }
56- funcInProgressSpec .Frame = & frameSpec
64+ funcInProgressSpec .Frame = frameSpec
5765 }
5866
5967 return funcInProgressSpec , outputType , nil
6068}
69+
70+ // initFrameFromAST initializes the spec based on tree.WindowFrame. It will
71+ // evaluate offset expressions if present in the frame.
72+ func initFrameFromAST (
73+ ctx context.Context ,
74+ spec * execinfrapb.WindowerSpec_Frame ,
75+ f * tree.WindowFrame ,
76+ evalCtx * eval.Context ,
77+ ) error {
78+ if err := initModeFromAST (& spec .Mode , f .Mode ); err != nil {
79+ return err
80+ }
81+ if err := initExclusionFromAST (& spec .Exclusion , f .Exclusion ); err != nil {
82+ return err
83+ }
84+ return initBoundsFromAST (ctx , & spec .Bounds , f .Bounds , f .Mode , evalCtx )
85+ }
86+
87+ func initModeFromAST (
88+ spec * execinfrapb.WindowerSpec_Frame_Mode , w treewindow.WindowFrameMode ,
89+ ) error {
90+ switch w {
91+ case treewindow .RANGE :
92+ * spec = execinfrapb .WindowerSpec_Frame_RANGE
93+ case treewindow .ROWS :
94+ * spec = execinfrapb .WindowerSpec_Frame_ROWS
95+ case treewindow .GROUPS :
96+ * spec = execinfrapb .WindowerSpec_Frame_GROUPS
97+ default :
98+ return errors .AssertionFailedf ("unexpected WindowFrameMode" )
99+ }
100+ return nil
101+ }
102+
103+ func initExclusionFromAST (
104+ spec * execinfrapb.WindowerSpec_Frame_Exclusion , e treewindow.WindowFrameExclusion ,
105+ ) error {
106+ switch e {
107+ case treewindow .NoExclusion :
108+ * spec = execinfrapb .WindowerSpec_Frame_NO_EXCLUSION
109+ case treewindow .ExcludeCurrentRow :
110+ * spec = execinfrapb .WindowerSpec_Frame_EXCLUDE_CURRENT_ROW
111+ case treewindow .ExcludeGroup :
112+ * spec = execinfrapb .WindowerSpec_Frame_EXCLUDE_GROUP
113+ case treewindow .ExcludeTies :
114+ * spec = execinfrapb .WindowerSpec_Frame_EXCLUDE_TIES
115+ default :
116+ return errors .AssertionFailedf ("unexpected WindowerFrameExclusion" )
117+ }
118+ return nil
119+ }
120+
121+ func initBoundTypeFromAST (
122+ spec * execinfrapb.WindowerSpec_Frame_BoundType , bt treewindow.WindowFrameBoundType ,
123+ ) error {
124+ switch bt {
125+ case treewindow .UnboundedPreceding :
126+ * spec = execinfrapb .WindowerSpec_Frame_UNBOUNDED_PRECEDING
127+ case treewindow .OffsetPreceding :
128+ * spec = execinfrapb .WindowerSpec_Frame_OFFSET_PRECEDING
129+ case treewindow .CurrentRow :
130+ * spec = execinfrapb .WindowerSpec_Frame_CURRENT_ROW
131+ case treewindow .OffsetFollowing :
132+ * spec = execinfrapb .WindowerSpec_Frame_OFFSET_FOLLOWING
133+ case treewindow .UnboundedFollowing :
134+ * spec = execinfrapb .WindowerSpec_Frame_UNBOUNDED_FOLLOWING
135+ default :
136+ return errors .AssertionFailedf ("unexpected WindowFrameBoundType" )
137+ }
138+ return nil
139+ }
140+
141+ // If offset exprs are present, we evaluate them and save the encoded results
142+ // in the spec.
143+ func initBoundsFromAST (
144+ ctx context.Context ,
145+ spec * execinfrapb.WindowerSpec_Frame_Bounds ,
146+ b tree.WindowFrameBounds ,
147+ m treewindow.WindowFrameMode ,
148+ evalCtx * eval.Context ,
149+ ) error {
150+ if b .StartBound == nil {
151+ return errors .Errorf ("unexpected: Start Bound is nil" )
152+ }
153+ spec .Start = execinfrapb.WindowerSpec_Frame_Bound {}
154+ if err := initBoundTypeFromAST (& spec .Start .BoundType , b .StartBound .BoundType ); err != nil {
155+ return err
156+ }
157+ if b .StartBound .HasOffset () {
158+ typedStartOffset := b .StartBound .OffsetExpr .(tree.TypedExpr )
159+ dStartOffset , err := eval .Expr (ctx , evalCtx , typedStartOffset )
160+ if err != nil {
161+ return err
162+ }
163+ if dStartOffset == tree .DNull {
164+ return pgerror .Newf (pgcode .NullValueNotAllowed , "frame starting offset must not be null" )
165+ }
166+ switch m {
167+ case treewindow .ROWS :
168+ startOffset := int64 (tree .MustBeDInt (dStartOffset ))
169+ if startOffset < 0 {
170+ return pgerror .Newf (pgcode .InvalidWindowFrameOffset , "frame starting offset must not be negative" )
171+ }
172+ spec .Start .IntOffset = uint64 (startOffset )
173+ case treewindow .RANGE :
174+ if neg , err := isNegative (ctx , evalCtx , dStartOffset ); err != nil {
175+ return err
176+ } else if neg {
177+ return pgerror .Newf (pgcode .InvalidWindowFrameOffset , "invalid preceding or following size in window function" )
178+ }
179+ typ := dStartOffset .ResolvedType ()
180+ spec .Start .OffsetType = execinfrapb.DatumInfo {Encoding : catenumpb .DatumEncoding_VALUE , Type : typ }
181+ var buf []byte
182+ var a tree.DatumAlloc
183+ datum := rowenc .DatumToEncDatum (typ , dStartOffset )
184+ buf , err = datum .Encode (typ , & a , catenumpb .DatumEncoding_VALUE , buf )
185+ if err != nil {
186+ return err
187+ }
188+ spec .Start .TypedOffset = buf
189+ case treewindow .GROUPS :
190+ startOffset := int64 (tree .MustBeDInt (dStartOffset ))
191+ if startOffset < 0 {
192+ return pgerror .Newf (pgcode .InvalidWindowFrameOffset , "frame starting offset must not be negative" )
193+ }
194+ spec .Start .IntOffset = uint64 (startOffset )
195+ }
196+ }
197+
198+ if b .EndBound != nil {
199+ spec .End = & execinfrapb.WindowerSpec_Frame_Bound {}
200+ if err := initBoundTypeFromAST (& spec .End .BoundType , b .EndBound .BoundType ); err != nil {
201+ return err
202+ }
203+ if b .EndBound .HasOffset () {
204+ typedEndOffset := b .EndBound .OffsetExpr .(tree.TypedExpr )
205+ dEndOffset , err := eval .Expr (ctx , evalCtx , typedEndOffset )
206+ if err != nil {
207+ return err
208+ }
209+ if dEndOffset == tree .DNull {
210+ return pgerror .Newf (pgcode .NullValueNotAllowed , "frame ending offset must not be null" )
211+ }
212+ switch m {
213+ case treewindow .ROWS :
214+ endOffset := int64 (tree .MustBeDInt (dEndOffset ))
215+ if endOffset < 0 {
216+ return pgerror .Newf (pgcode .InvalidWindowFrameOffset , "frame ending offset must not be negative" )
217+ }
218+ spec .End .IntOffset = uint64 (endOffset )
219+ case treewindow .RANGE :
220+ if neg , err := isNegative (ctx , evalCtx , dEndOffset ); err != nil {
221+ return err
222+ } else if neg {
223+ return pgerror .Newf (pgcode .InvalidWindowFrameOffset , "invalid preceding or following size in window function" )
224+ }
225+ typ := dEndOffset .ResolvedType ()
226+ spec .End .OffsetType = execinfrapb.DatumInfo {Encoding : catenumpb .DatumEncoding_VALUE , Type : typ }
227+ var buf []byte
228+ var a tree.DatumAlloc
229+ datum := rowenc .DatumToEncDatum (typ , dEndOffset )
230+ buf , err = datum .Encode (typ , & a , catenumpb .DatumEncoding_VALUE , buf )
231+ if err != nil {
232+ return err
233+ }
234+ spec .End .TypedOffset = buf
235+ case treewindow .GROUPS :
236+ endOffset := int64 (tree .MustBeDInt (dEndOffset ))
237+ if endOffset < 0 {
238+ return pgerror .Newf (pgcode .InvalidWindowFrameOffset , "frame ending offset must not be negative" )
239+ }
240+ spec .End .IntOffset = uint64 (endOffset )
241+ }
242+ }
243+ }
244+
245+ return nil
246+ }
247+
248+ // isNegative returns whether offset is negative.
249+ func isNegative (ctx context.Context , evalCtx * eval.Context , offset tree.Datum ) (bool , error ) {
250+ switch o := offset .(type ) {
251+ case * tree.DInt :
252+ return * o < 0 , nil
253+ case * tree.DDecimal :
254+ return o .Negative , nil
255+ case * tree.DFloat :
256+ return * o < 0 , nil
257+ case * tree.DInterval :
258+ cmp , err := o .Compare (ctx , evalCtx , & tree.DInterval {Duration : duration.Duration {}})
259+ return cmp < 0 , err
260+ default :
261+ panic ("unexpected offset type" )
262+ }
263+ }
![craig[bot]](https://github.githubassets.com/images/gravatars/gravatar-user-420.png?size=40){kind=avatar}
0 commit comments