elasticsearch-ri/exercise3.py at main · DiegoMfer/elasticsearch-ri · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from elasticsearch import Elasticsearch, helpers
from exercise1 import save_results


def retrieve_data(client, index, big_query):
    data = helpers.scan(client, index=index, query=big_query)

    result_data = {}

    for piece in data:
        source = piece["_source"]
        result_data[source["user_id_str"]] = {
            "date": source["created_at"],
            "text": source["text"]
        }

    return result_data


def big_query(topics, first, language):
    return {
        "query": {
            "query_string": {
                "query": construct_query(topics) + first + " AND lang:" + language
            }
        }
    }

def construct_query(topics):
    query = "text:"
    for topic in topics:
        query += topic + " OR "

    return query

def initial_query(client, index, topic, language, metric, size):

    res = client.search(
        index = index,
        body = get_query(topic, language, metric, size)
    )
    return [bucket["key"] for bucket in res["aggregations"]["Trending topics"]["buckets"]]


def get_query(topic, language, metric, size):
    query = {
            "size": 0,
            "query": {
                "query_string": {
                    "query": "text:\"{}\" AND lang:{}".format(topic, language)
                }
            },
            "aggs": {
                "Trending topics": {
                    "significant_terms": {
                        "field": "text",
                        "size": size,
                        metric: {}
                    }
                }
            }
        }
    return query

def retrieve_docs(doc_size, client, index, big_query):
    big_query["size"] = doc_size
    data = client.search(index=index, body=big_query)

    res = {}
    for hit in data["hits"]["hits"]:
        source = hit["_source"]
        res[source["user_id_str"]] = {
            "created_at": source["created_at"],
            "text": source["text"]
        }

    return res


def main():
    client = Elasticsearch("http://localhost:9200")

    # Customize here the search script
    index = "tweets-20090624-20090626-en_es-10percent-v2"
    language = "en"
    metric = "gnd" # Metric to experiment for exercise 3
    size = 5

    # Exercise 3 parameters
    topic = "iran" # Change topic as desired
    file_name = "output3-" + metric + ".json"
    doc_size = 20 # Increment number of documents
    file_20docs_name = "output3-20docs-" + metric +".json"

    topics = initial_query(client, index, topic, language, metric, size)

    query = big_query(topics, topic, language)

    results = retrieve_data(client, index, query)

    save_results(results, file_name)

    docs = retrieve_docs(doc_size, client, index, query)

    save_results(docs, file_20docs_name )

if __name__ == '__main__':
    main()