diff --git a/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java b/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java
index 253e03a80c8..54f7a112d48 100644
--- a/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java
+++ b/client/src/main/java/com/vaadin/client/ui/treegrid/TreeGridConnector.java
@@ -208,7 +208,7 @@ public void dataRemoved(int firstRowIndex, int numberOfRows) {
                             .getVisibleRowRange();
                     getDataSource().ensureAvailability(
                             visibleRowRange.getStart(),
-                            visibleRowRange.getEnd());
+                            visibleRowRange.length());
                 }
                 checkExpand();
             }
@@ -222,7 +222,7 @@ public void dataAdded(int firstRowIndex, int numberOfRows) {
                             .getVisibleRowRange();
                     getDataSource().ensureAvailability(
                             visibleRowRange.getStart(),
-                            visibleRowRange.getEnd());
+                            visibleRowRange.length());
                 }
                 checkExpand();
             }