Added map batch loaders so that you can be more natural in getting less results than asked for

Author: bbakerman

README.md

@@ -194,6 +194,42 @@ credentials or the database connection parameters.  You can do this by implement
 The batch loading code will now receive this context object and it can be used to get to data layers or
 to connect to other systems. 
 
+### Returning a Map of results from your batch loader
+
+Often there is not a 1:1 mapping of your batch loaded keys to the values returned.
+
+For example, let's assume you want to load users from a database, you could probably use a query that looks like this:
+
+```sql
+  SELECT * FROM User WHERE id IN (keys)
+```
+ 
+ Given say 10 user id keys you might only get 7 results back.  This can be more naturally represented in a map
+ than in an order list of values when returning values from the batch loader function.
+ 
+ You can use `org.dataloader.MapBatchLoader` for this purpose. 
+ 
+ When the map is processed by the `DataLoader` code, any keys that are missing in the map
+ will be replaced with null values.  The semantics that the number of `DataLoader.load` requests
+ are matched with values is kept.
+ 
+ Your keys provided MUST be first class keys since they will be used to examine the returned map and
+ create the list of results, with nulls filling in for missing values.
+ 
+```java
+        MapBatchLoader<Long, User> mapBatchLoader = new MapBatchLoader<Long, User>() {
+            @Override
+            public CompletionStage<Map<Long, User>> load(List<Long> userIds, Object context) {
+                SecurityCtx callCtx = (SecurityCtx) context;
+                return CompletableFuture.supplyAsync(() -> userManager.loadMapOfUsersById(callCtx, userIds));
+            }
+        };
+
+        DataLoader<Long, User> userLoader = DataLoader.newDataLoader(mapBatchLoader);
+
+        // ...
+```
+
 ### Error object is not a thing in a type safe Java world
 
 In the reference JS implementation if the batch loader returns an `Error` object back from the `load()` promise is rejected

src/main/java/org/dataloader/DataLoader.java

@@ -24,6 +24,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.stream.Collectors;
@@ -63,6 +64,7 @@
 public class DataLoader<K, V> {
 
     private final BatchLoader<K, V> batchLoadFunction;
+    private final MapBatchLoader<K, V> mapBatchLoadFunction;
     private final DataLoaderOptions loaderOptions;
     private final CacheMap<Object, CompletableFuture<V>> futureCache;
     private final List<SimpleImmutableEntry<K, CompletableFuture<V>>> loaderQueue;
@@ -96,6 +98,35 @@ public static <K, V> DataLoader<K, V> newDataLoader(BatchLoader<K, V> batchLoadF
         return new DataLoader<>(batchLoadFunction, options);
     }
 
+    /**
+     * Creates new DataLoader with the specified map batch loader function and default options
+     * (batching, caching and unlimited batch size).
+     *
+     * @param mapBatchLoaderFunction the batch load function to use
+     * @param <K>                    the key type
+     * @param <V>                    the value type
+     *
+     * @return a new DataLoader
+     */
+    public static <K, V> DataLoader<K, V> newDataLoader(MapBatchLoader<K, V> mapBatchLoaderFunction) {
+        return newDataLoader(mapBatchLoaderFunction, null);
+    }
+
+    /**
+     * Creates new DataLoader with the specified map batch loader function with the provided options
+     *
+     * @param mapBatchLoaderFunction the batch load function to use
+     * @param options                the options to use
+     * @param <K>                    the key type
+     * @param <V>                    the value type
+     *
+     * @return a new DataLoader
+     */
+    public static <K, V> DataLoader<K, V> newDataLoader(MapBatchLoader<K, V> mapBatchLoaderFunction, DataLoaderOptions options) {
+        return new DataLoader<>(mapBatchLoaderFunction, options);
+    }
+
+
     /**
      * Creates new DataLoader with the specified batch loader function and default options
      * (batching, caching and unlimited batch size) where the batch loader function returns a list of
@@ -134,6 +165,43 @@ public static <K, V> DataLoader<K, V> newDataLoaderWithTry(BatchLoader<K, Try<V>
         return new DataLoader<>((BatchLoader<K, V>) batchLoadFunction, options);
     }
 
+    /**
+     * Creates new DataLoader with the specified map abatch loader function and default options
+     * (batching, caching and unlimited batch size) where the batch loader function returns a list of
+     * {@link org.dataloader.Try} objects.
+     *
+     * This allows you to capture both the value that might be returned and also whether exception that might have occurred getting that individual value.  If its important you to
+     * know gather exact status of each item in a batch call and whether it threw exceptions when fetched then
+     * you can use this form to create the data loader.
+     *
+     * @param mapBatchLoaderFunction the map batch load function to use that uses {@link org.dataloader.Try} objects
+     * @param <K>                    the key type
+     * @param <V>                    the value type
+     *
+     * @return a new DataLoader
+     */
+    public static <K, V> DataLoader<K, V> newDataLoaderWithTry(MapBatchLoader<K, Try<V>> mapBatchLoaderFunction) {
+        return newDataLoaderWithTry(mapBatchLoaderFunction, null);
+    }
+
+    /**
+     * Creates new DataLoader with the specified map batch loader function and with the provided options
+     * where the batch loader function returns a list of
+     * {@link org.dataloader.Try} objects.
+     *
+     * @param mapBatchLoaderFunction the map batch load function to use that uses {@link org.dataloader.Try} objects
+     * @param options                the options to use
+     * @param <K>                    the key type
+     * @param <V>                    the value type
+     *
+     * @return a new DataLoader
+     *
+     * @see #newDataLoaderWithTry(MapBatchLoader)
+     */
+    @SuppressWarnings("unchecked")
+    public static <K, V> DataLoader<K, V> newDataLoaderWithTry(MapBatchLoader<K, Try<V>> mapBatchLoaderFunction, DataLoaderOptions options) {
+        return new DataLoader<>((MapBatchLoader<K, V>) mapBatchLoaderFunction, options);
+    }
 
     /**
      * Creates a new data loader with the provided batch load function, and default options.
@@ -144,19 +212,53 @@ public DataLoader(BatchLoader<K, V> batchLoadFunction) {
         this(batchLoadFunction, null);
     }
 
+    /**
+     * Creates a new data loader with the provided map batch load function.
+     *
+     * @param mapBatchLoadFunction the map batch load function to use
+     */
+    public DataLoader(MapBatchLoader<K, V> mapBatchLoadFunction) {
+        this(mapBatchLoadFunction, null);
+    }
+
     /**
      * Creates a new data loader with the provided batch load function and options.
      *
      * @param batchLoadFunction the batch load function to use
      * @param options           the batch load options
      */
     public DataLoader(BatchLoader<K, V> batchLoadFunction, DataLoaderOptions options) {
-        this.batchLoadFunction = nonNull(batchLoadFunction);
-        this.loaderOptions = options == null ? new DataLoaderOptions() : options;
+        this.batchLoadFunction = batchLoadFunction;
+        this.mapBatchLoadFunction = null;
+        this.loaderOptions = determineOptions(options);
         this.futureCache = determineCacheMap(loaderOptions);
         // order of keys matter in data loader
         this.loaderQueue = new ArrayList<>();
-        this.stats = nonNull(this.loaderOptions.getStatisticsCollector());
+        this.stats = determineCollector(this.loaderOptions);
+    }
+
+    /**
+     * Creates a new data loader with the provided map batch load function and options.
+     *
+     * @param mapBatchLoadFunction the map batch load function to use
+     * @param options              the batch load options
+     */
+    public DataLoader(MapBatchLoader<K, V> mapBatchLoadFunction, DataLoaderOptions options) {
+        this.batchLoadFunction = null;
+        this.mapBatchLoadFunction = mapBatchLoadFunction;
+        this.loaderOptions = determineOptions(options);
+        this.futureCache = determineCacheMap(loaderOptions);
+        // order of keys matter in data loader
+        this.loaderQueue = new ArrayList<>();
+        this.stats = determineCollector(this.loaderOptions);
+    }
+
+    private StatisticsCollector determineCollector(DataLoaderOptions loaderOptions) {
+        return nonNull(loaderOptions.getStatisticsCollector());
+    }
+
+    private DataLoaderOptions determineOptions(DataLoaderOptions options) {
+        return options == null ? new DataLoaderOptions() : options;
     }
 
     @SuppressWarnings("unchecked")
@@ -197,11 +299,7 @@ public CompletableFuture<V> load(K key) {
                 stats.incrementBatchLoadCountBy(1);
                 // immediate execution of batch function
                 Object context = loaderOptions.getBatchContextProvider().get();
-                CompletableFuture<List<V>> batchedLoad = batchLoadFunction
-                        .load(singletonList(key), context)
-                        .toCompletableFuture();
-                future = batchedLoad
-                        .thenApply(list -> list.get(0));
+                future = invokeLoaderImmediately(key, context);
             }
             if (cachingEnabled) {
                 futureCache.set(cacheKey, future);
@@ -210,6 +308,7 @@ public CompletableFuture<V> load(K key) {
         }
     }
 
+
     /**
      * Requests to load the list of data provided by the specified keys asynchronously, and returns a composite future
      * of the resulting values.
@@ -232,6 +331,25 @@ public CompletableFuture<List<V>> loadMany(List<K> keys) {
         }
     }
 
+    private CompletableFuture<V> invokeLoaderImmediately(K key, Object context) {
+        List<K> keys = singletonList(key);
+        CompletionStage<V> singleLoadCall;
+        if (isMapLoader()) {
+            singleLoadCall = mapBatchLoadFunction
+                    .load(keys, context)
+                    .thenApply(map -> map.get(key));
+        } else {
+            singleLoadCall = batchLoadFunction
+                    .load(keys, context)
+                    .thenApply(list -> list.get(0));
+        }
+        return singleLoadCall.toCompletableFuture();
+    }
+
+    private boolean isMapLoader() {
+        return mapBatchLoadFunction != null;
+    }
+
     /**
      * Dispatches the queued load requests to the batch execution function and returns a promise of the result.
      * <p>
@@ -302,17 +420,11 @@ private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<
     @SuppressWarnings("unchecked")
     private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<CompletableFuture<V>> queuedFutures) {
         stats.incrementBatchLoadCountBy(keys.size());
-        CompletionStage<List<V>> batchLoad;
-        try {
-            Object context = loaderOptions.getBatchContextProvider().get();
-            batchLoad = nonNull(batchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise");
-        } catch (Exception e) {
-            batchLoad = CompletableFutureKit.failedFuture(e);
-        }
+        CompletionStage<List<V>> batchLoad = invokeBatchFunction(keys);
         return batchLoad
                 .toCompletableFuture()
                 .thenApply(values -> {
-                    assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list");
+                    assertResultSize(keys, values);
 
                     for (int idx = 0; idx < queuedFutures.size(); idx++) {
                         Object value = values.get(idx);
@@ -351,6 +463,45 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Complet
                 });
     }
 
+    private CompletionStage<List<V>> invokeBatchFunction(List<K> keys) {
+        CompletionStage<List<V>> batchLoad;
+        try {
+            Object context = loaderOptions.getBatchContextProvider().get();
+            if (isMapLoader()) {
+                batchLoad = invokeMapBatchLoader(keys, context);
+            } else {
+                batchLoad = invokeListBatchLoader(keys, context);
+            }
+        } catch (Exception e) {
+            batchLoad = CompletableFutureKit.failedFuture(e);
+        }
+        return batchLoad;
+    }
+
+    private CompletionStage<List<V>> invokeListBatchLoader(List<K> keys, Object context) {
+        return nonNull(batchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise");
+    }
+
+    /*
+     * Turns a map of results that MAY be smaller than the key list back into a list by mapping null
+     * to missing elements.
+     */
+    private CompletionStage<List<V>> invokeMapBatchLoader(List<K> keys, Object context) {
+        CompletionStage<Map<K, V>> mapBatchLoad = nonNull(mapBatchLoadFunction.load(keys, context), "Your batch loader function MUST return a non null CompletionStage promise");
+        return mapBatchLoad.thenApply(map -> {
+            List<V> values = new ArrayList<>();
+            for (K key : keys) {
+                V value = map.get(key);
+                values.add(value);
+            }
+            return values;
+        });
+    }
+
+    private void assertResultSize(List<K> keys, List<V> values) {
+        assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list");
+    }
+
     /**
      * Normally {@link #dispatch()} is an asynchronous operation but this version will 'join' on the
      * results if dispatch and wait for them to complete.  If the {@link CompletableFuture} callbacks make more

src/main/java/org/dataloader/MapBatchLoader.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 2016 The original author or authors
+ *
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Eclipse Public License v1.0
+ * and Apache License v2.0 which accompanies this distribution.
+ *
+ *      The Eclipse Public License is available at
+ *      http://www.eclipse.org/legal/epl-v10.html
+ *
+ *      The Apache License v2.0 is available at
+ *      http://www.opensource.org/licenses/apache2.0.php
+ *
+ * You may elect to redistribute this code under either of these licenses.
+ */
+
+package org.dataloader;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletionStage;
+
+/**
+ * A function that is invoked for batch loading a map of of data values indicated by the provided list of keys. The
+ * function returns a promise of a map of results of individual load requests.
+ *
+ * There are a few constraints that must be upheld:
+ * <ul>
+ * <li>The keys MUST be able to be first class keys in a Java map.  Get your equals() and hashCode() methods in order</li>
+ * <li>The caller of the {@link org.dataloader.DataLoader} that uses this batch loader function MUSt be able to cope with
+ * null values coming back as results
+ * </li>
+ * <li>The function MUST be resilient to the same key being presented twice.</li>
+ * </ul>
+ *
+ * This form is useful when you don't have a 1:1 mapping of keys to values or when null is an acceptable value for a missing value.
+ *
+ * For example, let's assume you want to load users from a database, you could probably use a query that looks like this:
+ *
+ * <pre>
+ *    SELECT * FROM User WHERE id IN (keys)
+ * </pre>
+ *
+ * Given say 10 user id keys you might only get 7 results back.  This can be more naturally represented in a map
+ * than in an order list of values when returning values from the batch loader function.
+ *
+ * When the map is processed by the {@link org.dataloader.DataLoader} code, any keys that are missing in the map
+ * will be replaced with null values.  The semantics that the number of {@link org.dataloader.DataLoader#load(Object)} requests
+ * are matched with values is kept.
+ *
+ * This means that if 10 keys are asked for then {@link DataLoader#dispatch()} will return a promise 10 value results and each
+ * of the {@link org.dataloader.DataLoader#load(Object)} will complete with a value, null or an exception.
+ *
+ * When caching is disabled, its possible for the same key to be presented in the list of keys more than once.  You map
+ * batch loader function needs to be resilient to this situation.
+ *
+ * @param <K> type parameter indicating the type of keys to use for data load requests.
+ * @param <V> type parameter indicating the type of values returned
+ *
+ * @author <a href="https://github.com/aschrijver/">Arnold Schrijver</a>
+ * @author <a href="https://github.com/bbakerman/">Brad Baker</a>
+ */
+@FunctionalInterface
+public interface MapBatchLoader<K, V> {
+
+    /**
+     * Called to batch load the provided keys and return a promise to a map of values.  It can be given a context object to
+     * that maybe be useful during the call.  A typical use case is passing in security credentials or database connection details say.
+     *
+     * @param keys    the collection of keys to load
+     * @param context a context object that can help with the call
+     *
+     * @return a promise to a map of values for those keys
+     */
+    @SuppressWarnings("unused")
+    CompletionStage<Map<K, V>> load(List<K> keys, Object context);
+}