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