util: create MapE utility function

Epic: None

Release note: None

Author: kev-cao

pkg/util/BUILD.bazel

@@ -56,6 +56,7 @@ go_test(
     deps = [
         "//pkg/util/randutil",
         "//pkg/util/timeutil",
+        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/util/slices.go

@@ -80,6 +80,21 @@ func Map[T, K any](collection []T, fn func(T) K) []K {
 	return out
 }
 
+// MapE returns a new slice containing the results of fn for each element
+// within a collection. If fn returns an error for any element, the function
+// exits early and returns the error.
+func MapE[T, K any](collection []T, fn func(T) (K, error)) ([]K, error) {
+	out := make([]K, len(collection))
+	for i, el := range collection {
+		value, err := fn(el)
+		if err != nil {
+			return nil, err
+		}
+		out[i] = value
+	}
+	return out, nil
+}
+
 // MapFrom returns a map populated with keys and values returned by fn.
 // Usage:
 //

pkg/util/slices_test.go

@@ -8,6 +8,7 @@ package util
 import (
 	"testing"
 
+	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 )
 
@@ -123,6 +124,46 @@ func TestMap(t *testing.T) {
 	}))
 }
 
+func TestMapErr(t *testing.T) {
+	t.Run("empty slice", func(t *testing.T) {
+		out, err := MapE(nil, func(i int) (int, error) {
+			require.Fail(t, "should not be called")
+			return 0, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, []int{}, out)
+	})
+
+	t.Run("map to same type", func(t *testing.T) {
+		out, err := MapE([]int{1, 2, 3}, func(i int) (int, error) {
+			return i * 2, nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, []int{2, 4, 6}, out)
+	})
+
+	t.Run("map to different type", func(t *testing.T) {
+		out, err := MapE([]int{1, 2, 3}, func(i int) (string, error) {
+			return string(rune('a' + i - 1)), nil
+		})
+		require.NoError(t, err)
+		require.Equal(t, []string{"a", "b", "c"}, out)
+	})
+
+	t.Run("error case", func(t *testing.T) {
+		out, err := MapE([]int{1, 2, 3}, func(i int) (int, error) {
+			if i == 2 {
+				return 0, errors.New("error on 2")
+			} else if i == 3 {
+				require.Fail(t, "should not be called")
+			}
+			return i * 2, nil
+		})
+		require.Error(t, err)
+		require.Nil(t, out)
+	})
+}
+
 func TestMapFrom(t *testing.T) {
 	require.Equal(t, map[int]bool{}, MapFrom(nil, func(i int) (int, bool) {
 		require.Fail(t, "should not be called")