from wmfdata.spark import create_session
from pyspark.sql import functions as F

spark = create_session(app_name='shh')
q = """SELECT mediawiki.database AS wiki, experiment.assigned AS group, action, action_subtype, action_source, action_context
FROM event.product_metrics_web_base
WHERE datacenter = 'eqiad'
AND year = 2026
AND month = 5
AND day >= 19
AND experiment.enrolled = 'share-highlight'
"""
ddf = spark.sql(q)

# Inspect ddf

q = """SELECT mediawiki.database AS wiki, action, action_source, action_context, action_subtype, experiment.assigned AS group
FROM event.product_metrics_web_base
WHERE year = 2026
AND month = 5
AND day >= 19
AND experiment.enrolled = 'share-highlight'
AND mediawiki.database IN ('zhwiki', 'viwiki')
AND action_source = 'selection_share'
"""
ddf = spark.sql(q)
selection_share = ddf.groupBy('wiki', 'action_source', 'action_context').agg(F.collect_set('action').alias('action'), F.collect_set('action_subtype').alias('action_subtype'), F.collect_set('group').alias('group'), F.count('*').alias('events')).toPandas()
selection_share.sort_values('wiki')

     wiki    action_source       action_context                              action action_subtype        group  events
3  viwiki  selection_share  download_share_card                             [click]             []  [treatment]       1
4  viwiki  selection_share                share                             [click]             []  [treatment]       1
6  viwiki  selection_share      share_initiated                             [click]             []  [treatment]       6
7  viwiki  selection_share                 None                   [share_abandoned]             []  [treatment]       4
0  zhwiki  selection_share                 None  [share_abandoned, share_completed]             []  [treatment]      23
1  zhwiki  selection_share                share                             [click]             []  [treatment]       1
2  zhwiki  selection_share  download_share_card                             [click]             []  [treatment]       1
5  zhwiki  selection_share      share_initiated                             [click]             []  [treatment]      26

q = """SELECT mediawiki.database AS wiki, action, action_source, action_context, action_subtype, experiment.assigned AS group
FROM event.product_metrics_web_base
WHERE year = 2026
AND month = 5
AND day >= 19
AND experiment.enrolled = 'share-highlight'
AND action_source = 'selection_share'
AND action_context = 'copy_share_link'
"""
ddf = spark.sql(q)
ddf.count()

0