tenshift-core 0.1.2

//! CSV file source.
//!
//! Reads a CSV file and produces one [`Sample`] per row. Each column becomes
//! a named tensor field. Numeric columns are stored as `f64` tensors; all
//! others as `Bytes` tensors containing the raw UTF-8 string value.
//!
//! # Feature gate
//!
//! Requires the `csv` feature:
//! ```toml
//! tenshift-core = { version = "0.1", features = ["csv"] }
//! ```
//!
//! # Example
//!
//! ```rust,ignore
//! use tenshift_core::sources::CsvSource;
//! use tenshift_core::Pipeline;
//!
//! let source = CsvSource::new("data/train.csv");
//! let pipeline = Pipeline::from_source(source)
//!     .batch(64)
//!     .prefetch(4);
//!
//! for batch in pipeline {
//!     // Each sample has fields named after CSV columns
//! }
//! ```

use std::path::{Path, PathBuf};

use crate::error::{Error, Result};
use crate::sample::{Sample, Tensor};
use crate::source::{Source, SourceIterator};

/// A source that reads a CSV file, one sample per row.
///
/// Column headers become field names. Values are parsed as `f64` when possible,
/// otherwise stored as raw byte tensors.
pub struct CsvSource {
    path: PathBuf,
    has_headers: bool,
    delimiter: u8,
}

impl CsvSource {
    /// Create a CSV source from a file path.
    pub fn new(path: impl AsRef<Path>) -> Self {
        Self {
            path: path.as_ref().to_path_buf(),
            has_headers: true,
            delimiter: b',',
        }
    }

    /// Set whether the first row contains column headers (default: true).
    pub fn headers(mut self, has_headers: bool) -> Self {
        self.has_headers = has_headers;
        self
    }

    /// Set the field delimiter (default: `,`).
    pub fn delimiter(mut self, delimiter: u8) -> Self {
        self.delimiter = delimiter;
        self
    }
}

impl Source for CsvSource {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        let mut reader = csv::ReaderBuilder::new()
            .has_headers(self.has_headers)
            .delimiter(self.delimiter)
            .from_path(&self.path)
            .map_err(|error| Error::ReadFailed {
                path: self.path.clone(),
                reason: error.to_string(),
            })?;

        let headers: Vec<String> = if self.has_headers {
            reader
                .headers()
                .map_err(|error| Error::ReadFailed {
                    path: self.path.clone(),
                    reason: format!("failed to read CSV headers: {error}"),
                })?
                .iter()
                .map(String::from)
                .collect()
        } else {
            Vec::new()
        };

        Ok(Box::new(CsvIterator {
            reader,
            headers,
            path: self.path.clone(),
            index: 0,
        }))
    }

    fn name(&self) -> &str {
        self.path.to_str().map_or("csv", |s| s)
    }
}

/// Iterator over CSV rows.
struct CsvIterator {
    reader: csv::Reader<std::fs::File>,
    headers: Vec<String>,
    path: PathBuf,
    index: u64,
}

impl SourceIterator for CsvIterator {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        let mut record = csv::StringRecord::new();
        match self.reader.read_record(&mut record) {
            Ok(true) => {}
            Ok(false) => return None, // EOF
            Err(error) => {
                return Some(Err(Error::CorruptData {
                    path: self.path.clone(),
                    reason: format!("CSV parse error on row {}: {error}", self.index + 1),
                }));
            }
        }

        let index = self.index;
        self.index += 1;

        let mut sample = Sample::new();

        while self.headers.len() < record.len() {
            self.headers.push(format!("col_{}", self.headers.len()));
        }

        for (col_idx, value) in record.iter().enumerate() {
            let field_name = self.headers[col_idx].as_str();
            let tensor = parse_csv_value(value);
            sample = sample.with(field_name, tensor);
        }

        Some(Ok(sample.with_metadata(self.path.to_string_lossy(), index)))
    }
}

/// Parse a CSV field value into a tensor.
///
/// Attempts `f64` parsing first; falls back to raw bytes.
fn parse_csv_value(value: &str) -> Tensor {
    if let Ok(number) = value.parse::<f64>() {
        Tensor::f64(&[number], vec![1])
    } else {
        Tensor::bytes(value.as_bytes().to_vec())
    }
}