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 load_meta_and_data_by_id(datamart_id: int, first_n_rows: int = None, constrains=None): qm = QueryManager(es_host=ES_HOST, es_port=ES_PORT, es_index=PRODUCTION_ES_INDEX) res = qm.get_by_id(datamart_id) if res and res.get('_source'): df = Utils.get_dataset(res['_source'], constrains=constrains) if first_n_rows: df = df.head(first_n_rows) return res['_source'], df return None, None
def test_match_key_value_pairs_list(self): query = QueryManager.match_key_value_pairs(key_value_pairs=[( "description", ["average", "temperature"]), ("title.keyword", "TAVG"), ("variables.datamart_id", 0)]) expected = { "query": { "bool": { "must": [{ "terms": { "description": ["average", "temperature"] } }, { "term": { "title.keyword": "TAVG" } }, { "term": { "variables.datamart_id": 0 } }] } } } self.assertEqual(json.dumps(expected), query)
def test_match_temporal_coverage(self): query = QueryManager.match_temporal_coverage(start="2018-09-23", end="2018-09-30T00:00:00") expected = { "query": { "nested": { "path": "variables", "inner_hits": {"_source": ["temporal_coverage"]}, "query": { "bool": { "must": [ { "range": { "variables.temporal_coverage.start": { "lte": "2018-09-23T00:00:00", "format": "yyyy-MM-dd'T'HH:mm:ss" } } }, { "range": { "variables.temporal_coverage.end": { "gte": "2018-09-30T00:00:00", "format": "yyyy-MM-dd'T'HH:mm:ss" } } } ] } } } } } self.assertEqual(json.dumps(expected), query)
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)
def test_match_any(self): query = QueryManager.match_any(query_string="los angeles average") expected = { "query_string": { "query": "los angeles average" } } self.assertEqual(expected, query)
def test_match_global_datamart_id(self): query = QueryManager.match_global_datamart_id(datamart_id=0) expected = { "term": { "datamart_id": 0 } } self.assertEqual(expected, query)
def test_match_variable_datamart_id(self): query = QueryManager.match_variable_datamart_id(datamart_id=0) expected = { "nested": { "path": "variables", "inner_hits": {"_source": ["datamart_id"]}, "query": {"bool": {"must": [{"term": {"variables.datamart_id": 0}}]}}} } self.assertEqual(expected, query)
def test_match_match_global_datamart_id(self): query = QueryManager.match_global_datamart_id(datamart_id=0) expected = { "query": { "bool": { "must": [{"term": {"datamart_id": 0}}] } } } self.assertEqual(json.dumps(expected), query)
def test_match_all(self): query = QueryManager.match_all() expected = { "query": { "bool": { "must": [{"match_all": {}}] } } } self.assertEqual(json.dumps(expected), query)
def test_match_key_value_pairs_list(self): query = QueryManager.match_key_value_pairs(key_value_pairs=[ ("description", ["average", "temperature"]), ("title.keyword", "TAVG"), ("datamart_id", 0) ]) expected = { "bool": { "must": [{"match": {"description": "average"}}, {"match": {"description": "temperature"}}, {"match": {"title.keyword": "TAVG"}}, {"match": {"datamart_id": 0}} ] } } self.assertEqual(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_variables_array(terms=terms) expected = { "nested": { "path": "variables", "inner_hits": { "_source": ["named_entity"], "highlight": {"fields": {"variables.named_entity": { "pre_tags": [""], "post_tags": [""], "number_of_fragments": 0}}} }, "query": { "bool": { "should": [ { "match_phrase": { "variables.named_entity": {"query": "los angeles", "_name": "los angeles"} } }, { "match_phrase": { "variables.named_entity": {"query": "new york", "_name": "new york"}} }, { "match_phrase": { "variables.named_entity": {"query": "akba", "_name": "akba"}} }, { "match_phrase": { "variables.named_entity": {"query": "shanghai", "_name": "shanghai"}} }, { "match_phrase": { "variables.named_entity": {"query": "tokyo", "_name": "tokyo"}} } ], "minimum_should_match": 3 } } } } self.assertEqual(expected, query)
def test_match_temporal_coverage_invalid(self): query = QueryManager.match_temporal_coverage(start="2222s", end="2018-09-30T00:00:00") expected = { "query": { "bool": { "must": [{ "range": { "variables.temporal_coverage.end": { "gte": "2018-09-30T00:00:00", "format": "yyyy-MM-dd'T'HH:mm:ss" } } }] } } } self.assertEqual(json.dumps(expected), query)
def test_match_some_terms_from_array(self): query = QueryManager.match_some_terms_from_variables_array(terms=["los angeles", "NEW YORK", "AKBA"], key="variables.named_entity", minimum_should_match=2) expected = { "query": { "nested": { "path": "variables", "inner_hits": { "_source": ["named_entity"] }, "query": { "bool": { "should": [ { "match_phrase": { "variables.named_entity": {"query": "los angeles", "_name": "los angeles"} } }, { "match_phrase": { "variables.named_entity": {"query": "new york", "_name": "new york"} } }, { "match_phrase": { "variables.named_entity": {"query": "akba", "_name": "akba"} } } ], "minimum_should_match": 2 } } } }, "highlight": {"fields": {"variables.named_entity": {"pre_tags": [""], "post_tags": [""]}}} } self.assertEqual(json.dumps(expected), query)
def test_match_key_value_pairs(self): query = QueryManager.match_key_value_pairs(key_value_pairs=[ ("description", "average"), ("title.keyword", "TAVG"), ("datamart_id", 0) ]) expected = { "query": { "bool": { "must": [ {"match": {"description": "average"}}, {"match": {"title.keyword": "TAVG"}}, {"match": {"datamart_id": 0}}, {"nested": { "path": "variables", "inner_hits": {"_source": []}, "query": {"bool": {"must": []}}} } ] } } } 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) self.joiners = dict() self.profiler = Profiler() def query(self, col: pd.Series = None, minimum_should_match_ratio_for_col: float = None, query_string: str = None, temporal_coverage_start: str = None, temporal_coverage_end: str = None, global_datamart_id: int = None, variable_datamart_id: int = None, key_value_pairs: typing.List[tuple] = None, **kwargs) -> typing.Optional[typing.List[dict]]: """Query metadata by a pandas Dataframe column Args: col: pandas Dataframe column. minimum_should_match_ratio_for_col: An float ranges from 0 to 1 indicating the ratio of unique value of the column to be matched query_string: string to query any field in metadata temporal_coverage_start: start of a temporal coverage temporal_coverage_end: end of a temporal coverage global_datamart_id: match a global metadata id variable_datamart_id: match a variable metadata id key_value_pairs: match key value pairs Returns: matching docs of metadata """ queries = list() if query_string: queries.append(self.qm.match_any(query_string=query_string)) if temporal_coverage_start or temporal_coverage_end: queries.append( self.qm.match_temporal_coverage(start=temporal_coverage_start, end=temporal_coverage_end)) if global_datamart_id: queries.append( self.qm.match_global_datamart_id( datamart_id=global_datamart_id)) if variable_datamart_id: queries.append( self.qm.match_variable_datamart_id( datamart_id=variable_datamart_id)) if key_value_pairs: queries.append( self.qm.match_key_value_pairs(key_value_pairs=key_value_pairs)) if col is not None: queries.append( self.qm.match_some_terms_from_variables_array( terms=col.unique().tolist(), minimum_should_match=minimum_should_match_ratio_for_col)) if not queries: return self._query_all() return self.qm.search(body=self.qm.form_conjunction_query(queries), **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) def _query_all(self, **kwargs) -> typing.Optional[typing.List[dict]]: """Query all metadata Args: Returns: matching docs of metadata """ return self.qm.search(body=self.qm.match_all(), **kwargs) def join(self, left_df: pd.DataFrame, right_df: pd.DataFrame, left_columns: typing.List[typing.List[int]], right_columns: typing.List[typing.List[int]], left_metadata: dict = None, right_metadata: dict = None, joiner: str = "default") -> typing.Optional[pd.DataFrame]: """Join two dataframes based on different joiner. Args: left_df: pandas Dataframe right_df: pandas Dataframe left_metadata: metadata of left dataframe right_metadata: metadata of right dataframe left_columns: list of integers from left df for join right_columns: list of integers from right df for join joiner: string of joiner, default to be "default" Returns: Dataframe """ if joiner not in self.joiners: self.joiners[joiner] = JoinerPrepare.prepare_joiner(joiner=joiner) if not self.joiners[joiner]: warnings.warn("No suitable joiner, return original dataframe") return left_df if not left_metadata: # Left df is the user provided one. # We will generate metadata just based on the data itself, profiling and so on left_metadata = Utils.generate_metadata_from_dataframe( data=left_df) left_metadata = Utils.calculate_dsbox_features(data=left_df, metadata=left_metadata) right_metadata = Utils.calculate_dsbox_features( data=right_df, metadata=right_metadata) return self.joiners[joiner].join( left_df=left_df, right_df=right_df, left_columns=left_columns, right_columns=right_columns, left_metadata=left_metadata, right_metadata=right_metadata, )
def check_existence(materialization: dict, es_index: str = PRODUCTION_ES_INDEX): """ Query ElasticSearch with materializer name and arguments. Currently, only works with "general_materializer" and "wikitables_materializer." return the datamart id if exists else return None :param url: :return: datamart_id or None """ materializer = materialization['python_path'] if materializer == "general_materializer": query = { "query": { "bool": { "must": [{ "match_phrase": { "materialization.python_path": materializer } }, { "match_phrase": { "materialization.arguments.url": materialization['arguments']['url'] } }, { "match_phrase": { "materialization.arguments.index": materialization['arguments']['index'] } }] } } } elif materializer == "wikitables_materializer": query = { "query": { "bool": { "must": [{ "match_phrase": { "materialization.python_path": materializer } }, { "match_phrase": { "materialization.arguments.url": materialization['arguments']['url'] } }, { "match_phrase": { "materialization.arguments.xpath": materialization['arguments']['xpath'] } }] } } } else: raise Exception( f'Do not know how to perform existence check for materializer: {materializer}' ) qm = QueryManager(es_host=ES_HOST, es_port=ES_PORT, es_index=es_index) res = qm.search(dumps(query)) # TODO: how about return many results, should raise warning if res and res[0]: return int(res[0].get('_id'))
def test_form_conjunction_query(self): query = QueryManager.form_conjunction_query( queries=[ {'query_string': {'query': 'Budget'}}, {'term': {'datamart_id': 125980000}}, {'bool': {'must': [{'match': {'title': 'Government'}}]}}, {'nested': { 'path': 'variables', 'inner_hits': {'_source': ['named_entity'], 'highlight': { 'fields': { 'variables.named_entity': { 'pre_tags': [ ''], 'post_tags': [ ''], 'number_of_fragments': 0} } } }, 'query': {'bool': { 'should': [ { 'match_phrase': { 'variables.named_entity': { 'query': 'russia', '_name': 'russia' } } }, { 'match_phrase': { 'variables.named_entity': { 'query': 'saudi arabia', '_name': 'saudi arabia'}}} ], 'minimum_should_match': 1 } } } } ] ) expected = {"query": {"bool": { "must": [ {"query_string": {"query": "Budget"}}, {"term": {"datamart_id": 125980000}}, {"bool": {"must": [{"match": {"title": "Government"}}]}}, {"nested": {"path": "variables", "inner_hits": {"_source": ["named_entity"], "highlight": { "fields": { "variables.named_entity": { "pre_tags": [ ""], "post_tags": [ ""], "number_of_fragments": 0}}}}, "query": {"bool": { "should": [{ "match_phrase": { "variables.named_entity": { "query": "russia", "_name": "russia"}}}, { "match_phrase": { "variables.named_entity": { "query": "saudi arabia", "_name": "saudi arabia"}}}], "minimum_should_match": 1}}}}]}}} self.assertEqual(json.dumps(expected), query)
class Augment(object): DEFAULT_START_DATE = "1900-01-01T00:00:00" 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) self.joiners = dict() 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_variables_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_variables_array( terms=named_entities, key="variables.named_entity", 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 """ if "date_range" in constrains: if not constrains["date_range"].get("start", None): constrains["date_range"]["start"] = Augment.DEFAULT_START_DATE if not constrains["date_range"].get("end", None): constrains["date_range"]["end"] = datetime.now().strftime( '%Y-%m-%dT%H:%M:%S') df = Utils.materialize(metadata=metadata, constrains=constrains) if variables: return df.iloc[:, variables] return df @staticmethod def get_metadata_intersection(*metadata_lst) -> list: """Get the intersect metadata list. Args: metadata_lst: all metadata list returned by multiple queries Returns: list of intersect metadata """ metadata_dict = dict() metadata_sets = [] for lst in metadata_lst: this_set = set() for x in lst: if x["_source"]["datamart_id"] not in metadata_dict: metadata_dict[x["_source"]["datamart_id"]] = x this_set.add(x["_source"]["datamart_id"]) metadata_sets.append(this_set) return [ metadata_dict[datamart_id] for datamart_id in metadata_sets[0].intersection( *metadata_sets[1:]) ] def join(self, left_df: pd.DataFrame, right_df: pd.DataFrame, left_columns: typing.List[int], right_columns: typing.List[int], left_metadata: dict = None, right_metadata: dict = None, joiner: str = "default") -> typing.Optional[pd.DataFrame]: """Join two dataframes based on different joiner. Args: left_df: pandas Dataframe right_df: pandas Dataframe left_metadata: metadata of left dataframe right_metadata: metadata of right dataframe left_columns: list of integers from left df for join right_columns: list of integers from right df for join joiner: string of joiner, default to be "default" Returns: Dataframe """ if joiner not in self.joiners: self.joiners[joiner] = JoinerPrepare.prepare_joiner(joiner=joiner) if not self.joiners[joiner]: warnings.warn("No suitable joiner, return original dataframe") return left_df return self.joiners[joiner].join( left_df=left_df, right_df=right_df, left_columns=left_columns, right_columns=right_columns, left_metadata=left_metadata, right_metadata=right_metadata, )
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)