Fix incorrect window function cancellation handling (#5818)

dbms/src/DataStreams/WindowBlockInputStream.cpp

@@ -75,15 +75,28 @@ void WindowBlockInputStream::initialWorkspaces()
     only_have_pure_window = onlyHaveRowNumberAndRank();
 }
 
+bool WindowBlockInputStream::returnIfCancelledOrKilled()
+{
+    if (isCancelledOrThrowIfKilled())
+    {
+        if (!window_blocks.empty())
+            window_blocks.erase(window_blocks.begin(), window_blocks.end());
+        input_is_finished = true;
+        return true;
+    }
+    return false;
+}
+
 Block WindowBlockInputStream::readImpl()
 {
     const auto & stream = children.back();
     while (!input_is_finished)
     {
+        if (returnIfCancelledOrKilled())
+            return {};
+
         if (Block output_block = tryGetOutputBlock())
-        {
             return output_block;
-        }
 
         Block block = stream->read();
         if (!block)
@@ -93,6 +106,8 @@ Block WindowBlockInputStream::readImpl()
         tryCalculate();
     }
 
+    if (returnIfCancelledOrKilled())
+        return {};
     // return last partition block, if already return then return null
     return tryGetOutputBlock();
 }
@@ -360,14 +375,6 @@ void WindowBlockInputStream::writeOutCurrentRow()
 
 Block WindowBlockInputStream::tryGetOutputBlock()
 {
-    if (isCancelledOrThrowIfKilled())
-    {
-        if (!window_blocks.empty())
-            window_blocks.erase(window_blocks.begin(), window_blocks.end());
-        input_is_finished = true;
-        return {};
-    }
-
     assert(first_not_ready_row.block >= first_block_number);
     // The first_not_ready_row might be past-the-end if we have already
     // calculated the window functions for all input rows. That's why the

dbms/src/DataStreams/WindowBlockInputStream.h

@@ -169,6 +169,11 @@ class WindowBlockInputStream : public IProfilingBlockInputStream
 
 protected:
     Block readImpl() override;
+<<<<<<< HEAD
+=======
+    void appendInfo(FmtBuffer & buffer) const override;
+    bool returnIfCancelledOrKilled();
+>>>>>>> 8ffb734c88 (Fix incorrect window function cancellation handling (#5818))
 
     LoggerPtr log;