add mapped dataloader fetch method

Author: CodingContraption

.gitignore

@@ -1 +1,29 @@
+# If you prefer the allow list template instead of the deny list, see community template:
+# https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore
+#
+# Binaries for programs and plugins
+*.exe
+*.exe~
+*.dll
+*.so
+*.dylib
+
+# Test binary, built with `go test -c`
+*.test
+
+# Output of the go coverage tool, specifically when used with LiteIDE
+*.out
+
+# Dependency directories (remove the comment below to include it)
+# vendor/
+
+# Go workspace file
+# go.work
+go.work.sum
+
+# env file
+.env
+
+# Editors config
 .vscode
+.idea

README.md

@@ -6,11 +6,14 @@
 
 It works as follows:
 * A Loader object is created per graphql request.
-* Each of many concurrently executing graphql resolver functions call Load() on the Loader object with different keys. Let's say K1, K2, K3
-* Each call to Load() with a new key is delayed slightly (a few milliseconds) so that the Loader can load them together.
-* The customizable fetch function of the loader takes a list of keys and loads data for all of them in a single batched request to the data storage layer. It might send `[K1,K2,K3]` and get back `[V1,V2,V3]`.
+* Each of many concurrently executing graphql resolver functions call `Load()` on the Loader object with different keys. Let's say `K1`, `K2`, `K3`
+* Each call to `Load()` with a new key is delayed slightly (a few milliseconds) so that the Loader can load them together.
+* The customizable `fetch` function of the loader takes a list of keys and loads data for all of them in a single batched request to the data storage layer. It might send `[K1,K2,K3]` and get back `[V1,V2,V3]`. Alternatively, a mappedFetch function of the loader takes a list of keys and returns a map instead of a list. It might send `[K1, K2, K3]` and get back `{K1: V1, K2: V2, K3: V3}`.
 * The Loader takes case of sending the right result to the right caller and the result is cached for the duration of the graphql request.
 
+> [!NOTE]
+> The `fetch` method expect the returned list to correspond to the provided keys in the same order. Alternatively  the `mappedFetch` can be used which allows returning a map, ensuring correct ordering and easy handling of nil values.
+
 Usage:
 
 ```sh
@@ -25,26 +28,37 @@ import (
 	"context"
 	"fmt"
 	"strconv"
-	"time"
 
 	"github.com/vikstrous/dataloadgen"
 )
 
-// fetchFn is shown as a function here, but it might work better as a method
+// fetchFn/mappedFetchFn is shown as a function here, but it might work better as a method
 // ctx is the context from the first call to Load for the current batch
 func fetchFn(ctx context.Context, keys []string) (ret []int, errs []error) {
     for _, key := range keys {
         num, err := strconv.ParseInt(key, 10, 32)
         ret = append(ret, int(num))
         errs = append(errs, err)
     }
-    return
+    return ret, errs
+}
+func mappedFetchFn(ctx context.Context, keys []string) (ret map[string]int, errs map[string]error) {
+    for _, key := range keys {
+        num, err := strconv.ParseInt(key, 10, 32)
+        ret[key] = int(num)
+        errs[key] = err
+    }
+    return ret, errs
 }
 
 func main() {
     ctx := context.Background()
-    // Per-request setup code:
+	
+    // Per-request setup code. Either:
     loader := dataloadgen.NewLoader(fetchFn)
+    // or
+    loader := dataloadgen.NewMappedLoader(mappedFetchFn)
+	
     // In every graphql resolver:
     result, err := loader.Load(ctx, "1")
     if err != nil {

benchmark/benchmarks_test.go

@@ -89,6 +89,25 @@ func newVikstrous() *dataloadgen.Loader[int, benchmark.User] {
 	)
 }
 
+func newVikstrousMapped() *dataloadgen.Loader[int, benchmark.User] {
+	return dataloadgen.NewMappedLoader(func(_ context.Context, keys []int) (map[int]benchmark.User, map[int]error) {
+		users := make(map[int]benchmark.User, len(keys))
+		errors := make(map[int]error, len(keys))
+
+		for _, key := range keys {
+			if key%100 == 1 {
+				errors[key] = fmt.Errorf("user not found")
+			} else {
+				users[key] = benchmark.User{ID: strconv.Itoa(key), Name: "user " + strconv.Itoa(key)}
+			}
+		}
+		return users, errors
+	},
+		dataloadgen.WithBatchCapacity(100),
+		dataloadgen.WithWait(500*time.Nanosecond),
+	)
+}
+
 func BenchmarkAll(b *testing.B) {
 	ctx := context.Background()
 
@@ -113,6 +132,11 @@ func BenchmarkAll(b *testing.B) {
 				newVikstrous()
 			}
 		})
+		b.Run("dataloadgen_mapped", func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				newVikstrousMapped()
+			}
+		})
 	})
 
 	b.Run("cached", func(b *testing.B) {
@@ -124,7 +148,7 @@ func BenchmarkAll(b *testing.B) {
 					thunks[i] = dataloaderDL.Load(ctx, 1)
 				}
 				for i := 0; i < 100; i++ {
-					thunks[i]()
+					_, _ = thunks[i]()
 				}
 			}
 		})
@@ -137,7 +161,7 @@ func BenchmarkAll(b *testing.B) {
 				}
 
 				for i := 0; i < 100; i++ {
-					thunks[i]()
+					_, _ = thunks[i]()
 				}
 			}
 		})
@@ -150,7 +174,7 @@ func BenchmarkAll(b *testing.B) {
 				}
 
 				for i := 0; i < 100; i++ {
-					thunks[i].Get(ctx)
+					_, _ = thunks[i].Get(ctx)
 				}
 			}
 		})
@@ -163,7 +187,20 @@ func BenchmarkAll(b *testing.B) {
 				}
 
 				for i := 0; i < 100; i++ {
-					thunks[i]()
+					_, _ = thunks[i]()
+				}
+			}
+		})
+		b.Run("dataloadgen_mapped", func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				vikstrousDL := newVikstrousMapped()
+				thunks := make([]func() (benchmark.User, error), 100)
+				for i := 0; i < 100; i++ {
+					thunks[i] = vikstrousDL.LoadThunk(ctx, 1)
+				}
+
+				for i := 0; i < 100; i++ {
+					_, _ = thunks[i]()
 				}
 			}
 		})
@@ -178,7 +215,7 @@ func BenchmarkAll(b *testing.B) {
 					thunks[i] = dataloaderDL.Load(ctx, i)
 				}
 				for i := 0; i < 100; i++ {
-					thunks[i]()
+					_, _ = thunks[i]()
 				}
 			}
 		})
@@ -191,7 +228,7 @@ func BenchmarkAll(b *testing.B) {
 				}
 
 				for i := 0; i < 100; i++ {
-					thunks[i]()
+					_, _ = thunks[i]()
 				}
 			}
 		})
@@ -204,7 +241,7 @@ func BenchmarkAll(b *testing.B) {
 				}
 
 				for i := 0; i < 100; i++ {
-					thunks[i].Get(ctx)
+					_, _ = thunks[i].Get(ctx)
 				}
 			}
 		})
@@ -217,7 +254,20 @@ func BenchmarkAll(b *testing.B) {
 				}
 
 				for i := 0; i < 100; i++ {
-					thunks[i]()
+					_, _ = thunks[i]()
+				}
+			}
+		})
+		b.Run("dataloadgen_mapped", func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				vikstrousDL := newVikstrousMapped()
+				thunks := make([]func() (benchmark.User, error), 100)
+				for i := 0; i < 100; i++ {
+					thunks[i] = vikstrousDL.LoadThunk(ctx, i)
+				}
+
+				for i := 0; i < 100; i++ {
+					_, _ = thunks[i]()
 				}
 			}
 		})
@@ -295,11 +345,29 @@ func BenchmarkAll(b *testing.B) {
 				wg.Wait()
 			}
 		})
+		b.Run("dataloadgen_mapped", func(b *testing.B) {
+			for n := 0; n < b.N*10; n++ {
+				vikstrousDL := newVikstrousMapped()
+				results := make([]benchmark.User, 10)
+				var wg sync.WaitGroup
+				for i := 0; i < 10; i++ {
+					wg.Add(1)
+					go func(i int) {
+						for j := 0; j < b.N; j++ {
+							u, _ := vikstrousDL.Load(ctx, i)
+							results[i] = u
+						}
+						wg.Done()
+					}(i)
+				}
+				wg.Wait()
+			}
+		})
 	})
 
 	b.Run("all in one request", func(b *testing.B) {
 		b.Run("dataloader", func(b *testing.B) {
-			keys := []int{}
+			var keys []int
 			for n := 0; n < 10000; n++ {
 				keys = append(keys, n)
 			}
@@ -310,7 +378,7 @@ func BenchmarkAll(b *testing.B) {
 			}
 		})
 		b.Run("dataloaden", func(b *testing.B) {
-			keys := []int{}
+			var keys []int
 			for n := 0; n < 10000; n++ {
 				keys = append(keys, n)
 			}
@@ -321,7 +389,7 @@ func BenchmarkAll(b *testing.B) {
 			}
 		})
 		b.Run("yckao_dataloader", func(b *testing.B) {
-			keys := []int{}
+			var keys []int
 			for n := 0; n < 10000; n++ {
 				keys = append(keys, n)
 			}
@@ -330,19 +398,30 @@ func BenchmarkAll(b *testing.B) {
 				yckaoDL := newYckao()
 				thunks := yckaoDL.LoadMany(ctx, keys)
 				for _, t := range thunks {
-					t.Get(ctx)
+					_, _ = t.Get(ctx)
 				}
 			}
 		})
 		b.Run("dataloadgen", func(b *testing.B) {
-			keys := []int{}
+			var keys []int
 			for n := 0; n < 10000; n++ {
 				keys = append(keys, n)
 			}
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				vikstrousDL := newVikstrous()
-				vikstrousDL.LoadAll(ctx, keys)
+				_, _ = vikstrousDL.LoadAll(ctx, keys)
+			}
+		})
+		b.Run("dataloadgen_mapped", func(b *testing.B) {
+			var keys []int
+			for n := 0; n < 10000; n++ {
+				keys = append(keys, n)
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				vikstrousDL := newVikstrousMapped()
+				_, _ = vikstrousDL.LoadAll(ctx, keys)
 			}
 		})
 	})

dataloadgen.go

@@ -7,7 +7,7 @@ import (
 	"sync"
 	"time"
 
-	trace "go.opentelemetry.io/otel/trace"
+	"go.opentelemetry.io/otel/trace"
 )
 
 // Option allows for configuration of loader fields.
@@ -51,6 +51,25 @@ func NewLoader[KeyT comparable, ValueT any](fetch func(ctx context.Context, keys
 	return l
 }
 
+// NewMappedLoader creates a new GenericLoader given a mappedFetch, wait and maxBatch
+func NewMappedLoader[KeyT comparable, ValueT any](mappedFetch func(ctx context.Context, keys []KeyT) (map[KeyT]ValueT, map[KeyT]error), options ...Option) *Loader[KeyT, ValueT] {
+	return NewLoader(convertMappedFetch(mappedFetch), options...)
+}
+
+// convertMappedFetch accepts a fetcher method that returns maps, and converts it to a fetcher that returns lists.
+func convertMappedFetch[KeyT comparable, ValueT any](mappedFetch func(ctx context.Context, keys []KeyT) (map[KeyT]ValueT, map[KeyT]error)) func(ctx context.Context, keys []KeyT) ([]ValueT, []error) {
+	return func(ctx context.Context, keys []KeyT) ([]ValueT, []error) {
+		mappedResults, mappedErrs := mappedFetch(ctx, keys)
+		var values = make([]ValueT, len(keys))
+		var errs = make([]error, len(keys))
+		for i, key := range keys {
+			values[i] = mappedResults[key]
+			errs[i] = mappedErrs[key]
+		}
+		return values, errs
+	}
+}
+
 type loaderConfig struct {
 	// how long to done before sending a batch
 	wait time.Duration
@@ -177,18 +196,18 @@ func (l *Loader[KeyT, ValueT]) LoadAll(ctx context.Context, keys []KeyT) ([]Valu
 	}
 
 	values := make([]ValueT, len(keys))
-	errors := make([]error, len(keys))
+	errs := make([]error, len(keys))
 	allNil := true
 	for i, thunk := range thunks {
-		values[i], errors[i] = thunk()
-		if errors[i] != nil {
+		values[i], errs[i] = thunk()
+		if errs[i] != nil {
 			allNil = false
 		}
 	}
 	if allNil {
 		return values, nil
 	}
-	return values, ErrorSlice(errors)
+	return values, ErrorSlice(errs)
 }
 
 // LoadAllThunk returns a function that when called will block waiting for a ValueT.
@@ -201,18 +220,18 @@ func (l *Loader[KeyT, ValueT]) LoadAllThunk(ctx context.Context, keys []KeyT) fu
 	}
 	return func() ([]ValueT, error) {
 		values := make([]ValueT, len(keys))
-		errors := make([]error, len(keys))
+		errs := make([]error, len(keys))
 		allNil := true
 		for i, thunk := range thunks {
-			values[i], errors[i] = thunk()
-			if allNil && errors[i] != nil {
+			values[i], errs[i] = thunk()
+			if allNil && errs[i] != nil {
 				allNil = false
 			}
 		}
 		if allNil {
 			return values, nil
 		}
-		return values, ErrorSlice(errors)
+		return values, ErrorSlice(errs)
 	}
 }
 
@@ -296,7 +315,7 @@ func (l *Loader[KeyT, ValueT]) safeFetch(ctx context.Context, keys []KeyT) (valu
 	return l.fetch(ctx, keys)
 }
 
-// addKeyToBatch will return the location of the key in the batch, if its not found
+// addKeyToBatch will return the location of the key in the batch, if it's not found
 // it will add the key to the batch
 func (l *Loader[KeyT, ValueT]) addKeyToBatch(b *loaderBatch[KeyT, ValueT], key KeyT) int {
 	pos := len(b.keys)