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

@@ -170,6 +170,8 @@ class WindowBlockInputStream : public IProfilingBlockInputStream
 protected:
     Block readImpl() override;
 
+    bool returnIfCancelledOrKilled();
+
     LoggerPtr log;
 
 public:

format-diff.py

@@ -89,16 +89,22 @@ def main():
 
         if args.check_formatted:
             diff_res = run_cmd('git diff --name-only')
-            if diff_res:
+            files_not_in_contrib = [f for f in diff_res if not f.startswith('contrib')]
+            files_contrib = [f for f in diff_res if f.startswith('contrib')]
+            if files_not_in_contrib:
+                print('')
                 print('Error: found files NOT formatted')
-                print(''.join(diff_res))
+                print(''.join(files_not_in_contrib))
                 exit(-1)
+            elif files_contrib:
+                print('')
+                print('Warn: found contrib changed')
+                print(''.join(files_contrib))
+                print('')
+                print(''.join(run_cmd('git status')))
             else:
                 print("Format check passed")
         else:
-            cmd = 'clang-format -i {}'.format(' '.join(files_to_format))
-            if subprocess.Popen(cmd, shell=True, cwd=tics_repo_path).wait():
-                exit(-1)
             print("Finish code format")
     else:
         print('No file to format')