github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/sem/tree/window_funcs.go (about) 1 // Copyright 2017 The Cockroach Authors. 2 // 3 // Use of this software is governed by the Business Source License 4 // included in the file licenses/BSL.txt. 5 // 6 // As of the Change Date specified in that file, in accordance with 7 // the Business Source License, use of this software will be governed 8 // by the Apache License, Version 2.0, included in the file 9 // licenses/APL.txt. 10 11 package tree 12 13 import ( 14 "context" 15 "sort" 16 17 "github.com/cockroachdb/cockroach/pkg/sql/types" 18 "github.com/cockroachdb/cockroach/pkg/util/encoding" 19 "github.com/cockroachdb/cockroach/pkg/util/log" 20 "github.com/cockroachdb/errors" 21 ) 22 23 // IndexedRows are rows with the corresponding indices. 24 type IndexedRows interface { 25 Len() int // returns number of rows 26 GetRow(ctx context.Context, idx int) (IndexedRow, error) // returns a row at the given index or an error 27 } 28 29 // IndexedRow is a row with a corresponding index. 30 type IndexedRow interface { 31 GetIdx() int // returns index of the row 32 GetDatum(idx int) (Datum, error) // returns a datum at the given index 33 GetDatums(startIdx, endIdx int) (Datums, error) // returns datums at indices [startIdx, endIdx) 34 } 35 36 // WindowFrameRun contains the runtime state of window frame during calculations. 37 type WindowFrameRun struct { 38 // constant for all calls to WindowFunc.Add 39 Rows IndexedRows 40 ArgsIdxs []uint32 // indices of the arguments to the window function 41 Frame *WindowFrame // If non-nil, Frame represents the frame specification of this window. If nil, default frame is used. 42 StartBoundOffset Datum 43 EndBoundOffset Datum 44 FilterColIdx int 45 OrdColIdx int // Column over which rows are ordered within the partition. It is only required in RANGE mode. 46 OrdDirection encoding.Direction // Direction of the ordering over OrdColIdx. 47 PlusOp, MinusOp *BinOp // Binary operators for addition and subtraction required only in RANGE mode. 48 PeerHelper PeerGroupsIndicesHelper 49 50 // Any error that occurred within methods that cannot return an error (like 51 // within a closure that is passed into sort.Search()). 52 err error 53 54 // changes for each peer group 55 CurRowPeerGroupNum int // the number of the current row's peer group 56 57 // changes for each row (each call to WindowFunc.Add) 58 RowIdx int // the current row index 59 } 60 61 // WindowFrameRangeOps allows for looking up an implementation of binary 62 // operators necessary for RANGE mode of framing. 63 type WindowFrameRangeOps struct{} 64 65 // LookupImpl looks up implementation of Plus and Minus binary operators for 66 // provided left and right types and returns them along with a boolean which 67 // indicates whether lookup is successful. 68 func (o WindowFrameRangeOps) LookupImpl(left, right *types.T) (*BinOp, *BinOp, bool) { 69 plusOverloads, minusOverloads := BinOps[Plus], BinOps[Minus] 70 plusOp, found := plusOverloads.lookupImpl(left, right) 71 if !found { 72 return nil, nil, false 73 } 74 minusOp, found := minusOverloads.lookupImpl(left, right) 75 if !found { 76 return nil, nil, false 77 } 78 return plusOp, minusOp, true 79 } 80 81 // getValueByOffset returns a datum calculated as the value of the current row 82 // in the column over which rows are ordered plus/minus logic offset, and an 83 // error if encountered. It should be used only in RANGE mode. 84 func (wfr *WindowFrameRun) getValueByOffset( 85 ctx context.Context, evalCtx *EvalContext, offset Datum, negative bool, 86 ) (Datum, error) { 87 if wfr.OrdDirection == encoding.Descending { 88 // If rows are in descending order, we want to perform the "opposite" 89 // addition/subtraction to default ascending order. 90 negative = !negative 91 } 92 var binOp *BinOp 93 if negative { 94 binOp = wfr.MinusOp 95 } else { 96 binOp = wfr.PlusOp 97 } 98 value, err := wfr.valueAt(ctx, wfr.RowIdx) 99 if err != nil { 100 return nil, err 101 } 102 if value == DNull { 103 return DNull, nil 104 } 105 return binOp.Fn(evalCtx, value, offset) 106 } 107 108 // FrameStartIdx returns the index of starting row in the frame (which is the first to be included). 109 func (wfr *WindowFrameRun) FrameStartIdx(ctx context.Context, evalCtx *EvalContext) (int, error) { 110 if wfr.Frame == nil { 111 return 0, nil 112 } 113 switch wfr.Frame.Mode { 114 case RANGE: 115 switch wfr.Frame.Bounds.StartBound.BoundType { 116 case UnboundedPreceding: 117 return 0, nil 118 case OffsetPreceding: 119 value, err := wfr.getValueByOffset(ctx, evalCtx, wfr.StartBoundOffset, true /* negative */) 120 if err != nil { 121 return 0, err 122 } 123 if wfr.OrdDirection == encoding.Descending { 124 // We use binary search on [0, wfr.RowIdx) interval to find the first row 125 // whose value is smaller or equal to 'value'. If such row is not found, 126 // then Search will correctly return wfr.RowIdx. 127 return sort.Search(wfr.RowIdx, func(i int) bool { 128 if wfr.err != nil { 129 return false 130 } 131 valueAt, err := wfr.valueAt(ctx, i) 132 if err != nil { 133 wfr.err = err 134 return false 135 } 136 return valueAt.Compare(evalCtx, value) <= 0 137 }), wfr.err 138 } 139 // We use binary search on [0, wfr.RowIdx) interval to find the first row 140 // whose value is greater or equal to 'value'. If such row is not found, 141 // then Search will correctly return wfr.RowIdx. 142 return sort.Search(wfr.RowIdx, func(i int) bool { 143 if wfr.err != nil { 144 return false 145 } 146 valueAt, err := wfr.valueAt(ctx, i) 147 if err != nil { 148 wfr.err = err 149 return false 150 } 151 return valueAt.Compare(evalCtx, value) >= 0 152 }), wfr.err 153 case CurrentRow: 154 // Spec: in RANGE mode CURRENT ROW means that the frame starts with the current row's first peer. 155 return wfr.PeerHelper.GetFirstPeerIdx(wfr.CurRowPeerGroupNum), nil 156 case OffsetFollowing: 157 value, err := wfr.getValueByOffset(ctx, evalCtx, wfr.StartBoundOffset, false /* negative */) 158 if err != nil { 159 return 0, err 160 } 161 if wfr.OrdDirection == encoding.Descending { 162 // We use binary search on [0, wfr.PartitionSize()) interval to find 163 // the first row whose value is smaller or equal to 'value'. 164 return sort.Search(wfr.PartitionSize(), func(i int) bool { 165 if wfr.err != nil { 166 return false 167 } 168 valueAt, err := wfr.valueAt(ctx, i) 169 if err != nil { 170 wfr.err = err 171 return false 172 } 173 return valueAt.Compare(evalCtx, value) <= 0 174 }), wfr.err 175 } 176 // We use binary search on [0, wfr.PartitionSize()) interval to find the 177 // first row whose value is greater or equal to 'value'. 178 return sort.Search(wfr.PartitionSize(), func(i int) bool { 179 if wfr.err != nil { 180 return false 181 } 182 valueAt, err := wfr.valueAt(ctx, i) 183 if err != nil { 184 wfr.err = err 185 return false 186 } 187 return valueAt.Compare(evalCtx, value) >= 0 188 }), wfr.err 189 default: 190 return 0, errors.AssertionFailedf( 191 "unexpected WindowFrameBoundType in RANGE mode: %d", 192 log.Safe(wfr.Frame.Bounds.StartBound.BoundType)) 193 } 194 case ROWS: 195 switch wfr.Frame.Bounds.StartBound.BoundType { 196 case UnboundedPreceding: 197 return 0, nil 198 case OffsetPreceding: 199 offset := MustBeDInt(wfr.StartBoundOffset) 200 idx := wfr.RowIdx - int(offset) 201 if idx < 0 { 202 idx = 0 203 } 204 return idx, nil 205 case CurrentRow: 206 return wfr.RowIdx, nil 207 case OffsetFollowing: 208 offset := MustBeDInt(wfr.StartBoundOffset) 209 idx := wfr.RowIdx + int(offset) 210 if idx >= wfr.PartitionSize() { 211 idx = wfr.unboundedFollowing() 212 } 213 return idx, nil 214 default: 215 return 0, errors.AssertionFailedf( 216 "unexpected WindowFrameBoundType in ROWS mode: %d", 217 log.Safe(wfr.Frame.Bounds.StartBound.BoundType)) 218 } 219 case GROUPS: 220 switch wfr.Frame.Bounds.StartBound.BoundType { 221 case UnboundedPreceding: 222 return 0, nil 223 case OffsetPreceding: 224 offset := MustBeDInt(wfr.StartBoundOffset) 225 peerGroupNum := wfr.CurRowPeerGroupNum - int(offset) 226 if peerGroupNum < 0 { 227 peerGroupNum = 0 228 } 229 return wfr.PeerHelper.GetFirstPeerIdx(peerGroupNum), nil 230 case CurrentRow: 231 // Spec: in GROUPS mode CURRENT ROW means that the frame starts with the current row's first peer. 232 return wfr.PeerHelper.GetFirstPeerIdx(wfr.CurRowPeerGroupNum), nil 233 case OffsetFollowing: 234 offset := MustBeDInt(wfr.StartBoundOffset) 235 peerGroupNum := wfr.CurRowPeerGroupNum + int(offset) 236 lastPeerGroupNum := wfr.PeerHelper.GetLastPeerGroupNum() 237 if peerGroupNum > lastPeerGroupNum { 238 // peerGroupNum is out of bounds, so we return the index of the first 239 // row after the partition. 240 return wfr.unboundedFollowing(), nil 241 } 242 return wfr.PeerHelper.GetFirstPeerIdx(peerGroupNum), nil 243 default: 244 return 0, errors.AssertionFailedf( 245 "unexpected WindowFrameBoundType in GROUPS mode: %d", 246 log.Safe(wfr.Frame.Bounds.StartBound.BoundType)) 247 } 248 default: 249 return 0, errors.AssertionFailedf("unexpected WindowFrameMode: %d", wfr.Frame.Mode) 250 } 251 } 252 253 // IsDefaultFrame returns whether a frame equivalent to the default frame 254 // is being used (default is RANGE UNBOUNDED PRECEDING). 255 func (f *WindowFrame) IsDefaultFrame() bool { 256 if f == nil { 257 return true 258 } 259 if f.Bounds.StartBound.BoundType == UnboundedPreceding { 260 return f.DefaultFrameExclusion() && f.Mode == RANGE && 261 (f.Bounds.EndBound == nil || f.Bounds.EndBound.BoundType == CurrentRow) 262 } 263 return false 264 } 265 266 // DefaultFrameExclusion returns true if optional frame exclusion is omitted. 267 func (f *WindowFrame) DefaultFrameExclusion() bool { 268 return f == nil || f.Exclusion == NoExclusion 269 } 270 271 // FrameEndIdx returns the index of the first row after the frame. 272 func (wfr *WindowFrameRun) FrameEndIdx(ctx context.Context, evalCtx *EvalContext) (int, error) { 273 if wfr.Frame == nil { 274 return wfr.DefaultFrameSize(), nil 275 } 276 switch wfr.Frame.Mode { 277 case RANGE: 278 if wfr.Frame.Bounds.EndBound == nil { 279 // We're using default value of CURRENT ROW when EndBound is omitted. 280 // Spec: in RANGE mode CURRENT ROW means that the frame ends with the current row's last peer. 281 return wfr.DefaultFrameSize(), nil 282 } 283 switch wfr.Frame.Bounds.EndBound.BoundType { 284 case OffsetPreceding: 285 value, err := wfr.getValueByOffset(ctx, evalCtx, wfr.EndBoundOffset, true /* negative */) 286 if err != nil { 287 return 0, err 288 } 289 if wfr.OrdDirection == encoding.Descending { 290 // We use binary search on [0, wfr.PartitionSize()) interval to find 291 // the first row whose value is smaller than 'value'. If such row is 292 // not found, then Search will correctly return wfr.PartitionSize(). 293 // Note that searching up to wfr.RowIdx is not correct in case of a 294 // zero offset (we need to include all peers of the current row). 295 return sort.Search(wfr.PartitionSize(), func(i int) bool { 296 if wfr.err != nil { 297 return false 298 } 299 valueAt, err := wfr.valueAt(ctx, i) 300 if err != nil { 301 wfr.err = err 302 return false 303 } 304 return valueAt.Compare(evalCtx, value) < 0 305 }), wfr.err 306 } 307 // We use binary search on [0, wfr.PartitionSize()) interval to find 308 // the first row whose value is smaller than 'value'. If such row is 309 // not found, then Search will correctly return wfr.PartitionSize(). 310 // Note that searching up to wfr.RowIdx is not correct in case of a 311 // zero offset (we need to include all peers of the current row). 312 return sort.Search(wfr.PartitionSize(), func(i int) bool { 313 if wfr.err != nil { 314 return false 315 } 316 valueAt, err := wfr.valueAt(ctx, i) 317 if err != nil { 318 wfr.err = err 319 return false 320 } 321 return valueAt.Compare(evalCtx, value) > 0 322 }), wfr.err 323 case CurrentRow: 324 // Spec: in RANGE mode CURRENT ROW means that the frame end with the current row's last peer. 325 return wfr.DefaultFrameSize(), nil 326 case OffsetFollowing: 327 value, err := wfr.getValueByOffset(ctx, evalCtx, wfr.EndBoundOffset, false /* negative */) 328 if err != nil { 329 return 0, err 330 } 331 if wfr.OrdDirection == encoding.Descending { 332 // We use binary search on [0, wfr.PartitionSize()) interval to find 333 // the first row whose value is smaller than 'value'. 334 return sort.Search(wfr.PartitionSize(), func(i int) bool { 335 if wfr.err != nil { 336 return false 337 } 338 valueAt, err := wfr.valueAt(ctx, i) 339 if err != nil { 340 wfr.err = err 341 return false 342 } 343 return valueAt.Compare(evalCtx, value) < 0 344 }), wfr.err 345 } 346 // We use binary search on [0, wfr.PartitionSize()) interval to find 347 // the first row whose value is smaller than 'value'. 348 return sort.Search(wfr.PartitionSize(), func(i int) bool { 349 if wfr.err != nil { 350 return false 351 } 352 valueAt, err := wfr.valueAt(ctx, i) 353 if err != nil { 354 wfr.err = err 355 return false 356 } 357 return valueAt.Compare(evalCtx, value) > 0 358 }), wfr.err 359 case UnboundedFollowing: 360 return wfr.unboundedFollowing(), nil 361 default: 362 return 0, errors.AssertionFailedf( 363 "unexpected WindowFrameBoundType in RANGE mode: %d", 364 log.Safe(wfr.Frame.Bounds.EndBound.BoundType)) 365 } 366 case ROWS: 367 if wfr.Frame.Bounds.EndBound == nil { 368 // We're using default value of CURRENT ROW when EndBound is omitted. 369 return wfr.RowIdx + 1, nil 370 } 371 switch wfr.Frame.Bounds.EndBound.BoundType { 372 case OffsetPreceding: 373 offset := MustBeDInt(wfr.EndBoundOffset) 374 idx := wfr.RowIdx - int(offset) + 1 375 if idx < 0 { 376 idx = 0 377 } 378 return idx, nil 379 case CurrentRow: 380 return wfr.RowIdx + 1, nil 381 case OffsetFollowing: 382 offset := MustBeDInt(wfr.EndBoundOffset) 383 idx := wfr.RowIdx + int(offset) + 1 384 if idx >= wfr.PartitionSize() { 385 idx = wfr.unboundedFollowing() 386 } 387 return idx, nil 388 case UnboundedFollowing: 389 return wfr.unboundedFollowing(), nil 390 default: 391 return 0, errors.AssertionFailedf( 392 "unexpected WindowFrameBoundType in ROWS mode: %d", 393 log.Safe(wfr.Frame.Bounds.EndBound.BoundType)) 394 } 395 case GROUPS: 396 if wfr.Frame.Bounds.EndBound == nil { 397 // We're using default value of CURRENT ROW when EndBound is omitted. 398 // Spec: in GROUPS mode CURRENT ROW means that the frame ends with the current row's last peer. 399 return wfr.DefaultFrameSize(), nil 400 } 401 switch wfr.Frame.Bounds.EndBound.BoundType { 402 case OffsetPreceding: 403 offset := MustBeDInt(wfr.EndBoundOffset) 404 peerGroupNum := wfr.CurRowPeerGroupNum - int(offset) 405 if peerGroupNum < 0 { 406 // EndBound's peer group is "outside" of the partition. 407 return 0, nil 408 } 409 return wfr.PeerHelper.GetFirstPeerIdx(peerGroupNum) + wfr.PeerHelper.GetRowCount(peerGroupNum), nil 410 case CurrentRow: 411 return wfr.DefaultFrameSize(), nil 412 case OffsetFollowing: 413 offset := MustBeDInt(wfr.EndBoundOffset) 414 peerGroupNum := wfr.CurRowPeerGroupNum + int(offset) 415 lastPeerGroupNum := wfr.PeerHelper.GetLastPeerGroupNum() 416 if peerGroupNum > lastPeerGroupNum { 417 // peerGroupNum is out of bounds, so we return the index of the first 418 // row after the partition. 419 return wfr.unboundedFollowing(), nil 420 } 421 return wfr.PeerHelper.GetFirstPeerIdx(peerGroupNum) + wfr.PeerHelper.GetRowCount(peerGroupNum), nil 422 case UnboundedFollowing: 423 return wfr.unboundedFollowing(), nil 424 default: 425 return 0, errors.AssertionFailedf( 426 "unexpected WindowFrameBoundType in GROUPS mode: %d", 427 log.Safe(wfr.Frame.Bounds.EndBound.BoundType)) 428 } 429 default: 430 return 0, errors.AssertionFailedf( 431 "unexpected WindowFrameMode: %d", log.Safe(wfr.Frame.Mode)) 432 } 433 } 434 435 // FrameSize returns the number of rows in the current frame (taking into 436 // account - if present - a filter and a frame exclusion). 437 func (wfr *WindowFrameRun) FrameSize(ctx context.Context, evalCtx *EvalContext) (int, error) { 438 if wfr.Frame == nil { 439 return wfr.DefaultFrameSize(), nil 440 } 441 frameEndIdx, err := wfr.FrameEndIdx(ctx, evalCtx) 442 if err != nil { 443 return 0, err 444 } 445 frameStartIdx, err := wfr.FrameStartIdx(ctx, evalCtx) 446 if err != nil { 447 return 0, err 448 } 449 size := frameEndIdx - frameStartIdx 450 if !wfr.noFilter() || !wfr.Frame.DefaultFrameExclusion() { 451 size = 0 452 for idx := frameStartIdx; idx < frameEndIdx; idx++ { 453 if skipped, err := wfr.IsRowSkipped(ctx, idx); err != nil { 454 return 0, err 455 } else if skipped { 456 continue 457 } 458 size++ 459 } 460 } 461 if size <= 0 { 462 size = 0 463 } 464 return size, nil 465 } 466 467 // Rank returns the rank of the current row. 468 func (wfr *WindowFrameRun) Rank() int { 469 return wfr.RowIdx + 1 470 } 471 472 // PartitionSize returns the number of rows in the current partition. 473 func (wfr *WindowFrameRun) PartitionSize() int { 474 return wfr.Rows.Len() 475 } 476 477 // unboundedFollowing returns the index of the "first row beyond" the partition 478 // so that current frame contains all the rows till the end of the partition. 479 func (wfr *WindowFrameRun) unboundedFollowing() int { 480 return wfr.PartitionSize() 481 } 482 483 // DefaultFrameSize returns the size of default window frame which contains 484 // the rows from the start of the partition through the last peer of the current row. 485 func (wfr *WindowFrameRun) DefaultFrameSize() int { 486 return wfr.PeerHelper.GetFirstPeerIdx(wfr.CurRowPeerGroupNum) + wfr.PeerHelper.GetRowCount(wfr.CurRowPeerGroupNum) 487 } 488 489 // FirstInPeerGroup returns if the current row is the first in its peer group. 490 func (wfr *WindowFrameRun) FirstInPeerGroup() bool { 491 return wfr.RowIdx == wfr.PeerHelper.GetFirstPeerIdx(wfr.CurRowPeerGroupNum) 492 } 493 494 // Args returns the current argument set in the window frame. 495 func (wfr *WindowFrameRun) Args(ctx context.Context) (Datums, error) { 496 return wfr.ArgsWithRowOffset(ctx, 0) 497 } 498 499 // ArgsWithRowOffset returns the argument set at the given offset in the window frame. 500 func (wfr *WindowFrameRun) ArgsWithRowOffset(ctx context.Context, offset int) (Datums, error) { 501 return wfr.ArgsByRowIdx(ctx, wfr.RowIdx+offset) 502 } 503 504 // ArgsByRowIdx returns the argument set of the row at idx. 505 func (wfr *WindowFrameRun) ArgsByRowIdx(ctx context.Context, idx int) (Datums, error) { 506 row, err := wfr.Rows.GetRow(ctx, idx) 507 if err != nil { 508 return nil, err 509 } 510 datums := make(Datums, len(wfr.ArgsIdxs)) 511 for i, argIdx := range wfr.ArgsIdxs { 512 datums[i], err = row.GetDatum(int(argIdx)) 513 if err != nil { 514 return nil, err 515 } 516 } 517 return datums, nil 518 } 519 520 // valueAt returns the first argument of the window function at the row idx. 521 func (wfr *WindowFrameRun) valueAt(ctx context.Context, idx int) (Datum, error) { 522 row, err := wfr.Rows.GetRow(ctx, idx) 523 if err != nil { 524 return nil, err 525 } 526 return row.GetDatum(wfr.OrdColIdx) 527 } 528 529 // RangeModeWithOffsets returns whether the frame is in RANGE mode with at least 530 // one of the bounds containing an offset. 531 func (wfr *WindowFrameRun) RangeModeWithOffsets() bool { 532 return wfr.Frame.Mode == RANGE && wfr.Frame.Bounds.HasOffset() 533 } 534 535 // FullPartitionIsInWindow checks whether we have such a window frame that all 536 // rows of the partition are inside of the window for each of the rows. 537 func (wfr *WindowFrameRun) FullPartitionIsInWindow() bool { 538 // Note that we do not need to check whether a filter is present because 539 // application of the filter to a row does not depend on the position of the 540 // row or whether it is inside of the window frame. 541 if wfr.Frame == nil || !wfr.Frame.DefaultFrameExclusion() { 542 return false 543 } 544 if wfr.Frame.Bounds.EndBound == nil { 545 // If the end bound is omitted, it is CURRENT ROW (the default value) which 546 // doesn't guarantee full partition in the window for all rows. 547 return false 548 } 549 // precedingConfirmed and followingConfirmed indicate whether, for every row, 550 // all preceding and following, respectively, rows are always in the window. 551 precedingConfirmed := wfr.Frame.Bounds.StartBound.BoundType == UnboundedPreceding 552 followingConfirmed := wfr.Frame.Bounds.EndBound.BoundType == UnboundedFollowing 553 if wfr.Frame.Mode == ROWS || wfr.Frame.Mode == GROUPS { 554 // Every peer group in GROUPS modealways contains at least one row, so 555 // treating GROUPS as ROWS here is a subset of the cases when we should 556 // return true. 557 if wfr.Frame.Bounds.StartBound.BoundType == OffsetPreceding { 558 // Both ROWS and GROUPS have an offset of integer type, so this type 559 // conversion is safe. 560 startOffset := wfr.StartBoundOffset.(*DInt) 561 // The idea of this conditional is that to confirm that all preceding 562 // rows will always be in the window, we only need to look at the last 563 // row: if startOffset is at least as large as the number of rows in the 564 // partition before the last one, then it will be true for the first to 565 // last, second to last, etc. 566 precedingConfirmed = precedingConfirmed || *startOffset >= DInt(wfr.Rows.Len()-1) 567 } 568 if wfr.Frame.Bounds.EndBound.BoundType == OffsetFollowing { 569 // Both ROWS and GROUPS have an offset of integer type, so this type 570 // conversion is safe. 571 endOffset := wfr.EndBoundOffset.(*DInt) 572 // The idea of this conditional is that to confirm that all following 573 // rows will always be in the window, we only need to look at the first 574 // row: if endOffset is at least as large as the number of rows in the 575 // partition after the first one, then it will be true for the second, 576 // third, etc rows as well. 577 followingConfirmed = followingConfirmed || *endOffset >= DInt(wfr.Rows.Len()-1) 578 } 579 } 580 return precedingConfirmed && followingConfirmed 581 } 582 583 const noFilterIdx = -1 584 585 // noFilter returns whether a filter is present. 586 func (wfr *WindowFrameRun) noFilter() bool { 587 return wfr.FilterColIdx == noFilterIdx 588 } 589 590 // isRowExcluded returns whether the row at index idx should be excluded from 591 // the window frame of the current row. 592 func (wfr *WindowFrameRun) isRowExcluded(idx int) (bool, error) { 593 if wfr.Frame.DefaultFrameExclusion() { 594 // By default, no rows are excluded. 595 return false, nil 596 } 597 switch wfr.Frame.Exclusion { 598 case ExcludeCurrentRow: 599 return idx == wfr.RowIdx, nil 600 case ExcludeGroup: 601 curRowFirstPeerIdx := wfr.PeerHelper.GetFirstPeerIdx(wfr.CurRowPeerGroupNum) 602 curRowPeerGroupRowCount := wfr.PeerHelper.GetRowCount(wfr.CurRowPeerGroupNum) 603 return curRowFirstPeerIdx <= idx && idx < curRowFirstPeerIdx+curRowPeerGroupRowCount, nil 604 case ExcludeTies: 605 curRowFirstPeerIdx := wfr.PeerHelper.GetFirstPeerIdx(wfr.CurRowPeerGroupNum) 606 curRowPeerGroupRowCount := wfr.PeerHelper.GetRowCount(wfr.CurRowPeerGroupNum) 607 return curRowFirstPeerIdx <= idx && idx < curRowFirstPeerIdx+curRowPeerGroupRowCount && idx != wfr.RowIdx, nil 608 default: 609 return false, errors.AssertionFailedf("unexpected WindowFrameExclusion") 610 } 611 } 612 613 // IsRowSkipped returns whether a row at index idx is skipped from the window 614 // frame (it can either be filtered out according to the filter clause or 615 // excluded according to the frame exclusion clause) and any error if it 616 // occurs. 617 func (wfr *WindowFrameRun) IsRowSkipped(ctx context.Context, idx int) (bool, error) { 618 if !wfr.noFilter() { 619 row, err := wfr.Rows.GetRow(ctx, idx) 620 if err != nil { 621 return false, err 622 } 623 d, err := row.GetDatum(wfr.FilterColIdx) 624 if err != nil { 625 return false, err 626 } 627 if d != DBoolTrue { 628 // Row idx is filtered out from the window frame, so it is skipped. 629 return true, nil 630 } 631 } 632 // If a row is excluded from the window frame, it is skipped. 633 return wfr.isRowExcluded(idx) 634 } 635 636 // WindowFunc performs a computation on each row using data from a provided *WindowFrameRun. 637 type WindowFunc interface { 638 // Compute computes the window function for the provided window frame, given the 639 // current state of WindowFunc. The method should be called sequentially for every 640 // row in a partition in turn with the desired ordering of the WindowFunc. This is 641 // because there is an implicit carried dependency between each row and all those 642 // that have come before it (like in an AggregateFunc). As such, this approach does 643 // not present any exploitable associativity/commutativity for optimization. 644 Compute(context.Context, *EvalContext, *WindowFrameRun) (Datum, error) 645 646 // Reset resets the window function which allows for reusing it when 647 // computing over different partitions. 648 Reset(context.Context) 649 650 // Close allows the window function to free any memory it requested during execution, 651 // such as during the execution of an aggregation like CONCAT_AGG or ARRAY_AGG. 652 Close(context.Context, *EvalContext) 653 }