use crate::column::Column;
use crate::error::{Result, VolasError};
use super::{ComputedMeta, DataFrame};
impl DataFrame {
pub fn set_computed(&mut self, name: &str, directive: String, lookback: usize) {
self.computed.insert(
name.to_string(),
ComputedMeta {
directive,
lookback,
valid_rows: self.height,
state: None,
origin: 0,
},
);
}
pub fn set_computed_state(&mut self, name: &str, state: Option<Vec<f64>>) {
if let Some(meta) = self.computed.get_mut(name) {
meta.state = state;
}
}
pub fn has_stale_computed(&self) -> bool {
self.computed.values().any(|m| m.valid_rows < self.height)
}
pub fn computed_columns(&self) -> Vec<(String, ComputedMeta)> {
self.computed
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
}
pub fn stale_computed_columns(&self, only: Option<&str>) -> Vec<(String, ComputedMeta)> {
self.computed
.iter()
.filter(|(name, meta)| {
meta.valid_rows < self.height && only.is_none_or(|target| target == name.as_str())
})
.map(|(name, meta)| (name.clone(), meta.clone()))
.collect()
}
pub fn computed_names(&self) -> Vec<String> {
self.computed.keys().cloned().collect()
}
pub fn update_computed_tail(&mut self, name: &str, from: usize, tail: &Column) -> Result<()> {
let pos = self
.name_to_idx
.get(name)
.copied()
.ok_or_else(|| VolasError::ColumnNotFound(name.to_string()))?;
match (&mut self.columns[pos], tail) {
(Column::F64(arc), Column::F64(t)) => {
let buf = arc.make_mut();
for (i, &v) in t.iter().enumerate() {
if from + i < buf.len() {
buf[from + i] = v;
}
}
}
(Column::Bool(arc, _), Column::Bool(t, _)) => {
let buf = arc.make_mut();
for (i, &v) in t.iter().enumerate() {
if from + i < buf.len() {
buf[from + i] = v;
}
}
}
(col, t) => {
return Err(VolasError::DType(format!(
"computed tail dtype {} does not match column \"{name}\" dtype {}",
t.dtype(),
col.dtype()
)))
}
}
if let Some(meta) = self.computed.get_mut(name) {
meta.valid_rows = self.height;
}
Ok(())
}
pub fn update_computed_f64_value(&mut self, name: &str, row: usize, value: f64) -> Result<()> {
let pos = self
.name_to_idx
.get(name)
.copied()
.ok_or_else(|| VolasError::ColumnNotFound(name.to_string()))?;
match &mut self.columns[pos] {
Column::F64(arc) => {
let buf = arc.make_mut();
if row < buf.len() {
buf[row] = value;
}
}
col => {
return Err(VolasError::DType(format!(
"computed scalar dtype F64 does not match column \"{name}\" dtype {}",
col.dtype()
)))
}
}
if let Some(meta) = self.computed.get_mut(name) {
meta.valid_rows = self.height;
}
Ok(())
}
fn drop_computed_at(&mut self, col: usize) {
if let Some(name) = self.names.get(col) {
self.computed.remove(name);
}
}
pub(super) fn invalidate_computed_on_write_at(&mut self, col: usize) {
self.drop_computed_at(col);
for meta in self.computed.values_mut() {
meta.valid_rows = 0;
meta.state = None;
}
}
pub fn invalidate_computed_on_write(&mut self, written_name: &str) {
self.computed.remove(written_name);
for meta in self.computed.values_mut() {
meta.valid_rows = 0;
meta.state = None;
}
}
}