Fix starving caused by GIL

Author: auxten

src/Processors/Sources/PythonSource.cpp

@@ -48,6 +48,7 @@ extern const int PY_EXCEPTION_OCCURED;
 
 PythonSource::PythonSource(
     py::object & data_source_,
+    bool isInheritsFromPyReader_,
     const Block & sample_block_,
     PyColumnVecPtr column_cache,
     size_t data_source_row_count,
@@ -56,6 +57,7 @@ PythonSource::PythonSource(
     size_t num_streams)
     : ISource(sample_block_.cloneEmpty())
     , data_source(data_source_)
+    , isInheritsFromPyReader(isInheritsFromPyReader_)
     , sample_block(sample_block_)
     , column_cache(column_cache)
     , data_source_row_count(data_source_row_count)
@@ -544,7 +546,7 @@ Chunk PythonSource::generate()
 
     try
     {
-        if (isInheritsFromPyReader(data_source))
+        if (isInheritsFromPyReader)
         {
             PyObjectVecPtr data;
             py::gil_scoped_acquire acquire;

src/Processors/Sources/PythonSource.h

@@ -26,6 +26,7 @@ class PythonSource : public ISource
 public:
     PythonSource(
         py::object & data_source_,
+        bool isInheritsFromPyReader_,
         const Block & sample_block_,
         PyColumnVecPtr column_cache,
         size_t data_source_row_count,
@@ -42,6 +43,7 @@ class PythonSource : public ISource
 
 private:
     py::object & data_source; // Do not own the reference
+    bool isInheritsFromPyReader; // If the data_source is a PyReader object
 
     Block sample_block;
     PyColumnVecPtr column_cache;

src/Storages/StoragePython.cpp

@@ -70,7 +70,8 @@ Pipe StoragePython::read(
 
     if (isInheritsFromPyReader(data_source))
     {
-        return Pipe(std::make_shared<PythonSource>(data_source, sample_block, column_cache, data_source_row_count, max_block_size, 0, 1));
+        return Pipe(
+            std::make_shared<PythonSource>(data_source, true, sample_block, column_cache, data_source_row_count, max_block_size, 0, 1));
     }
 
     prepareColumnCache(column_names, sample_block.getColumns(), sample_block);
@@ -79,7 +80,7 @@ Pipe StoragePython::read(
     // num_streams = 32; // for chdb testing
     for (size_t stream = 0; stream < num_streams; ++stream)
         pipes.emplace_back(std::make_shared<PythonSource>(
-            data_source, sample_block, column_cache, data_source_row_count, max_block_size, stream, num_streams));
+            data_source, false, sample_block, column_cache, data_source_row_count, max_block_size, stream, num_streams));
     return Pipe::unitePipes(std::move(pipes));
 }