Merge pull request #5 from BillTheMaker/gobrrr

set up to_pandas logic and rerun for failed task

src/lib.rs

@@ -9,9 +9,7 @@ use tokio;
 use pyo3::types::{PyModule, PyAny};
 use pyo3::Bound;
 use pyo3_arrow::PyRecordBatch;
-
-// Use the high-performance mimalloc for better multi-threaded memory allocation.
-// We conditionally compile it to avoid issues on MSVC targets.
+// Required for global allocator (mimalloc)
 #[cfg(not(target_env = "msvc"))]
 use mimalloc;
 
@@ -19,12 +17,14 @@ use mimalloc;
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
+
 // --- Internal Crates ---
 use crate::config::{load_and_validate_config, ConnectorConfig};
-use crate::parser::run_db_logic;
+// NOTE: run_profiler_logic added, requires definition in parser.rs
+use crate::parser::{run_db_logic, run_profiler_logic};
 
 
-// --- THE PYTHON-CALLABLE ENTRY POINT ---
+// --- THE PYTHON-CALLABLE ENTRY POINT (Load Data) ---
 
 #[pyfunction]
 #[pyo3(signature = (config_path, blast_radius=312500))]
@@ -37,7 +37,6 @@ fn load_data_from_config<'py>(
     blast_radius: i64,
 ) -> PyResult<Bound<'py, PyAny>> {
 
-    // --- Phase 1: Load and Validate Configuration ---
     let config: ConnectorConfig = match load_and_validate_config(&config_path) {
         Ok(c) => c,
         Err(e) => return Err(PyValueError::new_err(format!("Configuration Error: {:?}", e))),
@@ -47,7 +46,6 @@ fn load_data_from_config<'py>(
     println!("Database: {}", config.connection_string);
     println!("Columns (in order): {:?}", config.schema.iter().map(|c| &c.column_name).collect::<Vec<_>>());
 
-    // --- Phase 2: Run Core Logic ---
     let record_batch = py.allow_threads(|| {
         tokio::runtime::Builder::new_multi_thread()
             .enable_all()
@@ -58,16 +56,42 @@ fn load_data_from_config<'py>(
             })
     }).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?;
 
-    // --- Phase 3: Return Data to Python ---
     let py_record_batch = PyRecordBatch::new(record_batch);
     py_record_batch.into_pyarrow(py)
 }
 
+// --- NEW PYTHON-CALLABLE FUNCTION FOR PANDAS CONVERSION (FIXED SIGNATURE) ---
+#[pyfunction]
+#[pyo3(signature = (arrow_table))] // FIX: Removed 'py' from the signature macro
+fn to_pandas_dataframe<'py>(py: Python<'py>, arrow_table: Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
+    // This calls the 'to_pandas' method on the PyArrow object.
+    arrow_table.call_method0("to_pandas")
+}
+
+// --- NEW PYTHON-CALLABLE ENTRY POINT FOR SCHEMA PROFILING ---
+#[pyfunction]
+#[pyo3(signature = (config_path))]
+fn profile_data(config_path: String) -> PyResult<String> {
+    // Note: The profiler uses a small, current_thread tokio runtime since it's sequential I/O.
+    let output = std::thread::spawn(move || {
+        tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .unwrap()
+            .block_on(async {
+                run_profiler_logic(&config_path).await
+            })
+    }).join().unwrap().map_err(|e| PyValueError::new_err(format!("Profiling Error: {:?}", e)))?;
+
+    Ok(output)
+}
+
 
 // --- PYTHON MODULE EXPORT ---
-
 #[pymodule]
 fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
+    m.add_function(wrap_pyfunction!(to_pandas_dataframe, m)?)?;
+    m.add_function(wrap_pyfunction!(profile_data, m)?)?; // NEW: Schema profiler
     Ok(())
 }

src/parser.rs

@@ -1,8 +1,8 @@
 // --- External Crates ---
 use std::pin::Pin;
 use std::sync::Arc;
-// FIX: Replace direct tokio_postgres connections with deadpool
 use tokio_postgres::{NoTls, CopyOutStream, Config as PgConfig};
+// Required for the connection pool
 use deadpool_postgres::{Pool, Manager, Runtime};
 use anyhow::{Context, Result, anyhow};
 use arrow::array::{
@@ -19,14 +19,14 @@ use bytes::{Bytes, BytesMut, Buf};
 use byteorder::{BigEndian, ReadBytesExt};
 use std::io::{Cursor, Read};
 use std::str;
-use chrono::{NaiveDateTime, NaiveDate};
-use std::mem;
 use std::str::FromStr;
+use chrono::{NaiveDateTime, NaiveDate};
+use std::time::Instant;
 use tokio::task::JoinSet;
 use arrow::compute::concat_batches;
 
 // --- Internal Crates ---
-use crate::config::ConnectorConfig;
+use crate::config::{ConnectorConfig, load_and_validate_config};
 
 
 // --- CONSTANTS ---
@@ -36,25 +36,26 @@ const POSTGRES_EPOCH_MICROS_OFFSET: i64 = 946684800000000;
 
 
 // --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) ---
-// Note: We use blast_radius from Python config as the partitioning strategy
+// This is the fully optimized function using Connection Pooling and Static Dispatch.
 pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<RecordBatch> {
 
+    // Start overall timer
+    let start_total = Instant::now();
+    let start_phase1 = Instant::now();
+
     println!("UncheckedIO: Starting Query Planner (Blast Radius: {})...", blast_radius);
 
-    // --- PHASE 1: SETUP CONNECTION POOL ---
+    // --- PHASE 1: SETUP CONNECTION POOL & STATS ---
     let pg_config: tokio_postgres::Config = PgConfig::from_str(&config.connection_string)
         .context("Invalid connection string in config")?;
 
-    // Initialize the Manager and Pool
     let manager = Manager::new(pg_config.clone(), NoTls);
-    // Set pool size higher than the expected partition count to ensure connections are always available.
     let pool = Pool::builder(manager)
         .max_size(20)
         .runtime(Runtime::Tokio1)
         .build()
         .context("Failed to build connection pool")?;
 
-    // 2. Query for Table Bounds (using pool connection)
     let client = pool.get().await.context("Failed to get pool connection for stats query")?;
     let partition_key = "id";
 
@@ -62,16 +63,14 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
         .context("Failed to parse base query from config")?;
     let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")");
 
-    // NOTE: We rely on MIN/MAX here, assuming dense key for benchmark data.
     let stats_query = format!("SELECT MIN({}), MAX({}) FROM ({}) AS subquery", partition_key, partition_key, base_query_inner);
 
     let row = client.query_one(&stats_query, &[]).await?;
     let min_id: i64 = row.try_get(0).context("Failed to get MIN(id)")?;
     let max_id: i64 = row.try_get(1).context("Failed to get MAX(id)")?;
-    // Connection is returned to the pool when 'client' is dropped here.
+
     println!("UncheckedIO: ID Range: {} to {}", min_id, max_id);
 
-    // 4. Generate Partitioned Queries (Based on blast_radius from Python)
     struct PartitionTask {
         index: usize,
         query: String,
@@ -95,25 +94,27 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
         idx += 1;
     }
     println!("UncheckedIO: Generated {} parallel partitions.", partitions.len());
+    let duration_phase1 = start_phase1.elapsed();
+
+
+    // --- PHASE 2: PARALLEL EXECUTION (DATA TRANSFER + PARSING) ---
+    let start_phase2 = Instant::now();
 
-    // --- Phase 2: Parallel Execution ---
     let arrow_schema = Arc::new(build_arrow_schema(&config)?);
     let mut join_set = JoinSet::new();
 
     for task in partitions {
-        let worker_pool = pool.clone(); // Pass the pool handle
+        let worker_pool = pool.clone();
         let worker_schema = arrow_schema.clone();
 
         join_set.spawn(async move {
             let worker_logic = async {
-                // Get connection from pool (This is the speedup)
                 let worker_client = worker_pool.get().await
                     .context("Worker: Failed to get pool connection")?;
 
                 let copy_stream = worker_client.copy_out(task.query.as_str()).await?;
                 let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
 
-                // Call the static dispatch parser
                 parse_data_with_schema(pinned_stream, worker_schema).await
             };
 
@@ -122,7 +123,7 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
         });
     }
 
-    // --- Phase 3: Collect, Order, and Stitch ---
+    // Collect results
     let mut results: Vec<Option<RecordBatch>> = vec![None; idx];
 
     while let Some(join_result) = join_set.join_next().await {
@@ -133,25 +134,48 @@ pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<
                 results[index] = Some(batch.1);
             }
             Err(e) => {
-                // --- Self-Healing Placeholder ---
                 eprintln!("UncheckedIO: Partition {} failed! Error: {}. Falling back to NULLs (Self-Healing logic required here).", index, e);
                 let null_batch = create_null_batch(arrow_schema.clone(), expected_rows)?;
                 results[index] = Some(null_batch);
             }
         }
     }
+    let duration_phase2 = start_phase2.elapsed();
+
+
+    // --- PHASE 3: CONCATENATION AND FINALIZATION ---
+    let start_phase3 = Instant::now();
 
     let batches: Vec<RecordBatch> = results.into_iter()
         .filter_map(|b| b)
         .collect();
 
     if batches.is_empty() {
+        println!("UncheckedIO: All workers returned empty batches.");
+        let duration_total = start_total.elapsed();
+        println!("--- UncheckedIO Internal Timing ---");
+        println!("Phase 1 (Setup, Query): {:.2?}", duration_phase1);
+        println!("Phase 2 (I/O, Parsing): {:.2?}", duration_phase2);
+        println!("Phase 3 (Concatenation): {:.2?}", start_phase3.elapsed());
+        println!("Total Wall Time: {:.2?}", duration_total);
         return Ok(RecordBatch::new_empty(arrow_schema));
     }
 
     let final_batch = concat_batches(&arrow_schema, &batches)
         .context("Failed to concatenate parallel batches")?;
 
+    let duration_phase3 = start_phase3.elapsed();
+    let duration_total = start_total.elapsed();
+
+
+    // --- FINAL REPORTING ---
+    println!("--- UncheckedIO Internal Timing ---");
+    println!("Phase 1 (Setup, Query): {:.2?}", duration_phase1);
+    println!("Phase 2 (I/O, Parsing): {:.2?}", duration_phase2);
+    println!("Phase 3 (Concatenation): {:.2?}", duration_phase3);
+    println!("Total Wall Time: {:.2?}", duration_total);
+
+
     Ok(final_batch)
 }
 
@@ -165,7 +189,7 @@ fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> Result<RecordBatch
 /// Helper function to build the Arrow Schema from the config
 fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> {
     let schema_fields: Vec<Field> = config.schema.iter().map(|col_cfg| {
-        let nullable = col_cfg.column_name == "notes"; // Hack for MVP
+        let nullable = col_cfg.column_name == "notes";
 
         let arrow_type = match col_cfg.arrow_type.as_str() {
             "Int64" => DataType::Int64,
@@ -410,3 +434,74 @@ fn read_string_field(
 
     Ok(())
 }
+
+// --------------------------------------------------------------------------------
+// --- 3. PROFILER LOGIC (New Feature) ---
+// --------------------------------------------------------------------------------
+
+/// Maps a PostgreSQL internal type name to a standard Arrow Type string for config.yaml.
+fn map_postgres_to_arrow_type(pg_type_name: &str) -> Option<&'static str> {
+    match pg_type_name {
+        "int8" | "bigint" | "serial8" => Some("Int64"),
+        "int4" | "integer" | "serial" => Some("Int32"),
+        "float8" | "double precision" => Some("Float64"),
+        "float4" | "real" => Some("Float32"),
+        "varchar" | "text" | "uuid" => Some("Utf8"),
+        "bool" | "boolean" => Some("Boolean"),
+        "timestamptz" | "timestamp" => Some("Timestamp(Nanosecond, None)"),
+        "date" => Some("Date32"),
+        _ => None, // Returns None for unsupported types (like JSON, arrays, etc.)
+    }
+}
+
+pub async fn run_profiler_logic(config_path: &str) -> Result<String> {
+    // Phase 1: Load config to get connection string and query
+    let config: ConnectorConfig = load_and_validate_config(config_path)
+        .context("Failed to load and validate config for profiling")?;
+
+    // Use a non-COPY query to get metadata
+    let (base_query, _) = config.query.trim().split_once("TO STDOUT (FORMAT binary)")
+        .context("Query in config is malformed or not a COPY command")?;
+
+    // We only need the base query for the metadata query
+    let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")");
+
+    // Construct the metadata query (limit 0 is fastest)
+    let metadata_query = format!("SELECT * FROM ({}) AS subquery LIMIT 0", base_query_inner);
+
+    // Phase 2: Connect and execute the query
+    let pg_config: PgConfig = PgConfig::from_str(&config.connection_string)?;
+    let (client, connection) = pg_config.connect(NoTls).await
+        .context("Profiler: Failed to connect to PostgreSQL")?;
+
+    tokio::spawn(async move {
+        if let Err(e) = connection.await { eprintln!("Profiler connection error: {}", e); }
+    });
+
+    let statement = client.prepare(&metadata_query).await
+        .context("Profiler: Failed to prepare metadata query")?;
+
+    let mut output = String::from("schema:\n");
+
+    // Phase 3: Inspect the statement's columns for metadata
+    for column in statement.columns() {
+        let pg_type_name = column.type_().name().to_lowercase();
+        let arrow_type = map_postgres_to_arrow_type(&pg_type_name)
+            .unwrap_or("UNKNOWN (Review Manually)");
+
+        let column_entry = format!(
+            "- arrow_type: {}\n  column_name: {}\n",
+            arrow_type,
+            column.name()
+        );
+        output.push_str(&column_entry);
+    }
+
+    // Final instructions for the user
+    output.push_str("\n# NOTE: Paste the 'schema' block above into your config.yaml\n");
+    output.push_str(
+        "# REVIEW any UNKNOWN types. PostgreSQL types: (int8, float8, text, bool, timestamp, date, etc.)\n"
+    );
+
+    Ok(output)
+}