(&mut self, states: &[ArrayRef])
| 652 | } |
| 653 | |
| 654 | fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { |
| 655 | let count_arr = as_uint64_array(&states[0])?; |
| 656 | let mean_x_arr = as_float64_array(&states[1])?; |
| 657 | let mean_y_arr = as_float64_array(&states[2])?; |
| 658 | let m2_x_arr = as_float64_array(&states[3])?; |
| 659 | let m2_y_arr = as_float64_array(&states[4])?; |
| 660 | let algo_const_arr = as_float64_array(&states[5])?; |
| 661 | |
| 662 | for i in 0..count_arr.len() { |
| 663 | let count_b = count_arr.value(i); |
| 664 | if count_b == 0_u64 { |
| 665 | continue; |
| 666 | } |
| 667 | let (count_a, mean_x_a, mean_y_a, m2_x_a, m2_y_a, algo_const_a) = ( |
| 668 | self.count, |
| 669 | self.mean_x, |
| 670 | self.mean_y, |
| 671 | self.m2_x, |
| 672 | self.m2_y, |
| 673 | self.algo_const, |
| 674 | ); |
| 675 | let (count_b, mean_x_b, mean_y_b, m2_x_b, m2_y_b, algo_const_b) = ( |
| 676 | count_b, |
| 677 | mean_x_arr.value(i), |
| 678 | mean_y_arr.value(i), |
| 679 | m2_x_arr.value(i), |
| 680 | m2_y_arr.value(i), |
| 681 | algo_const_arr.value(i), |
| 682 | ); |
| 683 | |
| 684 | // Assuming two different batches of input have calculated the states: |
| 685 | // batch A of Y, X -> {count_a, mean_x_a, mean_y_a, m2_x_a, algo_const_a} |
| 686 | // batch B of Y, X -> {count_b, mean_x_b, mean_y_b, m2_x_b, algo_const_b} |
| 687 | // The merged states from A and B are {count_ab, mean_x_ab, mean_y_ab, m2_x_ab, |
| 688 | // algo_const_ab} |
| 689 | // |
| 690 | // Reference for the algorithm to merge states: |
| 691 | // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm |
| 692 | let count_ab = count_a + count_b; |
| 693 | let (count_a, count_b) = (count_a as f64, count_b as f64); |
| 694 | let d_x = mean_x_b - mean_x_a; |
| 695 | let d_y = mean_y_b - mean_y_a; |
| 696 | let mean_x_ab = mean_x_a + d_x * count_b / count_ab as f64; |
| 697 | let mean_y_ab = mean_y_a + d_y * count_b / count_ab as f64; |
| 698 | let m2_x_ab = |
| 699 | m2_x_a + m2_x_b + d_x * d_x * count_a * count_b / count_ab as f64; |
| 700 | let m2_y_ab = |
| 701 | m2_y_a + m2_y_b + d_y * d_y * count_a * count_b / count_ab as f64; |
| 702 | let algo_const_ab = algo_const_a |
| 703 | + algo_const_b |
| 704 | + d_x * d_y * count_a * count_b / count_ab as f64; |
| 705 | |
| 706 | self.count = count_ab; |
| 707 | self.mean_x = mean_x_ab; |
| 708 | self.mean_y = mean_y_ab; |
| 709 | self.m2_x = m2_x_ab; |
| 710 | self.m2_y = m2_y_ab; |
| 711 | self.algo_const = algo_const_ab; |
nothing calls this directly
no test coverage detected