Replace Python generator protocol with Arrow C Stream Interface

Replace the per-batch Python generator protocol with native Arrow C
Stream Interface for parquet input. Instead of calling back into Python
for every RecordBatch (GIL acquisition, tuple unpacking, PyCapsule import
per batch), Python now constructs RecordBatchReader objects that C++
imports via ArrowArrayStream and pulls batches natively.

Key changes:
- Python: Stream factory functions return RecordBatchReader objects
  backed by pyarrow.dataset scanners (parquet) or pa.ipc.open_stream
  (IPC), enabling GIL-free batch consumption in C++.
- C++: New RecordBatchStreamSource abstract interface replaces
  Generator<RecordBatchWithFlag>. PyRecordBatchStreamSource imports
  readers via arrow::ImportRecordBatchReader(ArrowArrayStream*).
- C++: getNextBatch() pulls from RecordBatchReader::ReadNext() directly,
  only calling back to Python at stream boundaries (file/directory).
- C++: start() collects all needed columns before calling the factory,
  enabling column projection at the scanner level.
- Schema changes detected by comparing schemas at stream boundaries.
- Basket readers imported separately and pre-read into single batches.

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

Author: arhamchopra

cpp/csp/adapters/arrow/RecordBatchRowProcessor.h

@@ -64,8 +64,6 @@ class RecordBatchRowProcessor
     // Get dispatcher for a column (nullptr if not found).
     ColumnDispatcher * getDispatcher( const std::string & name );
 
-    int64_t numRows() const { return m_numRows; }
-    int64_t currentRow() const { return m_currentRow; }
     bool hasMoreRows() const { return m_currentRow < m_numRows; }
 
     // Schema from last setupFromSchema call.

cpp/csp/adapters/parquet/CMakeLists.txt

@@ -12,7 +12,7 @@ set(PARQUET_HEADER_FILES
         ParquetOutputFilenameAdapter.h
         ParquetStatusUtils.h
         ParquetWriter.h
-        RecordBatchWithFlag.h
+        RecordBatchStreamSource.h
 )
 
 set(PARQUET_SOURCE_FILES

cpp/csp/adapters/parquet/ParquetInputAdapterManager.cpp

@@ -3,10 +3,9 @@
 
 #include <csp/adapters/arrow/ColumnDispatcher.h>
 #include <csp/adapters/arrow/RecordBatchRowProcessor.h>
-#include <csp/adapters/parquet/RecordBatchWithFlag.h>
+#include <csp/adapters/parquet/RecordBatchStreamSource.h>
 #include <csp/adapters/utils/StructAdapterInfo.h>
 #include <csp/adapters/utils/ValueDispatcher.h>
-#include <csp/core/Generator.h>
 #include <csp/engine/AdapterManager.h>
 #include <csp/engine/Dictionary.h>
 #include <csp/engine/Struct.h>
@@ -24,10 +23,10 @@ namespace csp::adapters::parquet
 class ParquetInputAdapterManager final : public csp::AdapterManager
 {
 public:
-    using RecordBatchGeneratorPtr = csp::Generator<RecordBatchWithFlag, csp::DateTime, csp::DateTime>::Ptr;
+    using RecordBatchStreamSourcePtr = RecordBatchStreamSource::Ptr;
 
     ParquetInputAdapterManager( csp::Engine *engine, const Dictionary &properties,
-                                RecordBatchGeneratorPtr rbGeneratorPtr );
+                                RecordBatchStreamSourcePtr streamSource );
 
     ~ParquetInputAdapterManager();
 
@@ -73,7 +72,6 @@ class ParquetInputAdapterManager final : public csp::AdapterManager
         DictionaryPtr                     m_fieldMap;
         std::vector<FieldSetter>          m_fieldSetters;
         ValueDispatcher                   m_valueDispatcher;
-        bool                              m_needsReset = false;
 
         void createFieldSetters( arrow::RecordBatchRowProcessor & processor,
                                  const std::shared_ptr<::arrow::Schema> & schema );
@@ -84,9 +82,12 @@ class ParquetInputAdapterManager final : public csp::AdapterManager
     struct DictBasketReaderRecord
     {
         std::string                                     m_basketName;
+        std::string                                     m_basketSymbolColumn;
         arrow::ColumnDispatcher *                        m_valueCountDispatcher = nullptr;
         std::unique_ptr<arrow::RecordBatchRowProcessor>  m_processor;
         std::vector<std::unique_ptr<StructSubscription>>  m_structSubscriptions;
+
+        uint16_t getValueCount() const;
     };
 
     using AdaptersBySymbol = std::unordered_map<utils::Symbol, AdaptersSingleSymbol>;
@@ -96,6 +97,10 @@ class ParquetInputAdapterManager final : public csp::AdapterManager
     // Returns false when no more data is available.
     bool getNextBatch();
 
+    // Collect all column names referenced by adapters in the given map.
+    static void collectAdapterColumns( const AdaptersBySymbol & adaptersBySymbol,
+                                       std::set<std::string> & columns );
+
     // Set up a processor from the current batch schema and subscribe adapters.
     void setupProcessor( arrow::RecordBatchRowProcessor & processor,
                          const std::shared_ptr<::arrow::Schema> & schema,
@@ -106,7 +111,12 @@ class ParquetInputAdapterManager final : public csp::AdapterManager
     // Subscribe adapters to a processor's dispatchers.
     void subscribeAdapters( arrow::RecordBatchRowProcessor & processor,
                             const AdaptersBySymbol & adaptersBySymbol,
-                            bool subscribeAllOnEmptySymbol );
+                            bool subscribeAllOnEmptySymbol,
+                            std::vector<std::unique_ptr<StructSubscription>> & structSubscriptions );
+
+    // Set up a dict basket processor: setupFromSchema + subscribe adapters + bind batch.
+    void setupBasketProcessor( DictBasketReaderRecord & record,
+                               const AdaptersBySymbol & adaptersBySymbol );
 
     // Read the next row into the main processor (handles batch boundaries).
     bool readNextRow();
@@ -139,7 +149,7 @@ class ParquetInputAdapterManager final : public csp::AdapterManager
     csp::DateTime                       m_startTime;
     csp::DateTime                       m_endTime;
     csp::TimeDelta                      m_time_shift;
-    RecordBatchGeneratorPtr             m_rbGenerator;
+    RecordBatchStreamSourcePtr          m_streamSource;
     std::string                         m_symbolColumn;
     std::string                         m_timeColumn;
     std::string                         m_defaultTimezone;
@@ -151,6 +161,7 @@ class ParquetInputAdapterManager final : public csp::AdapterManager
 
     // Processor-based reader state
     std::unique_ptr<arrow::RecordBatchRowProcessor>  m_processor;
+    std::shared_ptr<::arrow::RecordBatchReader>      m_mainReader;
     std::shared_ptr<::arrow::RecordBatch>            m_curBatch;
     std::shared_ptr<::arrow::Schema>                 m_curSchema;
     std::unordered_map<std::string, std::shared_ptr<::arrow::RecordBatch>> m_curBasketBatches;

cpp/csp/adapters/parquet/ParquetInputAdapterManager.h

@@ -0,0 +1,46 @@
+#ifndef _IN_CSP_ADAPTERS_PARQUET_RecordBatchStreamSource_H
+#define _IN_CSP_ADAPTERS_PARQUET_RecordBatchStreamSource_H
+
+#include <arrow/record_batch.h>
+#include <csp/core/Time.h>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_map>
+
+namespace csp::adapters::parquet
+{
+
+// Abstract interface for streaming RecordBatch data from Python.
+// The Python implementation wraps a "stream factory" callable that yields
+// (RecordBatchReader, {basket: RecordBatchReader}) tuples.
+// C++ imports readers via ArrowArrayStream and pulls batches natively.
+class RecordBatchStreamSource
+{
+public:
+    using Ptr = std::shared_ptr<RecordBatchStreamSource>;
+
+    virtual ~RecordBatchStreamSource() = default;
+
+    // Call the stream factory with (starttime, endtime, needed_columns).
+    // Must be called before nextStream().
+    virtual void init( csp::DateTime start, csp::DateTime end,
+                       const std::set<std::string> & neededColumns ) = 0;
+
+    // Advance to the next stream set (file/directory boundary).
+    // After this call, mainReader() and basketBatches() reflect the new data.
+    // Returns false when all data is exhausted.
+    virtual bool nextStream() = 0;
+
+    // The current main RecordBatchReader. Pull batches with ReadNext().
+    // Returns nullptr before the first nextStream() call.
+    virtual std::shared_ptr<::arrow::RecordBatchReader> mainReader() = 0;
+
+    // Pre-read basket batches for the current stream boundary.
+    // Empty when there are no dict baskets.
+    virtual const std::unordered_map<std::string, std::shared_ptr<::arrow::RecordBatch>> & basketBatches() const = 0;
+};
+
+} // namespace csp::adapters::parquet
+
+#endif

cpp/csp/adapters/parquet/RecordBatchStreamSource.h

@@ -3,7 +3,6 @@
 #include <csp/adapters/parquet/ParquetOutputAdapterManager.h>
 #include <csp/adapters/parquet/ParquetDictBasketOutputWriter.h>
 #include <csp/adapters/parquet/ParquetStatusUtils.h>
-#include <csp/core/Generator.h>
 #include <csp/engine/PushInputAdapter.h>
 #include <csp/python/Conversions.h>
 #include <csp/python/Exception.h>
@@ -81,30 +80,50 @@ REGISTER_CPPNODE( csp::cppnodes, parquet_dict_basket_writer );
 namespace
 {
 
-// Generator that wraps a Python generator yielding (RecordBatch, dict, bool) tuples.
-// The dict maps basket names to RecordBatches; the bool indicates schema change.
-class RecordBatchGenerator : public csp::Generator<RecordBatchWithFlag, csp::DateTime, csp::DateTime>
+// Generator that wraps a Python "stream factory" callable.
+// The factory signature: factory(starttime, endtime, needed_columns) -> iterator of (reader, basket_dict)
+// Each reader is a pyarrow.RecordBatchReader; basket_dict maps basket names to readers.
+// Readers are imported via ArrowArrayStream for GIL-free batch pulling in C++.
+class PyRecordBatchStreamSource : public csp::adapters::parquet::RecordBatchStreamSource
 {
 public:
-    RecordBatchGenerator( PyObject *wrappedGenerator )
-            : m_wrappedGenerator( csp::python::PyObjectPtr::incref( wrappedGenerator ) )
+    PyRecordBatchStreamSource( PyObject *factory )
+            : m_factory( csp::python::PyObjectPtr::incref( factory ) )
     {
     }
 
-    void init( csp::DateTime start, csp::DateTime end ) override
+    void init( csp::DateTime start, csp::DateTime end,
+               const std::set<std::string> & neededColumns ) override
     {
-        auto tp = csp::python::PyObjectPtr::own( PyTuple_New( 2 ) );
+        auto tp = csp::python::PyObjectPtr::own( PyTuple_New( 3 ) );
         if( !tp.get() )
             CSP_THROW( csp::python::PythonPassthrough, "" );
 
         PyTuple_SET_ITEM( tp.get(), 0, csp::python::toPython( start ) );
         PyTuple_SET_ITEM( tp.get(), 1, csp::python::toPython( end ) );
-        m_iter = csp::python::PyObjectPtr::check( PyObject_Call( m_wrappedGenerator.ptr(), tp.get(), nullptr ) );
+
+        if( neededColumns.empty() )
+        {
+            Py_INCREF( Py_None );
+            PyTuple_SET_ITEM( tp.get(), 2, Py_None );
+        }
+        else
+        {
+            auto pyList = csp::python::PyObjectPtr::own( PyList_New( neededColumns.size() ) );
+            if( !pyList.get() )
+                CSP_THROW( csp::python::PythonPassthrough, "" );
+            Py_ssize_t idx = 0;
+            for( auto & col : neededColumns )
+                PyList_SET_ITEM( pyList.get(), idx++, PyUnicode_FromStringAndSize( col.c_str(), col.size() ) );
+            PyTuple_SET_ITEM( tp.get(), 2, pyList.release() );
+        }
+
+        m_iter = csp::python::PyObjectPtr::check( PyObject_Call( m_factory.ptr(), tp.get(), nullptr ) );
         CSP_TRUE_OR_THROW( PyIter_Check( m_iter.ptr() ), csp::TypeError,
-                           "RecordBatch generator expected to return iterator" );
+                           "Stream factory expected to return iterator" );
     }
 
-    bool next( RecordBatchWithFlag &value ) override
+    bool nextStream() override
     {
         if( m_iter.ptr() == nullptr )
             return false;
@@ -115,21 +134,18 @@ class RecordBatchGenerator : public csp::Generator<RecordBatchWithFlag, csp::Dat
         if( nextVal.get() == nullptr )
             return false;
 
-        // Expect a tuple of (RecordBatch, dict, bool)
-        CSP_TRUE_OR_THROW( PyTuple_Check( nextVal.get() ) && PyTuple_GET_SIZE( nextVal.get() ) == 3,
-                           csp::TypeError, "RecordBatch generator expected to yield (batch, basket_batches, schema_changed) tuples" );
+        // Expect a tuple of (RecordBatchReader, dict)
+        CSP_TRUE_OR_THROW( PyTuple_Check( nextVal.get() ) && PyTuple_GET_SIZE( nextVal.get() ) == 2,
+                           csp::TypeError, "Stream factory expected to yield (reader, basket_dict) tuples" );
 
-        PyObject *pyBatch         = PyTuple_GET_ITEM( nextVal.get(), 0 );
-        PyObject *pyBasketDict    = PyTuple_GET_ITEM( nextVal.get(), 1 );
-        PyObject *pySchemaChanged = PyTuple_GET_ITEM( nextVal.get(), 2 );
+        PyObject *pyReader     = PyTuple_GET_ITEM( nextVal.get(), 0 );
+        PyObject *pyBasketDict = PyTuple_GET_ITEM( nextVal.get(), 1 );
 
-        value.schemaChanged = PyObject_IsTrue( pySchemaChanged );
+        // Import main reader via ArrowArrayStream
+        m_mainReader = importRecordBatchReader( pyReader );
 
-        // Import main RecordBatch via PyCapsule C Data Interface
-        value.batch = importRecordBatch( pyBatch );
-
-        // Import basket RecordBatches
-        value.basketBatches.clear();
+        // Import basket readers and read their single batch
+        m_basketBatches.clear();
         if( PyDict_Check( pyBasketDict ) )
         {
             PyObject *key, *val;
@@ -139,43 +155,55 @@ class RecordBatchGenerator : public csp::Generator<RecordBatchWithFlag, csp::Dat
                 const char *basketName = PyUnicode_AsUTF8( key );
                 if( !basketName )
                     CSP_THROW( csp::python::PythonPassthrough, "" );
-                value.basketBatches[ basketName ] = importRecordBatch( val );
+
+                auto basketReader = importRecordBatchReader( val );
+                std::shared_ptr<::arrow::RecordBatch> batch;
+                auto status = basketReader -> ReadNext( &batch );
+                if( !status.ok() )
+                    CSP_THROW( csp::ValueError, "Failed to read basket batch: " << status.ToString() );
+                if( batch )
+                    m_basketBatches[ basketName ] = batch;
             }
         }
 
         return true;
     }
 
-private:
-    static std::shared_ptr<::arrow::RecordBatch> importRecordBatch( PyObject *pyBatch )
+    std::shared_ptr<::arrow::RecordBatchReader> mainReader() override
     {
-        auto exportResult = csp::python::PyObjectPtr::own(
-            PyObject_CallMethod( pyBatch, "__arrow_c_array__", nullptr ) );
-        if( !exportResult.get() || PyErr_Occurred() )
-            CSP_THROW( csp::python::PythonPassthrough, "" );
-
-        CSP_TRUE_OR_THROW( PyTuple_Check( exportResult.get() ) && PyTuple_GET_SIZE( exportResult.get() ) == 2,
-                           csp::TypeError, "__arrow_c_array__ expected to return (schema_capsule, array_capsule)" );
+        return m_mainReader;
+    }
 
-        PyObject *pySchemaCapsule = PyTuple_GET_ITEM( exportResult.get(), 0 );
-        PyObject *pyArrayCapsule  = PyTuple_GET_ITEM( exportResult.get(), 1 );
+    const std::unordered_map<std::string, std::shared_ptr<::arrow::RecordBatch>> & basketBatches() const override
+    {
+        return m_basketBatches;
+    }
 
-        auto *c_schema = reinterpret_cast<struct ArrowSchema *>( PyCapsule_GetPointer( pySchemaCapsule, "arrow_schema" ) );
-        auto *c_array  = reinterpret_cast<struct ArrowArray *>( PyCapsule_GetPointer( pyArrayCapsule, "arrow_array" ) );
+private:
+    static std::shared_ptr<::arrow::RecordBatchReader> importRecordBatchReader( PyObject *pyReader )
+    {
+        // Call __arrow_c_stream__() to export as ArrowArrayStream PyCapsule
+        auto capsule = csp::python::PyObjectPtr::own(
+            PyObject_CallMethod( pyReader, "__arrow_c_stream__", nullptr ) );
+        if( !capsule.get() || PyErr_Occurred() )
+            CSP_THROW( csp::python::PythonPassthrough, "" );
 
-        auto schemaResult = arrow::ImportSchema( c_schema );
-        if( !schemaResult.ok() )
-            CSP_THROW( csp::ValueError, "Failed to import RecordBatch schema: " << schemaResult.status().ToString() );
+        auto *stream = reinterpret_cast<struct ArrowArrayStream *>(
+            PyCapsule_GetPointer( capsule.get(), "arrow_array_stream" ) );
+        if( !stream )
+            CSP_THROW( csp::ValueError, "Failed to get ArrowArrayStream from PyCapsule" );
 
-        auto batchResult = arrow::ImportRecordBatch( c_array, schemaResult.ValueUnsafe() );
-        if( !batchResult.ok() )
-            CSP_THROW( csp::ValueError, "Failed to import RecordBatch: " << batchResult.status().ToString() );
+        auto result = ::arrow::ImportRecordBatchReader( stream );
+        if( !result.ok() )
+            CSP_THROW( csp::ValueError, "Failed to import RecordBatchReader: " << result.status().ToString() );
 
-        return std::move( batchResult.ValueUnsafe() );
+        return result.ValueUnsafe();
     }
 
-    csp::python::PyObjectPtr m_wrappedGenerator;
+    csp::python::PyObjectPtr m_factory;
     csp::python::PyObjectPtr m_iter;
+    std::shared_ptr<::arrow::RecordBatchReader> m_mainReader;
+    std::unordered_map<std::string, std::shared_ptr<::arrow::RecordBatch>> m_basketBatches;
 };
 
 }
@@ -185,9 +213,9 @@ namespace csp::python
 
 //AdapterManager
 csp::AdapterManager *create_parquet_input_adapter_manager_impl( PyEngine *engine, const Dictionary &properties,
-                                                                RecordBatchGenerator::Ptr rbGenerator )
+                                                                ParquetInputAdapterManager::RecordBatchStreamSourcePtr streamSource )
 {
-    auto res = engine -> engine() -> createOwnedObject<ParquetInputAdapterManager>( properties, rbGenerator );
+    auto res = engine -> engine() -> createOwnedObject<ParquetInputAdapterManager>( properties, streamSource );
     return res;
 }
 
@@ -443,9 +471,9 @@ static PyObject *create_parquet_input_adapter_manager( PyObject *args )
                                &PyFunction_Type, &pyFileGenerator ) )
             CSP_THROW( PythonPassthrough, "" );
 
-        auto rbGenerator = std::make_shared<RecordBatchGenerator>( pyFileGenerator );
+        auto streamSource = std::make_shared<PyRecordBatchStreamSource>( pyFileGenerator );
         auto *adapterMgr = create_parquet_input_adapter_manager_impl( pyEngine, fromPython<Dictionary>( pyProperties ),
-                                                                      rbGenerator );
+                                                                      streamSource );
         auto res         = PyCapsule_New( adapterMgr, "adapterMgr", nullptr );
         return res;
     CSP_RETURN_NULL;