Merge pull request #504 from wudidapaopao/fix_multi_refer_df

Fix crash on multi refer Python(df)

Author: auxten

programs/local/StoragePython.cpp

@@ -151,7 +151,21 @@ void StoragePython::prepareColumnCache(
 #endif
 
     auto & data_source = data_source_wrapper->getDataSource();
-    if (column_cache == nullptr)
+
+    bool need_rebuild = (column_cache == nullptr) || (column_cache->size() != names.size());
+    if (!need_rebuild)
+    {
+        for (size_t i = 0; i < names.size(); ++i)
+        {
+            if ((*column_cache)[i].name != names[i])
+            {
+                need_rebuild = true;
+                break;
+            }
+        }
+    }
+
+    if (need_rebuild)
     {
         column_cache = std::make_shared<PyColumnVec>(names.size());
         for (size_t i = 0; i < names.size(); ++i)

tests/test_query_py.py

@@ -400,6 +400,19 @@ def test_parquet_row_number_order(self):
         finally:
             os.unlink(path)
 
+    def test_multiple_python_references_in_subquery(self):
+        df = pd.DataFrame({'id': [1, 2, 3], 'value': [10, 20, 30]})
+
+        result = chdb.query('''
+            SELECT id, value + (SELECT SUM(value) FROM Python(df)) AS total
+            FROM Python(df)
+            ORDER BY id
+        ''', "DataFrame")
+
+        self.assertEqual(len(result), 3)
+        self.assertEqual(result['id'].tolist(), [1, 2, 3])
+        self.assertEqual(result['total'].tolist(), [70, 80, 90])
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=3)