Provide in-memory sorting information in Query

Change-Id: Iebafff6079816c08e5a4d144f6891d1379751f12

9 files changed, 78 insertions, 54 deletions

diff --git a/server/src/main/java/com/vaadin/server/data/BackEndDataProvider.java b/server/src/main/java/com/vaadin/server/data/BackEndDataProvider.java
index 0f83710d90..9c6433bc4d 100644
--- a/server/src/main/java/com/vaadin/server/data/BackEndDataProvider.java
+++ b/server/src/main/java/com/vaadin/server/data/BackEndDataProvider.java
@@ -32,8 +32,8 @@ import com.vaadin.server.SerializableFunction;
  */
 public class BackEndDataProvider<T, F> extends AbstractDataProvider<T, F> {
 
-    private final SerializableFunction<Query<F>, Stream<T>> request;
-    private final SerializableFunction<Query<F>, Integer> sizeCallback;
+    private final SerializableFunction<Query<T, F>, Stream<T>> request;
+    private final SerializableFunction<Query<T, F>, Integer> sizeCallback;
 
     /**
      * Constructs a new DataProvider to request data from an arbitrary back end
@@ -45,8 +45,8 @@ public class BackEndDataProvider<T, F> extends AbstractDataProvider<T, F> {
      *            function that return the amount of data in back end for query
      */
     public BackEndDataProvider(
-            SerializableFunction<Query<F>, Stream<T>> request,
-            SerializableFunction<Query<F>, Integer> sizeCallback) {
+            SerializableFunction<Query<T, F>, Stream<T>> request,
+            SerializableFunction<Query<T, F>, Integer> sizeCallback) {
         Objects.requireNonNull(request, "Request function can't be null");
         Objects.requireNonNull(sizeCallback, "Size callback can't be null");
         this.request = request;
@@ -54,12 +54,12 @@ public class BackEndDataProvider<T, F> extends AbstractDataProvider<T, F> {
     }
 
     @Override
-    public Stream<T> fetch(Query<F> query) {
+    public Stream<T> fetch(Query<T, F> query) {
         return request.apply(query);
     }
 
     @Override
-    public int size(Query<F> query) {
+    public int size(Query<T, F> query) {
         return sizeCallback.apply(query);
     }
 
@@ -77,9 +77,9 @@ public class BackEndDataProvider<T, F> extends AbstractDataProvider<T, F> {
             List<SortOrder<String>> queryOrder = new ArrayList<>(
                     query.getSortOrders());
             queryOrder.addAll(sortOrders);
-            return request
-                    .apply(new Query<>(query.getLimit(), query.getOffset(),
-                            queryOrder, query.getFilter().orElse(null)));
+            return request.apply(new Query<>(query.getLimit(),
+                    query.getOffset(), queryOrder, query.getInMemorySorting(),
+                    query.getFilter().orElse(null)));
         }, sizeCallback);
     }
 
diff --git a/server/src/main/java/com/vaadin/server/data/DataCommunicator.java b/server/src/main/java/com/vaadin/server/data/DataCommunicator.java
index 2628c79058..61545ada83 100644
--- a/server/src/main/java/com/vaadin/server/data/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/server/data/DataCommunicator.java
@@ -237,20 +237,9 @@ public class DataCommunicator<T, F> extends AbstractExtension {
             int offset = pushRows.getStart();
             int limit = pushRows.length();
 
-            Stream<T> rowsToPush;
-
-            if (getDataProvider().isInMemory()) {
-                // TODO: Move in-memory sorting to Query.
-                // We can safely request all the data when in memory
-                rowsToPush = getDataProvider().fetch(new Query<>(filter));
-                if (inMemorySorting != null) {
-                    rowsToPush = rowsToPush.sorted(inMemorySorting);
-                }
-                rowsToPush = rowsToPush.skip(offset).limit(limit);
-            } else {
-                rowsToPush = getDataProvider().fetch(
-                        new Query<>(offset, limit, backEndSorting, filter));
-            }
+            Stream<T> rowsToPush = getDataProvider().fetch(new Query<>(offset,
+                    limit, backEndSorting, inMemorySorting, filter));
+
             pushData(offset, rowsToPush);
         }
 
diff --git a/server/src/main/java/com/vaadin/server/data/DataProvider.java b/server/src/main/java/com/vaadin/server/data/DataProvider.java
index 0975a02c94..a373128358 100644
--- a/server/src/main/java/com/vaadin/server/data/DataProvider.java
+++ b/server/src/main/java/com/vaadin/server/data/DataProvider.java
@@ -57,7 +57,7 @@ public interface DataProvider<T, F> extends Serializable {
      *            query with sorting and filtering
      * @return the size of the data provider
      */
-    int size(Query<F> t);
+    int size(Query<T, F> t);
 
     /**
      * Fetches data from this DataProvider using given {@code query}.
@@ -67,7 +67,7 @@ public interface DataProvider<T, F> extends Serializable {
      * @return the result of the query request: a stream of data objects, not
      *         {@code null}
      */
-    Stream<T> fetch(Query<F> query);
+    Stream<T> fetch(Query<T, F> query);
 
     /**
      * Refreshes all data based on currently available data in the underlying
diff --git a/server/src/main/java/com/vaadin/server/data/FilteringDataProviderWrapper.java b/server/src/main/java/com/vaadin/server/data/FilteringDataProviderWrapper.java
index 08a76d4904..b888fd79c7 100644
--- a/server/src/main/java/com/vaadin/server/data/FilteringDataProviderWrapper.java
+++ b/server/src/main/java/com/vaadin/server/data/FilteringDataProviderWrapper.java
@@ -100,15 +100,15 @@ public abstract class FilteringDataProviderWrapper<T, F, M>
     }
 
     @Override
-    public int size(Query<F> t) {
-        return dataProvider.size(new Query<M>(t.getOffset(), t.getLimit(),
-                t.getSortOrders(), getFilter(t)));
+    public int size(Query<T, F> t) {
+        return dataProvider.size(new Query<>(t.getOffset(), t.getLimit(),
+                t.getSortOrders(), t.getInMemorySorting(), getFilter(t)));
     }
 
     @Override
-    public Stream<T> fetch(Query<F> t) {
-        return dataProvider.fetch(new Query<M>(t.getOffset(), t.getLimit(),
-                t.getSortOrders(), getFilter(t)));
+    public Stream<T> fetch(Query<T, F> t) {
+        return dataProvider.fetch(new Query<>(t.getOffset(), t.getLimit(),
+                t.getSortOrders(), t.getInMemorySorting(), getFilter(t)));
     }
 
     /**
@@ -118,7 +118,7 @@ public abstract class FilteringDataProviderWrapper<T, F, M>
      *            the current query
      * @return filter for the modified Query
      */
-    protected abstract M getFilter(Query<F> query);
+    protected abstract M getFilter(Query<T, F> query);
 
     /**
      * Creates a data provider wrapper with a static filter set to each Query.
@@ -144,7 +144,7 @@ public abstract class FilteringDataProviderWrapper<T, F, M>
         return new FilteringDataProviderWrapper<T, Void, F>(dataProvider) {
 
             @Override
-            protected F getFilter(Query<Void> query) {
+            protected F getFilter(Query<T, Void> query) {
                 return filter;
             }
         };
@@ -176,7 +176,7 @@ public abstract class FilteringDataProviderWrapper<T, F, M>
         return new FilteringDataProviderWrapper<T, F, M>(dataProvider) {
 
             @Override
-            protected M getFilter(Query<F> query) {
+            protected M getFilter(Query<T, F> query) {
                 return query.getFilter().map(mapper).orElse(null);
             }
         };
@@ -203,7 +203,7 @@ public abstract class FilteringDataProviderWrapper<T, F, M>
         return new AppendableFilterDataProviderWrapper<T, F>(dataProvider) {
 
             @Override
-            protected F getFilter(Query<F> query) {
+            protected F getFilter(Query<T, F> query) {
                 return combineFilters(filter, query.getFilter());
             }
         };
diff --git a/server/src/main/java/com/vaadin/server/data/ListDataProvider.java b/server/src/main/java/com/vaadin/server/data/ListDataProvider.java
index 7d2fa5c1f3..5665ac5f29 100644
--- a/server/src/main/java/com/vaadin/server/data/ListDataProvider.java
+++ b/server/src/main/java/com/vaadin/server/data/ListDataProvider.java
@@ -18,6 +18,7 @@ package com.vaadin.server.data;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
@@ -34,7 +35,7 @@ public class ListDataProvider<T>
         extends AbstractDataProvider<T, SerializablePredicate<T>>
         implements AppendableFilterDataProvider<T, SerializablePredicate<T>> {
 
-    private Comparator<T> sortOrder;
+    private Comparator<T> sortOrder = null;
     private final Collection<T> backend;
 
     /**
@@ -67,13 +68,20 @@ public class ListDataProvider<T>
     }
 
     @Override
-    public Stream<T> fetch(Query<SerializablePredicate<T>> query) {
+    public Stream<T> fetch(Query<T, SerializablePredicate<T>> query) {
         Stream<T> stream = backend.stream()
                 .filter(t -> query.getFilter().orElse(p -> true).test(t));
-        if (sortOrder != null) {
-            stream = stream.sorted(sortOrder);
+
+        Optional<Comparator<T>> comparing = Stream
+                .of(sortOrder, query.getInMemorySorting())
+                .filter(c -> c != null)
+                .reduce((c1, c2) -> c1.thenComparing(c2));
+
+        if (comparing.isPresent()) {
+            stream = stream.sorted(comparing.get());
         }
-        return stream;
+
+        return stream.skip(query.getOffset()).limit(query.getLimit());
     }
 
     /**
@@ -116,7 +124,7 @@ public class ListDataProvider<T>
     }
 
     @Override
-    public int size(Query<SerializablePredicate<T>> query) {
+    public int size(Query<T, SerializablePredicate<T>> query) {
         return (int) backend.stream()
                 .filter(t -> query.getFilter().orElse(p -> true).test(t))
                 .count();
diff --git a/server/src/main/java/com/vaadin/server/data/Query.java b/server/src/main/java/com/vaadin/server/data/Query.java
index 3aeaac6acc..6415af0e2d 100644
--- a/server/src/main/java/com/vaadin/server/data/Query.java
+++ b/server/src/main/java/com/vaadin/server/data/Query.java
@@ -17,6 +17,7 @@ package com.vaadin.server.data;
 
 import java.io.Serializable;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
 
@@ -24,16 +25,19 @@ import java.util.Optional;
  * Immutable query object used to request data from a backend. Contains index
  * limits, sorting and filtering information.
  *
+ * @param <T>
+ *            bean type
  * @param <F>
  *            filter type
  *
  * @since 8.0
  */
-public class Query<F> implements Serializable {
+public class Query<T, F> implements Serializable {
 
     private final int offset;
     private final int limit;
     private final List<SortOrder<String>> sortOrders;
+    private final Comparator<T> inMemorySorting;
     private final F filter;
 
     /**
@@ -44,6 +48,7 @@ public class Query<F> implements Serializable {
         offset = 0;
         limit = Integer.MAX_VALUE;
         sortOrders = Collections.emptyList();
+        inMemorySorting = null;
         filter = null;
     }
 
@@ -59,6 +64,7 @@ public class Query<F> implements Serializable {
         offset = 0;
         limit = Integer.MAX_VALUE;
         sortOrders = Collections.emptyList();
+        inMemorySorting = null;
         this.filter = filter;
     }
 
@@ -71,15 +77,18 @@ public class Query<F> implements Serializable {
      * @param limit
      *            fetched item count
      * @param sortOrders
-     *            sorting order for fetching
+     *            sorting order for fetching; used for sorting backends
+     * @param inMemorySorting
+     *            comparator for sorting in-memory data
      * @param filter
      *            filtering for fetching; can be null
      */
     public Query(int offset, int limit, List<SortOrder<String>> sortOrders,
-            F filter) {
+            Comparator<T> inMemorySorting, F filter) {
         this.offset = offset;
         this.limit = limit;
         this.sortOrders = sortOrders;
+        this.inMemorySorting = inMemorySorting;
         this.filter = filter;
     }
 
@@ -105,7 +114,12 @@ public class Query<F> implements Serializable {
     }
 
     /**
-     * Gets the sorting for items to fetch.
+     * Gets the sorting for items to fetch. This list of sort orders is used for
+     * sorting backends.
+     * <p>
+     * <strong>Note: </strong> Sort orders and in-memory sorting are mutually
+     * exclusive. If the {@link DataProvider} handles one, it should ignore the
+     * other.
      *
      * @return list of sort orders
      */
@@ -121,4 +135,17 @@ public class Query<F> implements Serializable {
     public Optional<F> getFilter() {
         return Optional.ofNullable(filter);
     }
+
+    /**
+     * Gets the comparator for sorting in-memory data.
+     * <p>
+     * <strong>Note: </strong> Sort orders and in-memory sorting are mutually
+     * exclusive. If the {@link DataProvider} handles one, it should ignore the
+     * other.
+     *
+     * @return sorting comparator
+     */
+    public Comparator<T> getInMemorySorting() {
+        return inMemorySorting;
+    }
 }
diff --git a/server/src/test/java/com/vaadin/server/data/AbstractDataProviderTest.java b/server/src/test/java/com/vaadin/server/data/AbstractDataProviderTest.java
index db2f704e2d..b6033ad098 100644
--- a/server/src/test/java/com/vaadin/server/data/AbstractDataProviderTest.java
+++ b/server/src/test/java/com/vaadin/server/data/AbstractDataProviderTest.java
@@ -32,12 +32,12 @@ public class AbstractDataProviderTest {
     private static class TestDataProvider
             extends AbstractDataProvider<Object, Object> {
         @Override
-        public Stream<Object> fetch(Query<Object> t) {
+        public Stream<Object> fetch(Query<Object, Object> t) {
             return null;
         }
 
         @Override
-        public int size(Query<Object> t) {
+        public int size(Query<Object, Object> t) {
             return 0;
         }
 
diff --git a/server/src/test/java/com/vaadin/tests/server/component/combobox/ComboBoxFilteringTest.java b/server/src/test/java/com/vaadin/tests/server/component/combobox/ComboBoxFilteringTest.java
index 54ea367fca..de03684307 100644
--- a/server/src/test/java/com/vaadin/tests/server/component/combobox/ComboBoxFilteringTest.java
+++ b/server/src/test/java/com/vaadin/tests/server/component/combobox/ComboBoxFilteringTest.java
@@ -156,16 +156,15 @@ public class ComboBoxFilteringTest {
 
         Assert.assertEquals(
                 "ComboBox filtered out results with no filter applied",
-                totalMatches, dataProvider.size(new Query<String>()));
+                totalMatches, dataProvider.size(new Query<>()));
         Assert.assertEquals(
                 "ComboBox filtered out results with empty filter string",
-                totalMatches, dataProvider.size(new Query<String>("")));
+                totalMatches, dataProvider.size(new Query<>("")));
         Assert.assertEquals("ComboBox filtered out wrong number of results",
-                matchingResults,
-                dataProvider.size(new Query<String>(filterText)));
+                matchingResults, dataProvider.size(new Query<>(filterText)));
         Assert.assertEquals(
                 "ComboBox should have no results with a non-matching filter", 0,
-                dataProvider.size(new Query<String>(nonMatchingFilterText)));
+                dataProvider.size(new Query<>(nonMatchingFilterText)));
     }
 
     private List<Person> getPersonCollection() {
diff --git a/uitest/src/main/java/com/vaadin/tests/util/ItemDataProvider.java b/uitest/src/main/java/com/vaadin/tests/util/ItemDataProvider.java
index f050d1f454..5f02135591 100644
--- a/uitest/src/main/java/com/vaadin/tests/util/ItemDataProvider.java
+++ b/uitest/src/main/java/com/vaadin/tests/util/ItemDataProvider.java
@@ -19,7 +19,8 @@ public class ItemDataProvider extends BackEndDataProvider<String, String> {
                 q -> size(q, size));
     }
 
-    private static Stream<String> itemStream(Query<String> q, int size) {
+    private static Stream<String> itemStream(Query<String, String> q,
+            int size) {
         Stream<String> stream = IntStream.range(0, size)
                 .mapToObj(i -> "Item " + i);
         String filterText = q.getFilter().orElse("").toLowerCase(Locale.US);
@@ -32,7 +33,7 @@ public class ItemDataProvider extends BackEndDataProvider<String, String> {
                 text -> text.toLowerCase(Locale.US).contains(filterText));
     }
 
-    private static int size(Query<String> q, int size) {
+    private static int size(Query<String, String> q, int size) {
         if (!q.getFilter().orElse("").isEmpty()) {
             return (int) itemStream(q, size).count();
         }