improving efficiency and writing to file

src/lib.rs

@@ -1,18 +1,18 @@
 // --- Declare our new modules ---
-// This tells Rust to look for src/config.rs and src/parser.rs
 mod config;
 mod parser;
 
 // --- External Crates ---
 use pyo3::prelude::*;
 use pyo3::exceptions::PyValueError;
-use tokio; // We need the tokio runtime here
+use tokio;
 
 // --- Internal Crates ---
-// We 'use' the public functions from our new modules
 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,9 +20,9 @@ use crate::parser::run_db_logic;
 #[allow(unsafe_code)]
 #[allow(unsafe_op_in_unsafe_fn)]
 #[allow(rust_2024_compatibility)]
-fn load_data_from_config(py: Python, config_path: String) -> PyResult<()> {
+// FIX: Change the return type from () to PyObject
+fn load_data_from_config(py: Python, config_path: String) -> PyResult<PyObject> {
     // --- Phase 1: Load and Validate Configuration ---
-    // We call the function that now lives in src/config.rs
     let config: ConnectorConfig = match load_and_validate_config(&config_path) {
         Ok(c) => c,
         Err(e) => return Err(PyValueError::new_err(format!("Configuration Error: {:?}", e))),
@@ -34,24 +34,28 @@ fn load_data_from_config(py: Python, config_path: String) -> PyResult<()> {
 
     // --- Phase 2: Run Core Logic ---
     // We release the GIL and start our own Tokio runtime
-    py.allow_threads(|| {
+    let record_batch = py.allow_threads(|| {
         tokio::runtime::Builder::new_multi_thread()
             .enable_all()
             .build()
             .unwrap()
             .block_on(async {
-                // We call the function that now lives in src/parser.rs
                 run_db_logic(config).await
             })
     }).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?;
 
-    Ok(())
+    // --- 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)
 }
 
 
 // --- PYTHON MODULE EXPORT ---
 
 #[pymodule]
+// FIX: Add the attribute to enable Arrow <-> Python conversion
+#[pyo3(with_arrow)]
 fn unchecked_io(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
     Ok(())

src/parser.rs

@@ -3,13 +3,11 @@ use std::pin::Pin;
 use std::sync::Arc;
 use tokio_postgres::{NoTls, CopyOutStream};
 use anyhow::{Context, Result, anyhow};
-// FIX: Import new builders
 use arrow::array::{
     ArrayBuilder, ArrayRef,
     Int64Builder, Float64Builder, Float32Builder, StringBuilder, BooleanBuilder,
     TimestampNanosecondBuilder, Date32Builder, Int32Builder
 };
-// FIX: Import new types
 use arrow::datatypes::{
     DataType, Field, Schema,
     Float64Type, Float32Type, Int64Type, Int32Type, Utf8Type, BooleanType, TimestampNanosecondType,
@@ -28,7 +26,8 @@ use crate::config::ConnectorConfig;
 
 
 // --- 1. CORE DATABASE LOGIC ---
-pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> {
+// FIX: Change return type from Result<()> to Result<RecordBatch>
+pub async fn run_db_logic(config: ConnectorConfig) -> Result<RecordBatch> {
     // 1. Establish the connection
     println!("UncheckedIO: Attempting connection...");
 
@@ -42,7 +41,6 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> {
     });
 
     // 2. Execute the COPY TO STDOUT command
-    // FIX: We are now using the query from the config file!
     let copy_query = &config.query;
     println!("UncheckedIO: Executing user-defined query...");
 
@@ -54,18 +52,17 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> {
 
     // Build the Arrow Schema from the config
     let schema_fields: Vec<Field> = config.schema.iter().map(|col_cfg| {
-        // We now allow "nullable" to be controlled by the column name, a temporary "hack"
         let nullable = col_cfg.column_name == "notes";
 
         let arrow_type = match col_cfg.arrow_type.as_str() {
             "Int64" => DataType::Int64,
-            "Int32" => DataType::Int32, // NEW
-            "Float64" => DataType::Float64, // NEW
+            "Int32" => DataType::Int32,
+            "Float64" => DataType::Float64,
             "Float32" => DataType::Float32,
             "Utf8" | "String" => DataType::Utf8,
             "Boolean" => DataType::Boolean,
             "Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
-            "Date32" => DataType::Date32, // NEW
+            "Date32" => DataType::Date32,
             _ => panic!("Unsupported type in config: {}", col_cfg.arrow_type),
         };
         Field::new(&col_cfg.column_name, arrow_type, nullable)
@@ -82,7 +79,9 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> {
              record_batch.num_rows(), record_batch.num_columns());
 
     println!("UncheckedIO: Data transfer complete. We lived.");
-    Ok(())
+
+    // FIX: Return the finished RecordBatch
+    Ok(record_batch)
 }
 
 
@@ -91,13 +90,13 @@ pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> {
 // This enum will hold our different builder types
 enum DynamicBuilder {
     Int64(Box<Int64Builder>),
-    Int32(Box<Int32Builder>), // NEW
-    Float64(Box<Float64Builder>), // NEW
+    Int32(Box<Int32Builder>),
+    Float64(Box<Float64Builder>),
     Float32(Box<Float32Builder>),
     String(Box<StringBuilder>),
     Boolean(Box<BooleanBuilder>),
     Timestamp(Box<TimestampNanosecondBuilder>),
-    Date32(Box<Date32Builder>), // NEW
+    Date32(Box<Date32Builder>),
 }
 
 // Postgres Epoch for timestamps