Merge branch 'main' into ac/arrow_nodes

Author: arhamchopra

.github/workflows/build.yml

@@ -295,28 +295,28 @@ jobs:
           run: make dist-check
 
         - name: Upload Wheel
-          uses: actions/upload-artifact@v6
+          uses: actions/upload-artifact@v7
           with:
             name: csp-dist-${{ runner.os }}-${{ runner.arch }}-3.10
             path: dist/*.whl
           if: ${{ matrix.cibuildwheel == 'cp310' }}
 
         - name: Upload Wheel
-          uses: actions/upload-artifact@v6
+          uses: actions/upload-artifact@v7
           with:
             name: csp-dist-${{ runner.os }}-${{ runner.arch }}-3.11
             path: dist/*.whl
           if: ${{ matrix.cibuildwheel == 'cp311' }}
 
         - name: Upload Wheel
-          uses: actions/upload-artifact@v6
+          uses: actions/upload-artifact@v7
           with:
             name: csp-dist-${{ runner.os }}-${{ runner.arch }}-3.12
             path: dist/*.whl
           if: ${{ matrix.cibuildwheel == 'cp312' }}
 
         - name: Upload Wheel
-          uses: actions/upload-artifact@v6
+          uses: actions/upload-artifact@v7
           with:
             name: csp-dist-${{ runner.os }}-${{ runner.arch }}-3.13
             path: dist/*.whl
@@ -377,7 +377,7 @@ jobs:
           run: make dist-check
 
         - name: Upload SDist
-          uses: actions/upload-artifact@v6
+          uses: actions/upload-artifact@v7
           with:
               name: csp-sdist
               path: dist/*.tar.gz
@@ -462,7 +462,7 @@ jobs:
           run: make requirements
 
         - name: Download wheel
-          uses: actions/download-artifact@v7
+          uses: actions/download-artifact@v8
           with:
             name: csp-dist-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}
 
@@ -573,7 +573,7 @@ jobs:
         - name: Install requirements
           run: sudo make dependencies-debian
 
-        - uses: actions/download-artifact@v7
+        - uses: actions/download-artifact@v8
           with:
             name: csp-sdist
             path: dist/
@@ -657,7 +657,7 @@ jobs:
           run: sudo apt-get install graphviz
 
         - name: Download wheel
-          uses: actions/download-artifact@v7
+          uses: actions/download-artifact@v8
           with:
             name: csp-dist-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}
 
@@ -737,7 +737,7 @@ jobs:
 
       steps:
         - name: Download wheels and sdist
-          uses: actions/download-artifact@v7
+          uses: actions/download-artifact@v8
           with:
             name:
             merge-multiple: true

ci/scripts/requirements.py

@@ -9,8 +9,12 @@ def deps():
     return requires + develop
 
 def main():
-    os.system("python -m pip install toml")
-    os.system("python -m pip install " + " ".join(f'"{_}"' for _ in deps()))
+    ret = os.system("python -m pip install --prefer-binary toml")
+    if ret != 0:
+        raise ValueError("Python requirement install failed: see output for error")
+    ret = os.system("python -m pip install --prefer-binary " + " ".join(f'"{_}"' for _ in deps()))
+    if ret != 0:
+        raise ValueError("Python requirement install failed: see output for error")
 
 
 if __name__ == "__main__":

cpp/csp/python/adapters/ArrowInputAdapter.h

@@ -24,7 +24,8 @@ class RecordBatchIterator
 {
 public:
     RecordBatchIterator() {}
-    RecordBatchIterator( PyObjectPtr iter, std::shared_ptr<::arrow::Schema> schema ): m_iter( std::move( iter ) ), m_schema( schema )
+
+    RecordBatchIterator( PyObjectPtr iter, std::shared_ptr<::arrow::Schema> schema ): m_iter( std::move( iter ) ), m_schema( std::move( schema ) )
     {
     }
 
@@ -36,7 +37,7 @@ class RecordBatchIterator
 
         if( py_tuple.get() == NULL )
         {
-            // No more data in the input steam
+            // No more data in the input stream
             return nullptr;
         }
 
@@ -47,6 +48,20 @@ class RecordBatchIterator
         if( num_elems != 2 )
             CSP_THROW( csp::TypeError, "Invalid arrow data, expected tuple (using the PyCapsule C interface) with 2 elements got " << num_elems );
 
+        // Extract schema from first batch if not provided upfront
+        if( !m_schema )
+        {
+            PyObject * py_schema = PyTuple_GetItem( py_tuple.get(), 0 );
+            if( !PyCapsule_IsValid( py_schema, "arrow_schema" ) )
+                CSP_THROW( csp::TypeError, "Invalid arrow data, expected schema capsule from the PyCapsule C interface" );
+
+            ArrowSchema * c_schema = reinterpret_cast<ArrowSchema*>( PyCapsule_GetPointer( py_schema, "arrow_schema" ) );
+            auto schema_result = ::arrow::ImportSchema( c_schema );
+            if( !schema_result.ok() )
+                CSP_THROW( ValueError, "Failed to load schema through PyCapsule C Data interface: " << schema_result.status().ToString() );
+            m_schema = std::move( schema_result.ValueUnsafe() );
+        }
+
         // Extract the record batch
         PyObject * py_array = PyTuple_GetItem( py_tuple.get(), 1 );
         if( !PyCapsule_IsValid( py_array, "arrow_array" ) )
@@ -60,20 +75,22 @@ class RecordBatchIterator
         return result.ValueUnsafe();
     }
 
+    std::shared_ptr<::arrow::Schema> schema() const { return m_schema; }
+
 private:
     PyObjectPtr m_iter;
     std::shared_ptr<::arrow::Schema> m_schema;
 };
 
-void ReleaseArrowSchemaPyCapsule( PyObject * capsule )
+inline void ReleaseArrowSchemaPyCapsule( PyObject * capsule )
 {
     ArrowSchema * schema = reinterpret_cast<ArrowSchema*>( PyCapsule_GetPointer( capsule, "arrow_schema" ) );
     if ( schema -> release != NULL )
         schema -> release( schema );
     free( schema );
 }
 
-void ReleaseArrowArrayPyCapsule( PyObject * capsule )
+inline void ReleaseArrowArrayPyCapsule( PyObject * capsule )
 {
     ArrowArray * array = reinterpret_cast<ArrowArray*>( PyCapsule_GetPointer( capsule, "arrow_array" ) );
     if ( array -> release != NULL )
@@ -88,36 +105,33 @@ class RecordBatchInputAdapter: public PullInputAdapter<std::vector<DialectGeneri
         : PullInputAdapter<std::vector<DialectGenericType>>( engine, type, PushMode::LAST_VALUE ),
           m_tsColName( tsColName ),
           m_expectSmallBatches( expectSmallBatches ),
-          m_finished( false )
+          m_finished( false ),
+          m_initialized( false ),
+          m_multiplier( 0 ),
+          m_arrayLastTime( 0 ),
+          m_numRows( 0 ),
+          m_startTime( 0 ),
+          m_endTime( 0 ),
+          m_curBatchIdx( 0 )
     {
-        // Extract the arrow schema
-        ArrowSchema * c_schema = reinterpret_cast<ArrowSchema*>( PyCapsule_GetPointer( pySchema.get(), "arrow_schema" ) );
-        auto result = ::arrow::ImportSchema( c_schema );
-        if( !result.ok() )
-            CSP_THROW( ValueError, "Failed to load schema for record batches through the PyCapsule C Data interface: " << result.status().ToString() );
-        m_schema = std::move( result.ValueUnsafe() );
+        // Check if schema is provided or should be deferred to first batch
+        if( pySchema.get() != Py_None )
+        {
+            // Extract the arrow schema upfront
+            ArrowSchema * c_schema = reinterpret_cast<ArrowSchema*>( PyCapsule_GetPointer( pySchema.get(), "arrow_schema" ) );
+            auto result = ::arrow::ImportSchema( c_schema );
+            if( !result.ok() )
+                CSP_THROW( ValueError, "Failed to load schema for record batches through the PyCapsule C Data interface: " << result.status().ToString() );
 
-        auto tsField = m_schema -> GetFieldByName( m_tsColName );
-        auto timestampType = std::static_pointer_cast<::arrow::TimestampType>( tsField -> type() );
-        switch( timestampType -> unit() )
+            m_source = RecordBatchIterator( source, std::move( result.ValueUnsafe() ) );
+            initializeTimestampMultiplier();
+            m_initialized = true;
+        }
+        else
         {
-            case ::arrow::TimeUnit::SECOND:
-                m_multiplier = csp::NANOS_PER_SECOND;
-                break;
-            case ::arrow::TimeUnit::MILLI:
-                m_multiplier = csp::NANOS_PER_MILLISECOND;
-                break;
-            case ::arrow::TimeUnit::MICRO:
-                m_multiplier = csp::NANOS_PER_MICROSECOND;
-                break;
-            case ::arrow::TimeUnit::NANO:
-                m_multiplier = 1;
-                break;
-            default:
-                CSP_THROW( ValueError, "Unsupported unit type for arrow timestamp column" );
+            // Schema will be extracted lazily from first batch
+            m_source = RecordBatchIterator( source, nullptr );
         }
-
-        m_source = RecordBatchIterator( source, m_schema );
     }
 
     int64_t findFirstMatchingIndex()
@@ -174,6 +188,16 @@ class RecordBatchInputAdapter: public PullInputAdapter<std::vector<DialectGeneri
         while( ( rb = m_source.next() ) && ( rb -> num_rows() == 0 ) ) {}
         if( rb )
         {
+            // Lazy schema initialization on first batch (when schema was not provided upfront)
+            if( !m_initialized ) [[unlikely]]
+            {
+                if( !m_source.schema() )
+                    CSP_THROW( ValueError, "Failed to extract schema from first record batch" );
+                initializeTimestampMultiplier();
+                computeTimeRange();
+                m_initialized = true;
+            }
+
             auto array = rb -> GetColumnByName( m_tsColName );
             if( !array )
                 CSP_THROW( ValueError, "Failed to get timestamp column " << m_tsColName << " from record batch " << rb -> ToString() );
@@ -189,24 +213,10 @@ class RecordBatchInputAdapter: public PullInputAdapter<std::vector<DialectGeneri
 
     void start( DateTime start, DateTime end ) override
     {
-        // start and end as multiples of the unit in timestamp column
-        auto start_nanos = start.asNanoseconds();
-        m_startTime = ( start_nanos % m_multiplier == 0 ) ? start_nanos / m_multiplier : start_nanos / m_multiplier + 1;
-        m_endTime = end.asNanoseconds() / m_multiplier;
+        m_rawStartTime = start;
+        m_rawEndTime = end;
 
-        // Find the starting index where time >= start
-        while( !m_finished )
-        {
-            updateStateFromNextRecordBatch();
-            if( !m_curRecordBatch )
-            {
-                m_finished = true;
-                break;
-            }
-            m_curBatchIdx = findFirstMatchingIndex();
-            if( m_curBatchIdx < m_numRows )
-                break;
-        }
+        initializeStartPosition();
         PullInputAdapter<std::vector<DialectGenericType>>::start( start, end );
     }
 
@@ -215,6 +225,12 @@ class RecordBatchInputAdapter: public PullInputAdapter<std::vector<DialectGeneri
         ArrowSchema* rb_schema = ( ArrowSchema* )malloc( sizeof( ArrowSchema ) );
         ArrowArray* rb_array = ( ArrowArray* )malloc( sizeof( ArrowArray ) );
         ::arrow::Status st = ::arrow::ExportRecordBatch( *rb, rb_array, rb_schema );
+        if( !st.ok() )
+        {
+            free( rb_array );
+            free( rb_schema );
+            CSP_THROW( ValueError, "Failed to export record batch through C Data interface: " << st.ToString() );
+        }
         auto py_schema = csp::python::PyObjectPtr::own( PyCapsule_New( rb_schema, "arrow_schema", ReleaseArrowSchemaPyCapsule ) );
         auto py_array = csp::python::PyObjectPtr::own( PyCapsule_New( rb_array, "arrow_array", ReleaseArrowArrayPyCapsule ) );
         auto py_tuple = csp::python::PyObjectPtr::own( PyTuple_Pack( 2, py_schema.get(), py_array.get() ) );
@@ -269,17 +285,85 @@ class RecordBatchInputAdapter: public PullInputAdapter<std::vector<DialectGeneri
     }
 
 private:
+    void initializeTimestampMultiplier()
+    {
+        auto schema = m_source.schema();
+        auto tsField = schema -> GetFieldByName( m_tsColName );
+        if( !tsField )
+            CSP_THROW( ValueError, "Timestamp column '" << m_tsColName << "' not found in schema" );
+
+        if( tsField -> type() -> id() != ::arrow::Type::TIMESTAMP )
+            CSP_THROW( ValueError, "Column '" << m_tsColName << "' is not a valid timestamp column" );
+
+        auto timestampType = std::static_pointer_cast<::arrow::TimestampType>( tsField -> type() );
+        switch( timestampType -> unit() )
+        {
+            case ::arrow::TimeUnit::SECOND:
+                m_multiplier = csp::NANOS_PER_SECOND;
+                break;
+            case ::arrow::TimeUnit::MILLI:
+                m_multiplier = csp::NANOS_PER_MILLISECOND;
+                break;
+            case ::arrow::TimeUnit::MICRO:
+                m_multiplier = csp::NANOS_PER_MICROSECOND;
+                break;
+            case ::arrow::TimeUnit::NANO:
+                m_multiplier = 1;
+                break;
+            default:
+                CSP_THROW( ValueError, "Unsupported unit type for arrow timestamp column" );
+        }
+    }
+
+    void computeTimeRange()
+    {
+        auto start_nanos = m_rawStartTime.asNanoseconds();
+        m_startTime = ( start_nanos % m_multiplier == 0 )
+            ? start_nanos / m_multiplier
+            : start_nanos / m_multiplier + 1;
+        m_endTime = m_rawEndTime.asNanoseconds() / m_multiplier;
+    }
+
+    void initializeStartPosition()
+    {
+        // If schema was provided upfront, compute time range now
+        // Otherwise, time range will be computed lazily in updateStateFromNextRecordBatch()
+        if( m_initialized )
+            computeTimeRange();
+
+        // Find the starting index where time >= start
+        while( !m_finished )
+        {
+            updateStateFromNextRecordBatch();
+            if( !m_curRecordBatch )
+            {
+                m_finished = true;
+                break;
+            }
+            m_curBatchIdx = findFirstMatchingIndex();
+            if( m_curBatchIdx < m_numRows )
+                break;
+            m_curRecordBatch = nullptr;
+        }
+    }
+
     std::string m_tsColName;
     RecordBatchIterator m_source;
 
     int m_expectSmallBatches;
     bool m_finished;
-    std::shared_ptr<::arrow::Schema> m_schema;
+    bool m_initialized;  // Whether schema and timestamp multiplier have been resolved
+    int64_t m_multiplier;
     std::shared_ptr<::arrow::RecordBatch> m_curRecordBatch;
     std::shared_ptr<::arrow::TimestampArray> m_tsArray;
     ::arrow::TimestampArray::IteratorType m_endIt;
     int64_t m_arrayLastTime;
-    int64_t m_multiplier, m_numRows, m_startTime, m_endTime, m_curBatchIdx;
+    int64_t m_numRows;
+    int64_t m_startTime;
+    int64_t m_endTime;
+    int64_t m_curBatchIdx;
+    DateTime m_rawStartTime;
+    DateTime m_rawEndTime;
 };
 
 };