Skip to content

perf(bigquery): [T2-1] add bounded 1-page lookahead prefetching to TableResult.iterateAll() #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lqiu96 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(bigquery): [T2-1] add bounded 1-page lookahead prefetching to TableResult.iterateAll() #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 90 additions & 1 deletion java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableResult.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,17 +23,32 @@
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
public abstract class TableResult implements Page<FieldValueList>, 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);
Expand Down Expand Up @@ -206,14 +221,88 @@ public TableResult getNextPage() {

@Override
public Iterable<FieldValueList> iterateAll() {
return addSchema(getPageNoSchema().iterateAll());
return addSchema(prefetchIterateAll(getPageNoSchema()));
}

@Override
public Iterable<FieldValueList> getValues() {
return addSchema(getPageNoSchema().getValues());
}

private static Iterable<FieldValueList> prefetchIterateAll(final Page<FieldValueList> firstPage) {
return new Iterable<FieldValueList>() {
@Override
public Iterator<FieldValueList> iterator() {
return new PrefetchingPageIterator(firstPage);
}
};
}

private static final class PrefetchingPageIterator extends AbstractIterator<FieldValueList> {
private Page<FieldValueList> currentPage;
private Iterator<FieldValueList> currentPageIterator;
private @Nullable Future<Page<FieldValueList>> nextPageFuture;
private boolean prefetchStartedForCurrentPage;

PrefetchingPageIterator(Page<FieldValueList> 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<FieldValueList> pageToFetchFrom = currentPage;
nextPageFuture =
PREFETCH_EXECUTOR.submit(
new Callable<Page<FieldValueList>>() {
@Override
public Page<FieldValueList> call() {
return pageToFetchFrom.getNextPage();
}
});
}
}
return currentPageIterator.next();
}

private static Page<FieldValueList> getPrefetchedPage(Future<Page<FieldValueList>> 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<FieldValueList> addSchema(Iterable<FieldValueList> iter) {
if (getSchema() == null) {
return iter;
Expand Down
113 changes: 113 additions & 0 deletions ...gquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/TableResultTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,18 @@

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;
import com.google.cloud.bigquery.JobStatistics.QueryStatistics;
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 {
Expand Down Expand Up @@ -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<FieldValueList> page1 =
new PageImpl<>(null, null, ImmutableList.of(newFieldValueList("2")));
Page<FieldValueList> page0 =
new PageImpl<>(
new PageImpl.NextPageFetcher<FieldValueList>() {
@Override
public Page<FieldValueList> 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<FieldValueList> 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<FieldValueList> page2 =
new PageImpl<>(null, null, ImmutableList.of(newFieldValueList("2")));
final Page<FieldValueList> page1 =
new PageImpl<>(
new PageImpl.NextPageFetcher<FieldValueList>() {
@Override
public Page<FieldValueList> getNextPage() {
fetchedPages.incrementAndGet();
return page2;
}
},
"cursor-2",
ImmutableList.of(newFieldValueList("1")));
Page<FieldValueList> page0 =
new PageImpl<>(
new PageImpl.NextPageFetcher<FieldValueList>() {
@Override
public Page<FieldValueList> 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<FieldValueList> 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<FieldValueList> page0 =
new PageImpl<>(
new PageImpl.NextPageFetcher<FieldValueList>() {
@Override
public Page<FieldValueList> 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<FieldValueList> 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");
}
}
Loading