Skip to content

Commit

Permalink
fix(parallel): order of GroupBy and PartitionBy (#572)
Browse files Browse the repository at this point in the history
liyishuai authored and samber committed Jan 24, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent 5dff357 commit 1010778
Showing 1 changed file with 18 additions and 39 deletions.
57 changes: 18 additions & 39 deletions parallel/slice.go
Original file line number Diff line number Diff line change
@@ -66,65 +66,44 @@ 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 the 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 {
result[resultIndex] = append(result[resultIndex], item)
} else {
resultIndex = len(result)
seen[keys[i]] = resultIndex
result = append(result, Slice{item})
}
}

wg.Wait()

return result
}

0 comments on commit 1010778

Please sign in to comment.