added batch/interactive flag, removed samples (#6953)

broadinstitute · Jan 29, 2021 · 90f756e · 90f756e
1 parent 8634ac3
commit 90f756e
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 6 deletions.
diff --git a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeatures.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeatures.java
@@ -54,6 +54,12 @@ public class ExtractFeatures extends ExtractTool {
     )
     protected boolean trainingSitesOnly = false;
 
+    @Argument(
+        fullName = "use-batch-queries",
+        doc = "If true, use batch (rather than interactive) priority queries in BigQuery",
+        optional = true)
+    protected boolean useBatchQueries = true;
+
     @Override
     public boolean requiresIntervals() {
         return true; // TODO -- do I need to check the boolean flag on this?
@@ -64,10 +70,8 @@ protected void onStartup() {
         super.onStartup();
 
         TableReference sampleTableRef = new TableReference(sampleTableName, SchemaUtils.SAMPLE_FIELDS);
-        SampleList sampleList = new SampleList(sampleTableName, null, printDebugInformation);
-        Set<String> sampleNames = new HashSet<>(sampleList.getSampleNames());
 
-        VCFHeader header = CommonCode.generateVcfHeader(sampleNames, reference.getSequenceDictionary());
+        VCFHeader header = CommonCode.generateVcfHeader(new HashSet<>(), reference.getSequenceDictionary());
 
         engine = new ExtractFeaturesEngine(
             projectID,
@@ -82,6 +86,7 @@ protected void onStartup() {
             intervalArgumentCollection.getIntervals(getBestAvailableSequenceDictionary()),
             localSortMaxRecordsInRam,
             printDebugInformation,
+            useBatchQueries,
             progressMeter);
         vcfWriter.writeHeader(header);
 }

diff --git a/...ain/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeaturesEngine.java b/...ain/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeaturesEngine.java
@@ -11,7 +11,6 @@
 import org.broadinstitute.hellbender.engine.ProgressMeter;
 import org.broadinstitute.hellbender.engine.ReferenceDataSource;
 import org.broadinstitute.hellbender.exceptions.UserException;
-import org.broadinstitute.hellbender.tools.variantdb.CommonCode;
 import org.broadinstitute.hellbender.tools.variantdb.SchemaUtils;
 import org.broadinstitute.hellbender.tools.walkers.ReferenceConfidenceVariantContextMerger;
 import org.broadinstitute.hellbender.tools.walkers.annotator.FisherStrand;
@@ -52,6 +51,7 @@ public class ExtractFeaturesEngine {
     private final TableReference sampleListTable;
     private final ProgressMeter progressMeter;
     private List<SimpleInterval> userIntervals;
+    private final boolean useBatchQueries;
 
 //    /** Set of sample names seen in the variant data from BigQuery. */
 //    private final Set<String> sampleNames = new HashSet<>();
@@ -68,6 +68,7 @@ public ExtractFeaturesEngine(final String projectID,
                                final List<SimpleInterval> userIntervals,
                                final int localSortMaxRecordsInRam,
                                final boolean printDebugInformation,
+                               final boolean useBatchQueries,
                                final ProgressMeter progressMeter) {
         this.localSortMaxRecordsInRam = localSortMaxRecordsInRam;
 
@@ -79,6 +80,7 @@ public ExtractFeaturesEngine(final String projectID,
         this.vetTable = new TableReference(fqVetTable, SchemaUtils.VET_FIELDS);
         this.sampleListTable = sampleListTable;
         this.printDebugInformation = printDebugInformation;
+        this.useBatchQueries = useBatchQueries;
         this.progressMeter = progressMeter;
         this.userIntervals = userIntervals;
 
@@ -89,7 +91,7 @@ public void traverse() {
         userIntervals.forEach(interval -> {
             final String featureQueryString = ExtractFeaturesBQ.getVQSRFeatureExtractQueryString(vetTable, altAlleleTable, sampleListTable, interval, trainingSitesOnly);
             logger.info(featureQueryString);
-            final StorageAPIAvroReader storageAPIAvroReader = BigQueryUtils.executeQueryWithStorageAPI(featureQueryString, SchemaUtils.FEATURE_EXTRACT_FIELDS, projectID, true);
+            final StorageAPIAvroReader storageAPIAvroReader = BigQueryUtils.executeQueryWithStorageAPI(featureQueryString, SchemaUtils.FEATURE_EXTRACT_FIELDS, projectID, useBatchQueries);
 
             createVQSRInputFromTableResult(storageAPIAvroReader, interval.getContig());
         });

diff --git a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractTool.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractTool.java
@@ -52,7 +52,7 @@ public enum QueryMode {
     @Argument(
             fullName = "sample-table",
             doc = "Fully qualified name of a bigquery table containing a single column `sample` that describes the full list of samples to extract",
-            optional = true
+            optional = false
     )
     protected String sampleTableName = null;