Skip to content

Commit

Permalink
feat: Implement ChunkMap (samber#533)
Browse files Browse the repository at this point in the history
  • Loading branch information
oswaldom-code committed Sep 23, 2024
1 parent bc0037c commit 2f5f43e
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 0 deletions.
40 changes: 40 additions & 0 deletions slice.go
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,46 @@ func Chunk[T any, Slice ~[]T](collection Slice, size int) []Slice {
return result
}

// ChunkMap splits a map into an array of elements in groups of a length equal to its size. If the map cannot be split evenly,
// the final chunk will contain the remaining elements.
func ChunkMap[K comparable, V any](m map[K]V, size int) []map[K]V {
if size <= 0 {
panic("The chunk size must be greater than 0")
}

keys := make([]K, 0, len(m))
for key := range m {
keys = append(keys, key)
}

if len(keys) == 0 {
return []map[K]V{}
}

chunksNum := len(keys) / size
if len(keys)%size != 0 {
chunksNum += 1
}

result := make([]map[K]V, 0, chunksNum)

for i := 0; i < chunksNum; i++ {
start := i * size
end := (i + 1) * size
if end > len(keys) {
end = len(keys)
}

chunk := make(map[K]V)
for _, key := range keys[start:end] {
chunk[key] = m[key]
}
result = append(result, chunk)
}

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.
Expand Down
46 changes: 46 additions & 0 deletions slice_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,52 @@ func TestChunk(t *testing.T) {
is.Equal(originalArray, []int{0, 1, 2, 3, 4, 5})
}

func TestChunkMap(t *testing.T) {
t.Parallel()
is := assert.New(t)

result1 := ChunkMap(map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}, 2)
result2 := ChunkMap(map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}, 3)
result3 := ChunkMap(map[string]int{}, 2)
result4 := ChunkMap(map[string]int{"a": 1}, 2)
result5 := ChunkMap(map[string]int{"a": 1, "b": 2}, 1)

expectedCount1 := 3
expectedCount2 := 2
expectedCount3 := 0
expectedCount4 := 1
expectedCount5 := 2

is.Len(result1, expectedCount1)
is.Len(result2, expectedCount2)
is.Len(result3, expectedCount3)
is.Len(result4, expectedCount4)
is.Len(result5, expectedCount5)

is.PanicsWithValue("The chunk size must be greater than 0", func() {
ChunkMap(map[string]int{"a": 1}, 0)
})
is.PanicsWithValue("The chunk size must be greater than 0", func() {
ChunkMap(map[string]int{"a": 1}, -1)
})

type myStruct struct {
Name string
Value int
}

allStructs := []myStruct{{"one", 1}, {"two", 2}, {"three", 3}}
nonempty := ChunkMap(map[string]myStruct{"a": allStructs[0], "b": allStructs[1], "c": allStructs[2]}, 2)
is.Len(nonempty, 2)

originalMap := map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}
result6 := ChunkMap(originalMap, 2)
for k := range result6[0] {
result6[0][k] = 10
}
is.Equal(originalMap, map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5})
}

func TestPartitionBy(t *testing.T) {
t.Parallel()
is := assert.New(t)
Expand Down

0 comments on commit 2f5f43e

Please sign in to comment.