From 2f5f43efcf2c2125053275c7b21ba94cee561675 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Oswaldo=20Monta=C3=B1o?= <oswaldoeliezer876@gmail.com>
Date: Mon, 23 Sep 2024 19:02:26 +0200
Subject: [PATCH] feat: Implement ChunkMap (#533)

---
 slice.go      | 40 ++++++++++++++++++++++++++++++++++++++++
 slice_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/slice.go b/slice.go
index d2d3fd84..f1e9c6f5 100644
--- a/slice.go
+++ b/slice.go
@@ -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.
diff --git a/slice_test.go b/slice_test.go
index 9f923eea..ba5dc623 100644
--- a/slice_test.go
+++ b/slice_test.go
@@ -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)