Fix incorrect window function cancellation handling

dbms/src/DataStreams/WindowBlockInputStream.cpp

@@ -73,15 +73,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)
@@ -91,6 +104,8 @@ Block WindowBlockInputStream::readImpl()
         tryCalculate();
     }
 
+    if (returnIfCancelledOrKilled())
+        return {};
     // return last partition block, if already return then return null
     return tryGetOutputBlock();
 }
@@ -374,14 +389,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

@@ -166,6 +166,7 @@ class WindowBlockInputStream : public IProfilingBlockInputStream
 protected:
     Block readImpl() override;
     void appendInfo(FmtBuffer & buffer) const override;
+    bool returnIfCancelledOrKilled();
 
     LoggerPtr log;