Create ReadSession asynchronously (#928)

* Create ReadSession asynchronously * moving executor to private member
GoogleCloudDataproc · Mar 20, 2023 · ae2a5eb · ae2a5eb
1 parent 5f83d7c
commit ae2a5eb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/...uery-lib/src/main/java/com/google/cloud/spark/bigquery/v2/Spark31BigQueryScanBuilder.java b/...uery-lib/src/main/java/com/google/cloud/spark/bigquery/v2/Spark31BigQueryScanBuilder.java
@@ -63,6 +63,9 @@ public TableId getTableId() {
 
   @Override
   public Scan build() {
+    // start creating ReadSession asynchronously
+    // no further pushFilters / pruneColumns expected from this point
+    ctx.build();
     return this; // new BigQueryScan(ctx);
   }
 

diff --git a/...main/java/com/google/cloud/spark/bigquery/v2/context/BigQueryDataSourceReaderContext.java b/...main/java/com/google/cloud/spark/bigquery/v2/context/BigQueryDataSourceReaderContext.java
@@ -50,6 +50,8 @@
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -109,6 +111,7 @@ public OptionalLong numRows() {
   // "planBatchInputPartitionContexts" or
   // "planInputPartitionContexts". We will use this to get table statistics in estimateStatistics.
   private Supplier<ReadSessionResponse> readSessionResponse;
+  private final ExecutorService asyncReadSessionExecutor = Executors.newSingleThreadExecutor();
 
   public BigQueryDataSourceReaderContext(
       TableInfo table,
@@ -451,4 +454,10 @@ public BigQueryRDDFactory getBigQueryRddFactory() {
   public TableInfo getTableInfo() {
     return this.table;
   }
+
+  public void build() {
+    // Supplier provided by Suppliers.memoize is thread-safe
+    asyncReadSessionExecutor.submit(() -> readSessionResponse.get());
+    asyncReadSessionExecutor.shutdown();
+  }
 }