Skip to content

Commit 359dc0f

Browse files
oswaldom-codesamber
authored andcommitted
feat: Implement ChunkMap (#533)
1 parent 6b110b3 commit 359dc0f

File tree

2 files changed

+86
-0
lines changed

2 files changed

+86
-0
lines changed

slice.go

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,46 @@ func Chunk[T any, Slice ~[]T](collection Slice, size int) []Slice {
214214
return result
215215
}
216216

217+
// 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,
218+
// the final chunk will contain the remaining elements.
219+
func ChunkMap[K comparable, V any](m map[K]V, size int) []map[K]V {
220+
if size <= 0 {
221+
panic("The chunk size must be greater than 0")
222+
}
223+
224+
keys := make([]K, 0, len(m))
225+
for key := range m {
226+
keys = append(keys, key)
227+
}
228+
229+
if len(keys) == 0 {
230+
return []map[K]V{}
231+
}
232+
233+
chunksNum := len(keys) / size
234+
if len(keys)%size != 0 {
235+
chunksNum += 1
236+
}
237+
238+
result := make([]map[K]V, 0, chunksNum)
239+
240+
for i := 0; i < chunksNum; i++ {
241+
start := i * size
242+
end := (i + 1) * size
243+
if end > len(keys) {
244+
end = len(keys)
245+
}
246+
247+
chunk := make(map[K]V)
248+
for _, key := range keys[start:end] {
249+
chunk[key] = m[key]
250+
}
251+
result = append(result, chunk)
252+
}
253+
254+
return result
255+
}
256+
217257
// PartitionBy returns an array of elements split into groups. The order of grouped values is
218258
// determined by the order they occur in collection. The grouping is generated from the results
219259
// of running each element of collection through iteratee.

slice_test.go

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,52 @@ func TestChunk(t *testing.T) {
283283
is.Equal(originalArray, []int{0, 1, 2, 3, 4, 5})
284284
}
285285

286+
func TestChunkMap(t *testing.T) {
287+
t.Parallel()
288+
is := assert.New(t)
289+
290+
result1 := ChunkMap(map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}, 2)
291+
result2 := ChunkMap(map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}, 3)
292+
result3 := ChunkMap(map[string]int{}, 2)
293+
result4 := ChunkMap(map[string]int{"a": 1}, 2)
294+
result5 := ChunkMap(map[string]int{"a": 1, "b": 2}, 1)
295+
296+
expectedCount1 := 3
297+
expectedCount2 := 2
298+
expectedCount3 := 0
299+
expectedCount4 := 1
300+
expectedCount5 := 2
301+
302+
is.Len(result1, expectedCount1)
303+
is.Len(result2, expectedCount2)
304+
is.Len(result3, expectedCount3)
305+
is.Len(result4, expectedCount4)
306+
is.Len(result5, expectedCount5)
307+
308+
is.PanicsWithValue("The chunk size must be greater than 0", func() {
309+
ChunkMap(map[string]int{"a": 1}, 0)
310+
})
311+
is.PanicsWithValue("The chunk size must be greater than 0", func() {
312+
ChunkMap(map[string]int{"a": 1}, -1)
313+
})
314+
315+
type myStruct struct {
316+
Name string
317+
Value int
318+
}
319+
320+
allStructs := []myStruct{{"one", 1}, {"two", 2}, {"three", 3}}
321+
nonempty := ChunkMap(map[string]myStruct{"a": allStructs[0], "b": allStructs[1], "c": allStructs[2]}, 2)
322+
is.Len(nonempty, 2)
323+
324+
originalMap := map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}
325+
result6 := ChunkMap(originalMap, 2)
326+
for k := range result6[0] {
327+
result6[0][k] = 10
328+
}
329+
is.Equal(originalMap, map[string]int{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5})
330+
}
331+
286332
func TestPartitionBy(t *testing.T) {
287333
t.Parallel()
288334
is := assert.New(t)

0 commit comments

Comments
 (0)