pub mod column_transformer;
pub use column_transformer::ColumnTransformer;
use crate::traits::{Error, Fit, Result, Transform};
use polars::prelude::*;
pub trait DataFrameTransformer:
Fit<DataFrame, DataFrame, Output = ()> + Transform<DataFrame, Output = DataFrame>
{
}
impl<T> DataFrameTransformer for T where
T: Fit<DataFrame, DataFrame, Output = ()> + Transform<DataFrame, Output = DataFrame>
{
}
pub struct Pipeline {
steps: Vec<(String, Box<dyn DataFrameTransformer>)>,
}
impl Pipeline {
pub fn new(steps: Vec<(String, Box<dyn DataFrameTransformer>)>) -> Self {
assert!(
!steps.is_empty(),
"Pipeline::new: at least one step is required. \
Provide a non-empty Vec of (name, transformer) pairs."
);
Self { steps }
}
pub fn steps(&self) -> &[(String, Box<dyn DataFrameTransformer>)] {
&self.steps
}
}
impl Fit<DataFrame, DataFrame> for Pipeline {
type Output = ();
fn fit(&mut self, x: DataFrame, y: DataFrame) -> Result<()> {
if x.height() == 0 {
return Err(Error::InvalidInput(
"Pipeline.fit received a DataFrame with 0 rows.".into(),
));
}
let mut x_curr = x;
let y_curr = y;
let n = self.steps.len();
for (i, (name, transformer)) in self.steps.iter_mut().enumerate() {
let is_last = i == n - 1;
transformer
.fit(x_curr.clone(), y_curr.clone())
.map_err(|e| {
Error::Computation(format!(
"Pipeline: step {} ('{}') failed during fit: {}",
i, name, e
))
})?;
if !is_last {
x_curr = transformer.transform(x_curr).map_err(|e| {
Error::Computation(format!(
"Pipeline: step {} ('{}') failed during intermediate transform: {}",
i, name, e
))
})?;
}
}
Ok(())
}
}
impl Transform<DataFrame> for Pipeline {
type Output = DataFrame;
fn transform(&self, x: DataFrame) -> Result<DataFrame> {
let mut x_curr = x;
for (i, (name, transformer)) in self.steps.iter().enumerate() {
x_curr = transformer.transform(x_curr).map_err(|e| {
Error::Computation(format!(
"Pipeline: step {} ('{}') failed during transform: {}",
i, name, e
))
})?;
}
Ok(x_curr)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::preprocessing::scaler::StandardScaler;
fn make_test_df() -> DataFrame {
let a = Column::from(Series::new("a".into(), &[1.0f64, 3.0, 5.0]));
let b = Column::from(Series::new("b".into(), &[2.0f64, 4.0, 6.0]));
DataFrame::new(3, vec![a, b]).unwrap()
}
#[test]
fn test_pipeline_single_step() {
let scaler = StandardScaler::new();
let mut pipeline = Pipeline::new(vec![("scaler".into(), Box::new(scaler))]);
let df = make_test_df();
let y = df.clone();
pipeline.fit(df.clone(), y).unwrap();
let result = pipeline.transform(df).unwrap();
assert_eq!(result.width(), 2);
assert_eq!(result.height(), 3);
}
}