trs-dataframe 0.11.0

use super::ColumnFrame;
use crate::Key;
use data_value::DataValue;
use std::collections::HashMap;
use std::ops::{AddAssign, DivAssign, MulAssign, SubAssign};

macro_rules! op_assign {
    ($item: ident, $op: ident, $oop: tt) => {
        impl $item for ColumnFrame {
            fn $op(&mut self, rhs: Self) {
                if self.ncolumns() != rhs.ncolumns() || self.ncolumns() == 1 {
                    let indexes = self.index.select(rhs.keys());
                    for (right_idx, left_index) in rhs.index.indexes().into_iter().zip(indexes.indexes()) {
                        let right_col = &rhs.data_frame[right_idx];
                        if right_col.len() == 1 {
                            let right_value = right_col.get_or_null(0);
                            let nrows = self.nrows();
                            let left = &mut self.data_frame[left_index];
                            for i in 0..nrows {
                                let mut current = left.get_or_null(i);
                                current.$op(&right_value);
                                let _ = left.set(i, current);
                            }
                        } else {
                            let right_values: Vec<DataValue> = right_col.to_vec();
                            let left = &mut self.data_frame[left_index];
                            for (i, rhs_val) in right_values.iter().enumerate() {
                                let mut current = left.get_or_null(i);
                                current.$op(rhs_val);
                                let _ = left.set(i, current);
                            }
                        }
                    }
                } else {
                    for (self_col, rhs_col) in self.data_frame.iter_mut().zip(rhs.data_frame.iter()) {
                        let rhs_values: Vec<DataValue> = rhs_col.to_vec();
                        for (i, rhs_val) in rhs_values.iter().enumerate() {
                            let current = self_col.get_or_null(i);
                            let result = current.clone() $oop rhs_val.clone();
                            let _ = self_col.set(i, result);
                        }
                    }
                }
            }
        }

        impl $item<HashMap<String, DataValue>> for ColumnFrame {
            fn $op(&mut self, rhs: HashMap<String, DataValue>) {
                let keys = rhs.keys().cloned().map(Key::from).collect::<Vec<_>>();
                let indexes = self.index.select(keys.as_slice());
                let int_idx = indexes.indexes();
                for left_index in int_idx.iter() {
                    let key_right = self.index.get_key(*left_index).expect("BUG: Defined above!");
                    let right_value = rhs.get(key_right.name()).cloned().unwrap_or_default();
                    let nrows = self.nrows();
                    let left = &mut self.data_frame[*left_index];
                    for i in 0..nrows {
                        let mut current = left.get_or_null(i);
                        current.$op(&right_value);
                        let _ = left.set(i, current);
                    }
                }
            }
        }
    };
}

op_assign!(AddAssign, add_assign, +);
op_assign!(SubAssign, sub_assign, -);
op_assign!(MulAssign, mul_assign, *);
op_assign!(DivAssign, div_assign, /);

#[cfg(test)]
mod test {
    use crate::column_frame;

    use super::*;
    use data_value::DataValue;
    use rstest::*;
    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 4, 6])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![3, 4, 5], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![3, 5, 3], "b" => vec![1, 2, 3])
    )]
    fn test_add_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a += b;
        assert_eq!(a, expected);
    }
    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![0f32, 0f32, 0f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![-1f32, 0f32, 1f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![DataValue::from(-1f32), DataValue::from(-1f32), DataValue::from(3)], "b" => vec![1, 2, 3])
    )]
    fn test_sub_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a -= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1f32, 4f32, 9f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![DataValue::from(2f32), DataValue::from(6f32), DataValue::from(3)], "b" => vec![1, 2, 3])
    )]
    fn test_mul_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a *= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1f32, 1f32, 1f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![0.5f32, 1f32, (3f32/2f32)], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![DataValue::from(0.5f32), DataValue::from(2f32/3f32), DataValue::from(3)], "b" => vec![1, 2, 3])
    )]
    fn test_div_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a /= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![2f32, 4f32, 6f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2),
        column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![3f32, 6f32, 9f32])
    )]
    fn test_mul_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a *= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2),
        column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32], "b" => vec![1f32/3f32, 2f32/3f32, 1f32])
    )]
    fn test_div_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a /= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![3, 4, 5])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2),
        column_frame!("a" => vec![3, 4, 5], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![3, 4, 5], "b" => vec![4, 5, 6])
    )]
    fn test_add_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a += b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 1, "b" => 3),
        column_frame!("a" => vec![0f32, 1f32, 2f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 1),
        column_frame!("a" => vec![0f32, 1f32, 2f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![10, 20, 30], "b" => vec![10, 20, 30]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![8f32, 18f32, 28f32], "b" => vec![7f32, 17f32, 27f32])
    )]
    fn test_sub_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a -= b;
        assert_eq!(a, expected);
    }
}