Merge pull request #2191 from ebean-orm/feature/findEachBatch

ENH: Add Query findEach() with batch consumer

ebean-api/src/main/java/io/ebean/ExpressionList.java

@@ -307,6 +307,13 @@ public interface ExpressionList<T> {
    */
   void findEach(Consumer<T> consumer);
 
+  /**
+   * Execute findEach with a batch consumer.
+   *
+   * @see Query#findEach(int, Consumer)
+   */
+  void findEach(int batch, Consumer<List<T>> consumer);
+
   /**
    * Execute the query processing the beans one at a time with the ability to
    * stop processing before reading all the beans.

ebean-api/src/main/java/io/ebean/ExtendedServer.java

@@ -159,6 +159,13 @@ public interface ExtendedServer {
    */
   <T> void findEach(Query<T> query, Consumer<T> consumer, Transaction transaction);
 
+  /**
+   * Execute findEach with batch consumer.
+   *
+   * @see Query#findEach(int, Consumer)
+   */
+  <T> void findEach(Query<T> query, int batch, Consumer<List<T>> consumer, Transaction t);
+
   /**
    * Execute the query visiting the each bean one at a time.
    * <p>

ebean-api/src/main/java/io/ebean/Query.java

@@ -810,7 +810,7 @@ public interface Query<T> {
    * </p>
    * <p>
    * This method is functionally equivalent to findIterate() but instead of using an
-   * iterator uses the Consumer interface which is better suited to use with Java8 closures.
+   * iterator uses the Consumer interface which is better suited to use with closures.
    * </p>
    * <pre>{@code
    *
@@ -829,6 +829,21 @@ public interface Query<T> {
    */
   void findEach(Consumer<T> consumer);
 
+  /**
+   * Execute findEach streaming query batching the results for consuming.
+   * <p>
+   * This query execution will stream the results and is suited to consuming
+   * large numbers of results from the database.
+   * <p>
+   * Typically we use this batch consumer when we want to do further processing on
+   * the beans and want to do that processing in batch form, for example - 100 at
+   * a time.
+   *
+   * @param batch    The number of beans processed in the batch
+   * @param consumer Process the batch of beans
+   */
+  void findEach(int batch, Consumer<List<T>> consumer);
+
   /**
    * Execute the query using callbacks to a visitor to process the resulting
    * beans one at a time.
@@ -839,12 +854,12 @@ public interface Query<T> {
    * </p>
    * <p>
    * This method is functionally equivalent to findIterate() but instead of using an
-   * iterator uses the Predicate (SAM) interface which is better suited to use with Java8 closures.
+   * iterator uses the Predicate interface which is better suited to use with closures.
    * </p>
    * <pre>{@code
    *
    *  DB.find(Customer.class)
-   *     .fetch("contacts", FetchConfig.ofQuery(2))
+   *     .fetchQuery("contacts")
    *     .where().eq("status", Status.NEW)
    *     .order().asc("id")
    *     .setMaxRows(2000)

ebean-core/src/main/java/io/ebeaninternal/server/core/DefaultServer.java

@@ -1446,6 +1446,18 @@ public final class DefaultServer implements SpiServer, SpiEbeanServer {
     // no try finally - findEach guarantee's cleanup of the transaction if required
   }
 
+  @Override
+  public <T> void findEach(Query<T> query, int batch, Consumer<List<T>> consumer, Transaction t) {
+    SpiOrmQueryRequest<T> request = createQueryRequest(Type.ITERATE, query, t);
+//    if (request.isUseDocStore()) {
+//      docStore().findEach(request, consumer);
+//      return;
+//    }
+    request.initTransIfRequired();
+    request.findEach(batch, consumer);
+    // no try finally - findEach guarantee's cleanup of the transaction if required
+  }
+
   @Override
   public <T> void findEachWhile(Query<T> query, Predicate<T> consumer, Transaction t) {
     SpiOrmQueryRequest<T> request = createQueryRequest(Type.ITERATE, query, t);

ebean-core/src/main/java/io/ebeaninternal/server/core/OrmQueryRequest.java

@@ -430,6 +430,24 @@ public final class OrmQueryRequest<T> extends BeanRequest implements SpiOrmQuery
     }
   }
 
+  @Override
+  public void findEach(int batch, Consumer<List<T>> batchConsumer) {
+    final List<T> buffer = new ArrayList<>(batch);
+    try (QueryIterator<T> it = queryEngine.findIterate(this)) {
+      while (it.hasNext()) {
+        buffer.add(it.next());
+        if (buffer.size() >= batch) {
+          batchConsumer.accept(buffer);
+          buffer.clear();
+        }
+      }
+      if (!buffer.isEmpty()) {
+        // consume the remainder
+        batchConsumer.accept(buffer);
+      }
+    }
+  }
+
   @Override
   public void findEachWhile(Predicate<T> consumer) {
     try (QueryIterator<T> it = queryEngine.findIterate(this)) {

ebean-core/src/main/java/io/ebeaninternal/server/core/SpiOrmQueryRequest.java

@@ -80,10 +80,15 @@ public interface SpiOrmQueryRequest<T> extends BeanQueryRequest<T>, DocQueryRequ
   <A> List<A> findIds();
 
   /**
-   * Execute the find returning a QueryIterator and visitor pattern.
+   * Execute findEach iterating results one bean at a time.
    */
   void findEach(Consumer<T> consumer);
 
+  /**
+   * Execute findEach with a batch consumer.
+   */
+  void findEach(int batch, Consumer<List<T>> batchConsumer);
+
   /**
    * Execute the find returning a QueryIterator and visitor pattern.
    */

ebean-core/src/main/java/io/ebeaninternal/server/expression/DefaultExpressionList.java

@@ -436,6 +436,11 @@ public class DefaultExpressionList<T> implements SpiExpressionList<T> {
     query.findEach(consumer);
   }
 
+  @Override
+  public void findEach(int batch, Consumer<List<T>> consumer) {
+    query.findEach(batch, consumer);
+  }
+
   @Override
   public void findEachWhile(Predicate<T> consumer) {
     query.findEachWhile(consumer);

ebean-core/src/main/java/io/ebeaninternal/server/expression/JunctionExpression.java

@@ -446,6 +446,11 @@ class JunctionExpression<T> implements SpiJunction<T>, SpiExpression, Expression
     exprList.findEach(consumer);
   }
 
+  @Override
+  public void findEach(int batch, Consumer<List<T>> consumer) {
+    exprList.findEach(batch, consumer);
+  }
+
   @Override
   public void findEachWhile(Predicate<T> consumer) {
     exprList.findEachWhile(consumer);

ebean-core/src/main/java/io/ebeaninternal/server/query/DefaultFetchGroupQuery.java

@@ -263,6 +263,11 @@ class DefaultFetchGroupQuery<T> implements SpiFetchGroupQuery<T>, SpiQueryFetch
     throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
   }
 
+  @Override
+  public void findEach(int batch, Consumer<List<T>> consumer) {
+    throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");
+  }
+
   @Override
   public void findEachWhile(Predicate<T> consumer) {
     throw new RuntimeException("EB102: Only select() and fetch() clause is allowed on FetchGroup");

ebean-core/src/main/java/io/ebeaninternal/server/querydefn/DefaultOrmQuery.java

@@ -1551,6 +1551,11 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
     server.findEach(this, consumer, transaction);
   }
 
+  @Override
+  public void findEach(int batch, Consumer<List<T>> consumer) {
+    server.findEach(this, batch, consumer, transaction);
+  }
+
   @Override
   public QueryIterator<T> findIterate() {
     return server.findIterate(this, transaction);

ebean-core/src/test/java/io/ebeaninternal/api/TDSpiEbeanServer.java

@@ -675,6 +675,10 @@ public class TDSpiEbeanServer implements SpiEbeanServer {
   public <T> void findEach(Query<T> query, Consumer<T> consumer, Transaction transaction) {
   }
 
+  @Override
+  public <T> void findEach(Query<T> query, int batch, Consumer<List<T>> consumer, Transaction t) {
+  }
+
   @Override
   public <T> void findEachWhile(Query<T> query, Predicate<T> consumer, Transaction transaction) {
   }

ebean-querybean/src/main/java/io/ebean/typequery/TQRootBean.java

@@ -1794,24 +1794,19 @@ public abstract class TQRootBean<T, R> {
    * This method is appropriate to process very large query results as the
    * beans are consumed one at a time and do not need to be held in memory
    * (unlike #findList #findSet etc)
-   * </p>
    * <p>
    * Note that internally Ebean can inform the JDBC driver that it is expecting larger
    * resultSet and specifically for MySQL this hint is required to stop it's JDBC driver
    * from buffering the entire resultSet. As such, for smaller resultSets findList() is
    * generally preferable.
-   * </p>
    * <p>
    * Compared with #findEachWhile this will always process all the beans where as
    * #findEachWhile provides a way to stop processing the query result early before
    * all the beans have been read.
-   * </p>
    * <p>
    * This method is functionally equivalent to findIterate() but instead of using an
-   * iterator uses the QueryEachConsumer (SAM) interface which is better suited to use
-   * with Java8 closures.
-   * </p>
-   * <p>
+   * iterator uses the Consumer interface which is better suited to use with closures.
+   *
    * <pre>{@code
    *
    *  new QCustomer()
@@ -1831,16 +1826,30 @@ public abstract class TQRootBean<T, R> {
     query.findEach(consumer);
   }
 
+  /**
+   * Execute findEach streaming query batching the results for consuming.
+   * <p>
+   * This query execution will stream the results and is suited to consuming
+   * large numbers of results from the database.
+   * <p>
+   * Typically we use this batch consumer when we want to do further processing on
+   * the beans and want to do that processing in batch form, for example - 100 at
+   * a time.
+   *
+   * @param batch    The number of beans processed in the batch
+   * @param consumer Process the batch of beans
+   */
+  public void findEach(int batch, Consumer<List<T>> consumer) {
+    query.findEach(batch, consumer);
+  }
+
   /**
    * Execute the query using callbacks to a visitor to process the resulting
    * beans one at a time.
    * <p>
    * This method is functionally equivalent to findIterate() but instead of using an
-   * iterator uses the QueryEachWhileConsumer (SAM) interface which is better suited to use
-   * with Java8 closures.
-   * </p>
-   * <p>
-   * <p>
+   * iterator uses the Predicate interface which is better suited to use with closures.
+   *
    * <pre>{@code
    *
    *  new QCustomer()

ebean-querybean/src/test/java/org/querytest/QCustomerTest.java

@@ -38,6 +38,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.StringJoiner;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -82,14 +83,67 @@ public class QCustomerTest {
   public void findSingleAttribute() {
 
     List<String> names = new QCustomer()
-        .setDistinct(true)
-        .select(QCustomer.alias().name)
-        .status.equalTo(Customer.Status.BAD)
-        .findSingleAttributeList();
+      .setDistinct(true)
+      .select(QCustomer.alias().name)
+      .status.equalTo(Customer.Status.BAD)
+      .findSingleAttributeList();
 
     assertThat(names).isNotNull();
   }
 
+  @Test
+  public void findEachBatch() {
+
+    for (int i = 0; i < 22; i++) {
+      Customer customer = new Customer();
+      customer.setStatus(Customer.Status.MIDDLING);
+      customer.setName("findEachBatch_a_" + i);
+      customer.save();
+    }
+
+    final List<Integer> batchSizes = new ArrayList<>();
+
+    final AtomicInteger counter = new AtomicInteger();
+    new QCustomer()
+      .status.eq(Customer.Status.MIDDLING)
+      .name.startsWith("findEachBatch_a_")
+      .findEach(10, customers -> {
+        batchSizes.add(customers.size());
+        System.out.println("Batch " + counter.incrementAndGet() + " size:" + customers.size());
+      });
+
+    assertThat(batchSizes).hasSize(3);
+    assertThat(batchSizes.get(0)).isEqualTo(10);
+    assertThat(batchSizes.get(1)).isEqualTo(10);
+    assertThat(batchSizes.get(2)).isEqualTo(2);
+  }
+
+  @Test
+  public void findEachBatch_when_lastBatchEmpty() {
+
+    for (int i = 0; i < 18; i++) {
+      Customer customer = new Customer();
+      customer.setStatus(Customer.Status.MIDDLING);
+      customer.setName("findEachBatch_b_" + i);
+      customer.save();
+    }
+
+    final List<Integer> batchSizes = new ArrayList<>();
+
+    final AtomicInteger counter = new AtomicInteger();
+    new QCustomer()
+      .status.eq(Customer.Status.MIDDLING)
+      .name.startsWith("findEachBatch_b_")
+      .findEach(9, customers -> {
+        batchSizes.add(customers.size());
+        System.out.println("Batch " + counter.incrementAndGet() + " size:" + customers.size());
+      });
+
+    assertThat(batchSizes).hasSize(2);
+    assertThat(batchSizes.get(0)).isEqualTo(9);
+    assertThat(batchSizes.get(1)).isEqualTo(9);
+  }
+
   @Test
   public void findIterate() {
 
@@ -99,21 +153,21 @@ public class QCustomerTest {
     cust.save();
 
     List<Long> ids = new QCustomer()
-        .status.equalTo(Customer.Status.GOOD)
-        .findIds();
+      .status.equalTo(Customer.Status.GOOD)
+      .findIds();
 
     assertThat(ids).isNotEmpty();
 
 
     Map<List, Customer> map = new QCustomer()
-        .status.equalTo(Customer.Status.GOOD)
-        .findMap();
+      .status.equalTo(Customer.Status.GOOD)
+      .findMap();
 
     assertThat(map.size()).isEqualTo(ids.size());
 
     QueryIterator<Customer> iterate = new QCustomer()
-        .status.equalTo(Customer.Status.GOOD)
-        .findIterate();
+      .status.equalTo(Customer.Status.GOOD)
+      .findIterate();
 
     try {
       while (iterate.hasNext()) {
@@ -130,12 +184,12 @@ public class QCustomerTest {
   public void isEmpty() {
 
     new QCustomer()
-        .contacts.isEmpty()
-        .findList();
+      .contacts.isEmpty()
+      .findList();
 
     new QCustomer()
-        .contacts.isNotEmpty()
-        .findList();
+      .contacts.isNotEmpty()
+      .findList();
   }
 
   @Transactional
@@ -143,19 +197,19 @@ public class QCustomerTest {
   public void forUpdate() {
 
     new QCustomer()
-        .id.eq(42)
-        .forUpdate()
-        .findOne();
+      .id.eq(42)
+      .forUpdate()
+      .findOne();
 
     new QCustomer()
-        .id.eq(42)
-        .forUpdateNoWait()
-        .findOne();
+      .id.eq(42)
+      .forUpdateNoWait()
+      .findOne();
 
     new QCustomer()
-        .id.eq(42)
-        .forUpdateSkipLocked()
-        .findOne();
+      .id.eq(42)
+      .forUpdateSkipLocked()
+      .findOne();
   }
 
 
@@ -164,42 +218,42 @@ public class QCustomerTest {
   public void arrayContains() {
 
     new QContact()
-        .phoneNumbers.contains("4312")
-        .findList();
+      .phoneNumbers.contains("4312")
+      .findList();
 
     new QCustomer()
-        .contacts.phoneNumbers.contains("4312")
-        .findList();
+      .contacts.phoneNumbers.contains("4312")
+      .findList();
   }
 
   @Test
   public void setIncludeSoftDeletes() {
 
     new QCustomer()
-        .setIdIn(42L)
-        .setIncludeSoftDeletes()
-        .findList();
+      .setIdIn(42L)
+      .setIncludeSoftDeletes()
+      .findList();
   }
 
   @Test
   public void testIdIn() {
 
     new QCustomer()
-        .setIdIn("1", "2")
-        .findList();
+      .setIdIn("1", "2")
+      .findList();
 
     new QCustomer()
-        .id.in(1L, 2L, 3L)
-        .findList();
+      .id.in(1L, 2L, 3L)
+      .findList();
   }
 
   @Test
   public void testIn() {
     new QCustomer()
-        .id.in(34L, 33L)
-        .name.in("asd", "foo", "bar")
-        .registered.in(new Date())
-        .findList();
+      .id.in(34L, 33L)
+      .name.in("asd", "foo", "bar")
+      .registered.in(new Date())
+      .findList();
   }
 
   @Test
@@ -284,24 +338,24 @@ public class QCustomerTest {
   @Test
   public void testNotIn() {
     new QCustomer()
-        .id.isIn(34L, 33L)
-        .name.notIn("asd", "foo", "bar")
-        .registered.in(new Date())
-        .findList();
+      .id.isIn(34L, 33L)
+      .name.notIn("asd", "foo", "bar")
+      .registered.in(new Date())
+      .findList();
   }
 
   @Test
   public void testQueryBoolean() {
 
     new QCustomer()
-        .name.contains("rob")
-        //.setUseDocStore(true)
-        .setMaxRows(10)
-        .findPagedList();
+      .name.contains("rob")
+      //.setUseDocStore(true)
+      .setMaxRows(10)
+      .findPagedList();
 
     new QCustomer()
-        .inactive.isFalse()
-        .findList();
+      .inactive.isFalse()
+      .findList();
   }
 
   @Test
@@ -484,7 +538,7 @@ public class QCustomerTest {
     assertThat(billingAddressIds).hasSize(2);
 
 
-    Map<Long,Customer> map
+    Map<Long, Customer> map
       = new QCustomer()
       .billingAddress.id.asMapKey()
       .name.startsWith("asdBilling")
@@ -515,21 +569,21 @@ public class QCustomerTest {
       .findList();
 
     new QCustomer()
-      .currentInet.in(Inet.setOf("129.1.1.4","129.1.1.5"))
+      .currentInet.in(Inet.setOf("129.1.1.4", "129.1.1.5"))
       .findList();
 
     new QCustomer()
       .contacts.fetch("email")
       .orderBy()
-        .name.asc()
-        .contacts.email.asc()
+      .name.asc()
+      .contacts.email.asc()
       .findList();
 
     new QCustomer()
       .contacts.fetchQuery("email")
       .orderBy()
-        .name.asc()
-        .contacts.email.asc()
+      .name.asc()
+      .contacts.email.asc()
       .findList();
   }
 
@@ -608,8 +662,8 @@ public class QCustomerTest {
 
     boolean customerExists =
       new QCustomer()
-      .name.equalTo("DoesNotExistReally")
-      .exists();
+        .name.equalTo("DoesNotExistReally")
+        .exists();
 
     assertThat(customerExists).isFalse();
   }
@@ -621,37 +675,37 @@ public class QCustomerTest {
     QCustomer cust = QCustomer.alias();
 
     new QCustomer()
-        // tune query
-        .select(cust.name)
-        .status.isIn(Customer.Status.BAD, Customer.Status.BAD)
-        .contacts.fetch()
-        // predicates
-        .findList();
+      // tune query
+      .select(cust.name)
+      .status.isIn(Customer.Status.BAD, Customer.Status.BAD)
+      .contacts.fetch()
+      // predicates
+      .findList();
 
     new QCustomer()
-        // tune query
-        .select(cust.name)
-        .contacts.fetch()
-        // predicates
-        .findList();
+      // tune query
+      .select(cust.name)
+      .contacts.fetch()
+      // predicates
+      .findList();
 
     new QCustomer()
-        // tune query
-        .select(cust.id, cust.name)
-        .contacts.fetch(contact.firstName, contact.lastName, contact.email)
-        // predicates
-        .id.greaterThan(1)
-        .findList();
+      // tune query
+      .select(cust.id, cust.name)
+      .contacts.fetch(contact.firstName, contact.lastName, contact.email)
+      // predicates
+      .id.greaterThan(1)
+      .findList();
 
     PagedList<Customer> pagedList = new QCustomer()
-        // tune query
-        .select(cust.id, cust.name)
-        .contacts.fetch(contact.firstName, contact.lastName, contact.email)
-        // predicates
-        .id.greaterThan(1)
-        .setFirstRow(20)
-        .setMaxRows(10)
-        .findPagedList();
+      // tune query
+      .select(cust.id, cust.name)
+      .contacts.fetch(contact.firstName, contact.lastName, contact.email)
+      // predicates
+      .id.greaterThan(1)
+      .setFirstRow(20)
+      .setMaxRows(10)
+      .findPagedList();
 
     pagedList.getList();
     pagedList.getList();
@@ -765,7 +819,7 @@ public class QCustomerTest {
     cust.setRegistered(new Date());
     cust.save();
 
-    java.util.Date maxDate =  new QCustomer()
+    java.util.Date maxDate = new QCustomer()
       .select("max(registered)")
       .findSingleAttribute();
 
@@ -799,15 +853,15 @@ public class QCustomerTest {
     assertThat(new QCustomer()
       .name.eq(testName.getMethodName())
       .email.gt(new ValidEmail("foo2@example.org"))
-            .findOne()).isNull();
+      .findOne()).isNull();
     assertThat(new QCustomer()
       .name.eq(testName.getMethodName())
       .email.gt(new ValidEmail("foo1@example.org"))
-            .findOne()).isNotNull();
+      .findOne()).isNotNull();
     assertThat(new QCustomer()
       .name.eq(testName.getMethodName())
       .email.greaterOrEqualTo(new ValidEmail("foo2@example.org"))
-            .findOne()).isNotNull();
+      .findOne()).isNotNull();
   }