Fix PrefetchingResultSetIterator regression introduced in latest commit

Latest commit introduces a regression & doesn't handle all corner cases. Fix this by calling maybePrefetch in constructor & update unit tests accordingly.
criteo-forks · Dec 3, 2024 · 04489b7 · 04489b7
1 parent a63795a
commit 04489b7
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 10 deletions.
diff --git a/...src/main/scala/com/datastax/spark/connector/rdd/reader/PrefetchingResultSetIterator.scala b/...src/main/scala/com/datastax/spark/connector/rdd/reader/PrefetchingResultSetIterator.scala
@@ -37,9 +37,9 @@ class PrefetchingResultSetIterator(resultSet: AsyncResultSet, timer: Option[Time
   }
 
   private def maybePrefetch(): Unit = {
-    // It is sometimes possible to see pages in the middle have zero
-    // elements, hence iterating until we get at least one element is
-    // needed (and it needs to be done before calling .next(), too)
+    // It is sometimes possible to see pages have zero elements that are
+    // not the last, hence iterating until we get at least one element is
+    // needed
     while (!currentIterator.hasNext && currentResultSet.hasMorePages) {
       currentResultSet = Await.result(nextResultSet.get, Duration.Inf)
       currentIterator = currentResultSet.currentPage().iterator()
@@ -51,8 +51,16 @@ class PrefetchingResultSetIterator(resultSet: AsyncResultSet, timer: Option[Time
     currentIterator.hasNext || currentResultSet.hasMorePages
 
   override def next(): Row = {
-    maybePrefetch()
     val row = currentIterator.next()
+
+    // This must be called after the call to next() and not before,
+    // so that hasNext returns is a correct result
+    maybePrefetch()
     row
   }
+
+  // It is possible to have empty pages at the start of the query result,
+  // for example when using "filtering"; discard those and ensure hasNext
+  // always returns the correct result
+  maybePrefetch()
 }
diff --git a/...test/scala/com/datastax/spark/connector/rdd/reader/PrefetchingResultSetIteratorSpec.scala b/...test/scala/com/datastax/spark/connector/rdd/reader/PrefetchingResultSetIteratorSpec.scala
@@ -10,17 +10,61 @@ import scala.jdk.CollectionConverters._
 
 class PrefetchingResultSetIteratorSpec extends FlatSpec with Matchers with MockitoSugar {
 
-  "PrefetchingResultSetIterator" should "handle empty pages that are not the last with AsyncResultSet" in {
+  "PrefetchingResultSetIterator" should "handle empty pages that are not the last" in {
     val row1 = mock[Row]
     val row2 = mock[Row]
     val row3 = mock[Row]
 
     val asyncResultSet1 = mock[AsyncResultSet]
     val asyncResultSet2 = mock[AsyncResultSet]
     val asyncResultSet3 = mock[AsyncResultSet]
+    val asyncResultSet4 = mock[AsyncResultSet]
+    val asyncResultSet5 = mock[AsyncResultSet]
 
-    // First page has data
-    when(asyncResultSet1.currentPage()).thenReturn(Seq(row1).asJava)
+    // First page is empty
+    when(asyncResultSet1.currentPage()).thenReturn(Seq.empty[Row].asJava)
+    when(asyncResultSet1.hasMorePages()).thenReturn(true)
+    when(asyncResultSet1.fetchNextPage()).thenReturn(CompletableFuture.completedFuture(asyncResultSet2))
+
+    // Second page has data
+    when(asyncResultSet2.currentPage()).thenReturn(Seq(row1).asJava)
+    when(asyncResultSet2.hasMorePages()).thenReturn(true)
+    when(asyncResultSet2.fetchNextPage()).thenReturn(CompletableFuture.completedFuture(asyncResultSet3))
+
+    // Third page is empty
+    when(asyncResultSet3.currentPage()).thenReturn(Seq.empty[Row].asJava)
+    when(asyncResultSet3.hasMorePages()).thenReturn(true)
+    when(asyncResultSet3.fetchNextPage()).thenReturn(CompletableFuture.completedFuture(asyncResultSet4))
+
+    // Fourth page has data
+    when(asyncResultSet4.currentPage()).thenReturn(Seq(row2, row3).asJava)
+    when(asyncResultSet4.hasMorePages()).thenReturn(true)
+    when(asyncResultSet4.fetchNextPage()).thenReturn(CompletableFuture.completedFuture(asyncResultSet5))
+
+    // Last page is empty
+    when(asyncResultSet5.currentPage()).thenReturn(Seq.empty[Row].asJava)
+    when(asyncResultSet5.hasMorePages()).thenReturn(false)
+
+    val iterator = new PrefetchingResultSetIterator(asyncResultSet1)
+
+    val rows = iterator.toList
+
+    rows should contain theSameElementsInOrderAs Seq(row1, row2, row3)
+
+    verify(asyncResultSet1).fetchNextPage()
+    verify(asyncResultSet2).fetchNextPage()
+    verify(asyncResultSet3).fetchNextPage()
+    verify(asyncResultSet4).fetchNextPage()
+    verify(asyncResultSet5, never()).fetchNextPage()
+  }
+
+  it should "handle a result made of empty pages only" in {
+    val asyncResultSet1 = mock[AsyncResultSet]
+    val asyncResultSet2 = mock[AsyncResultSet]
+    val asyncResultSet3 = mock[AsyncResultSet]
+
+    // First page is empty
+    when(asyncResultSet1.currentPage()).thenReturn(Seq.empty[Row].asJava)
     when(asyncResultSet1.hasMorePages()).thenReturn(true)
     when(asyncResultSet1.fetchNextPage()).thenReturn(CompletableFuture.completedFuture(asyncResultSet2))
 
@@ -29,15 +73,15 @@ class PrefetchingResultSetIteratorSpec extends FlatSpec with Matchers with Mocki
     when(asyncResultSet2.hasMorePages()).thenReturn(true)
     when(asyncResultSet2.fetchNextPage()).thenReturn(CompletableFuture.completedFuture(asyncResultSet3))
 
-    // Third page has data
-    when(asyncResultSet3.currentPage()).thenReturn(Seq(row2, row3).asJava)
+    // Last page is empty
+    when(asyncResultSet3.currentPage()).thenReturn(Seq.empty[Row].asJava)
     when(asyncResultSet3.hasMorePages()).thenReturn(false)
 
     val iterator = new PrefetchingResultSetIterator(asyncResultSet1)
 
     val rows = iterator.toList
 
-    rows should contain theSameElementsInOrderAs Seq(row1, row2, row3)
+    rows should contain theSameElementsInOrderAs Seq.empty[Row]
 
     verify(asyncResultSet1).fetchNextPage()
     verify(asyncResultSet2).fetchNextPage()