From 664f2bc5152838168d4c45ddfeba425dd623779f Mon Sep 17 00:00:00 2001 From: Lawrence Qiu Date: Tue, 13 Oct 2026 03:50:45 +0000 Subject: [PATCH] perf(bigquery): [T2-1] add bounded 1-page lookahead prefetching to TableResult.iterateAll() --- .../google/cloud/bigquery/TableResult.java | 91 +++++++++++++- .../cloud/bigquery/TableResultTest.java | 113 ++++++++++++++++++ 2 files changed, 203 insertions(+), 1 deletion(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java index 1f5c56d44766..5052d4530d67 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java @@ -23,10 +23,18 @@ import com.google.cloud.bigquery.JobStatistics.SessionInfo; import com.google.common.base.Function; import com.google.common.base.MoreObjects; +import com.google.common.collect.AbstractIterator; import com.google.common.collect.Iterables; import com.google.common.collect.Iterators; +import com.google.common.util.concurrent.ThreadFactoryBuilder; import java.io.Serializable; +import java.util.Iterator; import java.util.Objects; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import org.jspecify.annotations.Nullable; @AutoValue @@ -34,6 +42,13 @@ public abstract class TableResult implements Page, Serializable private static final long serialVersionUID = 1L; + private static final ExecutorService PREFETCH_EXECUTOR = + Executors.newCachedThreadPool( + new ThreadFactoryBuilder() + .setDaemon(true) + .setNameFormat("bigquery-table-result-prefetch-%d") + .build()); + @AutoValue.Builder public abstract static class Builder { public abstract TableResult.Builder setSchema(Schema schema); @@ -206,7 +221,7 @@ public TableResult getNextPage() { @Override public Iterable iterateAll() { - return addSchema(getPageNoSchema().iterateAll()); + return addSchema(prefetchIterateAll(getPageNoSchema())); } @Override @@ -214,6 +229,80 @@ public Iterable getValues() { return addSchema(getPageNoSchema().getValues()); } + private static Iterable prefetchIterateAll(final Page firstPage) { + return new Iterable() { + @Override + public Iterator iterator() { + return new PrefetchingPageIterator(firstPage); + } + }; + } + + private static final class PrefetchingPageIterator extends AbstractIterator { + private Page currentPage; + private Iterator currentPageIterator; + private @Nullable Future> nextPageFuture; + private boolean prefetchStartedForCurrentPage; + + PrefetchingPageIterator(Page firstPage) { + this.currentPage = firstPage; + this.currentPageIterator = firstPage.getValues().iterator(); + } + + @Override + protected FieldValueList computeNext() { + while (!currentPageIterator.hasNext()) { + if (!prefetchStartedForCurrentPage) { + currentPage = currentPage.hasNextPage() ? currentPage.getNextPage() : null; + } else if (nextPageFuture != null) { + currentPage = getPrefetchedPage(nextPageFuture); + nextPageFuture = null; + } else { + currentPage = null; + } + prefetchStartedForCurrentPage = false; + if (currentPage == null) { + return endOfData(); + } + currentPageIterator = currentPage.getValues().iterator(); + } + if (!prefetchStartedForCurrentPage) { + prefetchStartedForCurrentPage = true; + if (currentPage.hasNextPage()) { + final Page pageToFetchFrom = currentPage; + nextPageFuture = + PREFETCH_EXECUTOR.submit( + new Callable>() { + @Override + public Page call() { + return pageToFetchFrom.getNextPage(); + } + }); + } + } + return currentPageIterator.next(); + } + + private static Page getPrefetchedPage(Future> future) { + try { + return future.get(); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new RuntimeException(cause); + } catch (InterruptedException e) { + future.cancel(true); + Thread.currentThread().interrupt(); + throw new BigQueryException(0, "Interrupted while prefetching next page of results", e); + } + } + } + private Iterable addSchema(Iterable iter) { if (getSchema() == null) { return iter; diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java index b830ee8e030f..510a1e173cfd 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java @@ -18,6 +18,7 @@ import static com.google.cloud.bigquery.FieldValue.Attribute.PRIMITIVE; import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; import com.google.api.gax.paging.Page; import com.google.cloud.PageImpl; @@ -25,6 +26,10 @@ import com.google.cloud.bigquery.JobStatistics.QueryStatistics.StatementType; import com.google.cloud.bigquery.JobStatistics.SessionInfo; import com.google.common.collect.ImmutableList; +import java.util.Iterator; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.Test; class TableResultTest { @@ -260,4 +265,112 @@ void testEqualsAndHashCode() { assertThat(result1.toString()).contains("cacheHit=true"); assertThat(result1.toString()).contains("sessionId=" + SESSION_ID_1); } + + @Test + void testIterateAllPrefetchesNextPageAsynchronously() throws InterruptedException { + final CountDownLatch prefetchStarted = new CountDownLatch(1); + final CountDownLatch allowPage1Return = new CountDownLatch(1); + final Page page1 = + new PageImpl<>(null, null, ImmutableList.of(newFieldValueList("2"))); + Page page0 = + new PageImpl<>( + new PageImpl.NextPageFetcher() { + @Override + public Page getNextPage() { + prefetchStarted.countDown(); + try { + allowPage1Return.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return page1; + } + }, + "cursor-1", + ImmutableList.of(newFieldValueList("0"), newFieldValueList("1"))); + + TableResult result = + TableResult.newBuilder() + .setSchema(SCHEMA) + .setTotalRows(3L) + .setPageNoSchema(page0) + .setRowsInPage(2L) + .build(); + + Iterator iterator = result.iterateAll().iterator(); + assertThat(iterator.next()).isEqualTo(newFieldValueList("0").withSchema(SCHEMA.getFields())); + // While still on page 0, page 1 prefetch should already have started in the background. + assertThat(prefetchStarted.await(5, TimeUnit.SECONDS)).isTrue(); + allowPage1Return.countDown(); + assertThat(iterator.next()).isEqualTo(newFieldValueList("1").withSchema(SCHEMA.getFields())); + assertThat(iterator.next()).isEqualTo(newFieldValueList("2").withSchema(SCHEMA.getFields())); + assertThat(iterator.hasNext()).isFalse(); + } + + @Test + void testIterateAllBoundsLookaheadToOnePage() throws InterruptedException { + final AtomicInteger fetchedPages = new AtomicInteger(0); + final CountDownLatch page1Fetched = new CountDownLatch(1); + final Page page2 = + new PageImpl<>(null, null, ImmutableList.of(newFieldValueList("2"))); + final Page page1 = + new PageImpl<>( + new PageImpl.NextPageFetcher() { + @Override + public Page getNextPage() { + fetchedPages.incrementAndGet(); + return page2; + } + }, + "cursor-2", + ImmutableList.of(newFieldValueList("1"))); + Page page0 = + new PageImpl<>( + new PageImpl.NextPageFetcher() { + @Override + public Page getNextPage() { + fetchedPages.incrementAndGet(); + page1Fetched.countDown(); + return page1; + } + }, + "cursor-1", + ImmutableList.of(newFieldValueList("0"))); + + TableResult result = + TableResult.newBuilder().setTotalRows(3L).setPageNoSchema(page0).setRowsInPage(1L).build(); + + Iterator iterator = result.iterateAll().iterator(); + assertThat(iterator.next()).isEqualTo(newFieldValueList("0")); + assertThat(page1Fetched.await(5, TimeUnit.SECONDS)).isTrue(); + // Only page 1 should be prefetched while consuming page 0. Page 2 must not be fetched yet. + assertThat(fetchedPages.get()).isEqualTo(1); + assertThat(iterator.next()).isEqualTo(newFieldValueList("1")); + assertThat(iterator.next()).isEqualTo(newFieldValueList("2")); + assertThat(iterator.hasNext()).isFalse(); + assertThat(fetchedPages.get()).isEqualTo(2); + } + + @Test + void testIterateAllPropagatesExceptionAtPageBoundary() { + Page page0 = + new PageImpl<>( + new PageImpl.NextPageFetcher() { + @Override + public Page getNextPage() { + throw new BigQueryException(500, "Failed to fetch page 1"); + } + }, + "cursor-1", + ImmutableList.of(newFieldValueList("0"))); + + TableResult result = + TableResult.newBuilder().setTotalRows(2L).setPageNoSchema(page0).setRowsInPage(1L).build(); + + Iterator iterator = result.iterateAll().iterator(); + assertThat(iterator.hasNext()).isTrue(); + assertThat(iterator.next()).isEqualTo(newFieldValueList("0")); + BigQueryException thrown = assertThrows(BigQueryException.class, iterator::hasNext); + assertThat(thrown).hasMessageThat().contains("Failed to fetch page 1"); + } }