improving efficiency and writing to file

src/lib.rs

@@ -6,13 +6,17 @@ mod parser;
 use pyo3::prelude::*;
 use pyo3::exceptions::PyValueError;
 use tokio;
+use pyo3::types::PyModule;
+use pyo3::Bound;
+
+// FIX 1: Import PyRecordBatch.
+use pyo3_arrow::PyRecordBatch;
+
 
 // --- Internal Crates ---
 use crate::config::{load_and_validate_config, ConnectorConfig};
 use crate::parser::run_db_logic;
 
-// FIX: Import the Arrow bridge trait
-use pyo3_arrow::PyArrowConvert;
 
 // --- THE PYTHON-CALLABLE ENTRY POINT ---
 
@@ -20,8 +24,7 @@ use pyo3_arrow::PyArrowConvert;
 #[allow(unsafe_code)]
 #[allow(unsafe_op_in_unsafe_fn)]
 #[allow(rust_2024_compatibility)]
-// FIX: Change the return type from () to PyObject
-fn load_data_from_config(py: Python, config_path: String) -> PyResult<PyObject> {
+fn load_data_from_config<'py>(py: Python<'py>, config_path: String) -> PyResult<Bound<'py, PyAny>> {
     // --- Phase 1: Load and Validate Configuration ---
     let config: ConnectorConfig = match load_and_validate_config(&config_path) {
         Ok(c) => c,
@@ -33,7 +36,6 @@ fn load_data_from_config(py: Python, config_path: String) -> PyResult<PyObject>
     println!("Columns (in order): {:?}", config.schema.iter().map(|c| &c.column_name).collect::<Vec<_>>());
 
     // --- Phase 2: Run Core Logic ---
-    // We release the GIL and start our own Tokio runtime
     let record_batch = py.allow_threads(|| {
         tokio::runtime::Builder::new_multi_thread()
             .enable_all()
@@ -45,18 +47,22 @@ fn load_data_from_config(py: Python, config_path: String) -> PyResult<PyObject>
     }).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?;
 
     // --- Phase 3: Return Data to Python ---
-    // FIX: Convert the Rust RecordBatch to a Python PyObject (a pyarrow.Table)
-    // This is a ZERO-COPY operation.
-    record_batch.to_pyarrow(py)
+    // FIX 2: Create PyRecordBatch wrapper using the standard PyO3 __new__ convention.
+    // let py_record_batch = PyRecordBatch::new(record_batch);
+    // // Replace: let py_record_batch = PyRecordBatch::new(py, record_batch).map_err(|e| ...)?;
+    // // With:
+    let py_record_batch = PyRecordBatch::new(record_batch);
+    // Return the PyRecordBatch wrapper object as a generic PyObject reference.
+    py_record_batch.into_pyarrow(py)
 }
 
 
 // --- PYTHON MODULE EXPORT ---
-
 #[pymodule]
-// FIX: Add the attribute to enable Arrow <-> Python conversion
-#[pyo3(with_arrow)]
-fn unchecked_io(_py: Python, m: &PyModule) -> PyResult<()> {
+fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // FIX: Use the two-argument version of wrap_pyfunction!
+    // This resolves the E0308 type mismatch.
     m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
+
     Ok(())
 }