Skip to content

Commit

Permalink
storage: Allow include and exclude when querying annotations #830
Browse files Browse the repository at this point in the history
  • Loading branch information
j-coll committed May 1, 2018
1 parent 945dc4e commit 5a3db38
Show file tree
Hide file tree
Showing 14 changed files with 104 additions and 25 deletions.
Expand Up @@ -398,6 +398,8 @@ private void queryAnnotation() throws CatalogException, IOException, StorageEngi
QueryOptions options = new QueryOptions();
options.put(QueryOptions.LIMIT, cliOptions.limit);
options.put(QueryOptions.SKIP, cliOptions.skip);
options.put(QueryOptions.INCLUDE, cliOptions.dataModelOptions.include);
options.put(QueryOptions.EXCLUDE, cliOptions.dataModelOptions.exclude);
options.putAll(cliOptions.commonOptions.params);

Query query = new Query();
Expand Down
Expand Up @@ -578,6 +578,9 @@ public class QueryAnnotationCommandOptions extends StorageVariantCommandOptions.
@ParametersDelegate
public GeneralCliOptions.CommonCommandOptions commonOptions = commonCommandOptions;

@ParametersDelegate
public GeneralCliOptions.DataModelOptions dataModelOptions = new DataModelOptions();

@Parameter(names = {"-p", "--project"}, description = PROJECT_DESC, arity = 1)
public String project;

Expand Down
Expand Up @@ -376,6 +376,9 @@ public QueryResponse<VariantAnnotation> queryAnnotation() throws IOException {
QueryOptions options = new QueryOptions();
options.put(QueryOptions.LIMIT, cliOptions.limit);
options.put(QueryOptions.SKIP, cliOptions.skip);
options.put(QueryOptions.INCLUDE, cliOptions.dataModelOptions.include);
options.put(QueryOptions.EXCLUDE, cliOptions.dataModelOptions.exclude);
options.put(QueryOptions.SKIP, cliOptions.skip);
options.putAll(cliOptions.commonOptions.params);

Query query = new Query();
Expand Down
Expand Up @@ -470,8 +470,8 @@ public Response getVariants(@ApiParam(name = "params", value = "Query parameters
@Path("/annotation/{name}/query")
@ApiOperation(value = "", position = 15, response = VariantAnnotation[].class)
@ApiImplicitParams({
// @ApiImplicitParam(name = QueryOptions.INCLUDE, value = "Fields included in the response, whole JSON path must be provided", example = "name,attributes", dataType = "string", paramType = "query"),
// @ApiImplicitParam(name = QueryOptions.EXCLUDE, value = "Fields excluded in the response, whole JSON path must be provided", example = "id,status", dataType = "string", paramType = "query"),
@ApiImplicitParam(name = QueryOptions.INCLUDE, value = "Fields included in the response, whole JSON path must be provided", example = "name,attributes", dataType = "string", paramType = "query"),
@ApiImplicitParam(name = QueryOptions.EXCLUDE, value = "Fields excluded in the response, whole JSON path must be provided", example = "id,status", dataType = "string", paramType = "query"),
@ApiImplicitParam(name = QueryOptions.LIMIT, value = "Number of results to be returned in the queries", dataType = "integer", paramType = "query"),
@ApiImplicitParam(name = QueryOptions.SKIP, value = "Number of results to skip in the queries", dataType = "integer", paramType = "query"),
// @ApiImplicitParam(name = QueryOptions.COUNT, value = "Total number of results", dataType = "boolean", paramType = "query")
Expand Down
Expand Up @@ -153,20 +153,14 @@ public class IndexCommandOptions {

}

public class QueryCommandOptions {
public class QueryCommandOptions extends DataModelOptions{

@Parameter(names = {"-o", "--output"}, description = "Output file. [STDOUT]", required = false, arity = 1)
public String output;

@Parameter(names = {"-d", "--database"}, description = "DataBase name", required = false, arity = 1)
public String dbName;

@Parameter(names = {"-i", "--include"}, description = "", required = false, arity = 1)
public String include;

@Parameter(names = {"-e", "--exclude"}, description = "", required = false, arity = 1)
public String exclude;

@Parameter(names = {"--skip"}, description = "Skip some number of elements.", required = false, arity = 1)
public int skip;

Expand All @@ -178,4 +172,15 @@ public class QueryCommandOptions {

}

public static class DataModelOptions {

@Parameter(names = {"-I", "--include"}, description = "Comma separated list of fields to be included in the response", arity = 1)
public String include;

@Parameter(names = {"-E", "--exclude"}, description = "Comma separated list of fields to be excluded from the response", arity = 1)
public String exclude;

}


}
Expand Up @@ -480,6 +480,8 @@ private void queryAnnotation() throws VariantAnnotatorException, StorageEngineEx
QueryOptions options = new QueryOptions();
options.put(QueryOptions.LIMIT, cliOptions.limit);
options.put(QueryOptions.SKIP, cliOptions.skip);
options.put(QueryOptions.INCLUDE, cliOptions.dataModelOptions.include);
options.put(QueryOptions.EXCLUDE, cliOptions.dataModelOptions.exclude);
options.putAll(cliOptions.commonOptions.params);

Query query = new Query();
Expand Down
Expand Up @@ -570,6 +570,9 @@ public class QueryAnnotationCommandOptions extends GenericQueryAnnotationCommand
@ParametersDelegate
public GeneralCliOptions.CommonOptions commonOptions = commonCommandOptions;

@ParametersDelegate
public GeneralCliOptions.DataModelOptions dataModelOptions = new GeneralCliOptions.DataModelOptions();

@Parameter(names = {"-d", "--database"}, description = "DataBase name", required = true, arity = 1)
public String dbName;

Expand Down
Expand Up @@ -279,7 +279,8 @@ public void deleteAnnotationSnapshot(String project, String annotationName, Obje

public QueryResult<VariantAnnotation> getAnnotation(String name, Query query, QueryOptions options, String sessionId)
throws StorageEngineException, CatalogException, IOException {
return secure(query, options, sessionId, (engine) -> engine.getAnnotation(name, query, options));
QueryOptions finalOptions = VariantQueryUtils.validateAnnotationQueryOptions(options);
return secure(query, finalOptions, sessionId, (engine) -> engine.getAnnotation(name, query, finalOptions));
}

public void stats(String study, List<String> cohorts, String outDir, ObjectMap config, String sessionId)
Expand Down
Expand Up @@ -131,10 +131,6 @@ public static VariantField get(String field) {
return getNamesMap().get(field);
}

public static Set<VariantField> getIncludeFields(QueryOptions options) {
return getIncludeFields(options, false);
}

/**
* Given a QueryOptions, reads the {@link QueryOptions#INCLUDE}, {@link QueryOptions#EXCLUDE} and {@link VariantField#SUMMARY}
* to determine which fields from Variant will be included.
Expand All @@ -144,10 +140,9 @@ public static Set<VariantField> getIncludeFields(QueryOptions options) {
* 3. {@link VariantField#SUMMARY}
*
* @param options Non null options
* @param prune Remove intermediate nodes some child is missing, or all children from a node if all are present
* @return List of fields to be included.
*/
public static Set<VariantField> getIncludeFields(QueryOptions options, boolean prune) {
public static Set<VariantField> getIncludeFields(QueryOptions options) {
Set<VariantField> includeFields;

if (options == null) {
Expand Down Expand Up @@ -188,11 +183,7 @@ public static Set<VariantField> getIncludeFields(QueryOptions options, boolean p
}
}

if (prune) {
return prune(includeFields);
} else {
return includeFields;
}
return includeFields;
}

/**
Expand Down
Expand Up @@ -128,6 +128,10 @@ public static VariantQueryException unknownVariantField(String projectionOp, Str
return new VariantQueryException("Found unknown variant field '" + field + "' in " + projectionOp.toLowerCase());
}

public static VariantQueryException unknownVariantAnnotationField(String projectionOp, String field) {
return new VariantQueryException("Found unknown variant annotation field '" + field + "' in " + projectionOp.toLowerCase());
}

public static VariantQueryException internalException(Exception e) {
return new VariantQueryException("Internal exception: " + e.getMessage(), e);
}
Expand Down
Expand Up @@ -159,6 +159,49 @@ public static void validateAnnotationQuery(Query query) {
}
}

public static QueryOptions validateAnnotationQueryOptions(QueryOptions queryOptions) {
if (queryOptions == null) {
return new QueryOptions(QueryOptions.INCLUDE, VariantField.ANNOTATION);
} else {
queryOptions = new QueryOptions(queryOptions);
}

boolean anyPresent = transformVariantAnnotationField(QueryOptions.INCLUDE, queryOptions);
anyPresent |= transformVariantAnnotationField(QueryOptions.EXCLUDE, queryOptions);

if (!anyPresent) {
queryOptions.add(QueryOptions.INCLUDE, VariantField.ANNOTATION);
}

return queryOptions;
}

private static boolean transformVariantAnnotationField(String key, QueryOptions queryOptions) {
StringBuilder sb = new StringBuilder();
final String annotation = VariantField.ANNOTATION.fieldName();
for (String field : queryOptions.getAsStringList(key)) {
String newField;
if (!field.startsWith(annotation + '.')) {
newField = annotation + '.' + field;
} else {
newField = field;
}

if (VariantField.get(newField) == null) {
throw VariantQueryException.unknownVariantAnnotationField(key, field);
}

sb.append(newField);
sb.append(',');
}
if (sb.length() > 0) {
queryOptions.put(key, sb.toString());
return true;
} else {
return false;
}
}

/**
* Determines if the filter is negated.
*
Expand Down
Expand Up @@ -2,6 +2,7 @@

import org.junit.Test;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.models.variant.avro.ConsequenceType;
import org.opencb.biodata.models.variant.avro.VariantAnnotation;
import org.opencb.commons.datastore.core.ObjectMap;
import org.opencb.commons.datastore.core.Query;
Expand All @@ -10,14 +11,17 @@
import org.opencb.opencga.storage.core.exceptions.StorageEngineException;
import org.opencb.opencga.storage.core.variant.VariantStorageBaseTest;
import org.opencb.opencga.storage.core.variant.VariantStorageEngine;
import org.opencb.opencga.storage.core.variant.adaptors.VariantField;
import org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam;
import org.opencb.opencga.storage.core.variant.annotation.annotators.VariantAnnotator;
import org.opencb.opencga.storage.core.variant.annotation.annotators.VariantAnnotatorFactory;

import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.opencb.opencga.storage.core.variant.annotation.VariantAnnotationManager.ANNOTATOR;
import static org.opencb.opencga.storage.core.variant.annotation.VariantAnnotationManager.VARIANT_ANNOTATOR_CLASSNAME;

Expand Down Expand Up @@ -74,6 +78,14 @@ public void testQueries(VariantStorageEngine variantStorageEngine) throws Storag
assertEquals(count, partialCount);
}

String consequenceTypes = VariantField.ANNOTATION_CONSEQUENCE_TYPES.fieldName().replace(VariantField.ANNOTATION.fieldName() + ".", "");
for (VariantAnnotation annotation : variantStorageEngine.getAnnotation("v2", null, new QueryOptions(QueryOptions.INCLUDE, consequenceTypes)).getResult()) {
assertEquals(1, annotation.getConsequenceTypes().size());
}
for (VariantAnnotation annotation : variantStorageEngine.getAnnotation("v2", null, new QueryOptions(QueryOptions.EXCLUDE, consequenceTypes)).getResult()) {
assertTrue(annotation.getConsequenceTypes() == null || annotation.getConsequenceTypes().isEmpty());
}


// Get annotations from a deleted snapshot
// FIXME: Should throw an exception?
Expand Down Expand Up @@ -109,6 +121,12 @@ public List<VariantAnnotation> annotate(List<Variant> variants) throws VariantAn
a.setReference(v.getReference());
a.setAlternate(v.getAlternate());
a.setId(key);
ConsequenceType ct = new ConsequenceType();
ct.setGeneName("a gene");
ct.setSequenceOntologyTerms(Collections.emptyList());
ct.setExonOverlap(Collections.emptyList());
ct.setTranscriptAnnotationFlags(Collections.emptyList());
a.setConsequenceTypes(Collections.singletonList(ct));
return a;
}).collect(Collectors.toList());
}
Expand Down
Expand Up @@ -294,13 +294,14 @@ public QueryResult<VariantAnnotation> getAnnotation(String name, Query query, Qu

public Iterator<VariantAnnotation> annotationIterator(String name, Query query, QueryOptions options) {
query = query == null ? new Query() : new Query(query);
options = options == null ? new QueryOptions() : new QueryOptions(options);
options = validateAnnotationQueryOptions(options);
validateAnnotationQuery(query);
if (name.equals("LATEST")) {
name = "FULL";
}
query.put(ANNOT_NAME.key(), name);
List<Scan> scans = hbaseQueryParser.parseQueryMultiRegion(query, options.append(QueryOptions.INCLUDE, VariantField.ANNOTATION));
SelectVariantElements selectElements = VariantQueryUtils.parseSelectElements(query, options, getStudyConfigurationManager());
List<Scan> scans = hbaseQueryParser.parseQueryMultiRegion(selectElements, query, options);

try {
Table table = getConnection().getTable(TableName.valueOf(variantTable));
Expand All @@ -311,7 +312,8 @@ public Iterator<VariantAnnotation> annotationIterator(String name, Query query,
throw VariantQueryException.internalException(e);
}
}).iterator();
HBaseToVariantAnnotationConverter converter = new HBaseToVariantAnnotationConverter(genomeHelper);
HBaseToVariantAnnotationConverter converter = new HBaseToVariantAnnotationConverter(genomeHelper)
.setIncludeFields(selectElements.getFields());
converter.setAnnotationColumn(Bytes.toBytes(VariantPhoenixHelper.getAnnotationSnapshotColumn(name)));
Iterator<Result> iterator = Iterators.concat(iterators);
int skip = options.getInt(QueryOptions.SKIP);
Expand Down
Expand Up @@ -482,7 +482,9 @@ public VariantQueryResult<Variant> getPhased(String varStr, String studyName, St
public QueryResult<VariantAnnotation> getAnnotation(String name, Query query, QueryOptions options) {
query = query == null ? new Query() : query;
validateAnnotationQuery(query);
options = validateAnnotationQueryOptions(options);
Document mongoQuery = queryParser.parseQuery(query);
Document projection = queryParser.createProjection(query, options);

MongoDBCollection annotationCollection;
if (name.equals("LATEST")) {
Expand All @@ -495,7 +497,7 @@ public QueryResult<VariantAnnotation> getAnnotation(String name, Query query, Qu
query, new QueryOptions(QueryOptions.INCLUDE, VariantField.ANNOTATION), studyConfigurationManager);

DocumentToVariantConverter converter = getDocumentToVariantConverter(new Query(), selectVariantElements);
QueryResult<Variant> result = annotationCollection.find(mongoQuery, converter, options);
QueryResult<Variant> result = annotationCollection.find(mongoQuery, projection, converter, options);

List<VariantAnnotation> annotations = result.getResult()
.stream()
Expand Down

0 comments on commit 5a3db38

Please sign in to comment.