1use super::ColumnFrame;
2use crate::Key;
3use data_value::DataValue;
4use std::collections::HashMap;
5use std::ops::{AddAssign, DivAssign, MulAssign, SubAssign};
6
7macro_rules! op_assign {
8 ($item: ident, $op: ident, $oop: tt) => {
9 impl $item for ColumnFrame {
10 fn $op(&mut self, rhs: Self) {
11 if self.data_frame.ncols() != rhs.data_frame.ncols() || self.data_frame.ncols() == 1 {
12 let indexes = self.index.select(rhs.keys());
13 for (right_idx, left_index) in rhs.index.indexes().into_iter().zip(indexes.indexes()) {
14 let right_col = rhs.data_frame.column(right_idx);
15 if right_col.len() == 1 {
16 let right_value = &right_col[0];
18 for i in 0..self.data_frame.nrows() {
19 self.data_frame[(i, left_index)].$op(right_value);
20 }
21 } else {
22 for (alpha, rhs) in self
24 .data_frame
25 .column_mut(left_index)
26 .iter_mut()
27 .zip(right_col.iter())
28 {
29 alpha.$op(rhs);
30 }
31 }
32 }
33 } else {
34 self.data_frame.zip_mut_with(&rhs.data_frame, |x, y| {
35 let z = x.clone();
36 *x = z $oop y.clone();
37 });
38 }
39 }
40 }
41
42 impl $item<HashMap<String, DataValue>> for ColumnFrame {
43 fn $op(&mut self, rhs: HashMap<String, DataValue>) {
44 let keys = rhs.keys().cloned().map(Key::from).collect::<Vec<_>>();
45 let indexes = self.index.select(keys.as_slice());
46 let int_idx =indexes.indexes();
47 for left_index in int_idx.iter(){
48 let key_right = self.index.get_key(*left_index).expect("BUG: Defined above!");
49 let right_value = rhs.get(key_right.name()).cloned().unwrap_or_default();
50 for i in 0..self.data_frame.nrows() {
51 self.data_frame[(i, *left_index)].$op(&right_value);
52 }
53
54 }
55 }
56 }
57 };
58}
59
60op_assign!(AddAssign, add_assign, +);
61op_assign!(SubAssign, sub_assign, -);
62op_assign!(MulAssign, mul_assign, *);
63op_assign!(DivAssign, div_assign, /);
64
65#[cfg(test)]
66mod test {
67 use crate::column_frame;
68
69 use super::*;
70 use data_value::DataValue;
71 use rstest::*;
72 #[rstest]
73 #[case(
74 column_frame!("a" => vec![1, 2, 3]),
75 column_frame!("a" => vec![1, 2, 3]),
76 column_frame!("a" => vec![2, 4, 6])
77 )]
78 #[case(
79 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
80 column_frame!("a" => vec![2]),
81 column_frame!("a" => vec![3, 4, 5], "b" => vec![1, 2, 3])
82 )]
83 #[case(
84 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
85 column_frame!("a" => vec![2, 3]),
86 column_frame!("a" => vec![3, 5, 3], "b" => vec![1, 2, 3])
87 )]
88 fn test_add_assign(
89 #[case] mut a: ColumnFrame,
90 #[case] b: ColumnFrame,
91 #[case] expected: ColumnFrame,
92 ) {
93 a += b;
94 assert_eq!(a, expected);
95 }
96 #[rstest]
97 #[case(
98 column_frame!("a" => vec![1, 2, 3]),
99 column_frame!("a" => vec![1, 2, 3]),
100 column_frame!("a" => vec![0f32, 0f32, 0f32])
101 )]
102 #[case(
103 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
104 column_frame!("a" => vec![2]),
105 column_frame!("a" => vec![-1f32, 0f32, 1f32], "b" => vec![1, 2, 3])
106 )]
107 #[case(
108 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
109 column_frame!("a" => vec![2, 3]),
110 column_frame!("a" => vec![DataValue::from(-1f32), DataValue::from(-1f32), DataValue::from(3)], "b" => vec![1, 2, 3])
111 )]
112 fn test_sub_assign(
113 #[case] mut a: ColumnFrame,
114 #[case] b: ColumnFrame,
115 #[case] expected: ColumnFrame,
116 ) {
117 a -= b;
118 assert_eq!(a, expected);
119 }
120
121 #[rstest]
122 #[case(
123 column_frame!("a" => vec![1, 2, 3]),
124 column_frame!("a" => vec![1, 2, 3]),
125 column_frame!("a" => vec![1f32, 4f32, 9f32])
126 )]
127 #[case(
128 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
129 column_frame!("a" => vec![2]),
130 column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![1, 2, 3])
131 )]
132 #[case(
133 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
134 column_frame!("a" => vec![2, 3]),
135 column_frame!("a" => vec![DataValue::from(2f32), DataValue::from(6f32), DataValue::from(3)], "b" => vec![1, 2, 3])
136 )]
137 fn test_mul_assign(
138 #[case] mut a: ColumnFrame,
139 #[case] b: ColumnFrame,
140 #[case] expected: ColumnFrame,
141 ) {
142 a *= b;
143 assert_eq!(a, expected);
144 }
145
146 #[rstest]
147 #[case(
148 column_frame!("a" => vec![1, 2, 3]),
149 column_frame!("a" => vec![1, 2, 3]),
150 column_frame!("a" => vec![1f32, 1f32, 1f32])
151 )]
152 #[case(
153 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
154 column_frame!("a" => vec![2]),
155 column_frame!("a" => vec![0.5f32, 1f32, (3f32/2f32)], "b" => vec![1, 2, 3])
156 )]
157 #[case(
158 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
159 column_frame!("a" => vec![2, 3]),
160 column_frame!("a" => vec![DataValue::from(0.5f32), DataValue::from(2f32/3f32), DataValue::from(3)], "b" => vec![1, 2, 3])
161 )]
162 fn test_div_assign(
163 #[case] mut a: ColumnFrame,
164 #[case] b: ColumnFrame,
165 #[case] expected: ColumnFrame,
166 ) {
167 a /= b;
168 assert_eq!(a, expected);
169 }
170
171 #[rstest]
172 #[case(
173 column_frame!("a" => vec![1, 2, 3]),
174 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
175 column_frame!("a" => vec![2f32, 4f32, 6f32])
176 )]
177 #[case(
178 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
179 crate::data_value::stdhashmap!("a" => 2),
180 column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![1, 2, 3])
181 )]
182 #[case(
183 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
184 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
185 column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![3f32, 6f32, 9f32])
186 )]
187 fn test_mul_assign_map(
188 #[case] mut a: ColumnFrame,
189 #[case] b: HashMap<String, DataValue>,
190 #[case] expected: ColumnFrame,
191 ) {
192 a *= b;
193 assert_eq!(a, expected);
194 }
195
196 #[rstest]
197 #[case(
198 column_frame!("a" => vec![1, 2, 3]),
199 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
200 column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32])
201 )]
202 #[case(
203 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
204 crate::data_value::stdhashmap!("a" => 2),
205 column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32], "b" => vec![1, 2, 3])
206 )]
207 #[case(
208 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
209 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
210 column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32], "b" => vec![1f32/3f32, 2f32/3f32, 1f32])
211 )]
212 fn test_div_assign_map(
213 #[case] mut a: ColumnFrame,
214 #[case] b: HashMap<String, DataValue>,
215 #[case] expected: ColumnFrame,
216 ) {
217 a /= b;
218 assert_eq!(a, expected);
219 }
220
221 #[rstest]
222 #[case(
223 column_frame!("a" => vec![1, 2, 3]),
224 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
225 column_frame!("a" => vec![3, 4, 5])
226 )]
227 #[case(
228 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
229 crate::data_value::stdhashmap!("a" => 2),
230 column_frame!("a" => vec![3, 4, 5], "b" => vec![1, 2, 3])
231 )]
232 #[case(
233 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
234 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
235 column_frame!("a" => vec![3, 4, 5], "b" => vec![4, 5, 6])
236 )]
237 fn test_add_assign_map(
238 #[case] mut a: ColumnFrame,
239 #[case] b: HashMap<String, DataValue>,
240 #[case] expected: ColumnFrame,
241 ) {
242 a += b;
243 assert_eq!(a, expected);
244 }
245
246 #[rstest]
247 #[case(
248 column_frame!("a" => vec![1, 2, 3]),
249 crate::data_value::stdhashmap!("a" => 1, "b" => 3),
250 column_frame!("a" => vec![0f32, 1f32, 2f32])
251 )]
252 #[case(
253 column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
254 crate::data_value::stdhashmap!("a" => 1),
255 column_frame!("a" => vec![0f32, 1f32, 2f32], "b" => vec![1, 2, 3])
256 )]
257 #[case(
258 column_frame!("a" => vec![10, 20, 30], "b" => vec![10, 20, 30]),
259 crate::data_value::stdhashmap!("a" => 2, "b" => 3),
260 column_frame!("a" => vec![8f32, 18f32, 28f32], "b" => vec![7f32, 17f32, 27f32])
261 )]
262 fn test_sub_assign_map(
263 #[case] mut a: ColumnFrame,
264 #[case] b: HashMap<String, DataValue>,
265 #[case] expected: ColumnFrame,
266 ) {
267 a -= b;
268 assert_eq!(a, expected);
269 }
270}