Skip to content

Commit 229109d

Browse files
bbakermanbbakerman
authored
Merge pull request #166 from Hc747/master
Add support for batch loading of a Map of key-value pairs.
2 parents 466e482 + d3c13c5 commit 229109d

File tree

5 files changed

+168
-42
lines changed

5 files changed

+168
-42
lines changed

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

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,7 @@
2525
import java.time.Clock;
2626
import java.time.Duration;
2727
import java.time.Instant;
28-
import java.util.ArrayList;
29-
import java.util.Collections;
30-
import java.util.List;
31-
import java.util.Optional;
28+
import java.util.*;
3229
import java.util.concurrent.CompletableFuture;
3330
import java.util.function.BiConsumer;
3431

@@ -574,6 +571,35 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
574571
}
575572
}
576573

574+
/**
575+
* Requests to load the map of data provided by the specified keys asynchronously, and returns a composite future
576+
* of the resulting values.
577+
* <p>
578+
* If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
579+
* start batch execution. If you forget this call the future will never be completed (unless already completed,
580+
* and returned from cache).
581+
* <p>
582+
* The key context object may be useful in the batch loader interfaces such as {@link org.dataloader.BatchLoaderWithContext} or
583+
* {@link org.dataloader.MappedBatchLoaderWithContext} to help retrieve data.
584+
*
585+
* @param keysAndContexts the map of keys to their respective contexts
586+
*
587+
* @return the composite future of the map of keys and values
588+
*/
589+
public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
590+
nonNull(keysAndContexts);
591+
592+
synchronized (this) {
593+
Map<K, CompletableFuture<V>> collect = new HashMap<>(keysAndContexts.size());
594+
for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
595+
K key = entry.getKey();
596+
Object keyContext = entry.getValue();
597+
collect.put(key, load(key, keyContext));
598+
}
599+
return CompletableFutureKit.allOf(collect);
600+
}
601+
}
602+
577603
/**
578604
* Dispatches the queued load requests to the batch execution function and returns a promise of the result.
579605
* <p>

‎src/main/java/org/dataloader/impl/CompletableFutureKit.java‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33
import org.dataloader.annotations.Internal;
44

55
import java.util.List;
6+
import java.util.Map;
67
import java.util.concurrent.CompletableFuture;
78
import java.util.concurrent.ExecutionException;
9+
import java.util.stream.Collectors;
810

911
import static java.util.stream.Collectors.toList;
1012

@@ -48,10 +50,21 @@ public static <V> boolean failed(CompletableFuture<V> future) {
4850
}
4951

5052
public static <T> CompletableFuture<List<T>> allOf(List<CompletableFuture<T>> cfs) {
51-
return CompletableFuture.allOf(cfs.toArray(new CompletableFuture[0]))
53+
return CompletableFuture.allOf(cfs.toArray(CompletableFuture[]::new))
5254
.thenApply(v -> cfs.stream()
5355
.map(CompletableFuture::join)
5456
.collect(toList())
5557
);
5658
}
59+
60+
public static <K, V> CompletableFuture<Map<K, V>> allOf(Map<K, CompletableFuture<V>> cfs) {
61+
return CompletableFuture.allOf(cfs.values().toArray(CompletableFuture[]::new))
62+
.thenApply(v -> cfs.entrySet().stream()
63+
.collect(
64+
Collectors.toMap(
65+
Map.Entry::getKey,
66+
task -> task.getValue().join())
67+
)
68+
);
69+
}
5770
}

‎src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java‎

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,7 @@
22

33
import org.junit.jupiter.api.Test;
44

5-
import java.util.ArrayList;
6-
import java.util.HashMap;
7-
import java.util.List;
8-
import java.util.Map;
5+
import java.util.*;
96
import java.util.concurrent.CompletableFuture;
107
import java.util.concurrent.atomic.AtomicInteger;
118
import java.util.stream.Collectors;
@@ -50,10 +47,14 @@ public void context_is_passed_to_batch_loader_function() {
5047
loader.load("A");
5148
loader.load("B");
5249
loader.loadMany(asList("C", "D"));
50+
Map<String, ?> keysAndContexts = new LinkedHashMap<>();
51+
keysAndContexts.put("E", null);
52+
keysAndContexts.put("F", null);
53+
loader.loadMany(keysAndContexts);
5354

5455
List<String> results = loader.dispatchAndJoin();
5556

56-
assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx")));
57+
assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx", "E-ctx", "F-ctx")));
5758
}
5859

5960
@Test
@@ -66,10 +67,14 @@ public void key_contexts_are_passed_to_batch_loader_function() {
6667
loader.load("A", "aCtx");
6768
loader.load("B", "bCtx");
6869
loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx"));
70+
Map<String, String> keysAndContexts = new LinkedHashMap<>();
71+
keysAndContexts.put("E", "eCtx");
72+
keysAndContexts.put("F", "fCtx");
73+
loader.loadMany(keysAndContexts);
6974

7075
List<String> results = loader.dispatchAndJoin();
7176

72-
assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx")));
77+
assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:fCtx-l:fCtx")));
7378
}
7479

7580
@Test
@@ -82,12 +87,17 @@ public void key_contexts_are_passed_to_batch_loader_function_when_batching_disab
8287

8388
CompletableFuture<String> aLoad = loader.load("A", "aCtx");
8489
CompletableFuture<String> bLoad = loader.load("B", "bCtx");
85-
CompletableFuture<List<String>> canDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx"));
90+
CompletableFuture<List<String>> cAndDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx"));
91+
Map<String, String> keysAndContexts = new LinkedHashMap<>();
92+
keysAndContexts.put("E", "eCtx");
93+
keysAndContexts.put("F", "fCtx");
94+
CompletableFuture<Map<String, String>> eAndFLoad = loader.loadMany(keysAndContexts);
8695

8796
List<String> results = new ArrayList<>(asList(aLoad.join(), bLoad.join()));
88-
results.addAll(canDLoad.join());
97+
results.addAll(cAndDLoad.join());
98+
results.addAll(eAndFLoad.join().values());
8999

90-
assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx")));
100+
assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:fCtx-l:fCtx")));
91101
}
92102

93103
@Test
@@ -101,9 +111,14 @@ public void missing_key_contexts_are_passed_to_batch_loader_function() {
101111
loader.load("B");
102112
loader.loadMany(asList("C", "D"), singletonList("cCtx"));
103113

114+
Map<String, String> keysAndContexts = new LinkedHashMap<>();
115+
keysAndContexts.put("E", "eCtx");
116+
keysAndContexts.put("F", null);
117+
loader.loadMany(keysAndContexts);
118+
104119
List<String> results = loader.dispatchAndJoin();
105120

106-
assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null")));
121+
assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:null-l:null")));
107122
}
108123

109124
@Test
@@ -125,9 +140,14 @@ public void context_is_passed_to_map_batch_loader_function() {
125140
loader.load("B");
126141
loader.loadMany(asList("C", "D"), singletonList("cCtx"));
127142

143+
Map<String, String> keysAndContexts = new LinkedHashMap<>();
144+
keysAndContexts.put("E", "eCtx");
145+
keysAndContexts.put("F", null);
146+
loader.loadMany(keysAndContexts);
147+
128148
List<String> results = loader.dispatchAndJoin();
129149

130-
assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null")));
150+
assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null", "E-ctx-eCtx", "F-ctx-null")));
131151
}
132152

133153
@Test
@@ -142,9 +162,14 @@ public void null_is_passed_as_context_if_you_do_nothing() {
142162
loader.load("B");
143163
loader.loadMany(asList("C", "D"));
144164

165+
Map<String, String> keysAndContexts = new LinkedHashMap<>();
166+
keysAndContexts.put("E", null);
167+
keysAndContexts.put("F", null);
168+
loader.loadMany(keysAndContexts);
169+
145170
List<String> results = loader.dispatchAndJoin();
146171

147-
assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null")));
172+
assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null", "E-null", "F-null")));
148173
}
149174

150175
@Test
@@ -160,9 +185,14 @@ public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() {
160185
loader.load("B");
161186
loader.loadMany(asList("C", "D"));
162187

188+
Map<String, String> keysAndContexts = new LinkedHashMap<>();
189+
keysAndContexts.put("E", null);
190+
keysAndContexts.put("F", null);
191+
loader.loadMany(keysAndContexts);
192+
163193
List<String> results = loader.dispatchAndJoin();
164194

165-
assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null")));
195+
assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null", "E-null", "F-null")));
166196
}
167197

168198
@Test

‎src/test/java/org/dataloader/DataLoaderStatsTest.java‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313

1414
import java.util.ArrayList;
1515
import java.util.List;
16+
import java.util.Map;
1617
import java.util.concurrent.CompletableFuture;
1718

1819
import static java.util.Arrays.asList;
@@ -118,19 +119,20 @@ public void stats_are_collected_with_caching_disabled() {
118119
loader.load("A");
119120
loader.load("B");
120121
loader.loadMany(asList("C", "D"));
122+
loader.loadMany(Map.of("E", "E", "F", "F"));
121123

122124
Statistics stats = loader.getStatistics();
123-
assertThat(stats.getLoadCount(), equalTo(4L));
125+
assertThat(stats.getLoadCount(), equalTo(6L));
124126
assertThat(stats.getBatchInvokeCount(), equalTo(0L));
125127
assertThat(stats.getBatchLoadCount(), equalTo(0L));
126128
assertThat(stats.getCacheHitCount(), equalTo(0L));
127129

128130
loader.dispatch();
129131

130132
stats = loader.getStatistics();
131-
assertThat(stats.getLoadCount(), equalTo(4L));
133+
assertThat(stats.getLoadCount(), equalTo(6L));
132134
assertThat(stats.getBatchInvokeCount(), equalTo(1L));
133-
assertThat(stats.getBatchLoadCount(), equalTo(4L));
135+
assertThat(stats.getBatchLoadCount(), equalTo(6L));
134136
assertThat(stats.getCacheHitCount(), equalTo(0L));
135137

136138
loader.load("A");
@@ -139,9 +141,9 @@ public void stats_are_collected_with_caching_disabled() {
139141
loader.dispatch();
140142

141143
stats = loader.getStatistics();
142-
assertThat(stats.getLoadCount(), equalTo(6L));
144+
assertThat(stats.getLoadCount(), equalTo(8L));
143145
assertThat(stats.getBatchInvokeCount(), equalTo(2L));
144-
assertThat(stats.getBatchLoadCount(), equalTo(6L));
146+
assertThat(stats.getBatchLoadCount(), equalTo(8L));
145147
assertThat(stats.getCacheHitCount(), equalTo(0L));
146148
}
147149

0 commit comments

Comments
 (0)