def delete_from_datastore(project, pipeline_options, run_locally):
"""Creates a pipeline that reads entities from Cloud Datastore."""
p = beam.Pipeline(options=pipeline_options)
# Create a query to read entities from datastore.
client = datastore.Client()
if run_locally:
pass
#q.add_filter('category', '=', 'BEBOP')
q = client.query(kind='PRDebugAttendee')
query.order = ['-created_date']
results = list(q.fetch(1))
if not results:
logging.error('No PRDebugAttendee objects found')
return
newest_date = results[0]['created_date']
logging.info('Deleting elements older than %s', newest_date)
q1 = client.query(kind='PRDebugAttendee')
q2 = client.query(kind='PRCityCategory')
datastore_1 = p | 'read PRDebugAttendee from datastore' >> ReadFromDatastore(
project, query._pb_from_query(q1), num_splits=400)
datastore_2 = p | 'read PRCityCategory from datastore' >> ReadFromDatastore(
project, query._pb_from_query(q2), num_splits=400)
# Set up our map/reduce pipeline
output = (
(datastore_1, datastore_2) | beam.Flatten()
| 'convert to entity' >> beam.Map(ConvertToEntity)
# Find the events we want to count, and expand all the admins/attendees
| 'find old rankings' >> beam.FlatMap(OldPRRecord, newest_date)
# And save it all back to the database
)
if not run_locally:
output | 'delete from datastore' >> beam.ParDo(DeleteFromDatastore())
"""
(output
| 'convert from entity' >> beam.Map(ConvertFromEntity)
| 'write to datastore' >> WriteToDatastore(client.project)
)
"""
# Actually run the pipeline (all operations above are deferred).
result = p.run()
# Wait until completion, main thread would access post-completion job results.
result.wait_until_finish()
return result
def check_estimated_size_bytes(self, entity_bytes, timestamp, namespace=None):
"""A helper method to test get_estimated_size_bytes"""
timestamp_req = helper.make_request(
self._PROJECT, namespace, helper.make_latest_timestamp_query(namespace))
timestamp_resp = self.make_stats_response(
{'timestamp': datastore_helper.from_timestamp(timestamp)})
kind_stat_req = helper.make_request(
self._PROJECT, namespace, helper.make_kind_stats_query(
namespace, self._query.kind[0].name,
datastore_helper.micros_from_timestamp(timestamp)))
kind_stat_resp = self.make_stats_response(
{'entity_bytes': entity_bytes})
def fake_run_query(req):
if req == timestamp_req:
return timestamp_resp
elif req == kind_stat_req:
return kind_stat_resp
else:
print kind_stat_req
raise ValueError("Unknown req: %s" % req)
self._mock_datastore.run_query.side_effect = fake_run_query
self.assertEqual(entity_bytes, ReadFromDatastore.get_estimated_size_bytes(
self._PROJECT, namespace, self._query, self._mock_datastore))
self.assertEqual(self._mock_datastore.run_query.call_args_list,
[call(timestamp_req), call(kind_stat_req)])
def test_SplitQueryFn_without_num_splits(self):
with patch.object(helper,
'get_datastore',
return_value=self._mock_datastore):
# Force SplitQueryFn to compute the number of query splits
num_splits = 0
expected_num_splits = 23
entity_bytes = (expected_num_splits *
ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES)
with patch.object(ReadFromDatastore,
'get_estimated_size_bytes',
return_value=entity_bytes):
def fake_get_splits(datastore,
query,
num_splits,
partition=None):
return self.split_query(query, num_splits)
with patch.object(query_splitter,
'get_splits',
side_effect=fake_get_splits):
split_query_fn = ReadFromDatastore.SplitQueryFn(
self._PROJECT, self._query, None, num_splits)
split_query_fn.start_bundle()
returned_split_queries = []
for split_query in split_query_fn.process(self._query):
returned_split_queries.append(split_query)
self.assertEqual(len(returned_split_queries),
expected_num_splits)
self.assertEqual(
0, len(self._mock_datastore.run_query.call_args_list))
self.verify_unique_keys(returned_split_queries)
def test_SplitQueryFn_with_exception(self):
"""A test that verifies that no split is performed when failures occur."""
with patch.object(helper,
'get_datastore',
return_value=self._mock_datastore):
# Force SplitQueryFn to compute the number of query splits
num_splits = 0
expected_num_splits = 1
entity_bytes = (expected_num_splits *
ReadFromDatastore._DEFAULT_BUNDLE_SIZE_BYTES)
with patch.object(ReadFromDatastore,
'get_estimated_size_bytes',
return_value=entity_bytes):
with patch.object(
query_splitter,
'get_splits',
side_effect=ValueError("Testing query split error")):
split_query_fn = ReadFromDatastore.SplitQueryFn(
self._PROJECT, self._query, None, num_splits)
split_query_fn.start_bundle()
returned_split_queries = []
for split_query in split_query_fn.process(self._query):
returned_split_queries.append(split_query)
self.assertEqual(len(returned_split_queries),
expected_num_splits)
self.assertEqual(returned_split_queries[0][1], self._query)
self.assertEqual(
0, len(self._mock_datastore.run_query.call_args_list))
self.verify_unique_keys(returned_split_queries)
def read_from_datastore(project, user_options, pipeline_options):
"""Creates a pipeline that reads entities from Cloud Datastore."""
p = beam.Pipeline(options=pipeline_options)
# Create a query to read entities from datastore.
query = make_ancestor_query(user_options.kind, user_options.namespace,
user_options.ancestor)
# Read entities from Cloud Datastore into a PCollection.
lines = p | 'read from datastore' >> ReadFromDatastore(
project, query, user_options.namespace)
# Count the occurrences of each word.
counts = (lines
| 'split' >>
(beam.ParDo(WordExtractingDoFn()).with_output_types(unicode))
| 'pair_with_one' >> beam.Map(lambda x: (x, 1))
| 'group' >> beam.GroupByKey()
| 'count' >> beam.Map(lambda (word, ones): (word, sum(ones))))
# Format the counts into a PCollection of strings.
output = counts | 'format' >> beam.Map(lambda (word, c): '%s: %s' %
(word, c))
# Write the output using a "Write" transform that has side effects.
# pylint: disable=expression-not-assigned
output | 'write' >> beam.io.WriteToText(
file_path_prefix=user_options.output,
num_shards=user_options.num_shards)
result = p.run()
# Wait until completion, main thread would access post-completion job results.
result.wait_until_finish()
return result
def model_datastoreio():
"""Using a Read and Write transform to read/write to Cloud Datastore."""
import uuid
from google.cloud.proto.datastore.v1 import entity_pb2
from google.cloud.proto.datastore.v1 import query_pb2
import googledatastore
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.io.gcp.datastore.v1.datastoreio import ReadFromDatastore
from apache_beam.io.gcp.datastore.v1.datastoreio import WriteToDatastore
project = 'my_project'
kind = 'my_kind'
query = query_pb2.Query()
query.kind.add().name = kind
# [START model_datastoreio_read]
p = beam.Pipeline(options=PipelineOptions())
entities = p | 'Read From Datastore' >> ReadFromDatastore(project, query)
# [END model_datastoreio_read]
# [START model_datastoreio_write]
p = beam.Pipeline(options=PipelineOptions())
musicians = p | 'Musicians' >> beam.Create(
['Mozart', 'Chopin', 'Beethoven', 'Vivaldi'])
def to_entity(content):
entity = entity_pb2.Entity()
googledatastore.helper.add_key_path(entity.key, kind, str(uuid.uuid4()))
googledatastore.helper.add_properties(entity, {'content': unicode(content)})
return entity
entities = musicians | 'To Entity' >> beam.Map(to_entity)
entities | 'Write To Datastore' >> WriteToDatastore(project)
def check_estimated_size_bytes(self, entity_bytes, timestamp, namespace=None):
"""A helper method to test get_estimated_size_bytes"""
timestamp_req = helper.make_request(
self._PROJECT, namespace, helper.make_latest_timestamp_query(namespace))
timestamp_resp = self.make_stats_response(
{'timestamp': datastore_helper.from_timestamp(timestamp)})
kind_stat_req = helper.make_request(
self._PROJECT, namespace, helper.make_kind_stats_query(
namespace, self._query.kind[0].name,
datastore_helper.micros_from_timestamp(timestamp)))
kind_stat_resp = self.make_stats_response(
{'entity_bytes': entity_bytes})
def fake_run_query(req):
if req == timestamp_req:
return timestamp_resp
elif req == kind_stat_req:
return kind_stat_resp
else:
print kind_stat_req
raise ValueError("Unknown req: %s" % req)
self._mock_datastore.run_query.side_effect = fake_run_query
self.assertEqual(entity_bytes, ReadFromDatastore.get_estimated_size_bytes(
self._PROJECT, namespace, self._query, self._mock_datastore))
self.assertEqual(self._mock_datastore.run_query.call_args_list,
[call(timestamp_req), call(kind_stat_req)])
def read_from_datastore(user_options, pipeline_options):
"""Creates a pipeline that reads entities from Cloud Datastore."""
p = beam.Pipeline(options=pipeline_options)
# Create a query to read entities from datastore.
query = make_ancestor_query(user_options.inputKind, user_options.namespace,
user_options.ancestor)
# Read entities from Cloud Datastore into a PCollection.
lines = p | 'read from datastore' >> ReadFromDatastore(
user_options.project, query, user_options.namespace)
# Count the occurrences of each word.
def count_ones(word_ones):
(word, ones) = word_ones
return (word, sum(ones))
processedTweets = (
lines
| 'split' >> (beam.ParDo(processTweet()))
| 'create entity' >> beam.Map(
EntityWrapper(user_options.namespace, user_options.outputKind,
user_options.ancestor).make_entity)
| 'write to datastore' >> WriteToDatastore(user_options.project))
result = p.run()
# Wait until completion, main thread would access post-completion job results.
result.wait_until_finish()
return result
def test_SplitQueryFn_with_query_limit(self):
"""A test that verifies no split is performed when the query has a limit."""
with patch.object(helper, 'get_datastore',
return_value=self._mock_datastore):
self._query.limit.value = 3
split_query_fn = ReadFromDatastore.SplitQueryFn(
self._PROJECT, self._query, None, 4)
split_query_fn.start_bundle()
returned_split_queries = []
for split_query in split_query_fn.process(self._query):
returned_split_queries.append(split_query)
self.assertEqual(1, len(returned_split_queries))
self.assertEqual(0, len(self._mock_datastore.method_calls))
def expand(self, pcoll):
query = query_pb2.Query()
query.kind.add().name = 'Tweet'
now = datetime.datetime.now()
# The 'earlier' var will be set to a static value on template creation.
# That is, because of the way that templates work, the value is defined
# at template compile time, not runtime.
# But defining a filter based on this value will still serve to make the
# query more efficient than if we didn't filter at all.
earlier = now - datetime.timedelta(days=self.days)
datastore_helper.set_property_filter(query.filter, 'created_at',
PropertyFilter.GREATER_THAN,
earlier)
return (pcoll
| 'read from datastore' >> ReadFromDatastore(
self.project, query, None))
def test_SplitQueryFn_with_num_splits(self):
with patch.object(helper, 'get_datastore',
return_value=self._mock_datastore):
num_splits = 23
def fake_get_splits(datastore, query, num_splits, partition=None):
return self.split_query(query, num_splits)
with patch.object(query_splitter, 'get_splits',
side_effect=fake_get_splits):
split_query_fn = ReadFromDatastore.SplitQueryFn(
self._PROJECT, self._query, None, num_splits)
split_query_fn.start_bundle()
returned_split_queries = []
for split_query in split_query_fn.process(self._query):
returned_split_queries.append(split_query)
self.assertEqual(len(returned_split_queries), num_splits)
self.assertEqual(0, len(self._mock_datastore.run_query.call_args_list))
self.verify_unique_keys(returned_split_queries)
def run(argv=None):
"""Main entry point."""
parser = argparse.ArgumentParser()
parser.add_argument('--kind',
dest='kind',
default='writereadtest',
help='Datastore Kind')
parser.add_argument('--num_entities',
dest='num_entities',
type=int,
required=True,
help='Number of entities to write')
parser.add_argument('--limit',
dest='limit',
type=int,
help='Limit of number of entities to write')
known_args, pipeline_args = parser.parse_known_args(argv)
pipeline_options = PipelineOptions(pipeline_args)
pipeline_options.view_as(SetupOptions).save_main_session = True
gcloud_options = pipeline_options.view_as(GoogleCloudOptions)
job_name = gcloud_options.job_name
kind = known_args.kind
num_entities = known_args.num_entities
project = gcloud_options.project
# a random ancesor key
ancestor = str(uuid.uuid4())
query = make_ancestor_query(kind, None, ancestor)
# Pipeline 1: Create and write the specified number of Entities to the
# Cloud Datastore.
logging.info('Writing %s entities to %s', num_entities, project)
p = new_pipeline_with_job_name(pipeline_options, job_name, '-write')
# pylint: disable=expression-not-assigned
(p
| 'Input' >> beam.Create(list(range(known_args.num_entities)))
| 'To String' >> beam.Map(str)
|
'To Entity' >> beam.Map(EntityWrapper(kind, None, ancestor).make_entity)
| 'Write to Datastore' >> WriteToDatastore(project))
p.run()
# Optional Pipeline 2: If a read limit was provided, read it and confirm
# that the expected entities were read.
if known_args.limit is not None:
logging.info(
'Querying a limited set of %s entities and verifying count.',
known_args.limit)
p = new_pipeline_with_job_name(pipeline_options, job_name,
'-verify-limit')
query_with_limit = query_pb2.Query()
query_with_limit.CopyFrom(query)
query_with_limit.limit.value = known_args.limit
entities = p | 'read from datastore' >> ReadFromDatastore(
project, query_with_limit)
assert_that(entities | beam.combiners.Count.Globally(),
equal_to([known_args.limit]))
p.run()
# Pipeline 3: Query the written Entities and verify result.
logging.info('Querying entities, asserting they match.')
p = new_pipeline_with_job_name(pipeline_options, job_name, '-verify')
entities = p | 'read from datastore' >> ReadFromDatastore(project, query)
assert_that(entities | beam.combiners.Count.Globally(),
equal_to([num_entities]))
p.run()
# Pipeline 4: Delete Entities.
logging.info('Deleting entities.')
p = new_pipeline_with_job_name(pipeline_options, job_name, '-delete')
entities = p | 'read from datastore' >> ReadFromDatastore(project, query)
# pylint: disable=expression-not-assigned
(entities
| 'To Keys' >> beam.Map(lambda entity: entity.key)
| 'Delete keys' >> DeleteFromDatastore(project))
p.run()
# Pipeline 5: Query the written Entities, verify no results.
logging.info(
'Querying for the entities to make sure there are none present.')
p = new_pipeline_with_job_name(pipeline_options, job_name,
'-verify-deleted')
entities = p | 'read from datastore' >> ReadFromDatastore(project, query)
assert_that(entities | beam.combiners.Count.Globally(), equal_to([0]))
p.run()
def run():
import pickle
import sys
import math
import numpy as np
reload(sys)
sys.setdefaultencoding('utf8')
from gensim.models import KeyedVectors
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions, GoogleCloudOptions, StandardOptions, SetupOptions
from apache_beam.io.gcp.datastore.v1.datastoreio import ReadFromDatastore
from google.cloud.proto.datastore.v1 import query_pb2
from apache_beam.io.textio import WriteToText
import nltk.data
import re
import uuid
import perceptron
_sentence_tokenizer = nltk.data.load("./tokenizer/punkt_turkish.pickle")
abbreviations = set()
with open("./tokenizer/abbreviations-long.txt") as f:
for l in f:
abbreviations.add(l.split(':')[0])
_sentence_tokenizer._params.abbrev_types = abbreviations
model_file = "perceptron_word2vec_stemmed_normalized.pickle"
with open(model_file, 'rb') as model:
w, b = pickle.load(model)
def sentences_from_text(text):
return _sentence_tokenizer.tokenize(text.strip())
def tokens_from_sentence(sentence):
return sentence.split(" ") # nltk.word_tokenize(sentence)
def ngrams(obj, n):
tokens = []
sentences = (
sentences_from_text(obj["title"]) +
sentences_from_text(obj["description"]) +
sentences_from_text(obj["content"])
)
for sentence in sentences:
tokens += tokens_from_sentence(sentence)
pairs = nltk.ngrams(tokens, n)
return [" ".join(pair) for pair in pairs]
def convertToObject(jsonObj):
x = jsonObj
link = x.properties.get('link', None)
link = link.string_value if link else ""
title = x.properties.get('title', None)
title = title.string_value if title else ""
description = x.properties.get("description", None)
description = description.string_value if description else ""
content = x.properties.get("text", "")
content = content.string_value if content else ""
published = x.properties.get("published")
published = published.string_value if published else ""
obj = {
"link": link,
"title": title,
"description": description,
"content": content,
"published": published
}
obj["key"] = obj["link"] if obj["link"] else str(uuid.uuid4())
return obj
# https://stackoverflow.com/questions/9662346/python-code-to-remove-html-tags-from-a-string
def cleanhtml(raw_html):
cleanr = re.compile('<.*?>')
cleantext = re.sub(cleanr, '', raw_html)
return cleantext
def removeHTMLFromStrings(obj):
for key in obj.keys():
obj[key] = cleanhtml(obj[key])
return obj
def tokenize_to_sentences(obj):
obj["sentences"] = (
sentences_from_text(obj["title"]) +
sentences_from_text(obj["description"]) +
sentences_from_text(obj["content"])
)
return obj
def tokenize_to_words(obj):
obj["tokens"] = []
for sentence in obj["sentences"]:
obj["tokens"] += tokens_from_sentence(sentence)
for token in obj["tokens"]:
yield (obj["key"], token)
options = PipelineOptions()
google_cloud_options = options.view_as(GoogleCloudOptions)
google_cloud_options.project = 'news-197916'
google_cloud_options.job_name = 'sentiment-analysis'
google_cloud_options.staging_location = 'gs://news-197916.appspot.com/word_count/'
google_cloud_options.temp_location = 'gs://news-197916.appspot.com/df_tmp'
options.view_as(StandardOptions).runner = 'DataflowRunner'
setup_options = options.view_as(SetupOptions)
setup_options.requirements_file = "requirements.txt"
setup_options.save_main_session = True
p = beam.Pipeline(options=options)
query = query_pb2.Query()
query.kind.add().name = "News_Entry"
pairs = (p
| 'Read From Datastore' >> ReadFromDatastore(project = google_cloud_options.project, query=query)
# | "Read From Text" >> ReadFromText("news.json", coder=beam.coders.coders.StrUtf8Coder()) # line by line
# | "Convert to Json Object" >> beam.Map(convertToJsonObj)
| "Convert to Python Object" >> beam.Map(convertToObject)
| "Remove HTML Tags From Strings (Normalization 1)" >> beam.Map(removeHTMLFromStrings)
)
tokens_1gram = (pairs
| 'Sentence Tokenization' >> beam.Map(tokenize_to_sentences)
| 'Word Tokenization' >> beam.FlatMap(tokenize_to_words) # also convert to key value pairs
)
"""
tokens_2gram = (pairs
| "Create 2-grams" >> beam.FlatMap(lambda obj: [(obj["key"], token) for token in ngrams(obj, 2)])
)
"""
tokens = tokens_1gram
"""
vocabulary = (tokens
| "Get words only" >> beam.Values()
| "Remove duplicate words" >> beam.RemoveDuplicates()
)
vocabulary_size = (vocabulary
| "Count Vocabulary elements" >> beam.combiners.Count.Globally()
)
doc_total_words = (tokens
| "Count Words of Doc" >> beam.combiners.Count.PerKey()
)
"""
tokens_paired_with_1 = (tokens
| "Pair with 1" >> beam.Map(lambda (doc, token): ((doc, token), 1))
)
"""
token_counts_per_doc = (tokens_paired_with_1
| "Group by Doc,Word" >> beam.GroupByKey()
| "Count ones" >> beam.Map(lambda ((doc, token), counts): (doc, (token, sum(counts))))
| "Group by Doc" >> beam.GroupByKey()
)
num_docs = (token_counts_per_doc
| "Get Docs" >> beam.Keys()
| "Count Docs" >> beam.combiners.Count.Globally()
)
word_tf_pre = (
{ 'total_tokens': doc_total_words, 'token_counts_per_doc': token_counts_per_doc }
| "CoGroup By Document" >> beam.CoGroupByKey()
)
def calc_tf((doc, count)):
[token_count] = count['token_counts_per_doc']
[tokens_total] = count['total_tokens']
for token, cnt in token_count:
yield token, (doc, float(cnt) / tokens_total)
doc_word_tf = (word_tf_pre
| "Compute Term Frequencies" >> beam.FlatMap(calc_tf)
)
word_occurrences = (tokens
| "Remove Multiple occurrences per doc" >> beam.RemoveDuplicates()
| "Pair with 1s" >> beam.Map(lambda (doc, word): (word, 1))
| "Group by Word" >> beam.GroupByKey()
| "Sum 1s" >> beam.Map(lambda (word, counts): (word, sum(counts)))
)
token_df = (
word_occurrences
| "Compute Document Frequency">> beam.Map(lambda (token, count), total: (token, float(count) / total), AsSingleton(num_docs)))
token_tf_df = (
{ 'term_frequency': doc_word_tf, 'document_frequency': token_df}
| "CoGroup By Token" >> beam.CoGroupByKey())
def calc_tfidf((token, tfdf)):
[df] = tfdf['document_frequency']
for doc, tf in tfdf['term_frequency']:
yield (doc, token), tf * math.log(1.0 / df)
token_tf_idf = (token_tf_df
| "Calculate TF-IDF Scores" >> beam.FlatMap(calc_tfidf)
)
"""
word2vec = KeyedVectors.load_word2vec_format('tr_word2vec', binary=True)
def get_vec(word2vec, token):
if word2vec is None:
word2vec = KeyedVectors.load_word2vec_format('tr_word2vec', binary=True)
try:
x = word2vec.get_vector(token)
x = x.reshape(400)
except:
x = np.zeros(400)
return x
def analyze_sentiment(x):
res = perceptron.f(x, w, b)
return res
doc_sentiment = (tokens_paired_with_1
| "Create Word2Vec Vector" >> beam.Map(lambda ((doc, token), cnt): (doc, get_vec(word2vec, token)))
| "Group Word2Vec Vectors By Document" >> beam.GroupByKey()
| "Sum Word2Vec Vectors" >> beam.Map(
lambda (doc, vecs): (doc, analyze_sentiment(np.sum(vecs, axis=0))[0]))
)
result = (doc_sentiment |
"Format Results" >> beam.Map(lambda (doc, tokens): '%s %s' % (doc, tokens))
)
(result
| "Write Results" >> WriteToText("sentiments")
)
p.run()
def process_datastore_tweets(project, dataset, pipeline_options):
"""Creates a pipeline that reads tweets from Cloud Datastore from the last
N days. The pipeline finds the top most-used words, the top most-tweeted
URLs, ranks word co-occurrences by an 'interestingness' metric (similar to
on tf* idf).
"""
ts = str(datetime.datetime.utcnow())
p = beam.Pipeline(options=pipeline_options)
# Create a query to read entities from datastore.
query = make_query('Tweet')
# Read entities from Cloud Datastore into a PCollection.
lines = (p
|
'read from datastore' >> ReadFromDatastore(project, query, None))
global_count = AsSingleton(
lines
| 'global count' >> beam.combiners.Count.Globally())
# Count the occurrences of each word.
percents = (lines
| 'split' >>
(beam.ParDo(WordExtractingDoFn()).with_output_types(unicode))
| 'pair_with_one' >> beam.Map(lambda x: (x, 1))
| 'group' >> beam.GroupByKey()
| 'count' >> beam.Map(lambda (word, ones): (word, sum(ones)))
| 'in tweets percent' >> beam.Map(
lambda (word, wsum), gc:
(word, float(wsum) / gc), global_count))
top_percents = (
percents
| 'top 500' >> combiners.Top.Of(500, lambda x, y: x[1] < y[1]))
# Count the occurrences of each expanded url in the tweets
url_counts = (
lines
| 'geturls' >>
(beam.ParDo(URLExtractingDoFn()).with_output_types(unicode))
| 'urls_pair_with_one' >> beam.Map(lambda x: (x, 1))
| 'urls_group' >> beam.GroupByKey()
| 'urls_count' >> beam.Map(lambda (word, ones): (word, sum(ones)))
| 'urls top 300' >> combiners.Top.Of(300, lambda x, y: x[1] < y[1]))
# Define some inline helper functions.
def join_cinfo(cooccur, percents):
"""Calculate a co-occurence ranking."""
import math
word1 = cooccur[0][0]
word2 = cooccur[0][1]
try:
word1_percent = percents[word1]
weight1 = 1 / word1_percent
word2_percent = percents[word2]
weight2 = 1 / word2_percent
return (cooccur[0], cooccur[1],
cooccur[1] * math.log(min(weight1, weight2)))
except:
return 0
def generate_cooccur_schema():
"""BigQuery schema for the word co-occurrence table."""
json_str = json.dumps({
'fields': [{
'name': 'w1',
'type': 'STRING',
'mode': 'NULLABLE'
}, {
'name': 'w2',
'type': 'STRING',
'mode': 'NULLABLE'
}, {
'name': 'count',
'type': 'INTEGER',
'mode': 'NULLABLE'
}, {
'name': 'log_weight',
'type': 'FLOAT',
'mode': 'NULLABLE'
}, {
'name': 'ts',
'type': 'TIMESTAMP',
'mode': 'NULLABLE'
}]
})
return parse_table_schema_from_json(json_str)
def generate_url_schema():
"""BigQuery schema for the urls count table."""
json_str = json.dumps({
'fields': [{
'name': 'url',
'type': 'STRING',
'mode': 'NULLABLE'
}, {
'name': 'count',
'type': 'INTEGER',
'mode': 'NULLABLE'
}, {
'name': 'ts',
'type': 'TIMESTAMP',
'mode': 'NULLABLE'
}]
})
return parse_table_schema_from_json(json_str)
def generate_wc_schema():
"""BigQuery schema for the word count table."""
json_str = json.dumps({
'fields': [{
'name': 'word',
'type': 'STRING',
'mode': 'NULLABLE'
}, {
'name': 'percent',
'type': 'FLOAT',
'mode': 'NULLABLE'
}, {
'name': 'ts',
'type': 'TIMESTAMP',
'mode': 'NULLABLE'
}]
})
return parse_table_schema_from_json(json_str)
# Now build the rest of the pipeline.
# Calculate the word co-occurence scores.
cooccur_rankings = (
lines
| 'getcooccur' >> (beam.ParDo(CoOccurExtractingDoFn()))
| 'co_pair_with_one' >> beam.Map(lambda x: (x, 1))
| 'co_group' >> beam.GroupByKey()
| 'co_count' >> beam.Map(lambda (wordts, ones): (wordts, sum(ones)))
| 'weights' >> beam.Map(join_cinfo, AsDict(percents))
| 'co top 300' >> combiners.Top.Of(300, lambda x, y: x[2] < y[2]))
# Format the counts into a PCollection of strings.
wc_records = top_percents | 'format' >> beam.FlatMap(
lambda x: [{
'word': xx[0],
'percent': xx[1],
'ts': ts
} for xx in x])
url_records = url_counts | 'urls_format' >> beam.FlatMap(
lambda x: [{
'url': xx[0],
'count': xx[1],
'ts': ts
} for xx in x])
co_records = cooccur_rankings | 'co_format' >> beam.FlatMap(
lambda x: [{
'w1': xx[0][0],
'w2': xx[0][1],
'count': xx[1],
'log_weight': xx[2],
'ts': ts
} for xx in x])
# Write the results to three BigQuery tables.
wc_records | 'wc_write_bq' >> beam.io.Write(
beam.io.BigQuerySink(
'%s:%s.word_counts' % (project, dataset),
schema=generate_wc_schema(),
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND))
url_records | 'urls_write_bq' >> beam.io.Write(
beam.io.BigQuerySink(
'%s:%s.urls' % (project, dataset),
schema=generate_url_schema(),
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND))
co_records | 'co_write_bq' >> beam.io.Write(
beam.io.BigQuerySink(
'%s:%s.word_cooccur' % (project, dataset),
schema=generate_cooccur_schema(),
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND))
# Actually run the pipeline.
return p.run()
def run_pipeline(project, pipeline_options, run_locally, debug_attendees):
"""Creates a pipeline that reads entities from Cloud Datastore."""
p = beam.Pipeline(options=pipeline_options)
# Create a query to read entities from datastore.
client = datastore.Client()
q = client.query(kind='DBEvent')
if run_locally:
q.key_filter(client.key('DBEvent', '999'), '>')
q.key_filter(client.key('DBEvent', 'A'), '<')
# Let's build a timestamp to save all our objects with
timestamp = datetime.datetime.now()
# Set up our map/reduce pipeline
produce_attendees = (
p | 'read from datastore' >> ReadFromDatastore(project, query._pb_from_query(q), num_splits=400) |
'convert to entity' >> beam.Map(ConvertToEntity)
# Find the events we want to count, and expand all the admins/attendees
| 'filter events' >> beam.FlatMap(CountableEvent) | 'load fb attending' >> beam.ParDo(GetEventAndAttending()) |
'export attendees' >> beam.FlatMap(ExportPeople)
)
top_attendee_lists = (
produce_attendees | 'map category -> person' >> beam.FlatMap(GroupPeopleByCategory) | 'group by category' >> beam.GroupByKey() |
'build top-people lists' >> beam.FlatMap(CountPeopleInfos)
)
if debug_attendees:
attendee_event_debugging = (
produce_attendees | 'map city-attendee -> event' >> beam.FlatMap(DebugExportEventPeopleForGrouping) |
'group by city-attendee' >> beam.GroupByKey() |
'within city-attendee, group event_ids by admin_hash' >> beam.FlatMap(DebugGroupEventIds)
)
exploded_top_attendees = (
top_attendee_lists |
'explode the top attendees into a mapping: category-attendee -> YES' >> beam.FlatMap(DebugExplodeAttendeeList)
# We don't deal with duplicates, since it requires the objects (ie our dicts) to be hashable
# Instead, we rely on DebugFilterForTopAttendee to filter out duplicates created by the above
# | 'remove duplicates from multiple overlapping attendee-lists' >> beam.RemoveDuplicates()
)
(
# These both have the same keys:
# key contains {person_type, city, category, person_id}
(attendee_event_debugging, exploded_top_attendees) | beam.Flatten()
# keys are {city, person_id}
| 'group the attendee-debug info with the is-it-a-top-attendee info' >> beam.GroupByKey() |
'filter for TOP_ATTENDEE' >> beam.FlatMap(DebugFilterForTopAttendee) |
'build PRDebugAttendee' >> beam.ParDo(DebugBuildPRDebugAttendee(), timestamp) |
'write PRDebugAttendee to datastore (unbatched)' >> beam.ParDo(WriteToDatastoreSingle(), actually_save=not run_locally)
)
(
top_attendee_lists |
'generate PRCityCategory database record' >> beam.ParDo(BuildPRCityCategory(), timestamp, 'PRCityCategory', TOP_ALL_N) |
'write PRCityCategory to datastore (unbatched)' >> beam.ParDo(WriteToDatastoreSingle(), actually_save=not run_locally)
)
"""
(output
| 'convert from entity' >> beam.Map(ConvertFromEntity)
| 'write to datastore' >> WriteToDatastore(client.project)
)
"""
# Actually run the pipeline (all operations above are deferred).
result = p.run()
# Wait until completion, main thread would access post-completion job results.
result.wait_until_finish()
return result
