def test_match_some_terms_from_array(self):
query = QueryManager.match_some_terms_from_array(
terms=["los angeles", "NEW YORK", "AKBA"],
key="variables.named_entity",
minimum_should_match=2)
expected = {
"query": {
"bool": {
"should": [{
"term": {
"variables.named_entity": "los angeles"
}
}, {
"term": {
"variables.named_entity": "new york"
}
}, {
"term": {
"variables.named_entity": "akba"
}
}],
"minimum_should_match":
2
}
}
}
self.assertEqual(json.dumps(expected), query)
def test_match_some_terms_from_array_default(self):
terms = ["los angeles", "NEW YORK", "AKBA", "shanghai", "TOKyo"]
query = QueryManager.match_some_terms_from_array(terms=terms)
expected = {
"query": {
"bool": {
"should": [{
"term": {
"variables.named_entity.keyword": "los angeles"
}
}, {
"term": {
"variables.named_entity.keyword": "new york"
}
}, {
"term": {
"variables.named_entity.keyword": "akba"
}
}, {
"term": {
"variables.named_entity.keyword": "shanghai"
}
}, {
"term": {
"variables.named_entity.keyword": "tokyo"
}
}],
"minimum_should_match":
math.ceil(len(terms) / 2)
}
}
}
self.assertEqual(json.dumps(expected), query)
class Augment(object):
def __init__(self, es_index: str, es_host: str = "dsbox02.isi.edu", es_port: int = 9200) -> None:
"""Init method of QuerySystem, set up connection to elastic search.
Args:
es_index: elastic search index.
es_host: es_host.
es_port: es_port.
Returns:
"""
self.qm = QueryManager(es_host=es_host, es_port=es_port, es_index=es_index)
def query_by_column(self,
col: pd.Series,
minimum_should_match: int = None,
**kwargs
) -> typing.Optional[typing.List[dict]]:
"""Query metadata by a pandas Dataframe column
Args:
col: pandas Dataframe column.
minimum_should_match: An integer ranges from 0 to length of unique value in col.
Represent the minimum number of terms should match.
Returns:
matching docs of metadata
"""
body = self.qm.match_some_terms_from_array(terms=col.unique().tolist(),
minimum_should_match=minimum_should_match)
return self.qm.search(body=body, **kwargs)
def query_by_named_entities(self,
named_entities: list,
minimum_should_match: int = None,
**kwargs
) -> typing.Optional[typing.List[dict]]:
"""Query metadata by a pandas Dataframe column
Args:
named_entities: list of named entities
minimum_should_match: An integer ranges from 0 to length of named entities list.
Represent the minimum number of terms should match.
Returns:
matching docs of metadata
"""
body = self.qm.match_some_terms_from_array(terms=named_entities,
key="variables.named_entity.keyword",
minimum_should_match=minimum_should_match)
return self.qm.search(body=body, **kwargs)
def query_by_temporal_coverage(self, start=None, end=None, **kwargs) -> typing.Optional[typing.List[dict]]:
"""Query metadata by a temporal coverage of column
Args:
start: dataset should cover date time earlier than the start date.
end: dataset should cover date time later than the end date.
Returns:
matching docs of metadata
"""
body = self.qm.match_temporal_coverage(start=start, end=end)
return self.qm.search(body=body, **kwargs)
def query_by_datamart_id(self, datamart_id: int, **kwargs) -> typing.Optional[typing.List[dict]]:
"""Query metadata by datamart id
Args:
datamart_id: int
Returns:
matching docs of metadata
"""
global_body = self.qm.match_global_datamart_id(datamart_id=datamart_id)
variable_body = self.qm.match_variable_datamart_id(datamart_id=datamart_id)
return self.qm.search(body=global_body, **kwargs) or self.qm.search(body=variable_body, **kwargs)
def query_by_key_value_pairs(self,
key_value_pairs: typing.List[tuple],
**kwargs
) -> typing.Optional[typing.List[dict]]:
"""Query metadata by datamart id
Args:
key_value_pairs: list of key value tuple
Returns:
matching docs of metadata
"""
body = self.qm.match_key_value_pairs(key_value_pairs=key_value_pairs)
return self.qm.search(body=body, **kwargs)
def query_any_field_with_string(self, query_string, **kwargs) -> typing.Optional[typing.List[dict]]:
"""Query any field of matadata with query_string
Args:
key_value_pairs: list of key value tuple
Returns:
matching docs of metadata
"""
body = self.qm.match_any(query_string=query_string)
return self.qm.search(body=body, **kwargs)
def query_by_es_query(self, body: str, **kwargs) -> typing.Optional[typing.List[dict]]:
"""Query metadata by an elastic search query
Args:
body: query body
Returns:
matching docs of metadata
"""
return self.qm.search(body=body, **kwargs)
@staticmethod
def get_dataset(metadata: dict, variables: list = None, constrains: dict = None) -> typing.Optional[pd.DataFrame]:
"""Get the dataset with materializer.
Args:
metadata: metadata dict.
variables:
constrains:
Returns:
pandas dataframe
"""
return Utils.materialize(metadata=metadata, variables=variables, constrains=constrains)