From 334a3814bbd03e7af5ded876c70217818d601baa Mon Sep 17 00:00:00 2001
From: Yishuai Li <liyishuai.lys@alibaba-inc.com>
Date: Tue, 7 Jan 2025 16:59:36 +0800
Subject: [PATCH] fix: order of GroupBy and PartitionBy

---
 parallel/slice.go | 58 +++++++++++++++--------------------------------
 1 file changed, 18 insertions(+), 40 deletions(-)

diff --git a/parallel/slice.go b/parallel/slice.go
index a70fb70f..416935bc 100644
--- a/parallel/slice.go
+++ b/parallel/slice.go
@@ -66,65 +66,43 @@ func Times[T any](count int, iteratee func(index int) T) []T {
 }
 
 // GroupBy returns an object composed of keys generated from the results of running each element of collection through iteratee.
+// The order of grouped values is determined by the order they occur in collection.
 // `iteratee` is call in parallel.
 func GroupBy[T any, U comparable, Slice ~[]T](collection Slice, iteratee func(item T) U) map[U]Slice {
 	result := map[U]Slice{}
 
-	var mu sync.Mutex
-	var wg sync.WaitGroup
-	wg.Add(len(collection))
-
-	for _, item := range collection {
-		go func(_item T) {
-			key := iteratee(_item)
-
-			mu.Lock()
-
-			result[key] = append(result[key], _item)
+	keys := Map(collection, func(item T, _ int) U {
+		return iteratee(item)
+	})
 
-			mu.Unlock()
-			wg.Done()
-		}(item)
+	for i, item := range collection {
+		result[keys[i]] = append(result[keys[i]], item)
 	}
 
-	wg.Wait()
-
 	return result
 }
 
 // PartitionBy returns an array of elements split into groups. The order of grouped values is
 // determined by the order they occur in collection. The grouping is generated from the results
 // of running each element of collection through iteratee.
+// The order of groups is determined by their first appearance in the collection.
 // `iteratee` is call in parallel.
 func PartitionBy[T any, K comparable, Slice ~[]T](collection Slice, iteratee func(item T) K) []Slice {
 	result := []Slice{}
 	seen := map[K]int{}
 
-	var mu sync.Mutex
-	var wg sync.WaitGroup
-	wg.Add(len(collection))
-
-	for _, item := range collection {
-		go func(_item T) {
-			key := iteratee(_item)
-
-			mu.Lock()
-
-			resultIndex, ok := seen[key]
-			if !ok {
-				resultIndex = len(result)
-				seen[key] = resultIndex
-				result = append(result, []T{})
-			}
-
-			result[resultIndex] = append(result[resultIndex], _item)
+	keys := Map(collection, func(item T, _ int) K {
+		return iteratee(item)
+	})
 
-			mu.Unlock()
-			wg.Done()
-		}(item)
+	for i, item := range collection {
+		if resultIndex, ok := seen[keys[i]]; !ok {
+			resultIndex = len(result)
+			seen[keys[i]] = resultIndex
+			result = append(result, Slice{item})
+		} else {
+			result[resultIndex] = append(result[resultIndex], item)
+		}
 	}
-
-	wg.Wait()
-
 	return result
 }