use crate::{cli::Invocation, error::PhError, io};
use polars::prelude::*;
use polars::series::ops::NullBehavior;
pub fn run(df: DataFrame, inv: &Invocation) -> Result<Vec<u8>, PhError> {
let periods: i64 = inv
.option("periods")
.unwrap_or("1")
.parse()
.map_err(|_| PhError::new("diff --periods must be an integer."))?;
let axis: u8 = inv
.option("axis")
.unwrap_or("0")
.parse()
.map_err(|_| PhError::new("diff --axis must be 0 or 1."))?;
if axis > 1 {
return Err(PhError::new("diff --axis must be 0 or 1."));
}
let names: Vec<String> = if inv.args.is_empty() {
df.get_column_names()
.iter()
.map(|name| name.to_string())
.collect()
} else {
inv.args.clone()
};
for name in &names {
if df.column(name.as_str()).is_err() {
return Err(PhError::new(format!("ph diff: Unknown column {name}")));
}
}
let expressions: Vec<Expr> = if axis == 0 {
names
.iter()
.map(|name| {
col(name.as_str())
.cast(DataType::Float64)
.diff(lit(periods), NullBehavior::Ignore)
.alias(name.as_str())
})
.collect()
} else {
names
.iter()
.enumerate()
.map(|(i, name)| {
let j = i as i128 - periods as i128;
if j < 0 || j >= names.len() as i128 {
lit(NULL).cast(DataType::Float64).alias(name.as_str())
} else {
let rhs = &names[j as usize];
let lhs_dtype = df.column(name.as_str()).unwrap().dtype();
let rhs_dtype = df.column(rhs.as_str()).unwrap().dtype();
let dtype = if periods == 0
|| matches!(lhs_dtype, DataType::Float32 | DataType::Float64)
|| matches!(rhs_dtype, DataType::Float32 | DataType::Float64)
{
DataType::Float64
} else {
DataType::Int64
};
(col(name.as_str()).cast(dtype.clone()) - col(rhs.as_str()).cast(dtype))
.alias(name.as_str())
}
})
.collect()
};
if expressions.is_empty() {
return io::write_csv(df);
}
let output = df.lazy().with_columns(expressions).collect()?;
io::write_csv(output)
}