diff --git a/velox/exec/Window.cpp b/velox/exec/Window.cpp index 2681f18f3d12..4dbd531796e9 100644 --- a/velox/exec/Window.cpp +++ b/velox/exec/Window.cpp @@ -294,6 +294,15 @@ void updateKRowsOffsetsColumn( } } +// Get the start index that integer overflow can happen for kFollowing. +FOLLY_ALWAYS_INLINE int32_t getOverflowStart(int64_t constantOffset) { + if (constantOffset > (int64_t)std::numeric_limits::max()) { + return 0; + } else { + return std::numeric_limits::max() - constantOffset + 1; + } +} + } // namespace void Window::updateKRowsFrameBounds( @@ -306,14 +315,28 @@ void Window::updateKRowsFrameBounds( auto constantOffset = frameArg.constant.value(); auto startValue = (isKPreceding ? -constantOffset : constantOffset) + startRow; - for (int i = 0; i < numRows; i++) { - // Considers integer overflow case. - if (startValue != (int32_t)startValue) { - rawFrameBounds[i] = startValue < 0 ? 0 : numRows - 1; + + if (isKPreceding) { + // Considers a very large int64 constantOffset is used. + if (startValue < std::numeric_limits::min()) { + std::fill_n(rawFrameBounds, numRows, startRow); + } else { + // Integer overflow cannot happen. + std::iota(rawFrameBounds, rawFrameBounds + numRows, startValue); + } + } else { + auto overflowStart = getOverflowStart(constantOffset); + if (overflowStart >= 0 && overflowStart < numRows) { + std::iota(rawFrameBounds, rawFrameBounds + overflowStart, startValue); + // For remaining, set with the largest index for this partition. + std::fill_n( + rawFrameBounds + overflowStart, + numRows - overflowStart, + startRow + numRows - 1); } else { - rawFrameBounds[i] = startValue; + // Integer overflow cannot happen. + std::iota(rawFrameBounds, rawFrameBounds + numRows, startValue); } - startValue = startValue + 1; } } else { currentPartition_->extractColumn(