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 index f4b6add99cb..910f7488177 100644 --- 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 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/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeaturesEngine.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeaturesEngine.java index 572ee573d5b..26c3b2e63ab 100644 --- a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractFeaturesEngine.java +++ b/src/main/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 userIntervals; + private final boolean useBatchQueries; // /** Set of sample names seen in the variant data from BigQuery. */ // private final Set sampleNames = new HashSet<>(); @@ -68,6 +68,7 @@ public ExtractFeaturesEngine(final String projectID, final List 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 index 2272edfd26c..bbae65e093e 100644 --- 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;