Support for partitioning keys

[jhaber]

Sorry if this is already supported or has already been discussed, I wasn't able to find anything in the docs or in another GitHub issue. We have some APIs that support batching, but only if all of the inputs are uniform in some way. For example, you could have an image resizing API that supports batch resizing, but only if the desired dimensions are the same. We want to use java-dataloader for the batching and caching support, but handling these sorts of constraints ends up being a bit of a hassle. We get a batch of keys that might have mixed dimensions, so inside our load function we need to partition the keys, make multiple HTTP requests, and handle combining the futures and merging the data. Here is a simplified example of what this might involve:

public class ImageResizeDataLoader implements MappedBatchLoader<ImageResizeRequest, ImageResizeResponse> {

  @Override
  public CompletionStage<Map<ImageResizeRequest, ImageResizeResponse>> load(Set<ImageResizeRequest> keys) {
    Multimap<ResizeDimensions, ImageResizeRequest> dimensionsToKeys = keys
        .stream()
        .collect(Multimaps.toMultimap(key -> key.resizeDimensions, Function.identity(), HashMultimap::create));

    Map<ImageResizeRequest, ImageResizeResponse> results = new ConcurrentHashMap<>();
    List<CompletableFuture<?>> futures = new ArrayList<>();

    dimensionsToKeys.asMap().forEach((dimensions, batch) -> {
      futures.add(doLoad(dimensions, batch).thenAccept(results::putAll));
    });

    return CompletableFuture.allOf(futures.toArray(new CompletableFuture<?>[0])).thenApply(ignored -> results);
  }

  private CompletableFuture<Map<ImageResizeRequest, ImageResizeResponse>> doLoad(
      ResizeDimensions dimensions,
      Collection<ImageResizeRequest> batch
  ) {
    HttpRequest request = HttpRequest
        .newBuilder()
        .setMethod(StandardHttpMethod.POST)
        .setUri("https://myapi.com/resize")
        .addQueryParam("width", dimensions.width)
        .addQueryParam("height", dimensions.height)
        .setBody(batch.stream().map(i -> i.imageUrl).collect(Collectors.toSet()))
        .build();

    // make request, handle response
    return null;
  }

  public static class ResizeDimensions {
    public int width;
    public int height;

    // equals/hashCode
  }

  public static class ImageResizeRequest {
    public String imageUrl;
    public ResizeDimensions resizeDimensions;

    // equals/hashCode
  }

  public static class ImageResizeResponse {
    public String resizedUrl;
  }
}

Is there a simpler way to handle this that I'm missing?

[bbakerman]

Is there a simpler way to handle this that I'm missing?

No I dont think so - the batch loader is a simple design that is key -> value such that you end up with a list < ke y> -> list < value >

You have extended functionality where logically you want to batch N issues but your underlyign system requires you to group them (partion them) according to custom logic.

I am afraid you will need to write that logic and make N batch calls then combine them into a logical complete list< value > on the way out

[jhaber]

Makes sense, that's what we're doing currently. If there are no plans to expose a partition function via the DataLoaderOptions, then we can get most of the way there by making reusable subclasses of BatchLoader/MappedBatchLoader, for example:

public abstract class PartitionedMappedBatchLoader<K, V> {
  protected abstract Object partitionFor(K key);
  protected abstract CompletionStage<Map<K, V>> loadPartition(Set<K> partition);

  public final CompletionStage<Map<K, V>> load(Set<K> keys) {
    Multimap<Object, K> partitions = keys
        .stream()
        .collect(Multimaps.toMultimap(this::partitionFor, Function.identity(), HashMultimap::create));

    for (Collection<K> partition : partitions.asMap().values()) {
      CompletionStage<Map<K, V>> partitionResult = loadPartition(new HashSet<>(partition));
    }

    // merge all the partition results
  }
}