From 3b0d91c22bd3d4bb697c4402cc724cfc099b5b78 Mon Sep 17 00:00:00 2001
From: Arnaud Gourlay <arnaud.gourlay@gmail.com>
Date: Wed, 17 Apr 2024 13:53:59 +0200
Subject: [PATCH] Fix average sparse vector recommendation (#4047)

* Fix average sparse vector recommendation

* Revert "Fix average sparse vector recommendation"

This reverts commit 186c51209c95292a6adbb485bcfad7354484d985.

* sort within combine (Ivan's approach)

* Revert "sort within combine (Ivan's approach)"

This reverts commit 28d560cc569d5f52fcf3c0628c3b67e0376babca.

* simpler approach without mutation
---
 lib/sparse/src/common/sparse_vector.rs | 54 ++++++++++++++++++++------
 1 file changed, 42 insertions(+), 12 deletions(-)

diff --git a/lib/sparse/src/common/sparse_vector.rs b/lib/sparse/src/common/sparse_vector.rs
index 4fdf54a7e2..5d33daeb40 100644
--- a/lib/sparse/src/common/sparse_vector.rs
+++ b/lib/sparse/src/common/sparse_vector.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use itertools::Itertools;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -86,23 +88,42 @@ impl SparseVector {
         }
     }
 
-    /// Construct a new vector that is the result of performing all indices-wise operations
+    /// Construct a new vector that is the result of performing all indices-wise operations.
+    /// Automatically sort input vectors if necessary.
     pub fn combine_aggregate(
         &self,
         other: &SparseVector,
         op: impl Fn(DimWeight, DimWeight) -> DimWeight,
     ) -> Self {
-        debug_assert!(self.is_sorted());
-        debug_assert!(other.is_sorted());
+        // Copy and sort `self` vector if not already sorted
+        let this: Cow<SparseVector> = if !self.is_sorted() {
+            let mut this = self.clone();
+            this.sort_by_indices();
+            Cow::Owned(this)
+        } else {
+            Cow::Borrowed(self)
+        };
+        assert!(this.is_sorted());
+
+        // Copy and sort `other` vector if not already sorted
+        let cow_other: Cow<SparseVector> = if !other.is_sorted() {
+            let mut other = other.clone();
+            other.sort_by_indices();
+            Cow::Owned(other)
+        } else {
+            Cow::Borrowed(other)
+        };
+        let other = &cow_other;
+        assert!(other.is_sorted());
 
         let mut result = SparseVector::default();
         let mut i = 0;
         let mut j = 0;
-        while i < self.indices.len() && j < other.indices.len() {
-            match self.indices[i].cmp(&other.indices[j]) {
+        while i < this.indices.len() && j < other.indices.len() {
+            match this.indices[i].cmp(&other.indices[j]) {
                 std::cmp::Ordering::Less => {
-                    result.indices.push(self.indices[i]);
-                    result.values.push(op(self.values[i], 0.0));
+                    result.indices.push(this.indices[i]);
+                    result.values.push(op(this.values[i], 0.0));
                     i += 1;
                 }
                 std::cmp::Ordering::Greater => {
@@ -111,16 +132,16 @@ impl SparseVector {
                     j += 1;
                 }
                 std::cmp::Ordering::Equal => {
-                    result.indices.push(self.indices[i]);
-                    result.values.push(op(self.values[i], other.values[j]));
+                    result.indices.push(this.indices[i]);
+                    result.values.push(op(this.values[i], other.values[j]));
                     i += 1;
                     j += 1;
                 }
             }
         }
-        while i < self.indices.len() {
-            result.indices.push(self.indices[i]);
-            result.values.push(op(self.values[i], 0.0));
+        while i < this.indices.len() {
+            result.indices.push(this.indices[i]);
+            result.values.push(op(this.values[i], 0.0));
             i += 1;
         }
         while j < other.indices.len() {
@@ -249,14 +270,23 @@ mod tests {
 
     #[test]
     fn combine_aggregate_test() {
+        // Test with missing index
         let a = SparseVector::new(vec![1, 2, 3], vec![0.1, 0.2, 0.3]).unwrap();
         let b = SparseVector::new(vec![2, 3, 4], vec![2.0, 3.0, 4.0]).unwrap();
         let sum = a.combine_aggregate(&b, |x, y| x + 2.0 * y);
         assert_eq!(sum.indices, vec![1, 2, 3, 4]);
         assert_eq!(sum.values, vec![0.1, 4.2, 6.3, 8.0]);
 
+        // reverse arguments
         let sum = b.combine_aggregate(&a, |x, y| x + 2.0 * y);
         assert_eq!(sum.indices, vec![1, 2, 3, 4]);
         assert_eq!(sum.values, vec![0.2, 2.4, 3.6, 4.0]);
+
+        // Test with non-sorted input
+        let a = SparseVector::new(vec![1, 2, 3], vec![0.1, 0.2, 0.3]).unwrap();
+        let b = SparseVector::new(vec![4, 2, 3], vec![4.0, 2.0, 3.0]).unwrap();
+        let sum = a.combine_aggregate(&b, |x, y| x + 2.0 * y);
+        assert_eq!(sum.indices, vec![1, 2, 3, 4]);
+        assert_eq!(sum.values, vec![0.1, 4.2, 6.3, 8.0]);
     }
 }